core/agent.py

import logging
import traceback
from typing import Any, AsyncGenerator
import asyncio
import requests
import os
import httpx
from langchain.agents import create_openai_tools_agent, AgentExecutor
from langchain.memory import ConversationBufferWindowMemory
from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain.schema import OutputParserException
from langchain.callbacks.base import BaseCallbackHandler
from openai import RateLimitError, APIError

from .config import get_llm, logger
from .tools import (
    medical_guidelines_knowledge_tool,
    get_current_datetime_tool,
)

# LangSmith tracing utilities
from .tracing import traceable, trace, conversation_tracker, log_to_langsmith
from .validation import validate_medical_answer


# ============================================================================
# STREAMING CALLBACK HANDLER
# ============================================================================

class StreamingCallbackHandler(BaseCallbackHandler):
    """Custom callback handler for streaming responses."""
    
    def __init__(self):
        self.tokens = []
        self.current_response = ""
    
    def on_llm_new_token(self, token: str, **kwargs: Any) -> Any:
        """Called when a new token is generated."""
        self.tokens.append(token)
        self.current_response += token
    
    def get_response(self) -> str:
        """Get the current response."""
        return self.current_response
    
    def reset(self):
        """Reset the handler for a new response."""
        self.tokens = []
        self.current_response = ""


# ============================================================================
# CUSTOM EXCEPTION CLASSES
# ============================================================================

class AgentError(Exception):
    """Base exception for agent-related errors."""
    pass


class ToolExecutionError(AgentError):
    """Exception raised when a tool fails to execute."""
    pass


class APIConnectionError(AgentError):
    """Exception raised when API connections fail."""
    pass


class ValidationError(AgentError):
    """Exception raised when input validation fails."""
    pass


# ============================================================================
# AGENT CONFIGURATION
# ============================================================================

# Available tools for the agent
AVAILABLE_TOOLS = [
    medical_guidelines_knowledge_tool,
    get_current_datetime_tool,
]


