"""POST /pipeline/{bbb,eeg,mri} routes — thin dispatchers over the pipelines. Each route validates its request body via Pydantic, invokes the pipeline, reads back the produced Parquet to populate row/column counts, and returns a uniform PipelineResponse. Pipeline-domain errors map to standard HTTP codes: FileNotFoundError -> 404, ValueError -> 400, anything else -> 500. """ from __future__ import annotations import os import time from collections import deque from pathlib import Path from typing import Callable import mlflow import pandas as pd from fastapi import APIRouter, HTTPException from src.api.schemas import ( AgentRunRequest, AgentRunResponse, AgentToolTraceItem, BBBExplainRequest, BBBExplainResponse, BBBPredictRequest, BBBPredictResponse, BBBRequest, CalibrationContext, BBBPermeabilityMapRequest, BBBPermeabilityMapResponse, DrugDoseAdjustmentRequest, DrugDoseAdjustmentResponse, EEGClassProbability, EEGExplainRequest, EEGPredictRequest, EEGPredictResponse, FusionRequest, FusionResponse, EEGExplainResponse, EEGRequest, FeatureAttribution, HarmonizationRow, MLflowRunsResponse, MLflowRunSummary, ModelProvenance, MRIDiagnosticsRequest, MRIDiagnosticsResponse, MRIClassProbability, MRIExplainRequest, MRIExplainResponse, MRIPredictRequest, MRIPredictResponse, MRIRequest, PipelineResponse, RunDiffRequest, RunDiffResponse, RunDiffRow, ) from src.core.logger import get_logger from src.llm import explainer as llm_explainer from src.models import bbb_model, mri_model from src.pipelines import bbb_pipeline, eeg_pipeline, mri_pipeline logger = get_logger(__name__) router = APIRouter(prefix="/pipeline") predict_router = APIRouter(prefix="/predict") explain_router = APIRouter(prefix="/explain") experiments_router = APIRouter(prefix="/experiments") research_router = APIRouter(prefix="/research") def _wrap( experiment_name: str, output_path: Path, fn: Callable[[], None], ) -> PipelineResponse: """Run `fn()` (the pipeline call), gather metrics, return PipelineResponse.""" started = time.perf_counter() try: fn() except FileNotFoundError as e: raise HTTPException(status_code=404, detail=str(e)) except (ValueError, KeyError) as e: # KeyError: MRI pipeline raises this when sites_csv is missing a # site assignment for a subject — a user data problem, not a 500. raise HTTPException(status_code=400, detail=str(e)) duration_sec = time.perf_counter() - started df = pd.read_parquet(output_path) run_id = _latest_mlflow_run_id(experiment_name) return PipelineResponse( status="ok", output_path=str(output_path), rows=len(df), columns=df.shape[1], duration_sec=duration_sec, mlflow_run_id=run_id, ) def _latest_mlflow_run_id(experiment_name: str) -> str | None: """Return the newest MLflow run id, degrading to None when tracking is off.""" if os.environ.get("NEUROBRIDGE_DISABLE_MLFLOW") == "1": return None try: runs = mlflow.search_runs( experiment_names=[experiment_name], max_results=1, order_by=["start_time DESC"], ) except Exception as e: logger.warning("MLflow run lookup failed for %s: %s", experiment_name, e) return None return str(runs.iloc[0]["run_id"]) if len(runs) else None @router.post("/bbb", response_model=PipelineResponse) def run_bbb(req: BBBRequest) -> PipelineResponse: """Run the BBB pipeline; return rows/cols/duration + the MLflow run id.""" return _wrap( "bbb_pipeline", Path(req.output_path), lambda: bbb_pipeline.run_pipeline( input_path=Path(req.input_path), output_path=Path(req.output_path), smiles_col=req.smiles_col, n_bits=req.n_bits, radius=req.radius, ), ) @router.post("/eeg", response_model=PipelineResponse) def run_eeg(req: EEGRequest) -> PipelineResponse: """Run the EEG pipeline; return rows/cols/duration + the MLflow run id.""" return _wrap( "eeg_pipeline", Path(req.output_path), lambda: eeg_pipeline.run_pipeline( input_path=Path(req.input_path), output_path=Path(req.output_path), epoch_duration_s=req.epoch_duration_s, eog_ch_name=req.eog_ch_name, n_components=req.n_components, random_state=req.random_state, ), ) @router.post("/mri", response_model=PipelineResponse) def run_mri(req: MRIRequest) -> PipelineResponse: """Run the MRI pipeline; return rows/cols/duration + the MLflow run id.""" return _wrap( "mri_pipeline", Path(req.output_path), lambda: mri_pipeline.run_pipeline( input_dir=Path(req.input_dir), sites_csv=Path(req.sites_csv), output_path=Path(req.output_path), ), ) # Default artifact location. Overridable via BBB_MODEL_PATH env var so tests # can point at a tmp-built model without touching production paths. _DEFAULT_BBB_MODEL_PATH = Path("data/processed/bbb_model.joblib") _DEFAULT_MRI_MODEL_PATH = Path("data/processed/mri_model.onnx") def _bbb_model_path() -> Path: """Return the BBB model artifact path, overridable via BBB_MODEL_PATH env var.""" return Path(os.environ.get("BBB_MODEL_PATH", str(_DEFAULT_BBB_MODEL_PATH))) def _mri_model_path() -> Path: """Return the MRI ONNX model artifact path, overridable via MRI_MODEL_PATH.""" return Path(os.environ.get("MRI_MODEL_PATH", str(_DEFAULT_MRI_MODEL_PATH))) # Per-worker rolling window of recent prediction confidences. # Cleared on worker restart; multi-worker setups have independent windows. WORKER_CONFIDENCE_DEQUE: deque[float] = deque(maxlen=100) _DRIFT_MIN_SAMPLES = 10 def _compute_drift_z(model, confidence: float) -> tuple[float | None, int]: """Append `confidence` to the worker deque and compute the drift z-score. Returns (drift_z, rolling_n). drift_z is None until both: (1) the deque has at least `_DRIFT_MIN_SAMPLES` samples, AND (2) the model has `_neurobridge_train_stats` attached. z = (rolling_median - train_median) / max(train_std, 1e-9) """ import statistics WORKER_CONFIDENCE_DEQUE.append(float(confidence)) rolling_n = len(WORKER_CONFIDENCE_DEQUE) stats = getattr(model, "_neurobridge_train_stats", None) if rolling_n < _DRIFT_MIN_SAMPLES or stats is None: return None, rolling_n rolling_median = statistics.median(WORKER_CONFIDENCE_DEQUE) train_median = float(stats["median"]) train_std = max(float(stats["std"]), 1e-9) drift_z = (rolling_median - train_median) / train_std return float(drift_z), rolling_n _PROVENANCE_CACHE: ModelProvenance | None = None _MODEL_VERSION = "v1" # bump manually per train cycle def _build_provenance(model) -> ModelProvenance: """Look up the most recent BBB MLflow run; build a ModelProvenance. Cached at module level so we hit MLflow once per worker. Failures (no runs found, MLflow unreachable, NEUROBRIDGE_DISABLE_MLFLOW=1) all degrade to a partial ModelProvenance with mlflow_run_id=None — the badge still renders, just without a run id. """ global _PROVENANCE_CACHE if _PROVENANCE_CACHE is not None: # Refresh n_examples each call from the model (cheap lookup). n_train = None stats = getattr(model, "_neurobridge_train_stats", None) if stats is not None: n_train = int(stats.get("n_train", 0)) or None return _PROVENANCE_CACHE.model_copy(update={"n_examples": n_train}) run_id: str | None = None train_date: str | None = None if os.environ.get("NEUROBRIDGE_DISABLE_MLFLOW") != "1": try: runs = mlflow.search_runs( experiment_names=["bbb_pipeline"], max_results=1, order_by=["start_time DESC"], ) if len(runs): row = runs.iloc[0] run_id = str(row["run_id"]) ts = row.get("start_time") if ts is not None: train_date = str(pd.Timestamp(ts).isoformat()) except Exception as e: # broad: MLflow store unreachable, schema mismatch, etc. logger.warning("MLflow provenance lookup failed: %s", e) n_train = None stats = getattr(model, "_neurobridge_train_stats", None) if stats is not None: n_train = int(stats.get("n_train", 0)) or None _PROVENANCE_CACHE = ModelProvenance( mlflow_run_id=run_id, model_version=_MODEL_VERSION, train_date=train_date, n_examples=n_train, ) return _PROVENANCE_CACHE def _matching_calibration_bin(model, confidence: float) -> CalibrationContext | None: """Pick the highest-threshold bin whose threshold <= confidence. None if no match or no metadata.""" bins = getattr(model, "_neurobridge_calibration", None) if not bins: return None matched = None for bin_ in bins: if bin_["threshold"] <= confidence: matched = bin_ else: break if matched is None: return None return CalibrationContext( threshold=matched["threshold"], precision=matched["precision"], support=matched["support"], ) @predict_router.post("/bbb", response_model=BBBPredictResponse) def predict_bbb(req: BBBPredictRequest) -> BBBPredictResponse: """Predict BBB permeability + return SHAP attributions for one SMILES. Returns 503 if the model artifact is missing (operator hasn't run the trainer CLI yet); 400 on invalid SMILES; 200 with the decision payload on success. """ artifact = _bbb_model_path() if not artifact.exists(): raise HTTPException( status_code=503, detail=( f"BBB model artifact not available at {artifact}. " f"Run `python -m src.models.bbb_model` to train it." ), ) try: model = bbb_model.load(artifact) except FileNotFoundError as e: raise HTTPException(status_code=503, detail=str(e)) try: pred = bbb_model.predict_with_proba(model, req.smiles) attributions = bbb_model.explain_prediction(model, req.smiles, top_k=req.top_k) except ValueError as e: raise HTTPException(status_code=400, detail=str(e)) label_text = "permeable" if pred["label"] == 1 else "non-permeable" calibration = _matching_calibration_bin(model, pred["confidence"]) drift_z, rolling_n = _compute_drift_z(model, pred["confidence"]) provenance = _build_provenance(model) return BBBPredictResponse( label=pred["label"], label_text=label_text, confidence=pred["confidence"], top_features=[FeatureAttribution(**a) for a in attributions], calibration=calibration, drift_z=drift_z, rolling_n=rolling_n, provenance=provenance, ) @predict_router.post("/bbb_permeability_map", response_model=BBBPermeabilityMapResponse) def predict_bbb_permeability_map(req: BBBPermeabilityMapRequest) -> BBBPermeabilityMapResponse: """Compute a BBB permeability score from MRI input. Two modes: - heuristic_proxy (default): reuses the 2D resnet18 4-class classifier; score = 1 - P(NonDemented). Demo-ready today. - dce_onnx (real DCE artifact): loads an ONNX model trained on 4D DCE data; emits a Ktrans map normalised to [0, 1]. Stub — drop the real artifact at data/processed/bbb_permeability_dce.onnx (or set BBB_PERMEABILITY_DCE_PATH). Researcher-persona route — does NOT feed into the fusion engine. """ from src.models import bbb_permeability_map as bbb_perm try: result = bbb_perm.compute_permeability( input_path=Path(req.input_path), mode=req.mode, # type: ignore[arg-type] ) except FileNotFoundError as e: raise HTTPException(status_code=404, detail=str(e)) except ValueError as e: raise HTTPException(status_code=400, detail=str(e)) return BBBPermeabilityMapResponse( permeability_score=float(result["permeability_score"]), interpretation=str(result["interpretation"]), method=str(result["method"]), voxel_map_available=bool(result.get("voxel_map_available", False)), ) @research_router.post("/drug_dose_adjustment", response_model=DrugDoseAdjustmentResponse) def research_drug_dose_adjustment(req: DrugDoseAdjustmentRequest) -> DrugDoseAdjustmentResponse: """Suggest a revised drug dose given patient BBB permeability + drug profile. If `smiles` is supplied, the BBB classifier is consulted to populate `drug_bbb_permeable` (overriding any explicit value). Researcher-persona route — output is a research suggestion, NOT medical advice. """ from src.research import drug_dose_adjuster drug_permeable: bool | None = req.drug_bbb_permeable if req.smiles: try: artifact = _bbb_model_path() if artifact.exists(): model = bbb_model.load(artifact) bbb_pred = bbb_model.predict_with_proba(model, req.smiles) drug_permeable = bool(bbb_pred["label"] == 1) except (FileNotFoundError, ValueError, KeyError) as e: logger.warning("could not auto-resolve BBB permeability for smiles=%s: %s", req.smiles, e) try: adj = drug_dose_adjuster.adjust( baseline_dose_mg=req.baseline_dose_mg, bbb_permeability_score=req.bbb_permeability_score, drug_bbb_permeable=drug_permeable, ) except ValueError as e: raise HTTPException(status_code=400, detail=str(e)) return DrugDoseAdjustmentResponse( recommended_dose_mg=adj.recommended_dose_mg, adjustment_factor=adj.adjustment_factor, risk_level=adj.risk_level, rationale=adj.rationale, drug_bbb_permeable=drug_permeable, ) @predict_router.post("/eeg", response_model=EEGPredictResponse) def predict_eeg(req: EEGPredictRequest) -> EEGPredictResponse: """Predict from EEG features using an externally-trained sklearn classifier. Real artifact lands at data/processed/eeg_clf.joblib (override via EEG_CLF_ARTIFACT). For the demo a stub fixture (RandomForestClassifier on synthetic features) is acceptable — the response shape stays stable. """ import numpy as np from src.models import eeg_model artifact = Path(os.environ.get("EEG_CLF_ARTIFACT", "data/processed/eeg_clf.joblib")) if not artifact.exists(): raise HTTPException( status_code=503, detail=( f"EEG model artifact not available at {artifact}. " "Drop the trained joblib at this path or set EEG_CLF_ARTIFACT." ), ) try: clf = eeg_model.load(artifact) features = np.asarray(req.features, dtype=np.float32) out = eeg_model.predict_features(clf, features) except FileNotFoundError as e: raise HTTPException(status_code=404, detail=str(e)) except ValueError as e: raise HTTPException(status_code=400, detail=str(e)) return EEGPredictResponse( label=int(out["label"]), label_text=str(out["label_text"]), confidence=float(out["confidence"]), probabilities=[EEGClassProbability(**p) for p in out["probabilities"]], ) @predict_router.post("/mri", response_model=MRIPredictResponse) def predict_mri(req: MRIPredictRequest) -> MRIPredictResponse: """Predict from one MRI image. Backend selected by MRI_MODEL_KIND env. - `volumetric_onnx` (default): NIfTI volume + externally-trained ONNX. - `resnet18_2d`: 2D image (.png/.jpg) + PyTorch state_dict, 4-class Alzheimer's classifier (MildDemented/ModerateDemented/NonDemented/VeryMildDemented). """ from src.models import mri_selector kind = mri_selector.current_kind() if kind == "resnet18_2d": artifact = Path(os.environ.get( "MRI_MODEL_PATH_2D", "data/processed/mri_dl_2d/best_model.pt", )) else: artifact = _mri_model_path() if not artifact.exists(): raise HTTPException( status_code=503, detail=( f"MRI model artifact not available at {artifact} (kind={kind}). " "Drop the trained checkpoint at this path, or override the path " "via MRI_MODEL_PATH (3D ONNX) or MRI_MODEL_PATH_2D (2D resnet18)." ), ) try: if kind == "resnet18_2d": pred = mri_selector.predict( input_path=Path(req.input_path), checkpoint_path=artifact, ) else: pred = mri_selector.predict( input_path=Path(req.input_path), checkpoint_path=artifact, target_shape=tuple(req.target_shape), label_names=tuple(req.label_names) if req.label_names else None, ) except FileNotFoundError as e: raise HTTPException(status_code=404, detail=str(e)) except ValueError as e: raise HTTPException(status_code=400, detail=str(e)) return MRIPredictResponse( label=int(pred["label"]), label_text=str(pred["label_text"]), confidence=float(pred["confidence"]), probabilities=[ MRIClassProbability(**p) for p in pred["probabilities"] ], input_path=req.input_path, model_path=str(artifact), ) @router.post("/mri/diagnostics", response_model=MRIDiagnosticsResponse) def mri_diagnostics(req: MRIDiagnosticsRequest) -> MRIDiagnosticsResponse: """Run the MRI pipeline twice and return pre/post ComBat data + site-gap KPIs.""" input_dir = Path(req.input_dir) sites_csv = Path(req.sites_csv) try: df = mri_pipeline.compute_harmonization_diagnostics( input_dir=input_dir, sites_csv=sites_csv, ) except FileNotFoundError as e: raise HTTPException(status_code=404, detail=str(e)) except KeyError as e: raise HTTPException(status_code=400, detail=str(e)) if df.empty: return MRIDiagnosticsResponse( rows=[], site_gap_pre=0.0, site_gap_post=0.0, reduction_factor=0.0, ) # Site-gap KPI on the first feature, averaged per site feat = df["feature"].iloc[0] feat_df = df[df["feature"] == feat] pre_means = feat_df[feat_df["harmonization_state"] == "Pre-ComBat"].groupby( "site" )["feature_value"].mean() post_means = feat_df[feat_df["harmonization_state"] == "Post-ComBat"].groupby( "site" )["feature_value"].mean() site_gap_pre = float(pre_means.max() - pre_means.min()) site_gap_post = float(post_means.max() - post_means.min()) eps = 1e-9 reduction_factor = site_gap_pre / max(site_gap_post, eps) rows = [ HarmonizationRow(**rec) for rec in df.to_dict(orient="records") ] return MRIDiagnosticsResponse( rows=rows, site_gap_pre=site_gap_pre, site_gap_post=site_gap_post, reduction_factor=reduction_factor, ) @explain_router.post("/bbb", response_model=BBBExplainResponse) def explain_bbb(req: BBBExplainRequest) -> BBBExplainResponse: """Natural-language rationale for a single BBB prediction. Always returns 200 — the explainer is guaranteed to produce a rationale via deterministic-template fallback. Pydantic enforces a non-empty top_features list; an empty list returns 422 from FastAPI before this handler runs. """ payload: llm_explainer.ExplainPayload = { "smiles": req.smiles, "label": req.label, "label_text": req.label_text, "confidence": req.confidence, "top_features": [ {"feature": f.feature, "shap_value": f.shap_value} for f in req.top_features ], "calibration": ( None if req.calibration is None else { "threshold": req.calibration.threshold, "precision": req.calibration.precision, "support": req.calibration.support, } ), "drift_z": req.drift_z, "user_question": req.user_question or "", } result = llm_explainer.explain(payload) return BBBExplainResponse( rationale=result["rationale"], source=result["source"], model=result["model"], ) @explain_router.post("/eeg", response_model=EEGExplainResponse) def explain_eeg(req: EEGExplainRequest) -> EEGExplainResponse: """Natural-language rationale for an EEG pipeline run.""" payload = { "rows": req.rows, "columns": req.columns, "duration_sec": req.duration_sec, "mlflow_run_id": req.mlflow_run_id, "user_question": req.user_question or "", } result = llm_explainer.explain(payload, modality="eeg") return EEGExplainResponse( rationale=result["rationale"], source=result["source"], model=result["model"], ) @explain_router.post("/mri", response_model=MRIExplainResponse) def explain_mri(req: MRIExplainRequest) -> MRIExplainResponse: """Natural-language rationale for an MRI ComBat diagnostic run.""" payload = { "site_gap_pre": req.site_gap_pre, "site_gap_post": req.site_gap_post, "reduction_factor": req.reduction_factor, "n_subjects": req.n_subjects, "user_question": req.user_question or "", } result = llm_explainer.explain(payload, modality="mri") return MRIExplainResponse( rationale=result["rationale"], source=result["source"], model=result["model"], ) @experiments_router.get("/runs", response_model=MLflowRunsResponse) def list_runs(limit: int = 50) -> MLflowRunsResponse: """List recent MLflow runs across known experiments. Returns an empty list when MLflow is disabled or unreachable. """ if os.environ.get("NEUROBRIDGE_DISABLE_MLFLOW") == "1": return MLflowRunsResponse(runs=[]) summaries: list[MLflowRunSummary] = [] for exp_name in ("bbb_pipeline", "eeg_pipeline", "mri_pipeline"): try: df = mlflow.search_runs( experiment_names=[exp_name], max_results=limit, order_by=["start_time DESC"], ) except Exception as e: # broad: MLflow store unreachable / not found logger.warning("MLflow lookup failed for %s: %s", exp_name, e) continue for _, row in df.iterrows(): metrics = { col[len("metrics."):]: float(row[col]) for col in df.columns if col.startswith("metrics.") and pd.notna(row[col]) } params = { col[len("params."):]: str(row[col]) for col in df.columns if col.startswith("params.") and pd.notna(row[col]) } summaries.append( MLflowRunSummary( run_id=str(row["run_id"]), experiment_name=exp_name, start_time=str(pd.Timestamp(row["start_time"]).isoformat()) if pd.notna(row.get("start_time")) else "", status=str(row.get("status", "UNKNOWN")), metrics=metrics, params=params, ) ) summaries.sort(key=lambda s: s.start_time, reverse=True) return MLflowRunsResponse(runs=summaries[:limit]) @experiments_router.post("/diff", response_model=RunDiffResponse) def diff_runs(req: RunDiffRequest) -> RunDiffResponse: """Side-by-side diff of two MLflow runs (metrics + params). Returns 404 if either run id is not found in the local MLflow store. Returns 200 with an empty rows list when MLflow is disabled. """ if os.environ.get("NEUROBRIDGE_DISABLE_MLFLOW") == "1": return RunDiffResponse(rows=[]) try: run_a = mlflow.get_run(req.run_id_a) run_b = mlflow.get_run(req.run_id_b) except Exception as e: raise HTTPException(status_code=404, detail=f"Run not found: {e}") metrics_a = run_a.data.metrics metrics_b = run_b.data.metrics params_a = run_a.data.params params_b = run_b.data.params rows: list[RunDiffRow] = [] for key in sorted(set(metrics_a) | set(metrics_b)): va = metrics_a.get(key) vb = metrics_b.get(key) rows.append( RunDiffRow( key=key, kind="metric", value_a=None if va is None else f"{va:.6g}", value_b=None if vb is None else f"{vb:.6g}", differs=(va != vb), ) ) for key in sorted(set(params_a) | set(params_b)): va = params_a.get(key) vb = params_b.get(key) rows.append( RunDiffRow( key=key, kind="param", value_a=va, value_b=vb, differs=(va != vb), ) ) return RunDiffResponse(rows=rows) # --- Agent router ---------------------------------------------------------- agent_router = APIRouter(prefix="/agent") _DEFAULT_RAG_INDEX_DIR = Path("data/processed/faiss_index") _AGENT_MODEL_ENV = "NEUROBRIDGE_AGENT_MODEL" _AGENT_DEFAULT_MODEL = "openai/gpt-oss-20b:free" # Fallback chain probed at orchestrator-build time. First model returning a # non-404/429 ping wins. Override via NEUROBRIDGE_AGENT_MODEL env (single id) # or NEUROBRIDGE_AGENT_MODEL_CHAIN (comma-separated). _AGENT_FALLBACK_CHAIN: tuple[str, ...] = ( "openai/gpt-oss-20b:free", "minimax/minimax-m2.5:free", "tencent/hy3-preview:free", "inclusionai/ling-2.6-1t:free", "nvidia/nemotron-3-super-120b-a12b:free", "qwen/qwen3-next-80b-a3b-instruct:free", "google/gemma-4-31b-it:free", "meta-llama/llama-3.3-70b-instruct:free", ) # Cache the chosen model per process so we don't probe on every agent call. _AGENT_MODEL_CACHE: dict[str, str] = {} def _pick_working_agent_model(client: Any, candidates: tuple[str, ...]) -> str: """Return the first candidate that responds to a tiny ping; else last one. Cached per process — first /agent/run call probes once; subsequent calls reuse the picked model. To force a re-probe set NEUROBRIDGE_AGENT_MODEL_CHAIN or restart the worker. """ cache_key = "|".join(candidates) if cache_key in _AGENT_MODEL_CACHE: return _AGENT_MODEL_CACHE[cache_key] for m in candidates: try: client.chat.completions.create( model=m, messages=[{"role": "user", "content": "OK"}], max_tokens=4, temperature=0, ) logger.info("agent model selected: %s", m) _AGENT_MODEL_CACHE[cache_key] = m return m except Exception as e: logger.info("agent model unavailable: %s (%s)", m, type(e).__name__) fallback = candidates[-1] logger.warning("no agent model responded; falling back to %s", fallback) _AGENT_MODEL_CACHE[cache_key] = fallback return fallback def _build_orchestrator(): """Construct the default orchestrator. Patchable in tests.""" from openai import OpenAI from src.agents.orchestrator import Orchestrator from src.agents.prompts import ORCHESTRATOR_SYSTEM_PROMPT from src.agents.routing import build_retrieval_query, route_pipeline_input from src.agents.tools import build_default_tools api_key = os.environ.get("OPENROUTER_API_KEY") if not api_key: raise HTTPException( status_code=503, detail="OPENROUTER_API_KEY not set; agent surface unavailable.", ) client = OpenAI( base_url="https://openrouter.ai/api/v1", api_key=api_key, timeout=30.0, ) rag_dir = _DEFAULT_RAG_INDEX_DIR if _DEFAULT_RAG_INDEX_DIR.exists() else None clinical_idx = Path(os.environ.get( "CLINICAL_RAG_INDEX_PATH", "data/external_rag/index/rag_index.pkl", )) tools = build_default_tools( rag_index_dir=rag_dir, clinical_rag_index_path=clinical_idx if clinical_idx.exists() else None, ) # Resolve agent model. NEUROBRIDGE_AGENT_MODEL overrides; otherwise probe # the fallback chain (NEUROBRIDGE_AGENT_MODEL_CHAIN env to override the # candidate list) and pick the first one that responds. Demo robustness: # OpenRouter free-tier IDs churn; this avoids hard-coding a stale id. explicit = os.environ.get(_AGENT_MODEL_ENV) if explicit: model = explicit else: chain_raw = os.environ.get("NEUROBRIDGE_AGENT_MODEL_CHAIN") chain = ( tuple(s.strip() for s in chain_raw.split(",") if s.strip()) if chain_raw else _AGENT_FALLBACK_CHAIN ) model = _pick_working_agent_model(client, chain) return Orchestrator( llm_client=client, tools=tools, system_prompt=ORCHESTRATOR_SYSTEM_PROMPT, model=model, max_steps=5, enforce_workflow=True, workflow_pipeline_tools={ "run_bbb_pipeline", "run_eeg_pipeline", "run_mri_pipeline", }, workflow_retrieval_tool="retrieve_context", workflow_router=route_pipeline_input, workflow_query_builder=build_retrieval_query, ) @agent_router.post("/run", response_model=AgentRunResponse) def run_agent(req: AgentRunRequest) -> AgentRunResponse: """Run the orchestrator on `user_input`. Picks a pipeline + grounds via RAG.""" orch = _build_orchestrator() user_text = req.user_input if req.user_question: user_text = f"{req.user_input}\n\nUser question: {req.user_question}" result = orch.run(user_text, context={"sites_csv": req.sites_csv}) return AgentRunResponse( text=result.text, trace=[ AgentToolTraceItem(name=t.name, args=t.args, result=t.result, error=t.error) for t in result.trace ], model=result.model, finish_reason=result.finish_reason, ) # --- Fusion router --------------------------------------------------------- fusion_router = APIRouter(prefix="/fusion") @fusion_router.post("/predict", response_model=FusionResponse) def fusion_predict(req: FusionRequest) -> FusionResponse: """Combine MRI, EEG, and clinical scores into per-disease confidence.""" from src.fusion.engine import fuse as fuse_engine return fuse_engine(req)