|
|
| 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 ..state import CandidateCountry, TodoItem
|
| from .config import MAX_TODOS, RESEARCH_TOOL_NAMES
|
|
|
|
|
| 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()
|
| 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 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 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 split_todo_label(label: str) -> tuple[str, str]:
|
| if "—" in label:
|
| country, methods = label.split("—", 1)
|
| return country.strip(), methods.strip()
|
| if " - " in label:
|
| country, methods = label.split(" - ", 1)
|
| return country.strip(), methods.strip()
|
| return label.strip(), "Skilled migration pathway"
|
|
|
|
|
| def format_todo_label(todo: TodoItem | dict[str, Any]) -> str:
|
| country = str(todo.get("country") or "").strip()
|
| methods = str(todo.get("methods") or "").strip()
|
| if country and methods:
|
| return f"{country} — {methods}"
|
| return str(todo.get("title") or country or "Research task").strip()
|
|
|
|
|
| def todo_research_brief(todo: TodoItem, profile_summary: str) -> str:
|
| return (
|
| f"Country: {todo['country']}\n"
|
| f"Migration methods to research: {todo['methods']}\n\n"
|
| f"Research the best realistic skilled migration pathway to {todo['country']} "
|
| f"for this applicant. Focus on {todo['methods']}. Cover eligibility, required "
|
| f"documents, approximate costs, realistic timeline within 12-18 months, path to "
|
| f"permanent residence, and risks. Use official government or immigration authority "
|
| f"sources.\n\nApplicant profile: {profile_summary}"
|
| )
|
|
|
|
|
| def normalize_todo(
|
| raw: dict[str, Any],
|
| *,
|
| todo_id: int,
|
| profile_summary: str,
|
| candidate: CandidateCountry | None = None,
|
| ) -> TodoItem | None:
|
| country = str(raw.get("country") or "").strip()
|
| methods = str(raw.get("methods") or "").strip()
|
|
|
| if not country and raw.get("title"):
|
| country, parsed_methods = split_todo_label(str(raw["title"]))
|
| methods = methods or parsed_methods
|
|
|
| if candidate:
|
| country = country or candidate["name"]
|
| methods = methods or candidate["pathway_hint"]
|
|
|
| if not methods and raw.get("description"):
|
| methods = truncate(str(raw["description"]).strip(), 220)
|
|
|
| if not country:
|
| return None
|
| if not methods:
|
| methods = "Skilled migration pathway"
|
|
|
| return {"id": todo_id, "country": country, "methods": methods}
|
|
|
|
|
| def country_todo(
|
| candidate: CandidateCountry,
|
| profile_summary: str,
|
| *,
|
| todo_id: int,
|
| ) -> TodoItem:
|
| _ = profile_summary
|
| return {
|
| "id": todo_id,
|
| "country": candidate["name"],
|
| "methods": candidate["pathway_hint"],
|
| }
|
|
|
|
|
| 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:
|
| country = str(todo.get("country") or "").lower()
|
| methods = str(todo.get("methods") or "").lower()
|
| 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 or country in generic_titles:
|
| return True
|
| if methods == "research realistic migration options for this profile":
|
| 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
|
| candidate = shortlist[index] if index < len(shortlist) else None
|
| normalized = normalize_todo(
|
| raw,
|
| todo_id=len(todos) + 1,
|
| profile_summary=summary,
|
| candidate=candidate,
|
| )
|
| if normalized:
|
| todos.append(normalized)
|
|
|
| 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 _think_tags() -> tuple[str, str]:
|
| return "<" + "think" + ">", "</" + "think" + ">"
|
|
|
|
|
| _THINK_OPEN, _THINK_CLOSE = _think_tags()
|
| _THINK_BLOCK_PATTERN = re.compile(
|
| re.escape(_THINK_OPEN) + r".*?" + re.escape(_THINK_CLOSE),
|
| re.I | re.DOTALL,
|
| )
|
| _THINK_INNER_PATTERN = re.compile(
|
| re.escape(_THINK_OPEN) + r"(.*?)" + re.escape(_THINK_CLOSE),
|
| re.I | re.DOTALL,
|
| )
|
|
|
|
|
| def _flatten_content_blocks(content: Any) -> str:
|
| if content is None:
|
| return ""
|
| if isinstance(content, str):
|
| return content
|
| if isinstance(content, list):
|
| parts: list[str] = []
|
| for block in content:
|
| if isinstance(block, str):
|
| parts.append(block)
|
| continue
|
| if isinstance(block, dict):
|
| if block.get("type") == "text":
|
| parts.append(str(block.get("text") or ""))
|
| elif "text" in block:
|
| parts.append(str(block["text"]))
|
| return "\n".join(part for part in parts if part)
|
| return str(content)
|
|
|
|
|
| def _extract_reasoning(message: AIMessage) -> str:
|
| reasoning = getattr(message, "reasoning", None)
|
| if reasoning:
|
| return str(reasoning).strip()
|
|
|
| additional = getattr(message, "additional_kwargs", None) or {}
|
| for key in ("reasoning_content", "reasoning", "reasoning_text"):
|
| value = additional.get(key)
|
| if value:
|
| return str(value).strip()
|
|
|
| metadata = getattr(message, "response_metadata", None) or {}
|
| for key in ("reasoning_content", "reasoning"):
|
| value = metadata.get(key)
|
| if value:
|
| return str(value).strip()
|
| return ""
|
|
|
|
|
| def _strip_think_blocks(text: str) -> str:
|
| if not text.strip():
|
| return ""
|
| outside = _THINK_BLOCK_PATTERN.sub("", text).strip()
|
| if outside:
|
| return outside
|
| inner_parts = _THINK_INNER_PATTERN.findall(text)
|
| if inner_parts:
|
| return "\n".join(part.strip() for part in inner_parts if part.strip())
|
| return text.strip()
|
|
|
|
|
| def assistant_text_sources(message: AIMessage) -> tuple[str, str]:
|
| content_text = _flatten_content_blocks(message.content).strip()
|
| reasoning_text = _extract_reasoning(message)
|
| return content_text, reasoning_text
|
|
|
|
|
| def extract_assistant_text(message: AIMessage) -> str:
|
| """Extract user-visible assistant text from content, blocks, or reasoning."""
|
| content_text, reasoning_text = assistant_text_sources(message)
|
| for candidate in (content_text, reasoning_text):
|
| stripped = _strip_think_blocks(candidate)
|
| if stripped:
|
| return stripped
|
| return ""
|
|
|
|
|
| def extract_thinking_text(message: AIMessage) -> str:
|
| """Extract model reasoning / thinking text separate from the user-facing answer."""
|
| content_text, reasoning_text = assistant_text_sources(message)
|
| answer = extract_assistant_text(message)
|
| parts: list[str] = []
|
|
|
| for text in (reasoning_text, content_text):
|
| if not text.strip():
|
| continue
|
| inner_parts = _THINK_INNER_PATTERN.findall(text)
|
| if inner_parts:
|
| parts.extend(part.strip() for part in inner_parts if part.strip())
|
| continue
|
| outside = _THINK_BLOCK_PATTERN.sub("", text).strip()
|
| candidate = outside or text.strip()
|
| if candidate and candidate != answer:
|
| parts.append(candidate)
|
|
|
| deduped: list[str] = []
|
| seen: set[str] = set()
|
| for part in parts:
|
| if part not in seen:
|
| seen.add(part)
|
| deduped.append(part)
|
| return "\n\n".join(deduped).strip()
|
|
|
|
|
| def research_tool_calls(
|
| response: AIMessage,
|
| ) -> list[tuple[str, dict[str, Any], str]]:
|
| """Return normalized tool calls, including text-emitted calls from small models."""
|
| if response.tool_calls:
|
| return [
|
| (tool_call["name"], tool_call["args"] or {}, tool_call["id"])
|
| for tool_call in response.tool_calls
|
| ]
|
|
|
| content_text, reasoning_text = assistant_text_sources(response)
|
| parsed_calls = parse_text_tool_calls(content_text) or parse_text_tool_calls(
|
| reasoning_text
|
| )
|
| if not parsed_calls:
|
| return []
|
|
|
| normalized: list[tuple[str, dict[str, Any], str]] = []
|
| for tool_call in parsed_calls:
|
| function = tool_call.get("function") or {}
|
| tool_name = str(function.get("name") or "")
|
| if tool_name not in RESEARCH_TOOL_NAMES:
|
| continue
|
| normalized.append(
|
| (
|
| tool_name,
|
| _parse_arguments(str(function.get("arguments") or "")),
|
| str(tool_call.get("id") or uuid.uuid4().hex),
|
| )
|
| )
|
| 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"
|
| ),
|
| ]
|
|
|
