agents/research.py

from __future__ import annotations

import json
from typing import Generator

import anthropic

from agents.eligibility import bulk_parse_and_strip
from beacon_logging import get_logger
from config import RESEARCH_MODEL
from llm import cached_system, cached_tools
from models import PatientProfile
from prompts import RESEARCH_SYSTEM
from tools import RESEARCH_TOOLS
from translations import LANGUAGE_DIRECTIVE
from trials_api import search_trials_api, _flatten_and_rank
from _console import console

_logger = get_logger("agents.research")

_MAX_TRIALS_FOR_LLM = 15

# Fields with no synthesis value once eligibility is parsed; stripping them
# shrinks the tool-result payload significantly (nearest_sites alone is ~250 tokens/trial).
_STRIP_BEFORE_LLM = {
    "summary",       # LLM writes its own plain-language summary
    "conditions",    # patient already knows their disease
    "keywords",
    "min_age", "max_age", "sex", "healthy_volunteers",  # in parsed_criteria after bulk parse
    "std_ages",
    "eligibility",   # raw text; replaced by parsed_criteria for top-5
}


_PHASE_PRIORITY: dict[str, int] = {
    "PHASE4": 1,
    "PHASE3": 2,
    "PHASE2": 3,
    "PHASE1": 4,
    "EARLY_PHASE1": 5,
    "NA": 6,
}


def _phase_rank(trial: dict) -> int:
    """Lower = higher priority. Phase 4 > Phase 3 > ... > EAP > Observational."""
    study_type = trial.get("study_type", "")
    if study_type == "EXPANDED_ACCESS":
        return 7
    if study_type == "OBSERVATIONAL":
        return 8
    phase_str = trial.get("phase", "N/A")
    phases = [p.strip() for p in phase_str.replace(" ", "").split(",") if p.strip()]
    return min((_PHASE_PRIORITY.get(p, 6) for p in phases), default=6)


def _rank_and_slim(trials: list[dict]) -> list[dict]:
    """Sort by phase priority then distance, cap at _MAX_TRIALS_FOR_LLM, strip bloat."""
    ranked = sorted(
        trials,
        key=lambda t: (_phase_rank(t), t.get("closest_site_miles") or float("inf")),
    )
    slimmed = []
    for t in ranked[:_MAX_TRIALS_FOR_LLM]:
        t = {k: v for k, v in t.items() if k not in _STRIP_BEFORE_LLM}
        if "nearest_sites" in t:
            t["nearest_sites"] = t["nearest_sites"][:3]
        if "interventions" in t:
            t["interventions"] = [
                {"type": iv.get("type", ""), "name": iv.get("name", "")}
                for iv in t["interventions"]
            ]
        slimmed.append(t)
    return slimmed


def run_research_agent(client: anthropic.Anthropic, profile: PatientProfile) -> str:
    messages: list[anthropic.types.MessageParam] = [
        {
            "role": "user",
            "content": (
                f"Find clinical trials for this patient:\n\n{profile.summary()}\n\n"
                "Search within the specified radius and rank results by distance."
            ),
        }
    ]

    while True:
        response = client.messages.create(
            model=RESEARCH_MODEL,
            max_tokens=3000,
            system=cached_system(RESEARCH_SYSTEM),
            tools=cached_tools(RESEARCH_TOOLS),
            messages=messages,
        )

        messages.append({"role": "assistant", "content": response.content})

        if response.stop_reason == "end_turn":
            return next(
                (b.text for b in response.content if b.type == "text"),
                "No analysis produced.",
            )

        tool_results: list[anthropic.types.ToolResultBlockParam] = []
        for block in response.content:
            if block.type != "tool_use" or block.name != "search_clinical_trials":
                continue
            args = block.input
            radius = args.get("radius_miles", profile.radius_miles)
            study_type = args.get("study_type", "INTERVENTIONAL")
            # Enforce patient's phase preference; don't rely on LLM to repeat it.
            # Phase filters only apply to INTERVENTIONAL searches.
            phases = (profile.phases or None) if study_type == "INTERVENTIONAL" else None
            status_msg = (
                f"[cyan]Searching:[/cyan] '[bold]{args['condition']}[/bold]' | "
                f"radius=[bold]{radius}[/bold] mi | "
                f"type=[bold]{study_type}[/bold] | "
                f"phases=[bold]{phases or 'all'}[/bold]"
            )
            try:
                with console.status(status_msg, spinner="dots"):
                    studies = search_trials_api(
                        condition=args["condition"],
                        lat=args["lat"],
                        lon=args["lon"],
                        radius_miles=radius,
                        phases=phases,
                        study_type=study_type,
                    )
                    ranked = _flatten_and_rank(studies, profile.lat, profile.lon)
                    ranked = bulk_parse_and_strip(client, ranked, profile)
                console.print(f"  [green]✓[/green] {len(ranked)} trial(s) found.")
                content = json.dumps(_rank_and_slim(ranked))
                is_error = False
            except Exception as exc:
                console.print(f"[red bold]API error:[/red bold] {exc}")
                content = f"API request failed: {exc}. The ClinicalTrials.gov endpoint may be temporarily unavailable."
                is_error = True
            tool_results.append({
                "type": "tool_result",
                "tool_use_id": block.id,
                "content": content,
                "is_error": is_error,
            })

        messages.append({"role": "user", "content": tool_results})


