Research Assistant with Plan-Act-Reflect and a Structured Report in 10xGraph

"""Research assistant: plan, search, reflect, then write a structured report.

Claude plans, GPT-5 mini searches and reads pages, Gemini reviews, DeepSeek writes
the JSON report. Progress streams to the terminal and the report is saved as markdown.

Run from the repo root:
    python examples/use_cases/research_assistant/main.py "How do solid-state batteries work?"
"""

from __future__ import annotations

import asyncio
import sys
from pathlib import Path

from dotenv import load_dotenv
from graph import build_graph
from report import parse_report, save_markdown, to_markdown

from tenxgraph.core.state import Message
from tenxgraph.core.state.stream_chunks import StreamEvent


load_dotenv()

OUTPUT = Path(__file__).parent / "output" / "report.md"
DEFAULT_QUESTION = "What are the main trade-offs between solid-state and lithium-ion batteries?"


async def main() -> None:
    question = " ".join(sys.argv[1:]) or DEFAULT_QUESTION
    app = build_graph()
    node = ""
    deltas: list[str] = []  # REPORT text, collected instead of printed
    final_report = ""

    async for chunk in app.astream({"messages": [Message.text_message(question)]}):
        if chunk.event == StreamEvent.UPDATES and chunk.data.get("status") == "node_invoked":
            node = chunk.data["node"]
            print(f"\n--- {node} ---")
        elif chunk.event == StreamEvent.ERROR:
            print("error:", chunk.data)
        elif chunk.event == StreamEvent.MESSAGE and chunk.message is not None:
            msg = chunk.message
            if node == "REPORT":
                if msg.delta:
                    deltas.append(msg.text())
                elif msg.role == "assistant":
                    final_report = msg.text()
            elif msg.delta:
                print(msg.text(), end="", flush=True)
            elif msg.tools_calls:
                names = [c["function"]["name"] for c in msg.tools_calls]
                print(f"\n[tool calls] {', '.join(names)}")

    report = parse_report(final_report or "".join(deltas))
    path = save_markdown(report, OUTPUT)
    print("\n" + to_markdown(report))
    print(f"saved: {path}")
    await app.aclose()


if __name__ == "__main__":
    asyncio.run(main())

"""Plan -> Act -> Reflect -> Report graph with a different company's model per phase.

`PlanActReflectAgent` runs tools in a plain `ToolNode` and lets you set only the
plan and reflect models, so its ACT phase has no model of its own. To give ACT
its own model (OpenAI) and add a report phase (DeepSeek), this file wires the
same loop by hand with `StateGraph`.
"""

from __future__ import annotations

import os

import prompts
from tools import RESEARCH_TOOLS

from tenxgraph.core.graph import Agent, StateGraph, ToolNode
from tenxgraph.core.state import AgentState, Message
from tenxgraph.utils import END


MAX_ROUNDS = 2  # plan/act/reflect passes before the report is written anyway


class ResearchState(AgentState):
    rounds: int = 0


def stage(agent: Agent, task: str):
    """Run `agent` on a text-only view of the conversation plus a closing instruction.

    Tool calls and tool output are hidden from the view: they are long, and some
    providers reject a tool result they did not request. The real state is
    restored afterwards, so later phases still see everything.
    """

    async def node(state: ResearchState, config: dict):
        original = state.context
        visible = [m for m in original if m.role != "tool" and not m.tools_calls]
        state.context = [*visible, Message.text_message(task)]
        try:
            return await agent.execute(state, config)
        finally:
            state.context = original

    return node


def after_act(state: ResearchState) -> str:
    last = state.context[-1]
    return "TOOL" if last.role == "assistant" and last.tools_calls else "REFLECT"


def after_reflect(state: ResearchState) -> str:
    done = "[done]" in state.context[-1].text().lower()
    return "REPORT" if done or state.rounds + 1 >= MAX_ROUNDS else "NEXT_ROUND"


def next_round(state: ResearchState) -> list[Message]:
    state.rounds += 1
    return []


def build_graph():
    tool_node = ToolNode(RESEARCH_TOOLS)

    planner = Agent(model="claude-sonnet-5-5", system_prompt=prompts.PLAN)
    actor = Agent(model="gpt-5-mini", system_prompt=prompts.ACT, tool_node=tool_node)
    reflector = Agent(model="gemini-2.5-flash", system_prompt=prompts.REFLECT)
    writer = Agent(
        model="deepseek-chat",
        base_url="https://api.deepseek.com",
        api_key=os.environ["DEEPSEEK_API_KEY"],
        system_prompt=prompts.REPORT,
        reasoning_config=None,  # plain chat model; no reasoning parameters
    )

    graph = StateGraph(state=ResearchState())
    graph.add_node("PLAN", stage(planner, prompts.PLAN_TASK))
    graph.add_node("ACT", actor)
    graph.add_node("TOOL", tool_node)
    graph.add_node("REFLECT", stage(reflector, prompts.REFLECT_TASK))
    graph.add_node("NEXT_ROUND", next_round)
    graph.add_node("REPORT", stage(writer, prompts.REPORT_TASK))

    graph.set_entry_point("PLAN")
    graph.add_edge("PLAN", "ACT")
    graph.add_conditional_edges("ACT", after_act, {"TOOL": "TOOL", "REFLECT": "REFLECT"})
    graph.add_edge("TOOL", "ACT")
    graph.add_conditional_edges(
        "REFLECT", after_reflect, {"REPORT": "REPORT", "NEXT_ROUND": "NEXT_ROUND"}
    )
    graph.add_edge("NEXT_ROUND", "PLAN")
    graph.add_edge("REPORT", END)
    return graph.compile()

