| |
| from langchain_core.tools import BaseTool |
| from typing import Type, List, Dict, Any, Optional |
| from pydantic import BaseModel, Field |
|
|
| from services.logger import app_logger |
| from services.metrics import log_tool_usage |
|
|
| try: |
| from quantum.optimizer import optimize_treatment |
| except ImportError: |
| app_logger.warning("Actual 'quantum.optimizer.optimize_treatment' not found. Using mock function for QuantumTreatmentOptimizerTool.") |
| def optimize_treatment(patient_data: Dict[str, Any], current_treatments: List[str], conditions: List[str]) -> Dict[str, Any]: |
| import random, time |
| time.sleep(random.uniform(0.5,1.0)) |
| mock_actions = [f"Mock action for {conditions[0] if conditions else 'general health'} considering {patient_data.get('age', 'N/A')} years old."] |
| if current_treatments: mock_actions.append(f"Review interaction with {current_treatments[0]}.") |
| return { |
| "simulated_optimization_id": f"QO-MOCK-{random.randint(1000,9999)}", |
| "suggested_actions": mock_actions, |
| "primary_focus_condition": conditions[0] if conditions else "Overall Assessment", |
| "confidence_level_simulated": random.uniform(0.7, 0.9), |
| "summary_notes": "This is a simulated optimization result. Always consult with medical professionals for actual treatment decisions.", |
| } |
|
|
| class QuantumOptimizerInput(BaseModel): |
| patient_data: Dict[str, Any] = Field( |
| description=( |
| "A dictionary of relevant patient characteristics. " |
| "Examples: {'age': 55, 'gender': 'Male', 'relevant_labs': {'creatinine': 1.2, 'hbA1c': 7.5}, " |
| "'allergies': ['penicillin']}. This should be populated from the overall patient context." |
| ) |
| ) |
| current_treatments: List[str] = Field( |
| description="A list of current medications or therapies (e.g., ['Aspirin 81mg', 'Metformin 500mg OD'])." |
| ) |
| conditions: List[str] = Field( |
| description="A list of primary diagnosed conditions or symptoms to be addressed (e.g., ['Type 2 Diabetes', 'Hypertension'])." |
| ) |
|
|
| class QuantumTreatmentOptimizerTool(BaseTool): |
| name: str = "quantum_treatment_optimizer" |
| description: str = ( |
| "A specialized (simulated) tool that uses advanced algorithms to suggest optimized or alternative treatment plans " |
| "based on provided patient data, current treatments, and diagnosed conditions. " |
| "Use this when seeking novel therapeutic strategies, or to optimize complex polypharmacy. " |
| "You MUST provide detailed 'patient_data', 'current_treatments', and 'conditions'." |
| ) |
| args_schema: Type[BaseModel] = QuantumOptimizerInput |
|
|
| def _format_results_for_llm(self, optimization_output: Dict[str, Any]) -> str: |
| if not optimization_output or not isinstance(optimization_output, dict): |
| return "The optimizer did not return a structured result or the result was empty." |
| summary_lines = ["Quantum Treatment Optimizer Suggestions (Simulated):"] |
| actions = optimization_output.get("suggested_actions", []) |
| if actions: |
| summary_lines.append(" Key Suggested Actions/Considerations:") |
| for action_str in actions: summary_lines.append(f" - {action_str}") |
| focus = optimization_output.get("primary_focus_condition") |
| if focus: summary_lines.append(f" Primary Focus: Addressing {focus}.") |
| confidence = optimization_output.get("confidence_level_simulated") |
| if confidence is not None: summary_lines.append(f" Simulated Confidence Level: {confidence:.0%}") |
| notes = optimization_output.get("summary_notes") |
| if notes: summary_lines.append(f" Summary Notes: {notes}") |
| sim_id = optimization_output.get("simulated_optimization_id") |
| if sim_id: summary_lines.append(f" (Simulated Optimization ID: {sim_id})") |
| if len(summary_lines) == 1: |
| return f"The optimizer processed the request but provided no specific actionable suggestions. Raw data: {str(optimization_output)[:300]}" |
| return "\n".join(summary_lines) |
|
|
| def _run(self, patient_data: Dict[str, Any], current_treatments: List[str], conditions: List[str], **kwargs: Any) -> str: |
| app_logger.info( |
| f"Quantum Optimizer Tool called. Patient Data Keys: {list(patient_data.keys())}, " |
| f"Treatments: {current_treatments}, Conditions: {conditions}" |
| ) |
| log_tool_usage(self.name, {"conditions_count": len(conditions), "treatments_count": len(current_treatments)}) |
| if not patient_data or not conditions: |
| missing = [item for item, val in [("'patient_data'", patient_data), ("'conditions'", conditions)] if not val] |
| return f"Error: Insufficient information. Missing: {', '.join(missing)}. Provide comprehensive details." |
| try: |
| optimization_output: Dict[str, Any] = optimize_treatment( |
| patient_data=patient_data, current_treatments=current_treatments, conditions=conditions |
| ) |
| app_logger.info(f"Quantum optimizer raw output: {str(optimization_output)[:500]}...") |
| return self._format_results_for_llm(optimization_output) |
| except ImportError as ie: |
| app_logger.error(f"ImportError for quantum.optimizer: {ie}", exc_info=True) |
| return "Error: The core optimization module is currently unavailable." |
| except Exception as e: |
| app_logger.error(f"Unexpected error during quantum optimization: {e}", exc_info=True) |
| return f"Error in optimization process: {str(e)}. Ensure input data is correct." |
|
|
| async def _arun(self, patient_data: Dict[str, Any], current_treatments: List[str], conditions: List[str], **kwargs: Any) -> str: |
| return self._run(patient_data, current_treatments, conditions, **kwargs) |