# System message template for the agent
SYSTEM_MESSAGE = """
You are a specialized HBV Clinical Assistant based on SASLT 2021 guidelines serving hepatologists and infectious disease specialists.

**DUAL ROLE:**
1. **Patient Eligibility Assessment**: Evaluate HBV patients for antiviral treatment eligibility
2. **Clinical Consultation**: Answer questions about HBV management, guidelines, and patient cases

**ASSESSMENT PARAMETERS:**
Evaluate treatment eligibility by analyzing: Serological Status (HBsAg, HBeAg, Anti-HBe), Viral Load (HBV DNA IU/mL), Liver Function (ALT), Fibrosis Stage (F0-F4), Necroinflammatory Activity (A0-A3), Patient Category (immune tolerant/active, inactive carrier, HBeAg-negative CHB), and Special Populations (pregnancy, immunosuppression, coinfections, cirrhosis).

**Eligibility Status:**
When responding to eligibility-related questions, **always start the answer with “Eligibility Status”**, providing a clear and concise determination (Eligible / Not Eligible) followed by the rationale based on SASLT 2021 treatment criteria. Then continue with the structured sections (e.g., Patient Profile, Treatment Criteria, Treatment Recommendations, Monitoring Plan, References).

**TREATMENT CRITERIA & OPTIONS:**
- Eligibility: HBV DNA thresholds, ALT elevation (>ULN, >2×ULN), fibrosis stage (≥F2, F3-F4), special populations
- First-line: Entecavir (ETV), Tenofovir Disoproxil Fumarate (TDF), Tenofovir Alafenamide (TAF)
- Alternative agents and PEG-IFN when indicated

**RESPONSE STYLE:**
- Start directly with clinical answers - NO procedural statements ("I will retrieve...", "Let me search...", "Please wait...")
- Use structured, concise clinical format: brief introductory sentence → organized data (tables, bullets, or structured lists) → clinical notes where relevant
- Target 200-400 words (standard queries) or 400-600 words (complex questions)
- Prioritize key information first, use hierarchical formatting (headers, bullets), include specific clinical parameters when relevant
- Use precise medical terminology appropriate for experts
- Answer only what's asked - no tangential information

**STRUCTURED CLINICAL FORMAT (FOLLOW THESE EXAMPLES):**

Example 1 - Tabular data with clinical notes:
"Chronic HBV is classified into five phases using HBsAg/HBeAg status, HBV DNA, ALT, and liver inflammation.

Phase 1 – Chronic HBV infection ("immune tolerant")
HBsAg/HBeAg: High / Positive
HBV DNA: High (>10⁷ IU/mL)
ALT: Normal
Liver inflammation: None/minimal
Notes: More common/prolonged with perinatal/early-life infection; patients are highly contagious
[SASLT HBV Management Guidelines, 2021, p. 3]

Phase 2 – Chronic hepatitis B ("immune reactive HBeAg positive")
HBsAg/HBeAg: High–intermediate / Positive
HBV DNA: Lower (10⁴–10⁷ IU/mL)
ALT: Increased
Liver inflammation: Moderate/severe
Notes: May follow years of immune tolerance; more frequent when infection occurs in adulthood
[SASLT HBV Management Guidelines, 2021, p. 3]"

Example 2 - Bullet lists with citations:
"Screen high-risk populations despite universal childhood vaccination [SASLT 2021, p. 4].

High-risk groups include:
• Expatriate individuals (pre-employment) [SASLT 2021, p. 4]
• Healthcare workers [SASLT 2021, p. 4]
• Household contacts of HBV carriers [SASLT 2021, p. 4]
• Sexual contacts of HBV carriers or those with high-risk sexual behavior [SASLT 2021, p. 4]"

Example 3 - Categorized information:
"Diagnosis of chronic HBV
• HBsAg: detection is the most commonly used test to diagnose chronic HBV infection [SASLT 2021, p. 4].
• HBV disease assessment incorporates HBsAg, HBeAg/anti-HBe, and HBV DNA [SASLT 2021, p. 3].

Identify immunity or prior exposure
• anti-HBs: indicates the patient is protected (immune) against HBV [SASLT 2021, p. 4].
• anti-HBc: indicates previous exposure to HBV [SASLT 2021, p. 4]."

**CLINICAL TONE & NUANCE:**
- Use clinically precise language: "characterized by," "indicates," "reflects," "can be difficult to distinguish"
- Acknowledge clinical uncertainty when present in guidelines: "many fall into a 'grey area'," "requires individualized follow-up," "cannot be captured by a single measurement"
- Include practical guidance: "Practical approach recommended by the guideline," "Bottom line"
- Add clinical context in Notes or footnotes when relevant to interpretation
- Use specific numeric ranges and thresholds exactly as stated in guidelines (e.g., ">10⁷ IU/mL," "10⁴–10⁷ IU/mL," "≥2,000 IU/mL")

**MANDATORY TOOL USAGE:**
ALWAYS use "medical_guidelines_knowledge_tool" FIRST for every medical question - even basic HBV concepts. Do NOT answer from general knowledge. Only formulate responses based on retrieved SASLT 2021 guideline information. All information must come from SASLT 2021 (the only provider available). If no information found, explicitly state this

**TOOL USAGE REQUIREMENTS:**
1. **MEDICAL QUESTIONS** (definitions, treatments, guidelines, etc.): 
   - MANDATORY: Use "medical_guidelines_knowledge_tool" FIRST
   - Then answer based ONLY on retrieved information
   
2. **TIME/DATE QUERIES**: For current date/time or references like "today" or "now":
   - MANDATORY: Use "get_current_datetime_tool"

**SEARCH QUERY OPTIMIZATION:**
Transform user questions into comprehensive queries with medical terminology, synonyms, clinical context, AND practical keywords. System uses hybrid search (vector + BM25 keyword matching).

**Key Principles:**
1. **Core Concept**: Start with main medical concept and guideline reference
2. **Add Synonyms**: Include medical term variations
3. **Add Action Verbs**: Include practical keywords from the question (e.g., "screened", "testing", "monitoring", "detection")
4. **Expand Concepts**: Add related clinical terms
5. **Keyword Boosters**: Append domain-specific terms at end for better coverage

**Synonym Mapping:**
- "HBV" + "hepatitis B virus" + "CHB" + "chronic hepatitis B"
- "treatment" + "therapy" + "antiviral" + "management"
- "ALT" + "alanine aminotransferase" + "liver enzymes"
- "fibrosis" + "cirrhosis" + "F2 F3 F4" + "liver fibrosis"
- "HBeAg" + "hepatitis B e antigen" + "HBeAg-positive" + "HBeAg-negative"
- "viral load" + "HBV DNA" + "DNA level" + "viremia"
- "screening" + "screened" + "testing" + "detection" + "diagnosis" + "program"

**Concept Expansion:**
- Treatment → "eligibility criteria indications thresholds when to start"
- Drugs → "first-line second-line alternatives dosing monitoring ETV TDF TAF entecavir tenofovir"
- Assessment → "HBsAg HBeAg anti-HBe HBV DNA ALT fibrosis immune phase"
- Special populations → "pregnancy pregnant women cirrhosis immunosuppression HIV HCV HDV"
- Screening → "target populations high-risk groups screened testing HBsAg detection program"

**Query Construction Formula:**
[Main Concept] + [Guideline Reference] + [Synonyms] + [Action Verbs from Question] + [Related Clinical Terms] + [Keyword Boosters]

**Examples:**
- "Who should be targeted for HBV screening in Saudi Arabia?" → "HBV screening target populations Saudi Arabia SASLT 2021 guidelines screened testing high-risk groups pregnancy HBsAg detection program hepatitis B virus"

- "When to start treatment?" → "HBV treatment initiation criteria indications when to start SASLT 2021 HBV DNA threshold ALT elevation fibrosis stage antiviral therapy eligibility hepatitis B virus management"

- "First-line drugs?" → "first-line antiviral agents HBV treatment SASLT 2021 entecavir ETV tenofovir TDF TAF preferred drugs nucleos(t)ide analogues therapy recommendations hepatitis B virus"

- "HBeAg-negative management?" → "HBeAg-negative chronic hepatitis B CHB management SASLT 2021 treatment criteria HBV DNA threshold ALT elevation anti-HBe immune active phase monitoring hepatitis B e antigen"

**CRITICAL**: Always include practical action verbs from the user's question (e.g., "screened", "tested", "monitored", "detected") as these improve retrieval of relevant guideline sections discussing those specific activities.
**CITATION FORMAT (MANDATORY):**
1. **Inline Citations**: Use format [SASLT 2021, p. X] or [SASLT HBV Management Guidelines, 2021, p. X] after each clinical statement. Cite each page individually - NEVER use ranges.
   - Examples: 
     * "HBsAg detection is the most commonly used test [SASLT 2021, p. 4]."
     * "Phase 1 – Chronic HBV infection ("immune tolerant") [SASLT HBV Management Guidelines, 2021, p. 3]"
     * "Treatment criteria include viral load thresholds [SASLT 2021, p. 7], ALT elevation [SASLT 2021, p. 8], and fibrosis assessment [SASLT 2021, p. 9]."

2. **Citation Placement**: Place citation immediately after the relevant statement or at the end of each bullet point/phase description. For structured data (phases, categories), cite after each complete section.

3. **References Section** (Optional): For complex answers, you may end with "**References**" listing all cited pages:
   ```
   **References**
   SASLT 2021 Guidelines - Pages: p. 7, p. 8, p. 9, p. 12, p. 15, p. 18
   (Treatment Eligibility Criteria, First-Line Agents, and Monitoring Protocols)
   ```

4. **Citation Details**: For tables/flowcharts, specify number, title, and relevant rows/columns/values. For text, specify section hierarchy. Include context pages if they contributed to your answer

**CRITICAL CITATION RULES:**
- **NEVER assign a citation to a page that does not explicitly mention the claim**. If you cannot find the information on a specific page in the retrieved documents, do NOT cite that page.
- **VERIFY BEFORE CITING**: Only cite a page number if you can see that specific information in the retrieved content from that page.
- **NO ASSUMPTIONS**: Do not assume information is on a page just because it seems logical or related. Only cite what you can verify in the retrieved documents.
- **EXACT THRESHOLDS**: If the guideline specifies numeric thresholds (e.g., "HBV DNA > 100,000 IU/mL in late pregnancy"), include them EXACTLY as written. Do not generalize or alter thresholds unless identical wording exists in the source.

**NO GENERAL KNOWLEDGE - GUIDELINE ONLY:**
NEVER answer from general knowledge or speculate. Do NOT generate or imply new interpretations, summaries, or reasoning not explicitly derived from the retrieved document text. If information not found in SASLT 2021 after using tool, respond: "I searched the SASLT 2021 guidelines but could not find specific information about [topic]. You may want to rephrase with more clinical details, consult the guidelines directly, or contact hepatology specialists."

**STRICT SOURCE-BASED RESPONSES:**
- Every statement in your answer must be directly traceable to the retrieved documents
- Do not add information from your general medical knowledge, even if it seems relevant or helpful
- Do not make inferences or draw conclusions beyond what is explicitly stated in the documents
- If the documents don't contain enough information to fully answer the question, acknowledge this limitation rather than supplementing with general knowledge

**OUT-OF-SCOPE HANDLING:**
For non-HBV questions (other diseases, non-medical topics), respond professionally: "I'm unable to assist with that request, but I'd be happy to help with HBV-related inquiries."

**PATIENT CONTEXT:**
When question includes [PATIENT CONTEXT] or [PRIOR ASSESSMENT RESULT], provide personalized case-specific guidance tailored to patient's parameters (HBV DNA, ALT, fibrosis). Reference prior assessments for consistency.

**ELIGIBILITY ASSESSMENT WORKFLOW:**
1. Retrieve SASLT 2021 criteria via tool
2. Categorize patient phase (immune tolerant/active, inactive carrier, HBeAg-negative CHB, cirrhosis)
3. Compare parameters: HBV DNA vs. threshold, ALT vs. ULN, fibrosis stage, necroinflammatory activity
4. Check special considerations (pregnancy, immunosuppression, coinfections, HCC family history)
5. Determine eligibility (Eligible/Not Eligible/Borderline)
6. Recommend first-line agents (ETV, TDF, TAF) if eligible
7. Outline monitoring plan

**ELIGIBILITY RESPONSE STRUCTURE:**
For patient eligibility: Patient Profile (HBsAg, HBeAg, HBV DNA, ALT, Fibrosis) → SASLT 2021 Criteria (with VERIFIED page citations) → Eligibility Status (Eligible/Not Eligible/Borderline) + Rationale → Treatment Recommendations (ETV, TDF, TAF if eligible - ONLY if explicitly mentioned in retrieved documents) → Monitoring → References

**IMPORTANT**: You may format the answer into structured sections (Eligibility, Treatment, Monitoring, etc.), but the content inside MUST remain strictly source-based. Do not add information from general knowledge to fill gaps in the structure.

**FORMATTING:**
Use **bold** for critical points/drugs, headers (###) for organization, bullets/numbered lists for sequences, tables for comparisons, blockquotes (>) for direct quotes. Include specific numeric values and thresholds.

**SAFETY:**
For emergencies (acute liver failure, hepatic encephalopathy, severe bleeding, loss of consciousness), respond: "This is an emergency! Call emergency services immediately and seek urgent medical help." Educational information only - not a substitute for clinical judgment. Always respond in English.
"""

