todo: langgraph integration

requirements.txt

@@ -50,6 +50,9 @@ jiter==0.15.0
 joserfc==1.7.0
 jsonschema==4.26.0
 jsonschema-specifications==2025.9.1
+langchain-core==1.4.6
+langchain-openai==1.3.0
+langgraph==1.2.4
 markdown-it-py==4.2.0
 MarkupSafe==3.0.3
 matplotlib-inline==0.2.2

ui/agent/graph/__init__.py

@@ -0,0 +1,4 @@
+# ui/agent/graph/__init__.py
+from .respond import respond_with_graph
+
+__all__ = ["respond_with_graph"]

ui/agent/graph/llm.py

@@ -0,0 +1,26 @@
+# ui/agent/graph/llm.py
+from __future__ import annotations
+
+from typing import Any
+
+from langchain_core.runnables import RunnableConfig
+from langchain_openai import ChatOpenAI
+
+from ..config import MODEL_ID
+
+HF_ROUTER_BASE_URL = "https://router.huggingface.co/v1"
+
+
+def build_llm(config: RunnableConfig, **overrides: Any) -> ChatOpenAI:
+    """Build a chat model from the per-request configurable values."""
+    configurable = config.get("configurable", {})
+    params: dict[str, Any] = {
+        "model": MODEL_ID,
+        "api_key": configurable["hf_token"],
+        "base_url": HF_ROUTER_BASE_URL,
+        "max_tokens": configurable.get("max_tokens", 1800),
+        "temperature": configurable.get("temperature", 0.35),
+        "top_p": configurable.get("top_p", 0.9),
+    }
+    params.update(overrides)
+    return ChatOpenAI(**params)

ui/agent/graph/nodes.py

