Robust Planner Discovery Workflow

ui/agent/graph/nodes/__init__.py

@@ -1,10 +1,12 @@
 # ui/agent/graph/nodes/__init__.py
 from .consolidator import consolidator_node
+from .discovery import discovery_node
 from .planner import fan_out_research, planner_node
 from .researcher import researcher_node
 
 __all__ = [
     "consolidator_node",
+    "discovery_node",
     "fan_out_research",
     "planner_node",
     "researcher_node",

ui/agent/graph/nodes/discovery.py

@@ -0,0 +1,105 @@
+# ui/agent/graph/nodes/discovery.py
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from langchain_core.runnables import RunnableConfig
+from langgraph.config import get_stream_writer
+
+from apis.exa import search_immigration
+from ..state import AgentState, CandidateCountry
+from .helpers import (
+    candidates_from_search_text,
+    discovery_queries,
+    heuristic_candidate_countries,
+    merge_candidates,
+    profile_summary_from_text,
+    user_text,
+)
+
+
+def _format_discovery_summary(
+    candidates: list[CandidateCountry],
+    search_notes: list[str],
+) -> str:
+    lines = [
+        "Initial destination shortlist for parallel research:",
+        "",
+    ]
+    for item in candidates:
+        lines.append(
+            f"- {item['name']} ({item['iso2']}): {item['pathway_hint']}"
+        )
+    if search_notes:
+        lines.extend(["", "Discovery search notes:"])
+        lines.extend(f"- {note}" for note in search_notes[:4])
+    lines.append(
+        "\nThese are starting candidates only — each country will be researched "
+        "in depth before the final recommendation."
+    )
+    return "\n".join(lines)
+
+
+def discovery_node(state: AgentState, config: RunnableConfig) -> dict[str, Any]:
+    del config  # discovery uses direct API calls, not the LLM
+    writer = get_stream_writer()
+    profile_text = user_text(state["user_content"])
+    profile_summary = profile_summary_from_text(profile_text)
+    heuristic = heuristic_candidate_countries(profile_text)
+
+    search_notes: list[str] = []
+    search_blob = ""
+    search_log: dict[str, Any] = {"queries": [], "results": []}
+
+    for query in discovery_queries(profile_text):
+        search_log["queries"].append(query)
+        try:
+            result = search_immigration(query=query, num_results=6)
+            search_log["results"].append(result)
+            titles = [
+                str(item.get("title") or "")
+                for item in result.get("results", [])
+                if isinstance(item, dict)
+            ]
+            if titles:
+                search_notes.append(f"{query} → {titles[0]}")
+            search_blob += "\n".join(
+                f"{item.get('title', '')} {item.get('url', '')}"
+                for item in result.get("results", [])
+                if isinstance(item, dict)
+            )
+        except Exception as exc:
+            search_notes.append(f"{query} → search unavailable ({exc})")
+
+    from_search = candidates_from_search_text(search_blob)
+    candidates = merge_candidates(from_search, heuristic)
+    if not candidates:
+        candidates = heuristic
+
+    summary = _format_discovery_summary(candidates, search_notes)
+    writer(
+        {
+            "type": "discovery",
+            "summary": summary,
+            "candidates": candidates,
+            "log": {
+                "tool": "destination_discovery",
+                "arguments": {"queries": search_log["queries"]},
+                "result": truncate_discovery_log(search_log),
+            },
+        }
+    )
+
+    return {
+        "profile_summary": profile_summary,
+        "candidate_countries": candidates,
+        "discovery_summary": summary,
+    }
+
+
+def truncate_discovery_log(search_log: dict[str, Any]) -> Any:
+    serialized = json.dumps(search_log, default=str)
+    if len(serialized) <= 1500:
+        return search_log
+    return serialized[:1500] + "\n… (truncated)"

ui/agent/graph/nodes/helpers.py

@@ -2,14 +2,61 @@
 from __future__ import annotations
 
 import json
+import re
 import uuid
 from typing import Any
 
 from langchain_core.messages import AIMessage
 
+from apis.rest_countries import lookup_country
 from ...messages import parse_text_tool_calls
 from ...tools import _parse_arguments, truncate
-from .config import RESEARCH_TOOL_NAMES
+from ..state import CandidateCountry, TodoItem
+from .config import MAX_TODOS, RESEARCH_TOOL_NAMES
+
+# Conservative starting shortlist for skilled IT applicants (used when LLM/search fail).
+DEFAULT_SKILLED_IT_COUNTRIES: list[CandidateCountry] = [
+    {
+        "iso2": "CA",
+        "name": "Canada",
+        "pathway_hint": "Express Entry / Provincial Nominee skilled worker route",
+        "label": "Skilled worker - 12-18 mo",
+    },
+    {
+        "iso2": "DE",
+        "name": "Germany",
+        "pathway_hint": "EU Blue Card / skilled worker residence route",
+        "label": "EU Blue Card - 6-12 mo",
+    },
+    {
+        "iso2": "AU",
+        "name": "Australia",
+        "pathway_hint": "Skilled Independent / State nomination route",
+        "label": "Skilled migration - 12-18 mo",
+    },
+    {
+        "iso2": "IE",
+        "name": "Ireland",
+        "pathway_hint": "Critical Skills Employment Permit route",
+        "label": "Critical Skills - 6-12 mo",
+    },
+]
+
+_COUNTRY_NAME_TO_ISO2 = {
+    "canada": "CA",
+    "germany": "DE",
+    "australia": "AU",
+    "ireland": "IE",
+    "new zealand": "NZ",
+    "united kingdom": "GB",
+    "uk": "GB",
+    "portugal": "PT",
+    "netherlands": "NL",
+    "singapore": "SG",
+    "united states": "US",
+    "usa": "US",
+}
+
 
 def extract_json(text: str) -> dict[str, Any] | None:
     decoder = json.JSONDecoder()
@@ -25,24 +72,6 @@ def extract_json(text: str) -> dict[str, Any] | None:
     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
@@ -54,6 +83,226 @@ def user_text(user_content: str | list[dict[str, Any]]) -> str:
     return "\n".join(part for part in parts if part)
 
 
+def profile_summary_from_text(profile_text: str) -> str:
+    """Build a short profile summary without calling the LLM."""
+    lines = [line.strip() for line in profile_text.splitlines() if line.strip()]
+    if not lines:
+        return truncate(profile_text, 600)
+    headline = lines[0]
+    bullets = [line.lstrip("- ").strip() for line in lines[1:] if line.strip().startswith("-")]
+    if bullets:
+        return truncate(f"{headline}. Key constraints: {'; '.join(bullets[:6])}.", 600)
+    return truncate(profile_text, 600)
+
+
+def heuristic_candidate_countries(profile_text: str) -> list[CandidateCountry]:
+    """Deterministic shortlist when discovery/planner cannot produce one."""
+    text = profile_text.lower()
+    candidates = list(DEFAULT_SKILLED_IT_COUNTRIES)
+
+    if any(word in text for word in ("software", "it", "engineer", "developer", "tech")):
+        return candidates[:MAX_TODOS]
+
+    if any(word in text for word in ("study", "student", "university")):
+        return [
+            {
+                "iso2": "DE",
+                "name": "Germany",
+                "pathway_hint": "Student visa / post-study residence route",
+                "label": "Study route - 12-24 mo",
+            },
+            {
+                "iso2": "CA",
+                "name": "Canada",
+                "pathway_hint": "Study permit / PGWP pathway",
+                "label": "Study route - 12-24 mo",
+            },
+            {
+                "iso2": "IE",
+                "name": "Ireland",
+                "pathway_hint": "Study / graduate route",
+                "label": "Study route - 12-24 mo",
+            },
+            {
+                "iso2": "AU",
+                "name": "Australia",
+                "pathway_hint": "Student visa / skilled graduate route",
+                "label": "Study route - 12-24 mo",
+            },
+        ][:MAX_TODOS]
+
+    return candidates[:MAX_TODOS]
+
+
+def _candidate_from_iso2(iso2: str) -> CandidateCountry | None:
+    info = lookup_country(iso2)
+    if not info:
+        return None
+    default = next(
+        (item for item in DEFAULT_SKILLED_IT_COUNTRIES if item["iso2"] == iso2.upper()),
+        None,
+    )
+    if default:
+        return default
+    return {
+        "iso2": info["cca2"],
+        "name": str(info["name"]),
+        "pathway_hint": "Skilled worker / residence pathway",
+        "label": "Skilled route - 12-18 mo",
+    }
+
+
+def candidates_from_search_text(text: str) -> list[CandidateCountry]:
+    """Extract mentioned countries from search result text."""
+    lowered = text.lower()
+    found: list[CandidateCountry] = []
+    seen: set[str] = set()
+    for name, iso2 in _COUNTRY_NAME_TO_ISO2.items():
+        if name in lowered and iso2 not in seen:
+            candidate = _candidate_from_iso2(iso2)
+            if candidate:
+                found.append(candidate)
+                seen.add(iso2)
+    return found
+
+
+def merge_candidates(
+    primary: list[CandidateCountry],
+    secondary: list[CandidateCountry],
+) -> list[CandidateCountry]:
+    merged: list[CandidateCountry] = []
+    seen: set[str] = set()
+    for item in [*primary, *secondary]:
+        iso2 = item["iso2"].upper()
+        if iso2 in seen:
+            continue
+        merged.append(item)
+        seen.add(iso2)
+        if len(merged) >= MAX_TODOS:
+            break
+    return merged
+
+
+def country_todo(
+    candidate: CandidateCountry,
+    profile_summary: str,
+    *,
+    todo_id: int,
+) -> TodoItem:
+    return {
+        "id": todo_id,
+        "title": f"{candidate['name']} — {candidate['pathway_hint']}",
+        "description": (
+            f"Research the best realistic skilled migration pathway to {candidate['name']} "
+            f"for this applicant. Focus on {candidate['pathway_hint']}. Cover eligibility, "
+            f"required documents, approximate costs, realistic timeline within 12-18 months, "
+            f"path to permanent residence, and risks. Use official government or immigration "
+            f"authority sources.\n\nApplicant profile: {profile_summary}"
+        ),
+    }
+
+
+def plan_from_candidates(
+    candidates: list[CandidateCountry],
+    profile_text: str,
+    *,
+    thinking: str = "",
+) -> dict[str, Any]:
+    summary = profile_summary_from_text(profile_text)
+    todos = [
+        country_todo(candidate, summary, todo_id=index + 1)
+        for index, candidate in enumerate(candidates[:MAX_TODOS])
+    ]
+    return {
+        "thinking": thinking,
+        "countries": [item["iso2"] for item in candidates[:MAX_TODOS]],
+        "labels": [item["label"] for item in candidates[:MAX_TODOS]],
+        "profile_summary": summary,
+        "todos": todos,
+    }
+
+
+def fallback_plan(
+    profile_text: str,
+    candidates: list[CandidateCountry] | None = None,
+) -> dict[str, Any]:
+    shortlist = candidates or heuristic_candidate_countries(profile_text)
+    return plan_from_candidates(
+        shortlist,
+        profile_text,
+        thinking=(
+            "Using a conservative starting shortlist of skilled-worker destinations. "
+            "Each country will be researched in parallel."
+        ),
+    )
+
+
+def _is_generic_todo(todo: dict[str, Any]) -> bool:
+    title = str(todo.get("title") or "").lower()
+    description = str(todo.get("description") or "").lower()
+    generic_titles = {"research migration options", "research task"}
+    if title in generic_titles:
+        return True
+    if "research realistic migration options for this profile" in description:
+        return True
+    return len(description) > 800 and description.count("\n") >= 4
+
+
+def normalize_plan(
+    plan: dict[str, Any] | None,
+    profile_text: str,
+    candidates: list[CandidateCountry],
+) -> dict[str, Any]:
+    """Ensure the plan has 3-4 useful country-specific todos."""
+    summary = str((plan or {}).get("profile_summary") or "").strip() or profile_summary_from_text(
+        profile_text
+    )
+    shortlist = candidates[:MAX_TODOS] or heuristic_candidate_countries(profile_text)
+
+    if plan is None:
+        return fallback_plan(profile_text, shortlist)
+
+    raw_todos = plan.get("todos") or []
+    todos: list[TodoItem] = []
+    for index, raw in enumerate(raw_todos[:MAX_TODOS]):
+        if not isinstance(raw, dict):
+            continue
+        if _is_generic_todo(raw):
+            continue
+        title = str(raw.get("title") or "").strip()
+        description = str(raw.get("description") or title).strip()
+        if not title:
+            continue
+        todos.append({"id": len(todos) + 1, "title": title, "description": description})
+
+    if len(todos) < 3:
+        todos = [
+            country_todo(candidate, summary, todo_id=index + 1)
+            for index, candidate in enumerate(shortlist[:MAX_TODOS])
+        ]
+
+    countries = [str(code) for code in plan.get("countries") or [] if code]
+    labels = [str(label) for label in plan.get("labels") or [] if label]
+    if len(countries) != len(todos):
+        countries = [item["iso2"] for item in shortlist[: len(todos)]]
+        labels = [item["label"] for item in shortlist[: len(todos)]]
+
+    thinking = str(plan.get("thinking") or "").strip()
+    if not thinking and len(todos) >= 3:
+        thinking = (
+            f"Split research into {len(todos)} parallel country tasks based on the "
+            f"applicant profile and discovery shortlist."
+        )
+
+    return {
+        "thinking": thinking,
+        "countries": countries,
+        "labels": labels,
+        "profile_summary": summary,
+        "todos": todos,
+    }
+
+
 def research_tool_calls(
     response: AIMessage,
 ) -> list[tuple[str, dict[str, Any], str]]:
@@ -82,3 +331,19 @@ def research_tool_calls(
             )
         )
     return normalized