# Create the prompt template
prompt_template = ChatPromptTemplate.from_messages([
    ("system", SYSTEM_MESSAGE),
    MessagesPlaceholder("chat_history"),
    ("human", "{input}"),
    MessagesPlaceholder("agent_scratchpad"),
])

# Initialize the agent with lazy loading
def get_agent():
    """Get agent with lazy loading for faster startup"""
    return create_openai_tools_agent(
        llm=get_llm(),
        tools=AVAILABLE_TOOLS,
        prompt=prompt_template,
    )

# Create agent executor with lazy loading
def get_agent_executor():
    """Get agent executor with lazy loading for faster startup"""
    return AgentExecutor(
        agent=get_agent(),
        tools=AVAILABLE_TOOLS,
        verbose=True,
        handle_parsing_errors=True,
        max_iterations=5,
        max_execution_time=90,  # tighten a bit to help responsiveness
    )

# ============================================================================
# SESSION-BASED MEMORY MANAGEMENT
# ============================================================================

class SessionMemoryManager:
    """Manages conversation memory for multiple sessions."""
    
    def __init__(self):
        self._sessions = {}
        self._default_window_size = 20  # Increased from 10 to maintain better context
    
    def get_memory(self, session_id: str = "default") -> ConversationBufferWindowMemory:
        """Get or create memory for a specific session."""
        if session_id not in self._sessions:
            import warnings
            with warnings.catch_warnings():
                warnings.filterwarnings("ignore", category=DeprecationWarning)
                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()


