Refactor agent and tools for session-based memory management and side effect reporting. Removed medical answer validation tool, added session memory management class, and enhanced side effect reporting with LLM classification. Updated agent functions to support session IDs for improved conversation tracking.

core/agent.py

@@ -17,7 +17,6 @@ from .tools import (
     compare_providers_tool,
     get_current_datetime_tool,
     side_effect_recording_tool,
-    medical_answer_validation_tool,
 )
 
 # LangSmith tracing utilities
@@ -85,7 +84,6 @@ AVAILABLE_TOOLS = [
     compare_providers_tool,
     get_current_datetime_tool,
     side_effect_recording_tool,
-    medical_answer_validation_tool,
 ]
 
 
@@ -94,11 +92,16 @@ SYSTEM_MESSAGE = """
 You are an advanced Medical Advisor Chatbot for healthcare professionals.
 Your primary purpose is to answer clinical and medical questions strictly based on authoritative medical guidelines using the tool "medical_guidelines_knowledge_tool".
 
+Your answers must be concise, medically informative, evidence-based responses in an authoritative, precise, and clinical tone. 
+You will be responding to practicing medical professionals so adjust your answer and language accordingly.
+
 **INSTRUCTIONS:**
 - Always answer using only the information retrieved from medical guidelines via "medical_guidelines_knowledge_tool".
 - **SIDE EFFECT REPORTING**: When a healthcare professional reports an adverse drug reaction, side effect, or medication-related complication, ALWAYS use the "side_effect_recording_tool" first to document the information. Return the tool's response directly to the user without modification. DO NOT use validation or generate additional reports for side effect reporting queries.
 - Use the side effect recording tool when the input contains phrases like: "patient experienced", "side effect", "adverse reaction", "drug reaction", "medication caused", "developed after taking", etc.
 - When the side effect recording tool requests additional information, present the request exactly as provided by the tool.
+- **PROVIDER COMPARISON**: When the user asks to compare guidance between two providers (e.g., "compare NCCN vs ESMO on ..."), use the "compare_providers_tool" with appropriate `provider_a` and `provider_b` values to retrieve side-by-side, cited results.
+- **TIME/DATE QUERIES**: For any questions about the current date/time or references like "today" or "now", use the "get_current_datetime_tool". Treat this tool as the only reliable source of current time information.
 - For every answer, you MUST provide detailed citations including:
   * Source file name
   * Page number
@@ -166,12 +169,47 @@ def get_agent_executor():
         max_execution_time=90,  # tighten a bit to help responsiveness
     )
 
-# Initialize memory
-memory = ConversationBufferWindowMemory(
-    memory_key="chat_history",
-    return_messages=True,
-    max_window_size=10
-)
+# ============================================================================
+# SESSION-BASED MEMORY MANAGEMENT
+# ============================================================================
+
+class SessionMemoryManager:
+    """Manages conversation memory for multiple sessions."""
+    
+    def __init__(self):
+        self._sessions = {}
+        self._default_window_size = 10
+    
+    def get_memory(self, session_id: str = "default") -> ConversationBufferWindowMemory:
+        """Get or create memory for a specific session."""
+        if session_id not in self._sessions:
+            self._sessions[session_id] = ConversationBufferWindowMemory(
+                memory_key="chat_history",
+                return_messages=True,
+                max_window_size=self._default_window_size
+            )
+        return self._sessions[session_id]
+    
+    def clear_session(self, session_id: str) -> bool:
+        """Clear memory for a specific session."""
+        if session_id in self._sessions:
+            self._sessions[session_id].clear()
+            del self._sessions[session_id]
+            return True
+        return False
+    
+    def clear_all_sessions(self):
+        """Clear all session memories."""
+        for memory in self._sessions.values():
+            memory.clear()
+        self._sessions.clear()
+    
+    def get_active_sessions(self) -> list:
+        """Get list of active session IDs."""
+        return list(self._sessions.keys())
+
+# Global session memory manager
+_memory_manager = SessionMemoryManager()
 
 
 # ============================================================================
@@ -304,7 +342,7 @@ def _perform_automatic_validation(user_input: str, response: str) -> str:
 # ============================================================================
 
 # @traceable(name="run_agent_streaming")
-async def run_agent_streaming(user_input: str, max_retries: int = 3) -> AsyncGenerator[str, None]:
+async def run_agent_streaming(user_input: str, session_id: str = "default", max_retries: int = 3) -> AsyncGenerator[str, None]:
     """
     Run the agent with streaming support and comprehensive error handling.
     
