agents.py

"""Agent Architecture Implementations.

This module contains the implementation of different agent architectures
for customer service query processing.
"""

import random
import time
from typing import Any, List, Optional, TypedDict

from langchain_core.messages import BaseMessage, HumanMessage
from langgraph.graph import END, START, StateGraph
from llm_utils import call_llm, should_use_real_llm
from materializers.prompt import Prompt


# LangGraph State Definition
class CustomerServiceState(TypedDict):
    """State definition for LangGraph customer service workflow."""

    messages: List[BaseMessage]
    query_type: str
    confidence: float
    response_text: str


class AgentResponse:
    """Response from an agent with metadata."""

    def __init__(
        self, text: str, latency_ms: float, confidence: float, tokens_used: int
    ):
        """Initialize agent response.

        Args:
            text: Response text from the agent
            latency_ms: Response time in milliseconds
            confidence: Confidence score (0-1)
            tokens_used: Number of tokens consumed
        """
        self.text = text
        self.latency_ms = latency_ms
        self.confidence = confidence
        self.tokens_used = tokens_used


class BaseAgent:
    """Base class for all agent architectures."""

    def __init__(self, name: str, prompts: Optional[List[Prompt]] = None):
        """Initialize base agent.

        Args:
            name: Name of the agent architecture
            prompts: List of Prompt objects loaded as ZenML artifacts
        """
        self.name = name
        self.prompts = {p.name: p for p in (prompts or [])}

    def process_query(self, query: str) -> AgentResponse:
        """Process a single query and return response with metadata.

        Args:
            query: Customer service query text

        Raises:
            NotImplementedError: This method must be implemented by subclasses
        """
        raise NotImplementedError


class SingleAgentRAG(BaseAgent):
    """Simple RAG agent that handles all queries with one approach."""

    def __init__(self, prompts: Optional[List[Prompt]] = None) -> None:
        """Initialize SingleAgentRAG with knowledge base.

        Args:
            prompts: Optional list of Prompt objects loaded as ZenML artifacts
        """
        super().__init__("SingleAgentRAG", prompts)
        self.knowledge_base = {
            "return": "Items can be returned within 30 days with original receipt.",
            "refund": "Refunds are processed within 5-7 business days to original payment method.",
            "shipping": "Free shipping on orders over $50, otherwise $5.99 flat rate.",
            "support": "Customer support available 24/7 via chat, email, or phone.",
            "warranty": "All products come with 1-year manufacturer warranty.",
        }

    def process_query(self, query: str) -> AgentResponse:
        """Process query using LLM or simple keyword matching fallback.

        Args:
            query: Customer service query text

        Returns:
            AgentResponse with generated response and metadata
        """
        start_time = time.time()

        if should_use_real_llm():
            # Use real LLM for customer service response
            knowledge_context = "\n".join(
                [f"{k}: {v}" for k, v in self.knowledge_base.items()]
            )
            # Use loaded prompt template
            prompt_obj = self.prompts.get("single_agent_rag")
            if prompt_obj:
                prompt = prompt_obj.format(
                    knowledge_context=knowledge_context, query=query
                )
            else:
                # Fallback prompt
                prompt = f"You are a helpful customer service agent. Use this knowledge base to answer customer questions:\n\nKnowledge Base:\n{knowledge_context}\n\nCustomer Question: {query}\n\nProvide a helpful, professional response:"

            response_text = call_llm(
                prompt,
                model="gpt-3.5-turbo",
                metadata={
                    "agent_type": "SingleAgentRAG",
                    "query_length": len(query),
                    "knowledge_categories": list(self.knowledge_base.keys()),
                },
            )
            confidence = random.uniform(
                0.8, 0.95
            )  # Higher confidence for real LLM
            # Note: This is a rough approximation. For accurate token counting, use a tokenizer]
            #  like tiktoken or the tokenizer from the model provider
            tokens_used = len(prompt.split()) + len(response_text.split())
        else:
            # Fallback to mock keyword matching
            time.sleep(random.uniform(0.1, 0.3))  # Simulate processing time

            query_lower = query.lower()
            response_text = "I'd be happy to help you with that! "

            if "return" in query_lower:
                response_text += self.knowledge_base["return"]
            elif "refund" in query_lower:
                response_text += self.knowledge_base["refund"]
            elif "shipping" in query_lower or "ship" in query_lower:
                response_text += self.knowledge_base["shipping"]
            elif "warranty" in query_lower or "guarantee" in query_lower:
                response_text += self.knowledge_base["warranty"]
            else:
                response_text += "Let me connect you with a specialist who can help with your specific question."

            confidence = random.uniform(0.7, 0.9)
            tokens_used = random.randint(50, 150)

        latency_ms = (time.time() - start_time) * 1000

        return AgentResponse(
            text=response_text,
            latency_ms=latency_ms,
            confidence=confidence,
            tokens_used=tokens_used,
        )

    def get_graph_visualization(self) -> str:
        """Get a simple text representation of the SingleAgentRAG architecture.

        Returns:
            Text representation of the SingleAgentRAG workflow structure
        """
        return """
SingleAgentRAG Architecture:

INPUT (Customer Query) → RAG Knowledge Base Lookup → LLM/Fallback Response → OUTPUT

Architecture Details:
- Single unified approach for all query types
- Uses predefined knowledge base with 5 categories: return, refund, shipping, support, warranty
- LLM generates responses using knowledge context, or falls back to keyword matching
- No query routing or specialization - one size fits all

Knowledge Base Categories:
• return: Return policy information
• refund: Refund processing details  
• shipping: Shipping costs and policies
• support: Customer support availability
• warranty: Product warranty information
"""

    def get_mermaid_diagram(self) -> str:
        """Get a Mermaid diagram representation of the SingleAgentRAG workflow.

        Returns:
            HTML string containing interactive Mermaid diagram of the workflow
        """
        return """
<!DOCTYPE html>
<html>
<head>
    <title>SingleAgentRAG Architecture</title>
    <script src="https://cdn.jsdelivr.net/npm/mermaid/dist/mermaid.min.js"></script>
</head>
<body>
    <div class="mermaid">
        graph TD
            A[Customer Query] --> B[Knowledge Base Lookup]
            B --> C{LLM Available?}
            C -->|Yes| D[LLM Processing]
            C -->|No| E[Keyword Matching]
            D --> F[Generate Response]
            E --> F
            F --> G[Customer Response]

            B[Knowledge Base Lookup<br/>🔍 Search across 5 categories:<br/>returns, refunds, shipping,<br/>support, warranty]
            D[LLM Processing<br/>🤖 Use prompt template<br/>with knowledge context]
            E[Keyword Matching<br/>📝 Simple fallback logic<br/>based on query keywords]
            F[Generate Response<br/>💬 Unified response<br/>regardless of query type]

            style A fill:#e1f5fe
            style G fill:#e8f5e8
            style B fill:#fff3e0
            style C fill:#f3e5f5
            style D fill:#e0f2f1
            style E fill:#ffebee
            style F fill:#fff8e1
    </div>
    <script>
        mermaid.initialize({ startOnLoad: true });
    </script>
</body>
</html>
"""