def stream_research_agent(
    client: anthropic.Anthropic,
    profile: PatientProfile,
) -> Generator[tuple, None, None]:
    """
    Stream the research agent for a given patient profile.

    Yields:
      ("token", str)   — partial text chunk
      ("done", str)    — research complete; the full analysis text (may be empty)
    """
    messages: list[anthropic.types.MessageParam] = [{
        "role": "user",
        "content": (
            f"Find clinical trials for this patient:\n\n{profile.summary()}\n\n"
            "Search within the specified radius and rank results by distance."
        ),
    }]

    while True:
        stream_text = ""
        with client.messages.stream(
            model=RESEARCH_MODEL,
            max_tokens=3000,
            system=cached_system(LANGUAGE_DIRECTIVE[profile.lang] + RESEARCH_SYSTEM),
            tools=cached_tools(RESEARCH_TOOLS),
            messages=messages,
        ) as stream:
            for chunk in stream.text_stream:
                stream_text += chunk
                yield ("token", chunk)
            rresponse = stream.get_final_message()

        messages.append({"role": "assistant", "content": rresponse.content})

        if rresponse.stop_reason == "end_turn":
            yield ("done", stream_text)
            return

        tool_results: list[anthropic.types.ToolResultBlockParam] = []
        for block in rresponse.content:
            if block.type != "tool_use" or block.name != "search_clinical_trials":
                continue
            args = block.input
            radius = args.get("radius_miles", profile.radius_miles)
            study_type = args.get("study_type", "INTERVENTIONAL")
            # Enforce patient's phase preference; don't rely on LLM to repeat it.
            # Phase filters only apply to INTERVENTIONAL searches.
            phases = (profile.phases or None) if study_type == "INTERVENTIONAL" else None
            type_label = {"INTERVENTIONAL": "clinical trials", "EXPANDED_ACCESS": "expanded access programs", "OBSERVATIONAL": "observational studies"}.get(study_type, study_type.lower())
            yield ("status", f"Searching ClinicalTrials.gov for **{args['condition']}** ({type_label}, {radius} mi radius)…")
            try:
                studies = search_trials_api(
                    condition=args["condition"],
                    lat=args["lat"],
                    lon=args["lon"],
                    radius_miles=radius,
                    phases=phases,
                    study_type=study_type,
                )
                ranked = _flatten_and_rank(studies, profile.lat, profile.lon)
                n = len(ranked)
                yield ("status", f"Found **{n}** {type_label} — checking eligibility for the {min(5, n)} closest…")
                ranked = bulk_parse_and_strip(client, ranked, profile)
                yield ("status", "Eligibility analysis complete — generating your report…")
                content = json.dumps(_rank_and_slim(ranked))
                is_error = False
            except Exception as exc:
                content = f"API request failed: {exc}. The ClinicalTrials.gov endpoint may be temporarily unavailable."
                is_error = True
            tool_results.append({
                "type": "tool_result",
                "tool_use_id": block.id,
                "content": content,
                "is_error": is_error,
            })
        messages.append({"role": "user", "content": tool_results})