@@ -313,6 +351,7 @@ async def run_agent_streaming(user_input: str, max_retries: int = 3) -> AsyncGen
     
     Args:
         user_input (str): The user's input message to process
+        session_id (str, optional): Session identifier for conversation memory. Defaults to "default".
         max_retries (int, optional): Maximum number of retries for recoverable errors. 
                                    Defaults to 3.
     
@@ -343,7 +382,8 @@ async def run_agent_streaming(user_input: str, max_retries: int = 3) -> AsyncGen
             # Tracing for streaming disabled to avoid duplicate traces.
             # We keep tracing only for the AgentExecutor in run_agent().
             current_run_id = None
-            # Load conversation history from memory
+            # Load conversation history from session-specific memory
+            memory = _memory_manager.get_memory(session_id)
             chat_history = memory.load_memory_variables({})["chat_history"]
             
             logger.info(f"Processing user input (attempt {retry_count + 1}): {user_input[:50]}...")
@@ -547,7 +587,7 @@ async def run_agent_streaming(user_input: str, max_retries: int = 3) -> AsyncGen
     yield "Sorry, I was unable to process your request after several attempts. Please try again later."
 
 
-async def safe_run_agent_streaming(user_input: str) -> AsyncGenerator[str, None]:
+async def safe_run_agent_streaming(user_input: str, session_id: str = "default") -> AsyncGenerator[str, None]:
     """
     Streaming wrapper function with additional safety checks and input validation.
     
@@ -557,6 +597,7 @@ async def safe_run_agent_streaming(user_input: str) -> AsyncGenerator[str, None]
     
     Args:
         user_input (str): The user's input message to process
+        session_id (str, optional): Session identifier for conversation memory. Defaults to "default".
         
     Yields:
         str: Chunks of the agent's response as they are generated
@@ -585,7 +626,7 @@ async def safe_run_agent_streaming(user_input: str) -> AsyncGenerator[str, None]
             return
         
         # Stream the response through the main agent function
-        async for chunk in run_agent_streaming(user_input):
+        async for chunk in run_agent_streaming(user_input, session_id):
             yield chunk
         
     except Exception as e:
@@ -595,7 +636,7 @@ async def safe_run_agent_streaming(user_input: str) -> AsyncGenerator[str, None]
 
 
 @traceable(name="run_agent")
-async def run_agent(user_input: str, max_retries: int = 3) -> str:
+async def run_agent(user_input: str, session_id: str = "default", max_retries: int = 3) -> str:
     """
     Run the agent with comprehensive error handling and retry logic.
     
@@ -605,6 +646,7 @@ async def run_agent(user_input: str, max_retries: int = 3) -> str:
     
     Args:
         user_input (str): The user's input message to process
+        session_id (str, optional): Session identifier for conversation memory. Defaults to "default".
         max_retries (int, optional): Maximum number of retries for recoverable errors. 
                                    Defaults to 3.
     
@@ -626,7 +668,8 @@ async def run_agent(user_input: str, max_retries: int = 3) -> str:
     
     while retry_count <= max_retries:
         try:
-            # Load conversation history from memory
+            # Load conversation history from session-specific memory
+            memory = _memory_manager.get_memory(session_id)
             chat_history = memory.load_memory_variables({})["chat_history"]
 
             logger.info(f"Processing user input (attempt {retry_count + 1}): {user_input[:50]}...")
@@ -766,7 +809,7 @@ async def run_agent(user_input: str, max_retries: int = 3) -> str:
     return "Sorry, I was unable to process your request after several attempts. Please try again later."
 
 