"""Prompts for the four research phases and the stage instructions."""

from __future__ import annotations

import json

from models import ResearchReport


def _prompt(text: str) -> list[dict[str, str]]:
    return [{"role": "system", "content": text}]


PLAN = _prompt(
    "You are a research planner. Turn the user's question into 3 to 5 concrete research "
    "steps. Each step names what to find out and a suggested search query. If an earlier "
    "evaluation lists gaps, plan only for those gaps. Do not answer the question."
)

ACT = _prompt(
    "You are a research assistant who executes plans. Carry out the most recent plan: "
    "call web_search for each step and fetch_page to read promising sources. When you "
    "have enough, reply with plain-text notes: one short paragraph per step, each fact "
    "followed by its source URL. Never invent URLs or facts."
)

REFLECT = _prompt(
    "You are a strict reviewer. Check the research notes against the original question. "
    "If every part of the question is answered with sourced facts, end your reply with "
    "the exact token [DONE]. Otherwise list the specific gaps and do not write [DONE]."
)

# The braces in the JSON schema are doubled because system prompts are formatted
# with str.format() to fill in state fields.
_SCHEMA = json.dumps(ResearchReport.model_json_schema()).replace("{", "{{").replace("}", "}}")

REPORT = _prompt(
    "You write the final research report. Use only facts from the research notes. "
    "Reply with one JSON object and nothing else (no code fences) that matches this "
    f"JSON Schema:\n{_SCHEMA}"
)

# Added to the end of each phase's view of the conversation (see graph.stage).
PLAN_TASK = "Write the research plan now, or a revised plan if gaps were reported."
REFLECT_TASK = "Evaluate the research notes against the original question."
REPORT_TASK = "Write the final report as a JSON object."

"""Research tools for the ACT phase.

Both wrap prebuilt tools from `tenxgraph.prebuilt.tools` and only shrink the
output, so a few fetched pages do not fill the context window.

`google_web_search` calls Gemini with Google Search grounding through its own
google-genai client. That is independent of the model running the ACT phase, so
it works with an OpenAI model. It needs GEMINI_API_KEY or GOOGLE_API_KEY.
To use another search service, replace `web_search` with any async function
that returns text, for example one that calls the Tavily or Brave API.
"""

from __future__ import annotations

from tenxgraph.prebuilt.tools import fetch_url, google_web_search


async def web_search(query: str) -> str:
    """Search the web and return a grounded summary with sources as JSON."""
    return await google_web_search(query, max_chars=5000)


async def fetch_page(url: str) -> str:
    """Download a public web page and return its text (truncated) as JSON."""
    return await fetch_url(url, max_chars=5000)


RESEARCH_TOOLS = [web_search, fetch_page]

"""The structured report the research assistant produces."""

from __future__ import annotations

from pydantic import BaseModel, Field


class Finding(BaseModel):
    claim: str = Field(description="One specific, checkable statement.")
    evidence: str = Field(description="What in the sources supports the claim.")
    source_url: str = Field(description="URL of the page that supports the claim.")


class ResearchReport(BaseModel):
    title: str
    summary: str = Field(description="Three to five sentences answering the question.")
    key_findings: list[Finding]
    open_questions: list[str] = Field(default_factory=list)
    sources: list[str] = Field(default_factory=list, description="All URLs used.")

"""Parse the DeepSeek output into a `ResearchReport` and save it as markdown."""

from __future__ import annotations

import re
from pathlib import Path

from models import ResearchReport


def parse_report(text: str) -> ResearchReport:
    """Validate the model's JSON. Raises pydantic.ValidationError if it is malformed."""
    cleaned = re.sub(r"^```(?:json)?\s*|\s*```$", "", text.strip())
    return ResearchReport.model_validate_json(cleaned)


def to_markdown(report: ResearchReport) -> str:
    lines = [f"# {report.title}", "", report.summary, "", "## Key findings", ""]
    for i, f in enumerate(report.key_findings, 1):
        lines += [
            f"{i}. **{f.claim}**",
            f"   - Evidence: {f.evidence}",
            f"   - Source: {f.source_url}",
        ]
    if report.open_questions:
        lines += ["", "## Open questions", ""] + [f"- {q}" for q in report.open_questions]
    if report.sources:
        lines += ["", "## Sources", ""] + [f"- {s}" for s in report.sources]
    return "\n".join(lines) + "\n"


def save_markdown(report: ResearchReport, path: Path) -> Path:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(to_markdown(report), encoding="utf-8")
    return path