|
|
| from __future__ import annotations
|
|
|
| from typing import Any
|
|
|
| from langchain_core.runnables import RunnableConfig
|
| from langgraph.config import get_stream_writer
|
| from langgraph.types import Send
|
|
|
| from ..llm import build_llm
|
| from ..state import AgentState, CandidateCountry, TodoItem
|
| from .config import MAX_TODOS, PLANNER_MAX_TOKENS, PLANNER_TEMPERATURE
|
| from .helpers import (
|
| extract_assistant_text,
|
| extract_json,
|
| heuristic_candidate_countries,
|
| normalize_plan,
|
| user_text,
|
| )
|
| from .prompts import PLANNER_SYSTEM_PROMPT
|
|
|
|
|
| def _format_candidate_shortlist(candidates: list[CandidateCountry]) -> str:
|
| lines = ["Discovery shortlist (preferred starting countries):"]
|
| for item in candidates:
|
| lines.append(
|
| f"- {item['iso2']} {item['name']}: {item['pathway_hint']} ({item['label']})"
|
| )
|
| return "\n".join(lines)
|
|
|
|
|
| def planner_node(state: AgentState, config: RunnableConfig) -> dict[str, Any]:
|
| writer = get_stream_writer()
|
| llm = build_llm(
|
| config,
|
| max_tokens=PLANNER_MAX_TOKENS,
|
| temperature=PLANNER_TEMPERATURE,
|
| )
|
|
|
| profile_text = user_text(state["user_content"])
|
| candidates = state.get("candidate_countries") or heuristic_candidate_countries(
|
| profile_text
|
| )
|
| discovery_summary = str(state.get("discovery_summary") or "").strip()
|
| profile_summary = str(state.get("profile_summary") or "").strip()
|
|
|
| planner_context = "\n\n".join(
|
| part
|
| for part in [
|
| _format_candidate_shortlist(candidates),
|
| f"Discovery notes:\n{discovery_summary}" if discovery_summary else "",
|
| f"Profile summary:\n{profile_summary}" if profile_summary else "",
|
| ]
|
| if part
|
| )
|
|
|
| messages: list[Any] = [
|
| {"role": "system", "content": PLANNER_SYSTEM_PROMPT},
|
| *state.get("history_messages", []),
|
| {
|
| "role": "user",
|
| "content": (
|
| f"{planner_context}\n\n"
|
| f"Original user request:\n{profile_text}\n\n"
|
| "Produce the JSON research plan now."
|
| ),
|
| },
|
| ]
|
|
|
| raw_plan: dict[str, Any] | None = None
|
| for _ in range(2):
|
| response = llm.invoke(messages)
|
| raw_plan = extract_json(extract_assistant_text(response))
|
| if raw_plan and raw_plan.get("todos"):
|
| break
|
| raw_plan = None
|
|
|
| plan = normalize_plan(raw_plan, profile_text, candidates)
|
|
|
| thinking = str(plan.get("thinking") or "").strip()
|
| if thinking:
|
| writer({"type": "thinking", "text": thinking})
|
|
|
| todos: list[TodoItem] = list(plan.get("todos") or [])
|
| if not todos:
|
| plan = normalize_plan(None, profile_text, candidates)
|
| todos = plan["todos"]
|
|
|
| writer({"type": "plan", "todos": todos})
|
|
|
| countries = [str(code) for code in plan.get("countries") or [] if code]
|
| if countries:
|
| labels = [str(label) for label in plan.get("labels") or []]
|
| writer(
|
| {
|
| "type": "globe",
|
| "args": {"action": "show", "countries": countries, "labels": labels},
|
| }
|
| )
|
|
|
| return {
|
| "todos": todos,
|
| "profile_summary": str(plan.get("profile_summary") or profile_summary),
|
| }
|
|
|
|
|
| def fan_out_research(state: AgentState) -> list[Send]:
|
| return [
|
| Send(
|
| "researcher",
|
| {"todo": todo, "profile_summary": state.get("profile_summary", "")},
|
| )
|
| for todo in state["todos"]
|
| ]
|
|
|