+
+
+def discovery_queries(profile_text: str) -> list[str]:
+    text = re.sub(r"\s+", " ", profile_text).strip()
+    occupation = "software engineer" if re.search(r"software|it|developer", text, re.I) else "skilled worker"
+    origin = "India" if re.search(r"\bindia\b", text, re.I) else "applicant country"
+    return [
+        (
+            f"best skilled worker immigration pathways {occupation} {origin} "
+            "official government permanent residence"
+        ),
+        (
+            f"countries skilled worker visa path to permanent residence "
+            f"{occupation} Indian citizen official immigration"
+        ),
+    ]

ui/agent/graph/nodes/planner.py

@@ -7,14 +7,27 @@ from langchain_core.runnables import RunnableConfig
 from langgraph.config import get_stream_writer
 from langgraph.types import Send
 
-from ...tools import truncate
 from ..llm import build_llm
-from ..state import AgentState, TodoItem
+from ..state import AgentState, CandidateCountry, TodoItem
 from .config import MAX_TODOS, PLANNER_MAX_TOKENS, PLANNER_TEMPERATURE
-from .helpers import extract_json, fallback_plan, user_text
+from .helpers import (
+    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(
@@ -23,23 +36,45 @@ def planner_node(state: AgentState, config: RunnableConfig) -> dict[str, Any]:
         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": state["user_content"]},
+        {
+            "role": "user",
+            "content": (
+                f"{planner_context}\n\n"
+                f"Original user request:\n{profile_text}\n\n"
+                "Produce the JSON research plan now."
+            ),
+        },
     ]
 
-    plan: dict[str, Any] | None = None
+    raw_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"):
+        raw_plan = extract_json(str(response.content or ""))
+        if raw_plan and raw_plan.get("todos"):
             break
-        plan = None
+        raw_plan = None
 
-    profile_text = user_text(state["user_content"])
-    if plan is None:
-        plan = fallback_plan(profile_text)
+    plan = normalize_plan(raw_plan, profile_text, candidates)
 
     thinking = str(plan.get("thinking") or "").strip()
     if thinking:
@@ -52,9 +87,10 @@ def planner_node(state: AgentState, config: RunnableConfig) -> dict[str, Any]:
         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(profile_text)["todos"][0]
-        todos = [{"id": 1, "title": fallback["title"], "description": fallback["description"]}]
+        plan = normalize_plan(None, profile_text, candidates)
+        todos = plan["todos"]
 
     writer({"type": "plan", "todos": todos})
 
@@ -68,10 +104,10 @@ def planner_node(state: AgentState, config: RunnableConfig) -> dict[str, Any]:
             }
         )
 