-async def safe_run_agent(user_input: str) -> str:
+async def safe_run_agent(user_input: str, session_id: str = "default") -> str:
     """
     Wrapper function for run_agent with additional safety checks and input validation.
     
@@ -776,6 +819,7 @@ async def safe_run_agent(user_input: str) -> str:
     
     Args:
         user_input (str): The user's input message to process
+        session_id (str, optional): Session identifier for conversation memory. Defaults to "default".
         
     Returns:
         str: The agent's response or an appropriate error message in English
@@ -801,7 +845,7 @@ async def safe_run_agent(user_input: str) -> str:
             return "Sorry, I didn't receive any questions. Please enter your question or request."
         
         # Process the input through the main agent function
-        return await run_agent(user_input)
+        return await run_agent(user_input, session_id)
         
     except Exception as e:
         logger.critical(f"Critical error in safe_run_agent: {str(e)}")
@@ -817,23 +861,49 @@ def clear_memory() -> None:
     effectively starting a fresh conversation session.
     """
     try:
-        memory.clear()
+        _memory_manager.clear_all_sessions()
         logger.info("Conversation memory cleared successfully")
     except Exception as e:
         logger.error(f"Error clearing memory: {str(e)}")
 
 
-def get_memory_summary() -> str:
+def get_memory_summary(session_id: str = "default") -> str:
     """
-    Get a summary of the current conversation memory.
+    Get a summary of the conversation history for a specific session.
+    
+    Args:
+        session_id (str, optional): Session identifier. Defaults to "default".
     
     Returns:
         str: A summary of the conversation history stored in memory
     """
     try:
+        memory = _memory_manager.get_memory(session_id)
         memory_vars = memory.load_memory_variables({})
         return str(memory_vars.get("chat_history", "No conversation history available"))
     except Exception as e:
         logger.error(f"Error getting memory summary: {str(e)}")
-        return "Error retrieving memory summary"
+        return "Error retrieving conversation history"
 
+
+def clear_session_memory(session_id: str) -> bool:
+    """
+    Clear conversation memory for a specific session.
+    
+    Args:
+        session_id (str): Session identifier to clear
+    
+    Returns:
+        bool: True if session was cleared, False if session didn't exist
+    """
+    return _memory_manager.clear_session(session_id)
+
+
+def get_active_sessions() -> list:
+    """
+    Get list of all active session IDs.
+    
+    Returns:
+        list: List of active session identifiers
+    """
+    return _memory_manager.get_active_sessions()

core/tools.py

@@ -6,11 +6,12 @@ from datetime import datetime
 from typing import Optional, List
 
 import pytz
-from langchain.schema import Document
+from langchain.schema import Document, HumanMessage, SystemMessage
 from langchain.tools import tool
 from .retrievers import hybrid_search, vector_search, bm25_search
 from .validation import validate_medical_answer
 from .github_storage import get_github_storage
+from langchain_openai import ChatOpenAI
 
 CANONICAL_PROVIDERS = {"Manus", "ASCO", "NCCN", "ESMO", "NICE"}
 
@@ -28,6 +29,38 @@ def store_user_question(user_question: str):
     global _last_user_question
     _last_user_question = user_question
 