class MultiSpecialistAgents(BaseAgent):
    """Multiple specialized agents for different query types."""

    def __init__(self, prompts: Optional[List[Prompt]] = None) -> None:
        """Initialize MultiSpecialistAgents with specialist routing.

        Args:
            prompts: Optional list of Prompt objects loaded as ZenML artifacts
        """
        super().__init__("MultiSpecialistAgents", prompts)
        self.specialists = {
            "returns": "Returns Specialist: I handle all return and exchange requests.",
            "billing": "Billing Specialist: I can help with payment and billing questions.",
            "technical": "Technical Support: I assist with product setup and troubleshooting.",
            "general": "Customer Service: I'm here to help with general questions.",
        }

    def _route_query(self, query: str) -> str:
        """Route query to appropriate specialist.

        Args:
            query: Customer service query text to route

        Returns:
            str: Specialist category: 'returns', 'billing', 'technical', or 'general'
        """
        query_lower = query.lower()

        if any(
            word in query_lower for word in ["return", "exchange", "refund"]
        ):
            return "returns"
        elif any(
            word in query_lower
            for word in ["payment", "billing", "charge", "price"]
        ):
            return "billing"
        elif any(
            word in query_lower
            for word in ["setup", "install", "technical", "broken"]
        ):
            return "technical"
        else:
            return "general"

    def process_query(self, query: str) -> AgentResponse:
        """Process query using specialist routing with LLM or mock responses.

        Args:
            query: Customer service query text

        Returns:
            AgentResponse with generated response and metadata
        """
        start_time = time.time()

        specialist = self._route_query(query)

        if should_use_real_llm():
            # Use loaded specialist prompts
            specialist_prompt_obj = self.prompts.get(
                f"specialist_{specialist}"
            )
            if specialist_prompt_obj:
                prompt = specialist_prompt_obj.format(query=query)
            else:
                # Fallback prompts
                fallback_prompts = {
                    "returns": "You are a Returns Specialist. Help customers with returns, exchanges, and refunds professionally.\n\nCustomer Question: {query}\n\nProvide a helpful, specific response for this returns inquiry:",
                    "billing": "You are a Billing Specialist. Help customers with payment issues, billing questions, and account management.\n\nCustomer Question: {query}\n\nProvide a helpful, specific response for this billing inquiry:",
                    "technical": "You are a Technical Support Specialist. Help customers with product setup, troubleshooting, and technical issues.\n\nCustomer Question: {query}\n\nProvide a helpful, specific response for this technical inquiry:",
                    "general": "You are a General Customer Service Agent. Help customers with general questions and direct them to specialists when needed.\n\nCustomer Question: {query}\n\nProvide a helpful, specific response for this general inquiry:",
                }
                prompt = fallback_prompts[specialist].format(query=query)

            response_text = call_llm(
                prompt,
                model="gpt-3.5-turbo",
                metadata={
                    "agent_type": "MultiSpecialistAgents",
                    "specialist": specialist,
                    "query_length": len(query),
                    "routing_decision": specialist,
                },
            )
            confidence = random.uniform(
                0.85, 0.98
            )  # Higher confidence for specialized LLM
            # Note: This is a rough approximation. For accurate token counting, use a tokenizer]
            #  like tiktoken or the tokenizer from the model provider
            tokens_used = len(prompt.split()) + len(response_text.split())
        else:
            # Fallback to mock specialist responses
            time.sleep(
                random.uniform(0.15, 0.4)
            )  # Simulate routing and processing time

            base_response = self.specialists[specialist]

            # Add specific responses based on specialist
            if specialist == "returns":
                response_text = f"{base_response} Items can be returned within 30 days with receipt. Would you like me to start a return for you?"
            elif specialist == "billing":
                response_text = f"{base_response} I can help you review your account, process refunds, or update payment methods."
            elif specialist == "technical":
                response_text = f"{base_response} Let me walk you through the troubleshooting steps for your issue."
            else:
                response_text = f"{base_response} I can direct you to the right specialist or answer general questions about our policies."

            confidence = random.uniform(0.8, 0.95)
            tokens_used = random.randint(80, 200)

        latency_ms = (time.time() - start_time) * 1000

        return AgentResponse(
            text=response_text,
            latency_ms=latency_ms,
            confidence=confidence,
            tokens_used=tokens_used,
        )

    def get_graph_visualization(self) -> str:
        """Get a simple text representation of the MultiSpecialistAgents architecture.

        Returns:
            Text representation of the MultiSpecialistAgents workflow structure
        """
        return """
MultiSpecialistAgents Architecture:

INPUT (Customer Query) → Intent Router → Specialist Agent → LLM/Fallback Response → OUTPUT

Architecture Details:
- Multi-agent system with specialized experts
- Query routing based on keyword analysis
- Each specialist has domain expertise and specialized prompts
- Higher confidence due to specialization

Routing Logic:
• returns → Returns Specialist (return, exchange, refund keywords)
• billing → Billing Specialist (payment, billing, charge, price keywords)  
• technical → Technical Support (setup, install, technical, broken keywords)
• general → General Support (catch-all for other queries)

Specialist Capabilities:
• Returns Specialist: Handles returns, exchanges, refunds
• Billing Specialist: Payment issues, billing questions, account management
• Technical Support: Product setup, troubleshooting, technical issues
• General Support: Directs to other specialists, handles general inquiries
"""

    def get_mermaid_diagram(self) -> str:
        """Get a Mermaid diagram representation of the MultiSpecialistAgents workflow.

        Returns:
            HTML string containing interactive Mermaid diagram of the workflow
        """
        return """
<!DOCTYPE html>
<html>
<head>
    <title>MultiSpecialistAgents Architecture</title>
    <script src="https://cdn.jsdelivr.net/npm/mermaid/dist/mermaid.min.js"></script>
</head>
<body>
    <div class="mermaid">
        graph TD
            A[Customer Query] --> B[Intent Router]
            B --> C{Query Classification}
            C -->|return, exchange, refund| D[Returns Specialist]
            C -->|payment, billing, charge| E[Billing Specialist]
            C -->|setup, install, technical| F[Technical Support]
            C -->|other queries| G[General Support]
            
            D --> H{LLM Available?}
            E --> I{LLM Available?}
            F --> J{LLM Available?}
            G --> K{LLM Available?}
            
            H -->|Yes| L[Returns LLM]
            H -->|No| M[Returns Fallback]
            I -->|Yes| N[Billing LLM]
            I -->|No| O[Billing Fallback]
            J -->|Yes| P[Technical LLM]
            J -->|No| Q[Technical Fallback]
            K -->|Yes| R[General LLM]
            K -->|No| S[General Fallback]
            
            L --> T[Customer Response]
            M --> T
            N --> T
            O --> T
            P --> T
            Q --> T
            R --> T
            S --> T

            B[Intent Router<br/>🎯 Keyword-based routing<br/>to appropriate specialist]
            D[Returns Specialist<br/>↩️ Handles returns,<br/>exchanges, refunds]
            E[Billing Specialist<br/>💳 Payment issues,<br/>billing questions]
            F[Technical Support<br/>🔧 Setup, troubleshooting,<br/>technical issues]
            G[General Support<br/>📞 General inquiries,<br/>routing assistance]

            style A fill:#e1f5fe
            style T fill:#e8f5e8
            style B fill:#fff3e0
            style C fill:#f3e5f5
            style D fill:#e0f2f1
            style E fill:#fff8e1
            style F fill:#fce4ec
            style G fill:#f1f8e9
    </div>
    <script>
        mermaid.initialize({ startOnLoad: true });
    </script>
</body>
</html>
"""