@@ -0,0 +1,356 @@
+# ui/agent/graph/nodes.py
+from __future__ import annotations
+
+import json
+import time
+import uuid
+from typing import Any
+
+from langchain_core.messages import AIMessage, ToolMessage
+from langchain_core.runnables import RunnableConfig
+from langgraph.config import get_stream_writer
+from langgraph.types import Send
+
+from ..tool_schemas import (
+    crawl_web_site,
+    get_country_profile,
+    scrape_web_page,
+    search_immigration_info,
+)
+from ..tools import (
+    _done_tool_message,
+    _pending_tool_message,
+    _tool_log_metadata,
+    run_tool,
+    truncate,
+)
+from ..traces import record_tool_trace
+from .llm import build_llm
+from .state import AgentState, Finding, ResearchTask, TodoItem
+
+RESEARCH_TOOL_SCHEMAS = [
+    get_country_profile,
+    search_immigration_info,
+    scrape_web_page,
+    crawl_web_site,
+]
+
+MAX_TODOS = 5
+MAX_RESEARCH_ROUNDS = 4
+TOOL_RESULT_LIMIT = 4000
+
+PLANNER_SYSTEM_PROMPT = """
+You are the planning supervisor of Borderless, an immigration research agency.
+
+Read the user's profile and goals, then produce a focused research plan that a
+team of parallel research analysts will execute. Identify 3-5 plausible
+destination countries (prefer realistic fit over popular destinations) and
+break the research into 3-5 self-contained to-dos. Each to-do must be
+researchable independently — typically one to-do per recommended country
+covering its best visa pathway, eligibility, documents, costs, timelines, and
+risks. You may add one cross-cutting to-do (e.g. comparing costs or document
+preparation) when useful.
+
+Respond with ONLY a JSON object, no other text:
+{
+  "thinking": "brief reasoning about the user's profile and country choices",
+  "countries": ["ISO-2 codes of recommended countries, e.g. CA", "DE"],
+  "labels": ["short marker label per country, e.g. Skilled worker - 6-12 mo"],
+  "profile_summary": "2-3 sentence summary of the user's profile and constraints",
+  "todos": [
+    {"title": "short title", "description": "specific research instructions for the analyst"}
+  ]
+}
+""".strip()
+
+RESEARCHER_SYSTEM_PROMPT = """
+You are a research analyst at Borderless, an immigration research agency.
+You are assigned ONE research to-do. Use your tools to research it thoroughly:
+
+- Use search_immigration_info to find current visa rules, fees, documents,
+  processing times, and official government pages. Prefer official immigration
+  authority, government, and embassy sources.
+- Use scrape_web_page on the best official URLs before making concrete claims.
+- Use get_country_profile when country metadata helps.
+- Do not invent point scores, income thresholds, fees, or processing times.
+
+When you have enough information (after at most a few tool calls), STOP calling
+tools and write a dense findings report in markdown covering: visa pathway(s),
+eligibility, required documents, approximate costs, realistic timeline, risks,
+and the official source URLs you verified. Label unofficial sources clearly.
+Note anything you could not verify.
+""".strip()
+
+CONSOLIDATOR_SYSTEM_PROMPT = """
+You are Borderless, an agentic immigration research assistant. Your research
+team has completed parallel research on each to-do. Consolidate their findings
+into one final answer for the user. You are not a lawyer and must not present
+the answer as legal advice. Be practical, specific, and clear about
+uncertainty. Only state facts supported by the research findings; if something
+is missing, say what still needs verification.
+
+Final answer format:
+Start with a short, plain-English recommendation. Then use these sections:
+
+## Snapshot
+Summarize the user's profile and key constraints in 2-4 bullets.
+
+## Best-Fit Countries
+Use a compact table with country, recommended pathway, why it fits, main risk,
+and rough timeline.
+
+## Pathway Details
+For each recommended country, list the visa pathway, eligibility notes,
+required documents, approximate steps, timeline, and budget sensitivity.
+
+## Documents To Prepare
+Group common documents first, then country-specific documents.
+
+## Risks And Tradeoffs
+Mention language, job market, funds, age/points, credential recognition,
+family constraints, and policy uncertainty where relevant.
+
+## Official Sources
+List cited official URLs and what each source supports.
+
+## Next Steps
+Give 3-5 concrete actions the user can take next. End by reminding them to
+verify details on official sites or with a qualified immigration professional.
+""".strip()
+
+
+def _extract_json(text: str) -> dict[str, Any] | None:
+    decoder = json.JSONDecoder()
+    for start in range(len(text)):
+        if text[start] != "{":
+            continue
+        try:
+            parsed, _ = decoder.raw_decode(text[start:])
+        except json.JSONDecodeError:
+            continue
+        if isinstance(parsed, dict):
+            return parsed
+    return None
+
+
+def _fallback_plan(user_text: str) -> dict[str, Any]:
+    return {
+        "thinking": "",
+        "countries": [],
+        "labels": [],
+        "profile_summary": truncate(user_text, 600),
+        "todos": [
+            {
+                "title": "Research migration options",
+                "description": (
+                    "Research realistic migration options for this profile: "
+                    f"{truncate(user_text, 1200)}"
+                ),
+            }
+        ],
+    }
+
+
+def _user_text(user_content: str | list[dict[str, Any]]) -> str:
+    if isinstance(user_content, str):
+        return user_content
+    parts = [
+        item.get("text", "")
+        for item in user_content
+        if isinstance(item, dict) and item.get("type") == "text"
+    ]
+    return "\n".join(part for part in parts if part)
+
+
+def planner_node(state: AgentState, config: RunnableConfig) -> dict[str, Any]:
+    writer = get_stream_writer()
+    llm = build_llm(config)
+
+    messages: list[Any] = [
+        {"role": "system", "content": PLANNER_SYSTEM_PROMPT},
+        *state.get("history_messages", []),
+        {"role": "user", "content": state["user_content"]},
+    ]
+
+    plan: dict[str, Any] | None = None
+    for _ in range(2):
+        response = llm.invoke(messages)
+        plan = _extract_json(str(response.content or ""))
+        if plan and plan.get("todos"):
+            break
+        plan = None
+
+    user_text = _user_text(state["user_content"])
+    if plan is None:
+        plan = _fallback_plan(user_text)
+
+    thinking = str(plan.get("thinking") or "").strip()
+    if thinking:
+        writer({"type": "thinking", "text": thinking})
+
+    todos: list[TodoItem] = []
+    for index, raw in enumerate(plan.get("todos", [])[:MAX_TODOS]):
+        if not isinstance(raw, dict):
+            continue
+        title = str(raw.get("title") or f"Research task {index + 1}").strip()
+        description = str(raw.get("description") or title).strip()
+        todos.append({"id": index + 1, "title": title, "description": description})
+    if not todos:
+        fallback = _fallback_plan(user_text)["todos"][0]
+        todos = [{"id": 1, "title": fallback["title"], "description": fallback["description"]}]
+
+    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},
+            }
+        )
+
+    profile_summary = str(plan.get("profile_summary") or "").strip() or truncate(
+        user_text, 600
+    )
+    return {"todos": todos, "profile_summary": 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"]
+    ]
+
+
+def _execute_tool_call(
+    writer: Any,
+    todo_title: str,
+    tool_name: str,
+    args: dict[str, Any],
+) -> str:
+    event_id = uuid.uuid4().hex
+    title, pending_message = _pending_tool_message(tool_name, args)
+    writer(
+        {
+            "type": "tool_start",
+            "id": event_id,
+            "title": f"{title} · {todo_title}",
+            "message": pending_message,
+            "log": {"tool": tool_name, "arguments": args},
+        }
+    )
+
+    started = time.monotonic()
+    result, _ = run_tool(tool_name, json.dumps(args), globe_state=None)
+    duration = time.monotonic() - started
+    record_tool_trace(
+        tool_name=tool_name,
+        arguments=json.dumps(args),
+        result=result,
+        duration=duration,
+    )
+
+    writer(
+        {
+            "type": "tool_end",
+            "id": event_id,
+            "title": f"{title} · {todo_title}",
+            "message": _done_tool_message(tool_name, args, result),
+            "duration": duration,
+            "log": _tool_log_metadata(tool_name, args, result),
+        }
+    )
+    return result
+
+
+def researcher_node(task: ResearchTask, config: RunnableConfig) -> dict[str, Any]:
+    writer = get_stream_writer()
+    todo = task["todo"]
+    llm = build_llm(config).bind_tools(RESEARCH_TOOL_SCHEMAS)
+
+    messages: list[Any] = [
+        {"role": "system", "content": RESEARCHER_SYSTEM_PROMPT},
+        {
+            "role": "user",
+            "content": (
+                f"Applicant profile:\n{task.get('profile_summary') or 'Not provided.'}\n\n"
+                f"Your research to-do: {todo['title']}\n\n{todo['description']}"
+            ),
+        },
+    ]
+
+    summary = ""
+    for round_index in range(MAX_RESEARCH_ROUNDS):
+        response: AIMessage = llm.invoke(messages)
+        if not response.tool_calls:
+            summary = str(response.content or "").strip()
+            break
+
+        messages.append(response)
+        for tool_call in response.tool_calls:
+            result = _execute_tool_call(
+                writer, todo["title"], tool_call["name"], tool_call["args"] or {}
+            )
+            messages.append(
+                ToolMessage(
+                    content=truncate(result, TOOL_RESULT_LIMIT),
+                    tool_call_id=tool_call["id"],
+                )
+            )
+
+        if round_index == MAX_RESEARCH_ROUNDS - 1:
+            messages.append(
+                {
+                    "role": "user",
+                    "content": (
+                        "Stop researching now. Write your findings report based "
+                        "on the information gathered so far, labeling anything "
+                        "unverified."
+                    ),
+                }
+            )
+            final = build_llm(config).invoke(messages)
+            summary = str(final.content or "").strip()
+
+    if not summary:
+        summary = "No findings could be produced for this to-do."
+
+    finding: Finding = {
+        "todo_id": todo["id"],
+        "todo_title": todo["title"],
+        "summary": summary,
+    }
+    writer({"type": "finding", "todo_title": todo["title"], "summary": summary})
+    return {"findings": [finding]}
+
+
+def consolidator_node(state: AgentState, config: RunnableConfig) -> dict[str, Any]:
+    findings = sorted(state.get("findings", []), key=lambda item: item["todo_id"])
+    findings_text = "\n\n".join(
+        f"### Finding {item['todo_id']}: {item['todo_title']}\n{item['summary']}"
+        for item in findings
+    )
+
+    llm = build_llm(config, max_tokens=3000)
+    messages: list[Any] = [
+        {"role": "system", "content": CONSOLIDATOR_SYSTEM_PROMPT},
+        *state.get("history_messages", []),
+        {"role": "user", "content": state["user_content"]},
+        {
+            "role": "user",
+            "content": (
+                "Research team findings:\n\n"
+                f"{findings_text or 'No findings were produced.'}\n\n"
+                "Consolidate these findings into the final answer now, following "
+                "the required format."
+            ),
+        },
+    ]
+    response = llm.invoke(messages)
+    answer = str(response.content or "").strip()
+    return {"final_answer": answer}