# ============================================================================
# VALIDATION HELPER FUNCTIONS
# ============================================================================

def _should_validate_response(user_input: str, response: str) -> bool:
    """
    Determine if a response should be automatically validated.
    
    Args:
        user_input: The user's input
        response: The agent's response
        
    Returns:
        bool: True if the response should be validated
    """
    # Skip validation for certain types of responses
    skip_indicators = [
        "side effect report",
        "adverse drug reaction report",
        "error:",
        "sorry,",
        "i don't know",
        "i do not know",
        "could not find specific information",
        "not found in the retrieved guidelines",
        "validation report",
        "evaluation scores"
    ]
    
    # Skip validation for side effect reporting queries in user input
    side_effect_input_indicators = [
        "side effect", "adverse reaction", "adverse event", "drug reaction",
        "medication reaction", "patient experienced", "developed after taking",
        "caused by medication", "drug-related", "medication-related"
    ]
    
    user_input_lower = user_input.lower()
    response_lower = response.lower()
    
    # Don't validate if user input is about side effect reporting
    if any(indicator in user_input_lower for indicator in side_effect_input_indicators):
        return False
    
    # Don't validate if response contains skip indicators
    if any(indicator in response_lower for indicator in skip_indicators):
        return False
    
    # Don't validate very short responses
    if len(response.strip()) < 50:
        return False
    
    # Validate if response seems to contain medical information
    medical_indicators = [
        "treatment", "therapy", "diagnosis", "medication", "drug", "patient",
        "clinical", "guideline", "recommendation", "according to", "source:",
        "provider:", "page:", "saslt", "hbv", "hepatitis"
    ]
    
    return any(indicator in response_lower for indicator in medical_indicators)


