Spaces:
Sleeping
Sleeping
| """ | |
| Endpoints for running what-if supply chain simulations. | |
| """ | |
| import logging | |
| from datetime import datetime | |
| from typing import Annotated, Any, List, Literal | |
| from fastapi import APIRouter, Depends, HTTPException | |
| from pydantic import BaseModel, Field | |
| from ..deps import get_shock_propagator, get_graph_service, get_data_loader, get_gemini_analyst | |
| from ..services.shock_propagation import ShockPropagator | |
| from ..services.graph_service import GraphService | |
| from ..services.gemini_analyst import GeminiAnalyst | |
| from ..models.drug import Drug | |
| # Setup Logging | |
| logger = logging.getLogger("backend.simulate") | |
| router = APIRouter(prefix="/api/v1", tags=["simulate"]) | |
| class SimulationRequest(BaseModel): | |
| """Parameters for a supply chain shock simulation.""" | |
| province: str = Field(..., description="The ID of the province where the shock occurs") | |
| duration_days: int = Field(..., ge=1, le=180, description="Duration of the disruption in days") | |
| severity: Literal["warning", "partial_shutdown", "full_shutdown"] = Field(..., description="Intensity of the shock") | |
| class SimulationResult(BaseModel): | |
| """The outcome of a shock simulation.""" | |
| affected_drugs: List[Drug] = Field(..., description="List of top 10 drugs most affected by the shock") | |
| propagation_explanation: str = Field(..., description="AI or rule-based explanation of why these drugs are at risk") | |
| simulated_at: datetime = Field(default_factory=datetime.utcnow) | |
| async def post_simulate( | |
| request: SimulationRequest, | |
| propagator: Annotated[ShockPropagator, Depends(get_shock_propagator)], | |
| graph_service: Annotated[GraphService, Depends(get_graph_service)], | |
| data_loader: Annotated[Any, Depends(get_data_loader)], | |
| analyst: Annotated[GeminiAnalyst, Depends(get_gemini_analyst)] | |
| ) -> SimulationResult: | |
| """ | |
| Simulates a regional production shock and predicts its impact on the drug supply. | |
| Uses the GNN shock propagation model to estimate risk scores across the network. | |
| """ | |
| logger.info(f"Simulation Request: province={request.province}, dur={request.duration_days}, sev={request.severity}") | |
| # 1. Validate Province — be lenient: accept province name or ID | |
| node = graph_service.get_node(request.province) | |
| if not node or node["type"] != "province": | |
| # Also try matching by name (case-insensitive) | |
| all_provinces = graph_service.nodes_by_type("province") | |
| match = next( | |
| (p for p in all_provinces if p and p["name"].lower() == request.province.lower()), | |
| None | |
| ) | |
| if match: | |
| request = request.model_copy(update={"province": match["id"]}) | |
| else: | |
| available = [p["name"] for p in all_provinces if p] | |
| raise HTTPException( | |
| status_code=422, | |
| detail=f"Province '{request.province}' not found. Available: {', '.join(available[:10])}" | |
| ) | |
| # 2. Run Shock Propagation | |
| risk_map = propagator.simulate_shock( | |
| province=request.province, | |
| duration_days=request.duration_days, | |
| severity=request.severity | |
| ) | |
| # 3. Filter and Hydrate Top 10 | |
| filtered_risks = sorted( | |
| [(d_id, score) for d_id, score in risk_map.items() if score > 30], | |
| key=lambda x: x[1], | |
| reverse=True | |
| )[:10] | |
| hydrated_drugs = [] | |
| top_drug_names = [] | |
| for d_id, score in filtered_risks: | |
| drug = data_loader.get_drug(d_id) | |
| if drug: | |
| # Pydantic v2: use model_copy to update immutable fields | |
| drug = drug.model_copy(update={"current_risk": score}) | |
| hydrated_drugs.append(drug) | |
| top_drug_names.append(drug.name) | |
| # 4. Generate Explanation (falls back to rule-based if Gemini unavailable) | |
| explanation = propagator.propagate_explanation( | |
| province=request.province, | |
| duration_days=request.duration_days, | |
| severity=request.severity, | |
| top_affected=top_drug_names | |
| ) | |
| if top_drug_names: | |
| try: | |
| from ..config import settings | |
| explain_prompt = ( | |
| f"Briefly explain in 2-3 sentences why a {request.severity.replace('_', ' ')} " | |
| f"in {request.province} for {request.duration_days} days affects these drugs: " | |
| f"{', '.join(top_drug_names)}. Focus on shared API precursors and regional manufacturing concentration." | |
| ) | |
| model = analyst.genai.GenerativeModel(settings.GEMINI_FLASH_MODEL) | |
| explanation = model.generate_content(explain_prompt).text | |
| except Exception as e: | |
| logger.warning(f"Gemini explanation failed, using rule-based fallback: {e}") | |
| return SimulationResult( | |
| affected_drugs=hydrated_drugs, | |
| propagation_explanation=explanation, | |
| simulated_at=datetime.utcnow() | |
| ) | |