-    profile_summary = str(plan.get("profile_summary") or "").strip() or truncate(
-        profile_text, 600
-    )
-    return {"todos": todos, "profile_summary": profile_summary}
+    return {
+        "todos": todos,
+        "profile_summary": str(plan.get("profile_summary") or profile_summary),
+    }
 
 
 def fan_out_research(state: AgentState) -> list[Send]:

ui/agent/graph/nodes/prompts.py

@@ -4,23 +4,32 @@ from __future__ import annotations
 PLANNER_SYSTEM_PROMPT = """
 You are the planning supervisor of Borderless, an immigration research agency.
 
+A discovery step has already produced a candidate country shortlist. Use
+`candidate_countries` as your preferred starting point unless the profile clearly
+rules a country out.
+
 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.
+team of parallel research analysts will execute. Always produce 3-4
+country-specific to-dos unless the profile is too sparse to justify that many.
+Each to-do must be researchable independently — one to-do per country covering
+its best visa pathway, eligibility, documents, costs, timelines, and risks.
+
+Each todo title MUST include the country name and pathway direction, for example:
+- "Canada — Express Entry / PNP skilled worker route"
+- "Germany — EU Blue Card / skilled worker residence route"
+- "Australia — Skilled Independent / State nomination route"
+- "Ireland — Critical Skills Employment Permit route"
+
+Do NOT create one broad todo that repeats the full user prompt.
 
 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"],
+  "countries": ["ISO-2 codes matching the todos, e.g. CA", "DE"],
+  "labels": ["short marker label per country, e.g. Skilled worker - 12-18 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"}
+    {"title": "Country — pathway title", "description": "specific research instructions for the analyst"}
   ]
 }
 """.strip()