class LangGraphCustomerServiceAgent(BaseAgent):
    """LangGraph-based customer service agent with workflow visualization."""

    def __init__(self, prompts: Optional[List[Prompt]] = None) -> None:
        """Initialize LangGraph agent with workflow and knowledge base.

        Args:
            prompts: Optional list of Prompt objects loaded as ZenML artifacts
        """
        super().__init__("LangGraphCustomerServiceAgent", prompts)
        self.graph = self._build_graph()
        self.knowledge_base = {
            "return": "Items can be returned within 30 days with original receipt.",
            "refund": "Refunds are processed within 5-7 business days to original payment method.",
            "shipping": "Free shipping on orders over $50, otherwise $5.99 flat rate.",
            "support": "Customer support available 24/7 via chat, email, or phone.",
            "warranty": "All products come with 1-year manufacturer warranty.",
        }

    def _build_graph(self) -> Any:
        """Build the LangGraph workflow for customer service.

        Returns:
            Compiled StateGraph with customer service workflow nodes and edges
        """
        workflow = StateGraph(CustomerServiceState)

        # Add nodes
        workflow.add_node("analyze_query", self._analyze_query)
        workflow.add_node("classify_intent", self._classify_intent)
        workflow.add_node("generate_response", self._generate_response)
        workflow.add_node("validate_response", self._validate_response)

        # Add edges
        workflow.add_edge(START, "analyze_query")
        workflow.add_edge("analyze_query", "classify_intent")
        workflow.add_edge("classify_intent", "generate_response")
        workflow.add_edge("generate_response", "validate_response")
        workflow.add_edge("validate_response", END)

        return workflow.compile()

    def _analyze_query(
        self, state: CustomerServiceState
    ) -> CustomerServiceState:
        """Analyze the customer query.

        Args:
            state: Current customer service state with messages

        Returns:
            Updated state with confidence score based on query complexity
        """
        if state["messages"]:
            content = state["messages"][-1].content
            query = content if isinstance(content, str) else str(content)
            # Simple analysis
            complexity = len(query.split())
            state["confidence"] = 0.9 if complexity < 10 else 0.8
        return state

    def _classify_intent(
        self, state: CustomerServiceState
    ) -> CustomerServiceState:
        """Classify the customer's intent.

        Args:
            state: Current customer service state with messages

        Returns:
            Updated state with classified query_type (returns, billing, shipping, warranty, or general)
        """
        if state["messages"]:
            content = state["messages"][-1].content
            query = (
                content if isinstance(content, str) else str(content)
            ).lower()

            if any(word in query for word in ["return", "exchange", "refund"]):
                state["query_type"] = "returns"
            elif any(
                word in query for word in ["payment", "billing", "charge"]
            ):
                state["query_type"] = "billing"
            elif any(
                word in query for word in ["shipping", "delivery", "ship"]
            ):
                state["query_type"] = "shipping"
            elif any(
                word in query for word in ["warranty", "guarantee", "broken"]
            ):
                state["query_type"] = "warranty"
            else:
                state["query_type"] = "general"
        return state

    def _generate_response(
        self, state: CustomerServiceState
    ) -> CustomerServiceState:
        """Generate a response based on the classified intent using LLM or fallback.

        Args:
            state: Current customer service state with classified query_type

        Returns:
            Updated state with generated response_text
        """
        query_type = state.get("query_type", "general")
        query = state["messages"][-1].content if state["messages"] else ""

        if should_use_real_llm():
            # Use real LLM for response generation in LangGraph workflow
            knowledge_context = ""
            if query_type in self.knowledge_base:
                knowledge_context = (
                    f"Relevant knowledge: {self.knowledge_base[query_type]}"
                )

            # Use loaded LangGraph workflow prompt
            prompt_obj = self.prompts.get("langgraph_workflow")
            if prompt_obj:
                prompt = prompt_obj.format(
                    query_type=query_type,
                    knowledge_context=knowledge_context,
                    query=query,
                )
            else:
                # Fallback prompt
                prompt = f"You are a customer service agent in a structured workflow. \nYou have analyzed this query and classified it as: {query_type}\n\n{knowledge_context}\n\nCustomer Question: {query}\n\nGenerate a helpful, professional response. Be specific and actionable:"

            response = call_llm(
                prompt,
                model="gpt-3.5-turbo",
                metadata={
                    "agent_type": "LangGraphCustomerServiceAgent",
                    "query_type": query_type,
                    "workflow_step": "generate_response",
                    "query_length": len(str(query)),
                },
            )
        else:
            # Fallback to mock response generation
            if query_type in self.knowledge_base:
                base_response = self.knowledge_base[query_type]
                response = f"I understand you have a {query_type} question. {base_response} Is there anything specific about this I can help clarify?"
            else:
                response = "I'm here to help you with your question. Let me connect you with the right specialist who can provide detailed assistance."

        state["response_text"] = response
        return state

    def _validate_response(
        self, state: CustomerServiceState
    ) -> CustomerServiceState:
        """Validate and finalize the response.

        Args:
            state: Current customer service state with response_text

        Returns:
            Updated state with validated response_text and adjusted confidence
        """
        # Simple validation - ensure response is not empty and has reasonable length
        if len(state["response_text"]) < 20:
            state["response_text"] = (
                "I apologize, but I need more information to provide you with the best assistance. Could you please provide more details about your request?"
            )
            state["confidence"] = 0.6
        elif len(state["response_text"]) > 500:
            state["response_text"] = state["response_text"][:500] + "..."
            state["confidence"] = max(0.7, state["confidence"] - 0.1)

        return state

    def process_query(self, query: str) -> AgentResponse:
        """Process a query through the LangGraph workflow.

        Args:
            query: Customer service query text

        Returns:
            AgentResponse with generated response and metadata
        """
        start_time = time.time()

        # Initialize state
        initial_state = CustomerServiceState(
            messages=[HumanMessage(content=query)],
            query_type="",
            confidence=0.8,
            response_text="",
        )

        # Run through the graph
        final_state = self.graph.invoke(initial_state)

        latency_ms = (time.time() - start_time) * 1000

        return AgentResponse(
            text=final_state["response_text"],
            latency_ms=latency_ms,
            confidence=final_state["confidence"],
            tokens_used=random.randint(80, 180),  # Simulated for demo
        )

    def get_graph_visualization(self) -> str:
        """Get a simple text representation of the graph structure.

        Returns:
            Text representation of the LangGraph workflow structure
        """
        return """
LangGraph Customer Service Agent Workflow:

START → analyze_query → classify_intent → generate_response → validate_response → END

Node Details:
- analyze_query: Analyzes query complexity and sets initial confidence
- classify_intent: Classifies into returns, billing, shipping, warranty, or general
- generate_response: Generates appropriate response based on intent
- validate_response: Validates response quality and adjusts confidence
"""

    def get_mermaid_diagram(self) -> str:
        """Get a Mermaid diagram representation of the workflow.

        Returns:
            HTML string containing interactive Mermaid diagram of the workflow
        """
        return """
<!DOCTYPE html>
<html>
<head>
    <title>LangGraph Customer Service Agent Workflow</title>
    <script src="https://cdn.jsdelivr.net/npm/mermaid/dist/mermaid.min.js"></script>
</head>
<body>
    <div class="mermaid">
        graph TD
            A[START] --> B[analyze_query]
            B --> C[classify_intent]
            C --> D[generate_response]
            D --> E[validate_response]
            E --> F[END]

            B[analyze_query<br/>📊 Analyze query complexity<br/>Set initial confidence]
            C[classify_intent<br/>🏷️ Classify into:<br/>returns, billing, shipping,<br/>warranty, or general]
            D[generate_response<br/>💬 Generate response<br/>based on intent]
            E[validate_response<br/>✅ Validate quality<br/>Adjust confidence]

            style A fill:#e1f5fe
            style F fill:#e8f5e8
            style B fill:#fff3e0
            style C fill:#f3e5f5
            style D fill:#e0f2f1
            style E fill:#fff8e1
    </div>
    <script>
        mermaid.initialize({ startOnLoad: true });
    </script>
</body>
</html>
"""