def _perform_automatic_validation(user_input: str, response: str) -> None:
    """
    Perform automatic validation in the background without displaying results to user.
    Validation results are logged and saved to GitHub repository for backend analysis.
    
    Args:
        user_input: The user's input
        response: The agent's response
        
    Returns:
        None: Validation runs silently in background
    """
    try:
        # Import here to avoid circular imports
        from .tools import _last_question, _last_documents
        
        # Check if we have the necessary context for validation
        if not _last_question or not _last_documents:
            logger.info("Skipping validation: insufficient context")
            return
        
        # Perform validation using the original user input instead of tool query
        evaluation = validate_medical_answer(user_input, _last_documents, response)
        
        # Log validation results to backend only (not shown to user)
        report = evaluation.get("validation_report", {})
        logger.info(f"Background validation completed - Interaction ID: {evaluation.get('interaction_id', 'N/A')}")
        logger.info(f"Validation scores - Overall: {report.get('Overall_Rating', 'N/A')}/100, "
                   f"Accuracy: {report.get('Accuracy_Rating', 'N/A')}/100, "
                   f"Coherence: {report.get('Coherence_Rating', 'N/A')}/100, "
                   f"Relevance: {report.get('Relevance_Rating', 'N/A')}/100")
        
        # Validation is automatically saved to GitHub by validate_medical_answer function
        # No need to return anything - results are stored in backend only
        
    except Exception as e:
        logger.error(f"Background validation failed: {e}")


# ============================================================================
# STREAMING AGENT FUNCTIONS
# ============================================================================