ui/agent/graph/respond.py

@@ -0,0 +1,210 @@
+# ui/agent/graph/respond.py
+from __future__ import annotations
+
+import time
+from typing import Any
+
+import gradio as gr
+from gradio import ChatMessage
+
+from ui.globe_commands import apply_update_globe
+
+from ..messages import history_to_api_messages, multimodal_input_to_api_content
+from ..respond import (
+    AUTH_REQUIRED_MESSAGE,
+    SESSION_EXPIRED_MESSAGE,
+    _auth_error_message,
+)
+from ..streaming import yield_response
+from .workflow import build_workflow
+
+FAILURE_MESSAGE = (
+    "I gathered research but could not finish the full synthesis. "
+    "Please ask me to summarize the countries already found."
+)
+
+
+def _format_plan_content(todos: list[dict[str, Any]]) -> str:
+    lines = ["Research to-dos:"]
+    for todo in todos:
+        lines.append(f"{todo['id']}. {todo['title']} — {todo['description']}")
+    return "\n".join(lines)
+
+
+class _UiState:
+    """Maps custom graph events onto the running ChatMessage list."""
+
+    def __init__(self, globe_state: dict[str, Any]):
+        self.ui_messages: list[ChatMessage] = []
+        self.globe_state = globe_state
+        self._pending_by_id: dict[str, int] = {}
+
+    def handle(self, event: dict[str, Any]) -> bool:
+        kind = event.get("type")
+        if kind == "thinking":
+            self.ui_messages.append(
+                ChatMessage(
+                    role="assistant",
+                    content=str(event.get("text") or ""),
+                    metadata={"title": "Thinking", "status": "done"},
+                )
+            )
+            return True
+
+        if kind == "plan":
+            todos = event.get("todos") or []
+            self.ui_messages.append(
+                ChatMessage(
+                    role="assistant",
+                    content=_format_plan_content(todos),
+                    metadata={
+                        "title": "Research plan",
+                        "status": "done",
+                        "log": {"tool": "plan_research", "result": todos},
+                    },
+                )
+            )
+            return True
+
+        if kind == "globe":
+            args = event.get("args") or {}
+            result, self.globe_state = apply_update_globe(self.globe_state, args)
+            countries = ", ".join(args.get("countries") or [])
+            self.ui_messages.append(
+                ChatMessage(
+                    role="assistant",
+                    content=f"Marked {countries or 'recommended countries'} on the globe.",
+                    metadata={
+                        "title": "Updating globe",
+                        "status": "done",
+                        "log": {
+                            "tool": "update_globe",
+                            "arguments": args,
+                            "result": result,
+                        },
+                    },
+                )
+            )
+            return True
+
+        if kind == "tool_start":
+            self.ui_messages.append(
+                ChatMessage(
+                    role="assistant",
+                    content=str(event.get("message") or ""),
+                    metadata={
+                        "title": str(event.get("title") or "Tool"),
+                        "status": "pending",
+                        "log": event.get("log") or {},
+                    },
+                )
+            )
+            self._pending_by_id[str(event.get("id"))] = len(self.ui_messages) - 1
+            return True
+
+        if kind == "tool_end":
+            index = self._pending_by_id.pop(str(event.get("id")), None)
+            message = ChatMessage(
+                role="assistant",
+                content=str(event.get("message") or ""),
+                metadata={
+                    "title": str(event.get("title") or "Tool"),
+                    "status": "done",
+                    "duration": event.get("duration"),
+                    "log": event.get("log") or {},
+                },
+            )
+            if index is not None and index < len(self.ui_messages):
+                self.ui_messages[index] = message
+            else:
+                self.ui_messages.append(message)
+            return True
+
+        if kind == "finding":
+            self.ui_messages.append(
+                ChatMessage(
+                    role="assistant",
+                    content=str(event.get("summary") or ""),
+                    metadata={
+                        "title": f"Findings · {event.get('todo_title') or 'research'}",
+                        "status": "done",
+                    },
+                )
+            )
+            return True
+
+        return False
+
+
+def respond_with_graph(
+    message: str | dict[str, Any],
+    history: list[dict[str, Any]],
+    system_message: str,
+    max_tokens: int,
+    temperature: float,
+    top_p: float,
+    globe_state: dict[str, Any],
+    hf_token: gr.OAuthToken | None,
+):
+    """LangGraph-backed drop-in replacement for ui.agent.respond.respond."""
+    if hf_token is None:
+        yield from yield_response([], AUTH_REQUIRED_MESSAGE, globe_state)
+        return
+
+    if hf_token.expires_at <= time.time():
+        yield from yield_response([], SESSION_EXPIRED_MESSAGE, globe_state)
+        return
+
+    user_content = multimodal_input_to_api_content(message)
+    if not user_content:
+        yield from yield_response(
+            [],
+            "Please enter a message or attach a file.",
+            globe_state,
+        )
+        return
+
+    workflow = build_workflow()
+    ui_state = _UiState(globe_state)
+    final_answer = ""
+
+    config = {
+        "configurable": {
+            "hf_token": hf_token.token,
+            "max_tokens": max_tokens,
+            "temperature": temperature,
+            "top_p": top_p,
+        },
+        "recursion_limit": 50,
+    }
+    graph_input = {
+        "user_content": user_content,
+        "history_messages": history_to_api_messages(history),
+        "findings": [],
+    }
+
+    try:
+        for mode, payload in workflow.stream(
+            graph_input,
+            config=config,
+            stream_mode=["custom", "values"],
+        ):
+            if mode == "custom":
+                if ui_state.handle(payload):
+                    yield list(ui_state.ui_messages), ui_state.globe_state
+            elif mode == "values":
+                final_answer = payload.get("final_answer") or final_answer
+    except Exception as exc:
+        error_message = _auth_error_message(exc) or (
+            f"Sorry, something went wrong while generating a response: {exc}"
+        )
+        yield from yield_response(
+            ui_state.ui_messages, error_message, ui_state.globe_state
+        )
+        return
+
+    yield from yield_response(
+        ui_state.ui_messages,
+        final_answer or FAILURE_MESSAGE,
+        ui_state.globe_state,
+    )