ui/agent/graph/respond.py

@@ -51,6 +51,20 @@ class _UiState:
             )
             return True
 
+        if kind == "discovery":
+            self.ui_messages.append(
+                ChatMessage(
+                    role="assistant",
+                    content=str(event.get("summary") or ""),
+                    metadata={
+                        "title": "Destination discovery",
+                        "status": "done",
+                        "log": event.get("log") or {},
+                    },
+                )
+            )
+            return True
+
         if kind == "plan":
             todos = event.get("todos") or []
             self.ui_messages.append(

ui/agent/graph/state.py

@@ -17,12 +17,21 @@ class Finding(TypedDict):
     summary: str
 
 
+class CandidateCountry(TypedDict):
+    iso2: str
+    name: str
+    pathway_hint: str
+    label: 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
+    candidate_countries: list[CandidateCountry]
+    discovery_summary: str
     todos: list[TodoItem]
     findings: Annotated[list[Finding], operator.add]
     final_answer: str

ui/agent/graph/workflow.py

@@ -5,18 +5,26 @@ 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 .nodes import (
+    consolidator_node,
+    discovery_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("discovery", discovery_node)
     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_edge(START, "discovery")
+    builder.add_edge("discovery", "planner")
     builder.add_conditional_edges("planner", fan_out_research, ["researcher"])
     builder.add_edge("researcher", "consolidator")
     builder.add_edge("consolidator", END)