agent/graph.py

"""LangGraph ReAct loop: state, pure helpers, nodes, graph build, and a streaming runner."""

import json
import time
from typing import Annotated, Optional, TypedDict

from langgraph.config import get_stream_writer
from langgraph.graph import END, START, StateGraph
from google.genai import types

from agent import config, llm


def split_response(response):
    """Return (thought_text, answer_text, function_calls, model_content)."""
    content = response.candidates[0].content
    thoughts, answers, calls = [], [], []
    for part in (content.parts or []):
        fc = getattr(part, "function_call", None)
        if fc:
            calls.append(fc)
            continue
        text = getattr(part, "text", None)
        if not text:
            continue
        if getattr(part, "thought", False):
            thoughts.append(text)
        else:
            answers.append(text)
    return "\n".join(thoughts).strip(), "\n".join(answers).strip(), calls, content


def classify_weak(tool_name: str, observation: str) -> bool:
    """True when an observation should trigger reconsideration."""
    text = (observation or "").strip()
    if text.startswith("ERROR:") or text.startswith("NO_RESULTS"):
        return True
    if tool_name == "calculator":
        return False
    if len(text) < config.WEAK_OBS_MIN_LEN:
        return True
    low = text.lower()
    return any(marker in low for marker in config.WEAK_MARKERS)


def signature(call) -> tuple:
    """Stable identity of a tool call (name + sorted args) for re-query detection."""
    args = dict(getattr(call, "args", None) or {})
    return call.name, json.dumps(args, sort_keys=True, default=str)


class AgentState(TypedDict):
    contents: Annotated[list, lambda a, b: a + b]
    step: int
    last_observation_weak: bool
    weak_call_sigs: set
    pending_calls: list
    final_answer: Optional[str]
    stop_reason: Optional[str]


def build_graph(*, client, tool_fns, declarations, system_prompt,
                max_steps=config.MAX_STEPS, sleep=time.sleep):

    def reason(state: AgentState):
        writer = get_stream_writer()
        step = state["step"] + 1
        if step > max_steps:
            writer({"kind": "limit", "text": "Reached the step limit; stopping.", "step": step})
            return {"step": step, "stop_reason": "limit"}

        writer({"kind": "status", "text": "model is thinking..."})
        response = llm.generate(
            client, contents=state["contents"], tools=declarations,
            system_instruction=system_prompt, sleep=sleep,
        )
        thought, answer, calls, model_content = split_response(response)

        new_sigs = [signature(c) for c in calls]
        is_revision = (
            bool(calls)
            and state["last_observation_weak"]
            and not set(new_sigs).issubset(state["weak_call_sigs"])
        )
        if thought:
            writer({"kind": "thought", "text": thought, "revision": is_revision, "step": step})

        updates = {"step": step, "contents": [model_content]}
        if calls:
            for call in calls:
                writer({"kind": "tool_call", "tool": call.name,
                        "args": dict(call.args or {}), "step": step})
            updates["pending_calls"] = calls
        else:
            final = answer or "(the model produced no answer)"
            writer({"kind": "final", "text": final, "step": step})
            updates["final_answer"] = final
            updates["stop_reason"] = "answered"
        return updates

    def act(state: AgentState):
        writer = get_stream_writer()
        tool_contents, weak_sigs, any_weak = [], set(), False
        for call in state["pending_calls"]:
            name = call.name
            args = dict(call.args or {})
            writer({"kind": "status", "text": f"running {name}..."})
            fn = tool_fns.get(name)
            if fn is None:
                obs = f"ERROR: unknown tool '{name}'"
            else:
                try:
                    obs = str(fn(**args))
                except Exception as exc:
                    obs = f"ERROR: {type(exc).__name__}: {exc}"
            weak = classify_weak(name, obs)
            if weak:
                any_weak = True
                weak_sigs.add(signature(call))
            writer({"kind": "observation", "tool": name, "text": obs,
                    "weak": weak, "step": state["step"]})
            tool_contents.append(types.Content(
                role="tool",
                parts=[types.Part.from_function_response(name=name, response={"result": obs})],
            ))
        return {"contents": tool_contents, "last_observation_weak": any_weak,
                "weak_call_sigs": weak_sigs, "pending_calls": []}

    def route_after_reason(state: AgentState):
        if state.get("stop_reason"):
            return END
        if state.get("pending_calls"):
            return "act"
        return END

    graph = StateGraph(AgentState)
    graph.add_node("reason", reason)
    graph.add_node("act", act)
    graph.add_edge(START, "reason")
    graph.add_conditional_edges("reason", route_after_reason, {"act": "act", END: END})
    graph.add_edge("act", "reason")
    return graph.compile()


def stream_run(task, *, client, tool_fns, declarations, system_prompt,
               max_steps=config.MAX_STEPS, sleep=time.sleep):
    """Yield UI event dicts as the agent runs."""
    compiled = build_graph(
        client=client, tool_fns=tool_fns, declarations=declarations,
        system_prompt=system_prompt, max_steps=max_steps, sleep=sleep,
    )
    initial = {
        "contents": [types.Content(role="user", parts=[types.Part.from_text(text=task)])],
        "step": 0,
        "last_observation_weak": False,
        "weak_call_sigs": set(),
        "pending_calls": [],
        "final_answer": None,
        "stop_reason": None,
    }
    recursion_limit = max_steps * 2 + 5
    # langgraph's stream_mode="custom" yields each writer payload as a bare dict.
    for chunk in compiled.stream(initial, stream_mode="custom",
                                 config={"recursion_limit": recursion_limit}):
        if isinstance(chunk, dict) and "kind" in chunk:
            yield chunk