+def _get_llm_safe(temperature: float = 0.0, model: str = "gpt-4o"):
+    """Create a ChatOpenAI client if API key/config is available, else return None."""
+    try:
+        # ChatOpenAI will read OPENAI_API_KEY from env as in validation.py
+        return ChatOpenAI(model=model, temperature=temperature, max_tokens=512, request_timeout=30)
+    except Exception:
+        return None
+
+def _is_side_effect_report_llm(user_input: str) -> Optional[bool]:
+    """Use LLM to classify if input is an adverse drug reaction/side-effect report.
+    Returns True/False if confident, or None if unavailable/uncertain.
+    """
+    llm = _get_llm_safe()
+    if not llm:
+        return None
+    try:
+        system = SystemMessage(content=(
+            "You are a medical triage classifier. Decide if the user's text is a report of an adverse drug reaction (side effect) about a medication.\n"
+            "Criteria: mentions a medication/drug and symptoms or adverse effects experienced by a patient.\n"
+            "Respond with exactly one token: yes or no."
+        ))
+        human = HumanMessage(content=user_input[:1500])
+        resp = llm.invoke([system, human])
+        ans = (resp.content or "").strip().lower()
+        if ans.startswith("yes"):
+            return True
+        if ans.startswith("no"):
+            return False
+        return None
+    except Exception:
+        return None
+
 # Map lowercase variants and full names to canonical provider codes
 _PROVIDER_ALIASES = {
     # NCCN
@@ -289,7 +322,7 @@ def side_effect_recording_tool(user_input: str) -> str:
         str: Interactive form for collecting missing information or confirmation of data recording
     """
     try:
-        # Define keywords that indicate side effect reporting
+        # LLM classification (preferred), with keyword fallback to preserve behavior
         side_effect_keywords = [
             'side effect', 'adverse reaction', 'adverse event', 'drug reaction',
             'medication reaction', 'allergic reaction', 'complication', 'toxicity',
@@ -297,7 +330,6 @@ def side_effect_recording_tool(user_input: str) -> str:
             'overdose', 'poisoning', 'drug-induced', 'medication-induced',
             'experienced after taking', 'developed after', 'caused by medication',
             'drug-related', 'medication-related', 'pharmaceutical reaction',
-            # Add more comprehensive problem/symptom keywords
             'kidney problems', 'liver problems', 'heart problems', 'breathing problems',
             'skin problems', 'stomach problems', 'nausea', 'vomiting', 'diarrhea',
             'headache', 'dizziness', 'fatigue', 'weakness', 'rash', 'swelling',
@@ -305,8 +337,8 @@ def side_effect_recording_tool(user_input: str) -> str:
             'has these', 'has serious', 'causes', 'resulted in', 'led to',
             'problems with', 'issues with', 'complications from'
         ]
-        
         input_lower = user_input.lower().strip()
+        llm_decision = _is_side_effect_report_llm(user_input)
         
         # Check for special commands first
         if input_lower in ['save report', 'save', 'submit report', 'submit']:
@@ -324,8 +356,9 @@ def side_effect_recording_tool(user_input: str) -> str:
             extracted_data = _extract_side_effect_data(user_input)
             return _process_followup_response(user_input, extracted_data)
         
-        # Check if input contains side effect reporting indicators
-        contains_side_effect = any(keyword in input_lower for keyword in side_effect_keywords)
+        # Combine LLM decision with keyword fallback to avoid behavior regression
+        keyword_detected = any(keyword in input_lower for keyword in side_effect_keywords)
+        contains_side_effect = (llm_decision is True) or (llm_decision is not False and keyword_detected)
         
         if not contains_side_effect:
             return "This input does not appear to contain a side effect report. If you are reporting an adverse drug reaction, please include specific details about the medication and symptoms."
@@ -532,74 +565,23 @@ def _save_side_effect_report(extracted_data: dict) -> str:
                 # Ensure the value is properly formatted
                 extracted_data[field] = str(value).strip()
         
-        # Save to GitHub repository
-
-        
+        # Save to GitHub repository (fallback to local if needed)
         github_storage = get_github_storage()
-
-        
         success = github_storage.save_side_effects_report(extracted_data)
-
-        
-        
-
-        
         if not success:
-
-        
-            # Fallback to local storage if GitHub fails
-
-        
             csv_filename = "side_effects_reports.csv"
-
-        
             csv_path = os.path.join(os.getcwd(), csv_filename)
-
-        
-            
-
-        
             file_exists = os.path.exists(csv_path)
-
-        
-            
-
-        
             with open(csv_path, 'a', newline='', encoding='utf-8') as csvfile:
-
-        
                 writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
-
-        
-                
-
-        
                 if not file_exists:
-
-        
                     writer.writeheader()
-
-        
-                
-
-        
                 writer.writerow(extracted_data)
-
-        
-            
-
-        
             storage_location = "locally to side_effects_reports.csv (GitHub upload failed)"
-
-        
         else:
-
-        
             storage_location = "to GitHub cloud repository"
 
-        
-        
-# Generate confirmation message
+        # Generate confirmation message
         drug_name = extracted_data.get('drug_name', 'NaN')
         side_effects = extracted_data.get('side_effects', 'NaN')
         report_id = extracted_data['timestamp'].replace(':', '').replace('-', '').replace(' ', '_')
@@ -661,14 +643,10 @@ def _extract_side_effect_data_with_llm(user_input: str) -> dict:
     Returns:
         dict: Structured data extracted from the input
     """
-    from langchain.llms import OpenAI
-    from langchain.prompts import PromptTemplate
     import json
-    
     # Get current timestamp
     egypt_tz = pytz.timezone('Africa/Cairo')
     current_time = datetime.now(egypt_tz).strftime('%Y-%m-%d %H:%M:%S')
-    
     # Initialize extracted data with defaults
     extracted_data = {
         'timestamp': current_time,
@@ -682,65 +660,32 @@ def _extract_side_effect_data_with_llm(user_input: str) -> dict:
         'outcome': 'NaN',
         'additional_details': 'NaN',
         'reporter_info': 'NaN',
-        'raw_input': user_input[:500]  # Limit raw input length
+        'raw_input': user_input[:500]
     }
-    
-    # Create extraction prompt
-    extraction_prompt = PromptTemplate(
-        input_variables=["user_input"],
-        template="""Extract medical side effect information from the following text. Return ONLY a JSON object with these exact fields:
-
-{
-  "drug_name": "name of the medication/drug mentioned",
-  "side_effects": "list of side effects or symptoms described",
-  "patient_age": "patient's age if mentioned",
-  "patient_gender": "Male or Female if mentioned",
-  "dosage": "medication dosage if mentioned",
-  "duration": "treatment duration if mentioned (e.g., '3 months', '2 weeks')",
-  "severity": "mild, moderate, or severe if mentioned",
-  "outcome": "current status like ongoing, resolved, recovered if mentioned"
-}
-
-IMPORTANT RULES:
-1. If any information is not found or unclear, use "NaN" as the value
-2. For duration, look for phrases like "Treatment duration: 3 months", "for 2 weeks", "over 6 months", etc.
-3. Extract exact values as mentioned in the text
-4. Return ONLY the JSON object, no other text
-
-Text to analyze:
-{user_input}
-
-JSON:"""
-    )
-    
-    try:
-        # Try to use LLM extraction if available
+    llm = _get_llm_safe()
+    if llm:
         try:
-            # This would use the actual LLM - for now we'll use a fallback approach
-            # llm = OpenAI(temperature=0)
-            # prompt = extraction_prompt.format(user_input=user_input)
-            # response = llm(prompt)
-            # extracted_json = json.loads(response.strip())
-            
-            # Fallback to improved regex-based extraction with better duration handling
-            extracted_json = _extract_with_improved_regex(user_input)
-            
+            system = SystemMessage(content=(
+                "Extract medical side effect information. Return ONLY a JSON object with these exact fields: "
+                "drug_name, side_effects, patient_age, patient_gender, dosage, duration, severity, outcome. "
+                "If missing/unclear, use 'NaN'."
+            ))
+            human = HumanMessage(content=user_input[:2000])
+            response = llm.invoke([system, human])
+            text = (response.content or "").strip()
+            # Try parse; if fails, fallback regex
+            try:
+                extracted_json = json.loads(text)
+            except json.JSONDecodeError:
+                extracted_json = _extract_with_improved_regex(user_input)
         except Exception:
-            # Fallback to improved regex extraction
             extracted_json = _extract_with_improved_regex(user_input)
-        
-        # Update extracted_data with LLM results, keeping NaN for empty values
-        for key, value in extracted_json.items():
-            if key in extracted_data and value and str(value).strip() and str(value).strip().lower() != 'nan':
-                extracted_data[key] = str(value).strip()
-    
-    except Exception as e:
-        # If LLM extraction fails, use improved regex fallback
+    else:
         extracted_json = _extract_with_improved_regex(user_input)
-        for key, value in extracted_json.items():
-            if key in extracted_data and value and str(value).strip() and str(value).strip().lower() != 'nan':
-                extracted_data[key] = str(value).strip()
-    
+    # Update extracted_data
+    for key, value in extracted_json.items():
+        if key in extracted_data and value and str(value).strip() and str(value).strip().lower() != 'nan':
+            extracted_data[key] = str(value).strip()
     return extracted_data
 
 
@@ -885,89 +830,3 @@ def _extract_side_effect_data(user_input: str) -> dict:
     return _extract_side_effect_data_with_llm(user_input)
 
 
-@tool
-def medical_answer_validation_tool(
-    question: Optional[str] = None, 
-    retrieved_documents: Optional[List] = None, 
-    generated_answer: Optional[str] = None
-) -> str:
-    """
-    Validate a medical answer using the comprehensive validation system.
-    
-    This tool evaluates medical responses across 6 criteria: Accuracy, Coherence, 
-    Relevance, Completeness, Citations/Attribution, and Length.
-    
-    Args:
-        question: The original medical question (optional - uses stored context if not provided)
-        retrieved_documents: List of documents used for the answer (optional - uses stored context)
-        generated_answer: The AI-generated answer to validate (optional - uses stored context)
-        
-    Returns:
-        str: Formatted validation report with scores and improvement recommendations
-    """
-    global _last_question, _last_documents, _last_answer, _last_user_question
-    
-    try:
-        # Use provided parameters or fall back to stored context
-        # Prefer the original user question over the tool query
-        eval_question = question or _last_user_question or _last_question
-        eval_documents = retrieved_documents or _last_documents or []
-        eval_answer = generated_answer or _last_answer
-        
-        # Validate that we have the required information
-        if not eval_question:
-            return "Error: No question available for validation. Please provide a question or ensure medical_guidelines_knowledge_tool was used first."
-        
-        if not eval_answer:
-            return "Error: No answer available for validation. Please provide an answer to validate."
-        
-        if not eval_documents:
-            return "Warning: No retrieved documents available for validation. Validation will proceed with limited context."
-        
-        # Store the answer for future reference
-        if generated_answer:
-            _last_answer = generated_answer
-        
-        # Perform validation
-        evaluation = validate_medical_answer(eval_question, eval_documents, eval_answer)
-        
-        # Format the validation report for display
-        report = evaluation.get("validation_report", {})
-        
-        formatted_report = f"""
-**🔍 MEDICAL ANSWER VALIDATION REPORT**
-
-**Interaction ID:** {evaluation.get('interaction_id', 'N/A')}
-**Timestamp:** {evaluation.get('timestamp', 'N/A')}
-
-**Overall Score:** {report.get('Overall_Rating', 'N/A')}/100
-
-**Key Metrics:**
-
-**Accuracy:** {report.get('Accuracy_Rating', 'N/A')}/100
-{report.get('Accuracy_Comment', 'No comment available')}
-
-**Coherence:** {report.get('Coherence_Rating', 'N/A')}/100
-{report.get('Coherence_Comment', 'No comment available')}
-
-**Relevance:** {report.get('Relevance_Rating', 'N/A')}/100
-{report.get('Relevance_Comment', 'No comment available')}
-
-**Completeness:** {report.get('Completeness_Rating', 'N/A')}/100
-{report.get('Completeness_Comment', 'No comment available')}
-
-**Citations:** {report.get('Citations_Attribution_Rating', 'N/A')}/100
-{report.get('Citations_Attribution_Comment', 'No comment available')}
-
-**Length:** {report.get('Length_Rating', 'N/A')}/100
-{report.get('Length_Comment', 'No comment available')}
-
-**Assessment:** {report.get('Final_Summary_and_Improvement_Plan', 'No improvement plan available')}
-
-**📁 Data Storage:** Evaluation saved to evaluation_results.json
-"""
-        
-        return formatted_report.strip()
-        
-    except Exception as e:
-        return f"Validation error: {str(e)}. Please ensure all required parameters are provided or that context is available from previous tool usage."