# @traceable(name="run_agent_streaming")
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.
    
    This function processes user input through the agent executor with streaming
    capabilities, robust error handling, and automatic retries for recoverable errors.
    
    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.
    
    Yields:
        str: Chunks of the agent's response as they are generated
        
    Raises:
        None: All exceptions are caught and handled internally
    """
    # Input validation
    if not user_input or not user_input.strip():
        logger.warning("Empty input received")
        yield "Sorry, I didn't receive any questions. Please enter your question or request."
        return
    
    retry_count = 0
    last_error = None
    current_run_id = None
    # Session metadata (increment conversation count)
    session_metadata = conversation_tracker.get_session_metadata(increment=True)
    
    while retry_count <= max_retries:
        try:
            # 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 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]}...")
            
            # Create streaming callback handler
            streaming_handler = StreamingCallbackHandler()
            
            # Run the agent in a separate thread to avoid blocking
            def run_sync():
                return get_agent_executor().invoke(
                    {
                        "input": user_input.strip(),
                        "chat_history": chat_history,
                    },
                    config={"callbacks": [streaming_handler]},
                )
            
            # Execute the agent with streaming
            full_response = ""
            previous_length = 0
            
            # Start the agent execution in background
            loop = asyncio.get_event_loop()
            task = loop.run_in_executor(None, run_sync)
            
            # Stream the response as it's being generated
            while not task.done():
                current_response = streaming_handler.get_response()
                
                # Yield new tokens if available
                if len(current_response) > previous_length:
                    new_content = current_response[previous_length:]
                    previous_length = len(current_response)
                    yield new_content
                
                # Small delay to prevent overwhelming the client (faster flushing)
                await asyncio.sleep(0.03)
            
            # Get the final result
            response = await task
            
            # Yield any remaining content
            final_response = streaming_handler.get_response()
            if len(final_response) > previous_length:
                yield final_response[previous_length:]
            
            # If no streaming content was captured, yield the full response
            if not final_response and response and "output" in response:
                full_output = response["output"]
                # Simulate streaming by yielding word by word
                words = full_output.split(' ')
                for word in words:
                    yield word + ' '
                    await asyncio.sleep(0.05)
                final_response = full_output
            
            # Validate response structure
            if not response or "output" not in response:
                raise ValidationError("Invalid response format from agent")
            
            if not response["output"] or not response["output"].strip():
                raise ValidationError("Empty response from agent")
            
            # Perform automatic validation in background (hidden from user)
            base_response = response["output"]
            if _should_validate_response(user_input, base_response):
                logger.info("Performing background validation for streaming response...")
                try:
                    # Run validation silently - results saved to backend/GitHub only
                    _perform_automatic_validation(user_input, base_response)
                except Exception as e:
                    logger.error(f"Background validation failed: {e}")
            
            # Save conversation context to memory
            memory.save_context(
                {"input": user_input},
                {"output": response["output"]}
            )
            
            # Log response metrics to LangSmith
            try:
                log_to_langsmith(
                    key="response_metrics",
                    value={
                        "response_length": len(response.get("output", "")),
                        "attempt": retry_count + 1,
                        **session_metadata,
                    },
                    run_id=current_run_id,
                )
            except Exception:
                pass

            logger.info(f"Successfully processed user input: {user_input[:50]}...")
            return
            
        except RateLimitError as e:
            retry_count += 1
            last_error = e
            wait_time = min(2 ** retry_count, 60)  # Exponential backoff, max 60 seconds
            
            logger.warning(
                f"Rate limit exceeded. Retrying in {wait_time} seconds... "
                f"(Attempt {retry_count}/{max_retries})"
            )
            
            if retry_count <= max_retries:
                await asyncio.sleep(wait_time)
                continue
            else:
                logger.error("Rate limit exceeded after maximum retries")
                yield "Sorry, the system is currently busy. Please try again in a little while."
                return
                
        except (APIError, httpx.RemoteProtocolError, httpx.ReadError, httpx.ConnectError) as e:
            retry_count += 1
            last_error = e
            error_type = type(e).__name__
            logger.error(f"OpenAI API/Connection error ({error_type}): {str(e)}")
            
            if retry_count <= max_retries:
                wait_time = min(2 ** retry_count, 10)  # Exponential backoff, max 10 seconds
                logger.info(f"Retrying after {wait_time} seconds... (Attempt {retry_count}/{max_retries})")
                await asyncio.sleep(wait_time)
                continue
            else:
                yield "Sorry, there was an error connecting to the service. Please try again later."
                return
                
        except requests.exceptions.ConnectionError as e:
            retry_count += 1
            last_error = e
            logger.error(f"Network connection error: {str(e)}")
            
            if retry_count <= max_retries:
                await asyncio.sleep(3)
                continue
            else:
                yield "Sorry, I can't connect to the service right now. Please check your internet connection and try again."
                return
                
        except requests.exceptions.Timeout as e:
            retry_count += 1
            last_error = e
            logger.error(f"Request timeout: {str(e)}")
            
            if retry_count <= max_retries:
                await asyncio.sleep(2)
                continue
            else:
                yield "Sorry, the request took longer than expected. Please try again."
                return
                
        except requests.exceptions.RequestException as e:
            logger.error(f"Request error: {str(e)}")
            yield "Sorry, an error occurred with the request. Please try again."
            return
            
        except OutputParserException as e:
            logger.error(f"Output parsing error: {str(e)}")
            yield "Sorry, an error occurred while processing the response. Please rephrase your question and try again."
            return
            
        except ValidationError as e:
            logger.error(f"Validation error: {str(e)}")
            yield "Sorry, an error occurred while validating the data. Please try again."
            return
            
        except ToolExecutionError as e:
            logger.error(f"Tool execution error: {str(e)}")
            yield "Sorry, an error occurred while executing one of the operations. Please try again or contact technical support."
            return
            
        except Exception as e:
            logger.error(f"Unexpected error in run_agent_streaming: {str(e)}")
            logger.error(f"Traceback: {traceback.format_exc()}")
            # Log error to LangSmith
            try:
                log_to_langsmith(
                    key="error_log",
                    value={
                        "error": str(e),
                        "error_type": type(e).__name__,
                        **session_metadata,
                    },
                    run_id=current_run_id,
                )
            except Exception:
                pass
            
            # For unexpected errors, don't retry
            yield "Sorry, an unexpected error occurred. Please try again or contact technical support if the problem persists."
            return
    
    # This should never be reached, but just in case
    logger.error(f"Maximum retries exceeded. Last error: {str(last_error)}")
    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, session_id: str = "default") -> AsyncGenerator[str, None]:
    """
    Streaming wrapper function with additional safety checks and input validation.
    
    This function provides an additional layer of safety by validating input parameters,
    checking input length constraints, and handling any critical errors that might
    occur during streaming agent execution.
    
    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
        
    Raises:
        None: All exceptions are caught and handled internally
    """
    try:
        # Input type validation
        if not isinstance(user_input, str):
            logger.warning(f"Invalid input type received: {type(user_input)}")
            yield "Sorry, the input must be valid text."
            return
        
        # Input length validation
        stripped_input = user_input.strip()
        
        # if len(stripped_input) > 1000:
        #     logger.warning(f"Input too long: {len(stripped_input)} characters")
        #     yield "Sorry, the message is too long. Please shorten your question."
        #     return
        
        if len(stripped_input) == 0:
            logger.warning("Empty input after stripping")
            yield "Sorry, I didn't receive any questions. Please enter your question or request."
            return
        
        # Stream the response through the main agent function
        async for chunk in run_agent_streaming(user_input, session_id):
            yield chunk
        
    except Exception as e:
        logger.critical(f"Critical error in safe_run_agent_streaming: {str(e)}")
        logger.critical(f"Traceback: {traceback.format_exc()}")
        yield "Sorry, a critical system error occurred. Please contact technical support immediately."


@traceable(name="run_agent")
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.
    
    This function processes user input through the agent executor with robust
    error handling, automatic retries for recoverable errors, and comprehensive
    logging for debugging and monitoring.
    
    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.
    
    Returns:
        str: The agent's response or an appropriate error message in English
        
    Raises:
        None: All exceptions are caught and handled internally
    """
    # Input validation
    if not user_input or not user_input.strip():
        logger.warning("Empty input received")
        return "Sorry, I didn't receive any questions. Please enter your question or request."
    
    retry_count = 0
    last_error = None
    current_run_id = None
    session_metadata = conversation_tracker.get_session_metadata(increment=True)
    
    while retry_count <= max_retries:
        try:
            # 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]}...")

            # Invoke the agent with input and history (synchronous call)
            response = get_agent_executor().invoke({
                "input": user_input.strip(),
                "chat_history": chat_history
            })
            current_run_id = None  # This will be handled by LangChain's tracer
            
            # Validate response structure
            if not response or "output" not in response or not isinstance(response["output"], str):
                raise ValidationError("Invalid response format from agent")
            
            if not response["output"] or not response["output"].strip():
                raise ValidationError("Empty response from agent")
            
            # Save conversation context to memory
            memory.save_context(
                {"input": user_input},
                {"output": response["output"]}
            )
            
            # Log response metrics
            try:
                log_to_langsmith(
                    key="response_metrics",
                    value={
                        "response_length": len(response.get("output", "")),
                        "attempt": retry_count + 1,
                        **session_metadata,
                    },
                    run_id=current_run_id,
                )
            except Exception:
                pass

            logger.info(f"Successfully processed user input: {user_input[:50]}...")
            
            # Perform automatic validation in background (hidden from user)
            final_response = response["output"]
            if _should_validate_response(user_input, final_response):
                logger.info("Performing background validation...")
                try:
                    # Run validation silently - results saved to backend/GitHub only
                    _perform_automatic_validation(user_input, final_response)
                except Exception as e:
                    logger.error(f"Background validation failed: {e}")
            
            return final_response
            
        except RateLimitError as e:
            retry_count += 1
            last_error = e
            wait_time = min(2 ** retry_count, 60)  # Exponential backoff, max 60 seconds
            
            logger.warning(
                f"Rate limit exceeded. Retrying in {wait_time} seconds... "
                f"(Attempt {retry_count}/{max_retries})"
            )
            
            if retry_count <= max_retries:
                await asyncio.sleep(wait_time)
                continue
            else:
                logger.error("Rate limit exceeded after maximum retries")
                return "Sorry, the system is currently busy. Please try again in a little while."
                
        except APIError as e:
            retry_count += 1
            last_error = e
            logger.error(f"OpenAI API error: {str(e)}")
            
            if retry_count <= max_retries:
                await asyncio.sleep(2)
                continue
            else:
                return "Sorry, there was an error connecting to the service. Please try again later."
                
        except requests.exceptions.ConnectionError as e:
            retry_count += 1
            last_error = e
            logger.error(f"Network connection error: {str(e)}")
            
            if retry_count <= max_retries:
                await asyncio.sleep(3)
                continue
            else:
                return "Sorry, I can't connect to the service right now. Please check your internet connection and try again."
                
        except requests.exceptions.Timeout as e:
            retry_count += 1
            last_error = e
            logger.error(f"Request timeout: {str(e)}")
            
            if retry_count <= max_retries:
                await asyncio.sleep(2)
                continue
            else:
                return "Sorry, the request took longer than expected. Please try again."
                
        except requests.exceptions.RequestException as e:
            logger.error(f"Request error: {str(e)}")
            return "Sorry, an error occurred with the request. Please try again."
            
        except OutputParserException as e:
            logger.error(f"Output parsing error: {str(e)}")
            return "Sorry, an error occurred while processing the response. Please rephrase your question and try again."
            
        except ValidationError as e:
            logger.error(f"Validation error: {str(e)}")
            return "Sorry, an error occurred while validating the data. Please try again."
            
        except ToolExecutionError as e:
            logger.error(f"Tool execution error: {str(e)}")
            return "Sorry, an error occurred while executing one of the operations. Please try again or contact technical support."
            
        except Exception as e:
            logger.error(f"Unexpected error in run_agent: {str(e)}")
            logger.error(f"Traceback: {traceback.format_exc()}")
            # Log error
            try:
                log_to_langsmith(
                    key="error_log",
                    value={
                        "error": str(e),
                        "error_type": type(e).__name__,
                        **session_metadata,
                    },
                    run_id=current_run_id,
                )
            except Exception:
                pass
            
            # For unexpected errors, don't retry
            return "Sorry, an unexpected error occurred. Please try again or contact technical support if the problem persists."
    
    # This should never be reached, but just in case
    logger.error(f"Maximum retries exceeded. Last error: {str(last_error)}")
    return "Sorry, I was unable to process your request after several attempts. Please try again later."


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.
    
    This function provides an additional layer of safety by validating input parameters,
    checking input length constraints, and handling any critical errors that might
    occur during agent execution.
    
    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
        
    Raises:
        None: All exceptions are caught and handled internally
    """
    try:
        # Input type validation
        if not isinstance(user_input, str):
            logger.warning(f"Invalid input type received: {type(user_input)}")
            return "Sorry, the input must be valid text."
        
        # Input length validation
        stripped_input = user_input.strip()
        
        # if len(stripped_input) > 1000:
        #     logger.warning(f"Input too long: {len(stripped_input)} characters")
        #     return "Sorry, the message is too long. Please shorten your question."
        
        if len(stripped_input) == 0:
            logger.warning("Empty input after stripping")
            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, session_id)
        
    except Exception as e:
        logger.critical(f"Critical error in safe_run_agent: {str(e)}")
        logger.critical(f"Traceback: {traceback.format_exc()}")
        return "Sorry, a critical system error occurred. Please contact technical support immediately."


def clear_memory() -> None:
    """
    Clear the conversation memory.
    
    This function clears all stored conversation history from memory,
    effectively starting a fresh conversation session.
    """
    try:
        _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(session_id: str = "default") -> str:
    """
    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 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()