|
|
| from __future__ import annotations
|
|
|
| import os
|
| import time
|
| from typing import Any
|
|
|
| import gradio as gr
|
| from gradio import ChatMessage
|
|
|
| from ui.globe_commands import apply_update_globe
|
| from ui.globe_details import attach_finding_to_globe, attach_research_progress_to_globe
|
|
|
| from ..messages import history_to_api_messages, multimodal_input_to_api_content
|
| from ..config import hub_token_available
|
| from ..respond import (
|
| AUTH_REQUIRED_MESSAGE,
|
| SESSION_EXPIRED_MESSAGE,
|
| _auth_error_message,
|
| )
|
| from ..streaming import yield_response
|
| from ..synthesis import build_structured_final_answer
|
| from .workflow import build_workflow
|
|
|
|
|
| def _format_plan_summary(todos: list[dict[str, Any]]) -> str:
|
| count = len(todos)
|
| if count == 0:
|
| return "No countries queued for research."
|
| noun = "country" if count == 1 else "countries"
|
| return f"{count} {noun} queued for parallel research."
|
|
|
|
|
| def _research_task_country_methods(
|
| research_task: dict[str, Any],
|
| ) -> tuple[str, str] | None:
|
| country = str(research_task.get("country") or "").strip()
|
| methods = str(research_task.get("methods") or "").strip()
|
| if country and methods:
|
| return country, methods
|
|
|
| title = str(research_task.get("title") or "").strip()
|
| if not title:
|
| return None
|
| from ui.globe_details import split_country_title
|
|
|
| country, methods = split_country_title(title)
|
| if country:
|
| return country, methods or "Skilled migration pathway"
|
| return None
|
|
|
|
|
| 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 _attach_research_progress(self, event: dict[str, Any]) -> None:
|
| research_task = event.get("research_task")
|
| if not research_task:
|
| return
|
| country_methods = _research_task_country_methods(research_task)
|
| if not country_methods:
|
| return
|
| country, methods = country_methods
|
| log = event.get("log") or {}
|
| self.globe_state = attach_research_progress_to_globe(
|
| self.globe_state,
|
| country=country,
|
| methods=methods,
|
| tool_name=str(log.get("tool") or ""),
|
| args=dict(log.get("arguments") or {}),
|
| message=str(event.get("message") or ""),
|
| )
|
|
|
| def handle(self, event: dict[str, Any]) -> bool:
|
| kind = event.get("type")
|
| if kind == "thinking":
|
| thinking_text = str(event.get("text") or "").strip()
|
| self.ui_messages.append(
|
| ChatMessage(
|
| role="assistant",
|
| content=thinking_text,
|
| metadata={
|
| "title": "Thinking",
|
| "status": "done",
|
| "display": "thinking",
|
| "thinking": thinking_text,
|
| },
|
| )
|
| )
|
| return True
|
|
|
| if kind == "discovery":
|
| self.ui_messages.append(
|
| ChatMessage(
|
| role="assistant",
|
| content=str(event.get("summary") or ""),
|
| metadata={
|
| "title": "Destination discovery",
|
| "status": "done",
|
| "compact": True,
|
| "log": event.get("log") or {},
|
| },
|
| )
|
| )
|
| return True
|
|
|
| if kind == "plan":
|
| todos = event.get("todos") or []
|
| self.ui_messages.append(
|
| ChatMessage(
|
| role="assistant",
|
| content=_format_plan_summary(todos),
|
| metadata={
|
| "title": "Research plan",
|
| "status": "done",
|
| "display": "plan",
|
| "plan_todos": todos,
|
| },
|
| )
|
| )
|
| return True
|
|
|
| if kind == "globe":
|
| args = event.get("args") or {}
|
| countries = ", ".join(args.get("countries") or [])
|
| try:
|
| result, self.globe_state = apply_update_globe(self.globe_state, args)
|
| content = (
|
| f"Marked {countries or 'recommended countries'} on the globe."
|
| )
|
| status = "done"
|
| except Exception as exc:
|
| result = {"error": str(exc)}
|
| content = (
|
| f"Could not update globe for {countries or 'recommended countries'}: "
|
| f"{exc}"
|
| )
|
| status = "done"
|
| self.ui_messages.append(
|
| ChatMessage(
|
| role="assistant",
|
| content=content,
|
| metadata={
|
| "title": "Updating globe",
|
| "status": status,
|
| "log": {
|
| "tool": "update_globe",
|
| "arguments": args,
|
| "result": result,
|
| },
|
| },
|
| )
|
| )
|
| return True
|
|
|
| if kind == "tool_start":
|
| research_task = event.get("research_task")
|
| 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 {},
|
| "research_task": research_task,
|
| "compact": bool(research_task),
|
| },
|
| )
|
| )
|
| self._pending_by_id[str(event.get("id"))] = len(self.ui_messages) - 1
|
| self._attach_research_progress(event)
|
| return True
|
|
|
| if kind == "tool_end":
|
| index = self._pending_by_id.pop(str(event.get("id")), None)
|
| research_task = event.get("research_task")
|
| 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 {},
|
| "research_task": research_task,
|
| "compact": bool(research_task),
|
| },
|
| )
|
| if index is not None and index < len(self.ui_messages):
|
| self.ui_messages[index] = message
|
| else:
|
| self.ui_messages.append(message)
|
| self._attach_research_progress(event)
|
| return True
|
|
|
| if kind == "finding":
|
| research_task = event.get("research_task") or {}
|
| country_methods = _research_task_country_methods(research_task)
|
| if country_methods:
|
| country, methods = country_methods
|
| else:
|
| from ui.globe_details import split_country_title
|
|
|
| todo_title = str(event.get("todo_title") or "research")
|
| country, methods = split_country_title(todo_title)
|
| todo_label = f"{country} — {methods}"
|
| summary = str(event.get("summary") or "")
|
| self.globe_state = attach_finding_to_globe(
|
| self.globe_state,
|
| country=country,
|
| methods=methods,
|
| summary=summary,
|
| )
|
| self.ui_messages.append(
|
| ChatMessage(
|
| role="assistant",
|
| content=summary,
|
| metadata={
|
| "title": f"Findings · {todo_label}",
|
| "status": "done",
|
| "markdown": True,
|
| "research_task": research_task
|
| or {
|
| "id": event.get("todo_id"),
|
| "country": country,
|
| "methods": methods,
|
| },
|
| "research_finding": True,
|
| },
|
| )
|
| )
|
| 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 and not hub_token_available():
|
| yield from yield_response([], AUTH_REQUIRED_MESSAGE, globe_state)
|
| return
|
|
|
| if hf_token is not None and 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 = ""
|
| last_state: dict[str, Any] = {}
|
|
|
| config = {
|
| "configurable": {
|
| "hf_token": (
|
| hf_token.token
|
| if hf_token is not None
|
| else os.environ.get("HF_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
|
| last_state = payload
|
| 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
|
|
|
| if not final_answer:
|
| final_answer = build_structured_final_answer(
|
| profile_summary=str(last_state.get("profile_summary") or "").strip(),
|
| findings=last_state.get("findings") or [],
|
| todos=last_state.get("todos"),
|
| preamble=(
|
| "I could not complete model synthesis, but here is the structured "
|
| "research gathered so far."
|
| ),
|
| )
|
|
|
| yield from yield_response(
|
| ui_state.ui_messages,
|
| final_answer,
|
| ui_state.globe_state,
|
| assistant_metadata={
|
| "markdown": True,
|
| "thinking": str(last_state.get("consolidator_thinking") or "").strip(),
|
| },
|
| )
|
|
|