ui/agent/graph/state.py

@@ -0,0 +1,35 @@
+# ui/agent/graph/state.py
+from __future__ import annotations
+
+import operator
+from typing import Annotated, Any, TypedDict
+
+
+class TodoItem(TypedDict):
+    id: int
+    title: str
+    description: str
+
+
+class Finding(TypedDict):
+    todo_id: int
+    todo_title: str
+    summary: str
+
+
+class AgentState(TypedDict, total=False):
+    """Top-level workflow state."""
+
+    user_content: str | list[dict[str, Any]]
+    history_messages: list[dict[str, Any]]
+    profile_summary: str
+    todos: list[TodoItem]
+    findings: Annotated[list[Finding], operator.add]
+    final_answer: str
+
+
+class ResearchTask(TypedDict):
+    """Payload sent to each parallel researcher via Send."""
+
+    todo: TodoItem
+    profile_summary: str

ui/agent/graph/workflow.py

@@ -0,0 +1,23 @@
+# ui/agent/graph/workflow.py
+from __future__ import annotations
+
+from functools import lru_cache
+
+from langgraph.graph import END, START, StateGraph
+
+from .nodes import consolidator_node, fan_out_research, planner_node, researcher_node
+from .state import AgentState
+
+
+@lru_cache(maxsize=1)
+def build_workflow():
+    builder = StateGraph(AgentState)
+    builder.add_node("planner", planner_node)
+    builder.add_node("researcher", researcher_node)
+    builder.add_node("consolidator", consolidator_node)
+
+    builder.add_edge(START, "planner")
+    builder.add_conditional_edges("planner", fan_out_research, ["researcher"])
+    builder.add_edge("researcher", "consolidator")
+    builder.add_edge("consolidator", END)
+    return builder.compile()

ui/server_api.py

@@ -6,7 +6,7 @@ from typing import Any
 import gradio as gr
 from gradio import ChatMessage
 
-from ui.agent.respond import respond
+from ui.agent.graph import respond_with_graph as respond
 from ui.agent.system_prompt import BORDERLESS_SYSTEM_PROMPT
 from ui.chat.defaults import DEFAULT_MAX_TOKENS, DEFAULT_TEMPERATURE, DEFAULT_TOP_P
 from ui.globe_commands import empty_globe_state