|
|
| from __future__ import annotations
|
|
|
| import re
|
| import time
|
| from typing import Any
|
|
|
| from apis.rest_countries import lookup_country
|
|
|
| from ui.agent.synthesis import build_country_report_sections
|
|
|
| _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",
|
| }
|
|
|
| _INCOMPLETE_SUMMARIES = {
|
| "no findings could be produced for this to-do.",
|
| "no detailed findings were captured for this country yet.",
|
| }
|
|
|
| _TOOL_ACTIVITY_LABELS = {
|
| "search_immigration_info": "Searching immigration sources",
|
| "get_country_profile": "Loading country profile",
|
| "scrape_web_page": "Reading official page",
|
| "crawl_web_site": "Crawling official site",
|
| }
|
|
|
| _URL_PATTERN = re.compile(r"https?://[^\s\)\]>\"']+")
|
|
|
|
|
| def split_country_title(title: str) -> tuple[str, str]:
|
| if "—" in title:
|
| country, pathway = title.split("—", 1)
|
| return country.strip(), pathway.strip()
|
| if " - " in title:
|
| country, pathway = title.split(" - ", 1)
|
| return country.strip(), pathway.strip()
|
| return title.strip(), "Skilled migration pathway"
|
|
|
|
|
| def is_incomplete_summary(summary: str) -> bool:
|
| text = summary.strip().lower()
|
| if not text:
|
| return True
|
| return text in _INCOMPLETE_SUMMARIES
|
|
|
|
|
| def _extract_urls(text: str) -> list[str]:
|
| seen: set[str] = set()
|
| urls: list[str] = []
|
| for url in _URL_PATTERN.findall(text):
|
| cleaned = url.rstrip(".,;")
|
| if cleaned not in seen:
|
| seen.add(cleaned)
|
| urls.append(cleaned)
|
| return urls
|
|
|
|
|
| def iso2_from_country_name(country: str, marker_isos: set[str] | None = None) -> str | None:
|
| lowered = country.strip().lower()
|
| if lowered in _COUNTRY_NAME_TO_ISO2:
|
| code = _COUNTRY_NAME_TO_ISO2[lowered]
|
| if marker_isos is None or code in marker_isos:
|
| return code
|
|
|
| candidates = marker_isos or set(_COUNTRY_NAME_TO_ISO2.values())
|
| for iso2 in candidates:
|
| info = lookup_country(iso2)
|
| if info and str(info.get("name", "")).lower() == lowered:
|
| return iso2.upper()
|
| return None
|
|
|
|
|
| def iso2_from_todo_title(title: str, marker_isos: set[str] | None = None) -> str | None:
|
| country, _ = split_country_title(title)
|
| return iso2_from_country_name(country, marker_isos)
|
|
|
|
|
| def _preview_text(text: str, limit: int = 220) -> str:
|
| cleaned = re.sub(r"[#>*`_\-[\]()]", " ", text)
|
| cleaned = re.sub(r"\s+", " ", cleaned).strip()
|
| if len(cleaned) <= limit:
|
| return cleaned
|
| return f"{cleaned[: limit - 1]}…"
|
|
|
|
|
| def _activity_from_tool(tool_name: str, args: dict[str, Any], message: str) -> str:
|
| label = _TOOL_ACTIVITY_LABELS.get(tool_name, f"Running {tool_name}")
|
| if tool_name == "search_immigration_info":
|
| query = str(args.get("query") or "").strip()
|
| if query:
|
| return f"{label}: {query[:160]}"
|
| if tool_name == "scrape_web_page":
|
| url = str(args.get("url") or "").strip()
|
| if url:
|
| return f"{label}: {url[:120]}"
|
| if tool_name == "crawl_web_site":
|
| url = str(args.get("url") or "").strip()
|
| if url:
|
| return f"{label}: {url[:120]}"
|
| if message.strip():
|
| return message.strip()[:220]
|
| return label
|
|
|
|
|
| def _base_detail(country: str, methods: str) -> dict[str, Any]:
|
| label = f"{country} — {methods}"
|
| return {
|
| "country": country,
|
| "title": label,
|
| "pathway": methods,
|
| "summary": "",
|
| "preview": "",
|
| "sources": [],
|
| "activities": [],
|
| "current_activity": "",
|
| "sections": {},
|
| "updated_at": int(time.time()),
|
| }
|
|
|
|
|
| def _enrich_detail_with_iso2(
|
| detail: dict[str, Any],
|
| country: str,
|
| marker_isos: set[str],
|
| ) -> tuple[str | None, dict[str, Any]]:
|
| iso2 = iso2_from_country_name(country, marker_isos or None)
|
| if iso2:
|
| detail["iso2"] = iso2
|
| info = lookup_country(iso2)
|
| if info:
|
| detail["country"] = str(info.get("name") or detail["country"])
|
| key = iso2 or str(detail["country"])
|
| return iso2, detail if iso2 else detail | {"iso2": None, "key": key}
|
|
|
|
|
| def _merge_detail_into_state(
|
| state: dict[str, Any],
|
| *,
|
| iso2: str | None,
|
| country_key: str,
|
| detail: dict[str, Any],
|
| ) -> dict[str, Any]:
|
| country_details = dict(state.get("country_details") or {})
|
| country_details[country_key] = detail
|
|
|
| markers = []
|
| for marker in state.get("markers", []):
|
| marker_copy = dict(marker)
|
| if iso2 and marker_copy.get("iso2") == iso2:
|
| marker_copy["detail"] = detail
|
| markers.append(marker_copy)
|
|
|
| return {
|
| **state,
|
| "markers": markers,
|
| "country_details": country_details,
|
| "details_version": int(state.get("details_version", 0)) + 1,
|
| }
|
|
|
|
|
| def build_research_progress_detail(
|
| country: str,
|
| methods: str,
|
| *,
|
| tool_name: str,
|
| args: dict[str, Any],
|
| message: str,
|
| existing: dict[str, Any] | None = None,
|
| ) -> dict[str, Any]:
|
| detail = dict(existing or _base_detail(country, methods))
|
| activity = _activity_from_tool(tool_name, args, message)
|
| activities = list(detail.get("activities") or [])
|
| if activity and (not activities or activities[-1] != activity):
|
| activities.append(activity)
|
| activities = activities[-6:]
|
|
|
| detail.update(
|
| {
|
| "status": "researching",
|
| "current_activity": activity,
|
| "activities": activities,
|
| "preview": activity,
|
| "updated_at": int(time.time()),
|
| }
|
| )
|
| return detail
|
|
|
|
|
| def build_country_detail(country: str, methods: str, summary: str) -> dict[str, Any]:
|
| detail = _base_detail(country, methods)
|
|
|
| if is_incomplete_summary(summary):
|
| detail.update(
|
| {
|
| "status": "incomplete",
|
| "preview": (
|
| "Research could not produce verified findings for this country. "
|
| "Check the parallel research panel for tool activity."
|
| ),
|
| "sections": build_country_report_sections("", country=country, pathway=methods),
|
| }
|
| )
|
| return detail
|
|
|
| sections = build_country_report_sections(summary, country=country, pathway=methods)
|
| detail.update(
|
| {
|
| "status": "complete",
|
| "summary": summary.strip(),
|
| "preview": _preview_text(summary),
|
| "sources": _extract_urls(summary),
|
| "sections": sections,
|
| "current_activity": "Research complete",
|
| }
|
| )
|
| return detail
|
|
|
|
|
| def attach_research_progress_to_globe(
|
| state: dict[str, Any],
|
| *,
|
| country: str,
|
| methods: str,
|
| tool_name: str,
|
| args: dict[str, Any],
|
| message: str,
|
| todo_title: str | None = None,
|
| ) -> dict[str, Any]:
|
| if todo_title and not country:
|
| country, methods = split_country_title(todo_title)
|
|
|
| marker_isos = {str(marker.get("iso2", "")).upper() for marker in state.get("markers", [])}
|
| marker_isos.discard("")
|
|
|
| country_details = dict(state.get("country_details") or {})
|
| iso2 = iso2_from_country_name(country, marker_isos or None)
|
| key = iso2 or country
|
| existing = country_details.get(key)
|
| if existing and existing.get("status") == "complete":
|
| return state
|
|
|
| detail = build_research_progress_detail(
|
| country,
|
| methods,
|
| tool_name=tool_name,
|
| args=args,
|
| message=message,
|
| existing=existing,
|
| )
|
| iso2, detail = _enrich_detail_with_iso2(detail, country, marker_isos)
|
| return _merge_detail_into_state(
|
| state,
|
| iso2=iso2,
|
| country_key=iso2 or detail["country"],
|
| detail=detail,
|
| )
|
|
|
|
|
| def attach_finding_to_globe(
|
| state: dict[str, Any],
|
| *,
|
| country: str,
|
| methods: str,
|
| summary: str,
|
| todo_title: str | None = None,
|
| ) -> dict[str, Any]:
|
| """Merge researcher findings into globe_state for marker popups."""
|
| if todo_title and not country:
|
| country, methods = split_country_title(todo_title)
|
|
|
| marker_isos = {str(marker.get("iso2", "")).upper() for marker in state.get("markers", [])}
|
| marker_isos.discard("")
|
|
|
| detail = build_country_detail(country, methods, summary)
|
| iso2, detail = _enrich_detail_with_iso2(detail, country, marker_isos)
|
| return _merge_detail_into_state(
|
| state,
|
| iso2=iso2,
|
| country_key=iso2 or detail["country"],
|
| detail=detail,
|
| )
|
|
|
