assistant/agent/core/state.py

"""
Agent state definitions for LangGraph workflow
"""
from typing import Any, Optional, Literal
from typing_extensions import TypedDict, Annotated
from dataclasses import dataclass, field
from enum import Enum


class Intent(str, Enum):
    """User intent categories"""
    CUSTOMER_SERVICE = "customer_service"  # FAQ, general inquiries
    ORDER = "order"                         # Order query, tracking, modify, cancel
    AFTERSALE = "aftersale"                # Return, exchange, complaint
    PRODUCT = "product"                     # Search, recommend, quote
    HUMAN_HANDOFF = "human_handoff"        # Transfer to human agent
    UNKNOWN = "unknown"                     # Cannot determine intent


class ConversationState(str, Enum):
    """Conversation state machine"""
    INITIAL = "initial"
    CLASSIFYING = "classifying"
    PROCESSING = "processing"
    AWAITING_INFO = "awaiting_info"
    TOOL_CALLING = "tool_calling"
    GENERATING = "generating"
    HUMAN_REVIEW = "human_review"
    COMPLETED = "completed"
    ERROR = "error"


@dataclass
class Entity:
    """Extracted entity from user message"""
    type: str           # Entity type (order_id, product_id, date, etc.)
    value: Any          # Entity value
    confidence: float   # Extraction confidence (0-1)


@dataclass
class ToolCall:
    """MCP tool call record"""
    tool_name: str
    arguments: dict[str, Any]
    server: str         # MCP server name (strapi, order, aftersale, product)


@dataclass
class ToolResult:
    """MCP tool execution result"""
    tool_name: str
    success: bool
    data: Any
    error: Optional[str] = None


class AgentState(TypedDict):
    """Main agent state for LangGraph workflow
    
    This state is passed through all nodes in the workflow graph.
    """
    
    # ============ Session Information ============
    conversation_id: str              # Chatwoot conversation ID
    user_id: str                      # User identifier
    account_id: str                   # B2B account identifier
    
    # ============ Message Content ============
    messages: list[dict[str, Any]]    # Conversation history [{role, content}]
    current_message: str              # Current user message being processed
    
    # ============ Intent Recognition ============
    intent: Optional[str]             # Recognized intent (Intent enum value)
    intent_confidence: float          # Intent confidence score (0-1)
    sub_intent: Optional[str]         # Sub-intent for more specific routing
    
    # ============ Entity Extraction ============
    entities: dict[str, Any]          # Extracted entities {type: value}
    
    # ============ Agent Routing ============
    current_agent: Optional[str]      # Current processing agent name
    agent_history: list[str]          # History of agents involved
    
    # ============ Tool Calling ============
    tool_calls: list[dict[str, Any]]  # Pending tool calls
    tool_results: list[dict[str, Any]]  # Tool execution results
    
    # ============ Response Generation ============
    response: Optional[str]           # Generated response text
    response_type: str                # Response type (text, rich, action)
    
    # ============ Human Handoff ============
    requires_human: bool              # Whether human intervention is needed
    handoff_reason: Optional[str]     # Reason for human handoff
    
    # ============ Conversation Context ============
    context: dict[str, Any]           # Accumulated context (order details, etc.)
    
    # ============ State Control ============
    state: str                        # Current conversation state
    step_count: int                   # Number of steps taken
    max_steps: int                    # Maximum allowed steps
    error: Optional[str]              # Error message if any
    finished: bool                    # Whether processing is complete


def create_initial_state(
    conversation_id: str,
    user_id: str,
    account_id: str,
    current_message: str,
    messages: Optional[list[dict[str, Any]]] = None,
    context: Optional[dict[str, Any]] = None
) -> AgentState:
    """Create initial agent state for a new message
    
    Args:
        conversation_id: Chatwoot conversation ID
        user_id: User identifier
        account_id: B2B account identifier
        current_message: User's message to process
        messages: Previous conversation history
        context: Existing conversation context
        
    Returns:
        Initialized AgentState
    """
    return AgentState(
        # Session
        conversation_id=conversation_id,
        user_id=user_id,
        account_id=account_id,
        
        # Messages
        messages=messages or [],
        current_message=current_message,
        
        # Intent
        intent=None,
        intent_confidence=0.0,
        sub_intent=None,
        
        # Entities
        entities={},
        
        # Routing
        current_agent=None,
        agent_history=[],
        
        # Tools
        tool_calls=[],
        tool_results=[],
        
        # Response
        response=None,
        response_type="text",
        
        # Human handoff
        requires_human=False,
        handoff_reason=None,
        
        # Context
        context=context or {},
        
        # Control
        state=ConversationState.INITIAL.value,
        step_count=0,
        max_steps=10,
        error=None,
        finished=False
    )


# ============ State Update Helpers ============

def add_message(state: AgentState, role: str, content: str) -> AgentState:
    """Add a message to the conversation history"""
    state["messages"].append({"role": role, "content": content})
    return state


def set_intent(
    state: AgentState,
    intent: Intent,
    confidence: float,
    sub_intent: Optional[str] = None
) -> AgentState:
    """Set the recognized intent"""
    state["intent"] = intent.value
    state["intent_confidence"] = confidence
    state["sub_intent"] = sub_intent
    return state


def add_entity(state: AgentState, entity_type: str, value: Any) -> AgentState:
    """Add an extracted entity"""
    state["entities"][entity_type] = value
    return state


def add_tool_call(
    state: AgentState,
    tool_name: str,
    arguments: dict[str, Any],
    server: str
) -> AgentState:
    """Add a tool call to pending calls"""
    state["tool_calls"].append({
        "tool_name": tool_name,
        "arguments": arguments,
        "server": server
    })
    return state


def add_tool_result(
    state: AgentState,
    tool_name: str,
    success: bool,
    data: Any,
    error: Optional[str] = None
) -> AgentState:
    """Add a tool execution result"""
    state["tool_results"].append({
        "tool_name": tool_name,
        "success": success,
        "data": data,
        "error": error
    })
    return state


def set_response(
    state: AgentState,
    response: str,
    response_type: str = "text"
) -> AgentState:
    """Set the generated response"""
    state["response"] = response
    state["response_type"] = response_type
    return state


def request_human_handoff(
    state: AgentState,
    reason: str
) -> AgentState:
    """Request transfer to human agent"""
    state["requires_human"] = True
    state["handoff_reason"] = reason
    return state


def update_context(state: AgentState, updates: dict[str, Any]) -> AgentState:
    """Update conversation context"""
    state["context"].update(updates)
    return state


def set_error(state: AgentState, error: str) -> AgentState:
    """Set error state"""
    state["error"] = error
    state["state"] = ConversationState.ERROR.value
    return state


def mark_finished(state: AgentState) -> AgentState:
    """Mark processing as complete"""
    state["finished"] = True
    state["state"] = ConversationState.COMPLETED.value
    return state