import os import re import json import time import asyncio import functools from langchain_core.messages import SystemMessage, HumanMessage, AIMessage from src.utils.logger import setup_logger logger = setup_logger("Agents") def throttle_agent(func): @functools.wraps(func) async def wrapper(self, state, *args, **kwargs): agent_name = getattr(self, "agent_name", self.__class__.__name__) logger.info(f"[{agent_name}] >>> Start executing") start_time = time.time() try: res = await func(self, state, *args, **kwargs) elapsed = time.time() - start_time logger.info(f"[{agent_name}] <<< Execution completed in {elapsed:.3f} seconds.") if elapsed < 1.0: delay = 1.5 - elapsed logger.info(f"[{agent_name}] Execution was faster than 1.0s. Throttling: waiting {delay:.3f}s to reach 1.5s total time.") await asyncio.sleep(delay) logger.info(f"[{agent_name}] Throttling completed. Proceeding to next step.") return res except Exception as e: logger.error(f"[{agent_name}] Exception during execution: {e}") raise e return wrapper from src.agent_params import get_agent_params from src.core.model_manager import model_manager from src.core.state import AgentState from src.core.evidence_models import ClinicalOutputWithEvidence, EvidenceCitation from src.tools.web_tools import web_search_tool from src.tools.dietary_tools import search_guidelines, get_nutritional_data, page_indexed_retrieval from src.tools.patient_memory import save_patient_memory, get_patient_memory from src.agents.role_utils import classify_role from datetime import datetime class BaseAgent: def __init__(self, fallback_prompt: str, prompt_file: str = None, tools: list = None, agent_name: str = None): self.agent_name = agent_name or self.__class__.__name__ self.fallback_prompt = fallback_prompt self.prompt_file = prompt_file self.tools = tools or [] self.params = get_agent_params(self.agent_name) self.temperature = float(self.params.get("temperature", 0.0)) self.model_name = self.params.get("model_name") self._refresh_llm() def _refresh_llm(self): llm = model_manager.get_llm( temperature=self.temperature, model_name=self.model_name, ) if self.tools: llm = llm.bind_tools(self.tools) self.llm = llm def parse_json_response(self, response_text: str): if not response_text: return {} text = response_text.strip() try: return json.loads(text) except json.JSONDecodeError: match = re.search(r"\{.*\}", text, re.S) if not match: return {} try: return json.loads(match.group(0)) except json.JSONDecodeError: return {} @property def system_prompt(self) -> str: """Dynamically load prompt from file if available, otherwise use fallback.""" if self.prompt_file: current_dir = os.path.dirname(os.path.abspath(__file__)) prompt_path = os.path.abspath(os.path.join(current_dir, "..", "prompts", self.prompt_file)) try: if os.path.exists(prompt_path): with open(prompt_path, "r", encoding="utf-8") as handle: prompt = handle.read().strip() else: prompt = self.fallback_prompt except Exception: prompt = self.fallback_prompt else: prompt = self.fallback_prompt skip_confidence = self.agent_name in {"ResponseValidator", "SafetyCheck"} confidence_instruction = ( "\n\nAt the end of your response, include a confidence score from 0.0 to 1.0 " "in the format: Confidence: 0.8" ) if not skip_confidence and "confidence" not in prompt.lower(): prompt = f"{prompt}{confidence_instruction}" return prompt @throttle_agent async def run(self, state: AgentState, config=None): """Standard run method for graph nodes.""" messages = [SystemMessage(content=self.system_prompt)] + state["messages"] logger.info(f"--- Sending {len(messages)} messages to LLM ({self.agent_name}) ---") start_time = time.time() try: response = await self.llm.ainvoke(messages, config=config) end_time = time.time() tokens = 0 if hasattr(response, "usage_metadata") and response.usage_metadata: tokens = response.usage_metadata.get("total_tokens", 0) elif "token_usage" in response.response_metadata: tokens = response.response_metadata["token_usage"].get("total_tokens", 0) confidence = None if hasattr(response, "content") and isinstance(response.content, str): match = re.search(r"Confidence(?:\s+Score)?:\s*([0-9.]+)", response.content, re.IGNORECASE) if match: try: confidence = float(match.group(1)) except ValueError: pass metrics = { "agent": self.agent_name, "tokens": tokens, "time": round(end_time - start_time, 3), "confidence": confidence } return {"messages": [response], "metrics": [metrics]} except Exception as e: logger.error(f"Error in {self.agent_name}.run: {e}") raise class RoleClassifier(BaseAgent): def __init__(self): fallback_prompt = """You are a medical triage assistant. Classify the user input into one of five roles: 'patient', 'caregiver', 'clinician', 'researcher', or 'dietary'. Return only the name.""" super().__init__(fallback_prompt, "RoleClassifier.txt") @throttle_agent async def run(self, state: AgentState, config=None, **kwargs): messages = [SystemMessage(content=self.system_prompt)] + state["messages"] logger.info(f"--- RoleClassifier: Sending {len(messages)} messages to LLM ---") user_message = state["messages"][-1].content.lower() if state["messages"] else "" role = classify_role(user_message) start_time = time.time() try: if role is None: response = await self.llm.ainvoke(messages) end_time = time.time() raw = response.content.lower() if not raw.strip(): role = "patient" else: roles = ["patient", "caregiver", "clinician", "researcher", "dietary"] role = next((r for r in roles if r in raw), "patient") tokens = 0 if hasattr(response, "usage_metadata") and response.usage_metadata: tokens = response.usage_metadata.get("total_tokens", 0) elif "token_usage" in response.response_metadata: tokens = response.response_metadata["token_usage"].get("total_tokens", 0) duration = round(end_time - start_time, 3) else: end_time = time.time() tokens = 0 duration = round(end_time - start_time, 3) metrics = { "agent": "RoleClassifier", "tokens": tokens, "time": duration } return {"user_role": role, "metrics": [metrics]} except Exception as e: logger.error(f"Error in RoleClassifier.run: {e}") raise class PatientLLM(BaseAgent): def __init__(self): fallback_prompt = """You are a compassionate medical assistant for patients. Provide helpful, empathetic, and medically sound advice.""" super().__init__(fallback_prompt, "PatientLLM.txt", tools=[web_search_tool, save_patient_memory, get_patient_memory]) class CaregiverLLM(BaseAgent): def __init__(self): fallback_prompt = """You are a supportive caregiver assistant for a diabetes management platform. Help caregivers interpret symptoms, monitor treatment adherence, and know when to escalate to urgent care. Frame advice as practical proxy guidance for a patient while remaining clear and compassionate.""" super().__init__(fallback_prompt, "CaregiverLLM.txt", tools=[web_search_tool, save_patient_memory, get_patient_memory]) class ResponseValidator(BaseAgent): def __init__(self): fallback_prompt = """You are a medical response validator. Check if the last response is medically accurate and follows guidelines. Return JSON only.""" super().__init__(fallback_prompt, "ResponseValidator.txt") @throttle_agent async def run(self, state: AgentState, config=None, **kwargs): last_message = state["messages"][-1].content start_time = time.time() response = await self.llm.ainvoke([ SystemMessage(content=self.system_prompt), HumanMessage(content=f"Verify this response: {last_message}") ]) end_time = time.time() parsed = self.parse_json_response(response.content) decision = parsed.get("decision", "invalid").lower() is_valid = decision == "valid" if not parsed: lower_response = response.content.lower() if re.search(r"\binvalid\b", lower_response): is_valid = False elif re.search(r"\bvalid\b", lower_response): is_valid = True else: is_valid = False tokens = 0 if hasattr(response, "usage_metadata") and response.usage_metadata: tokens = response.usage_metadata.get("total_tokens", 0) elif "token_usage" in response.response_metadata: tokens = response.response_metadata["token_usage"].get("total_tokens", 0) metrics = { "agent": "ResponseValidator", "tokens": tokens, "time": round(end_time - start_time, 3) } return {"is_valid": is_valid, "metrics": [metrics]} class SafetyCheck(BaseAgent): def __init__(self): fallback_prompt = """You are a medical safety officer. Check if the response contains any dangerous advice or misinformation. Return JSON only.""" super().__init__(fallback_prompt, "SafetyCheck.txt") @throttle_agent async def run(self, state: AgentState, config=None, **kwargs): last_message = state["messages"][-1].content start_time = time.time() response = await self.llm.ainvoke([ SystemMessage(content=self.system_prompt), HumanMessage(content=f"Safety check on this: {last_message}") ]) end_time = time.time() parsed = self.parse_json_response(response.content) decision = parsed.get("decision", "unsafe").lower() is_safe = decision == "safe" if not parsed: lower_response = response.content.lower() if re.search(r"\bunsafe\b", lower_response): is_safe = False elif re.search(r"\bsafe\b", lower_response): is_safe = True else: is_safe = False tokens = 0 if hasattr(response, "usage_metadata") and response.usage_metadata: tokens = response.usage_metadata.get("total_tokens", 0) elif "token_usage" in response.response_metadata: tokens = response.response_metadata["token_usage"].get("total_tokens", 0) metrics = { "agent": "SafetyCheck", "tokens": tokens, "time": round(end_time - start_time, 3) } return {"is_safe": is_safe, "metrics": [metrics]} class IntentClassifier(BaseAgent): def __init__(self): fallback_prompt = """You are a clinical intent classifier. Classify into: 'diagnosis', 'treatment', 'monitoring', or 'general'.""" super().__init__(fallback_prompt, "IntentClassifier.txt") @throttle_agent async def run(self, state: AgentState, config=None, **kwargs): start_time = time.time() response = await self.llm.ainvoke([SystemMessage(content=self.system_prompt)] + state["messages"]) end_time = time.time() intent = response.content.lower().strip() tokens = 0 if hasattr(response, "usage_metadata") and response.usage_metadata: tokens = response.usage_metadata.get("total_tokens", 0) elif "token_usage" in response.response_metadata: tokens = response.response_metadata["token_usage"].get("total_tokens", 0) metrics = { "agent": "IntentClassifier", "tokens": tokens, "time": round(end_time - start_time, 3) } return {"intent_type": intent, "metrics": [metrics]} class ClinicalSpecialist(BaseAgent): """ Clinical specialist agent with structured output including evidence citations. Bug 12.3: Provides explainability through guideline sources and evidence levels. """ def __init__(self, specialty: str): fallback_prompt = f"You are a clinical specialist in {specialty}. Provide expert medical support." super().__init__(fallback_prompt, f"ClinicalSpecialist_{specialty}.txt") self.specialty = specialty @throttle_agent async def run(self, state: AgentState, config=None, **kwargs): """ Run clinical specialist with structured output requiring evidence citations. Returns both the text response and evidence citations in AgentState. """ messages = [SystemMessage(content=self.system_prompt)] + state["messages"] logger.info(f"--- ClinicalSpecialist ({self.specialty}): Running with evidence structure ---") start_time = time.time() try: # Use structured output with the LLM if available try: # Try to use with_structured_output for models that support it llm_with_output = self.llm.with_structured_output(ClinicalOutputWithEvidence) response = await llm_with_output.ainvoke(messages, config=config) except (AttributeError, NotImplementedError): # Fallback: regular invocation and manual extraction logger.warning(f"Model does not support structured output, using fallback") response = await self.llm.ainvoke(messages, config=config) # Create a basic ClinicalOutputWithEvidence from the response from src.core.evidence_models import Citation response = ClinicalOutputWithEvidence( recommendation=response.content[:200] if hasattr(response, 'content') else str(response), explanation=response.content if hasattr(response, 'content') else str(response), citations=[ Citation( source_document="Knowledge Base", evidence_level="C", section="General" ) ], confidence_score=0.7 ) end_time = time.time() # Extract tokens tokens = 0 if hasattr(response, "usage_metadata") and response.usage_metadata: tokens = response.usage_metadata.get("total_tokens", 0) elif isinstance(response, dict) and "usage_metadata" in response: tokens = response["usage_metadata"].get("total_tokens", 0) elif hasattr(response, "response_metadata") and "token_usage" in response.response_metadata: tokens = response.response_metadata["token_usage"].get("total_tokens", 0) confidence = 0.7 if isinstance(response, ClinicalOutputWithEvidence): output_content = response.recommendation citations = response.citations confidence = getattr(response, "confidence_score", 0.7) else: output_content = response.content if hasattr(response, 'content') else str(response) citations = [] if isinstance(output_content, str): match = re.search(r"Confidence(?:\s+Score)?:\s*([0-9.]+)", output_content, re.IGNORECASE) if match: try: confidence = float(match.group(1)) except ValueError: pass metrics = { "agent": f"ClinicalSpecialist({self.specialty})", "tokens": tokens, "time": round(end_time - start_time, 3), "confidence": confidence } # Build evidence citations from the structured output evidence_citations = [] for idx, citation in enumerate(citations): evidence_citation = { "recommendation_id": f"{self.specialty}_{idx}", "source_document": citation.source_document if hasattr(citation, 'source_document') else "Unknown", "page_number": getattr(citation, 'page_number', None), "evidence_level": getattr(citation, 'evidence_level', 'C'), "agent_name": f"ClinicalSpecialist({self.specialty})", "timestamp": datetime.utcnow().isoformat() } evidence_citations.append(evidence_citation) from langchain_core.messages import AIMessage return { "messages": [AIMessage(content=output_content)], "metrics": [metrics], "evidence_citations": evidence_citations } except Exception as e: logger.error(f"Error in ClinicalSpecialist({self.specialty}).run: {e}") raise class OutputMerger(BaseAgent): def __init__(self): fallback_prompt = "You are a clinical coordinator. Merge outputs into a single cohesive report." super().__init__(fallback_prompt, "OutputMerger.txt") @throttle_agent async def run(self, state: AgentState, config=None, **kwargs): latest_user_message = None for message in reversed(state["messages"]): if getattr(message, "type", None) == "human": latest_user_message = message.content break human_messages = [] if latest_user_message: human_messages.append(HumanMessage(content=f"Original user request:\n{latest_user_message}")) clinician_outputs = state.get("clinician_outputs") or [] if clinician_outputs: human_messages.append( HumanMessage(content="Latest specialist outputs:\n" + "\n\n".join(clinician_outputs)) ) messages = [SystemMessage(content=self.system_prompt)] + human_messages start_time = time.time() response = None full_content = "" async for chunk in self.llm.astream(messages, config=config): response = chunk if chunk and hasattr(chunk, "content") and isinstance(chunk.content, str): full_content += chunk.content if response is None: response = await self.llm.ainvoke(messages, config=config) if response and hasattr(response, "content") and isinstance(response.content, str): full_content = response.content end_time = time.time() tokens = 0 if hasattr(response, "usage_metadata") and response.usage_metadata: tokens = response.usage_metadata.get("total_tokens", 0) elif "token_usage" in response.response_metadata: tokens = response.response_metadata["token_usage"].get("total_tokens", 0) confidence = None if full_content: match = re.search(r"Confidence(?:\s+Score)?:\s*([0-9.]+)", full_content, re.IGNORECASE) if match: try: confidence = float(match.group(1)) except ValueError: pass metrics = { "agent": self.__class__.__name__, "tokens": tokens, "time": round(end_time - start_time, 3), "confidence": confidence } return {"messages": [response], "metrics": [metrics]} class ResearchAgent(BaseAgent): def __init__(self): fallback_prompt = """You are a medical research assistant. Provide detailed information for researchers.""" super().__init__(fallback_prompt, "ResearchAgent.txt", tools=[web_search_tool, page_indexed_retrieval]) class DietarySpecialist(BaseAgent): def __init__(self): fallback_prompt = """You are a certified dietary specialist. Provide advice based on guidelines.""" super().__init__(fallback_prompt, "DietarySpecialist.txt", tools=[search_guidelines, get_nutritional_data, page_indexed_retrieval, save_patient_memory, get_patient_memory])