File size: 2,201 Bytes
2276177 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 |
from langchain.tools import BaseTool
from typing import Type, List, Dict, Any
from pydantic import BaseModel, Field
from quantum.optimizer import optimize_treatment
from services.logger import app_logger
from services.metrics import log_tool_usage
class QuantumOptimizerInput(BaseModel):
patient_data: Dict[str, Any] = Field(description="Dictionary of relevant patient characteristics (e.g., {'age': 55, 'gender': 'male'}).")
current_treatments: List[str] = Field(description="List of current medications or therapies (e.g., ['Aspirin 81mg', 'Metformin 500mg']).")
conditions: List[str] = Field(description="List of diagnosed conditions (e.g., ['Type 2 Diabetes', 'Hypertension']).")
class QuantumTreatmentOptimizerTool(BaseTool):
name: str = "quantum_treatment_optimizer"
description: str = (
"A specialized tool that uses quantum-inspired algorithms to suggest optimized treatment plans. "
"Provide patient data, current treatments, and diagnosed conditions. "
"Use this when seeking novel therapeutic strategies or to optimize complex polypharmacy."
)
args_schema: Type[BaseModel] = QuantumOptimizerInput
def _run(self, patient_data: Dict[str, Any], current_treatments: List[str], conditions: List[str]) -> str:
app_logger.info(f"Quantum Optimizer Tool called with: {patient_data}, {current_treatments}, {conditions}")
log_tool_usage(self.name)
try:
result = optimize_treatment(patient_data, current_treatments, conditions)
# Format result for LLM
# Example: "Optimized suggestions: ..., Confidence: ..., Summary: ..."
# You might want to pretty-print the dict or convert to a string summary
return f"Quantum Optimizer Results: {result}"
except Exception as e:
app_logger.error(f"Error in QuantumTreatmentOptimizerTool: {e}")
return f"Error during quantum optimization: {str(e)}"
async def _arun(self, patient_data: Dict[str, Any], current_treatments: List[str], conditions: List[str]) -> str:
# For simplicity, using sync version for now
return self._run(patient_data, current_treatments, conditions) |