models.py

from typing import TypedDict, List, Optional, Annotated, Union
from operator import add
from pydantic import BaseModel, Field

# --- 1. API Input/Output Models (Strictly Problem Statement 2) ---

class Message(BaseModel):
    """Represents a single message in the conversation."""
    sender: str = Field(..., description="scammer or user")
    text: str = Field(..., description="Message content")
    # Flexible timestamp to handle both ISO-8601 strings (PDF) and Unix integers (Tester)
    timestamp: Union[str, int] = Field(..., description="ISO-8601 format or Unix timestamp")

class Metadata(BaseModel):
    """Optional metadata about the conversation channel."""
    channel: Optional[str] = Field(None, description="SMS / WhatsApp / Email / Chat")
    language: Optional[str] = Field(None, description="Language used")
    locale: Optional[str] = Field(None, description="Country or region")

class HoneypotRequest(BaseModel):
    """The incoming request body for the honeypot API."""
    sessionId: str = Field(..., description="Unique session ID")
    message: Message = Field(..., description="The latest incoming message")
    conversationHistory: List[Message] = Field(default_factory=list, description="All previous messages")
    metadata: Optional[Metadata] = None

class ExtractedIntelligence(BaseModel):
    """Structured data to be extracted from the conversation."""
    bankAccounts: List[str] = Field(default_factory=list)
    upiIds: List[str] = Field(default_factory=list) # Strictly 'upiIds' as per PDF
    phishingLinks: List[str] = Field(default_factory=list)
    phoneNumbers: List[str] = Field(default_factory=list)
    suspiciousKeywords: List[str] = Field(default_factory=list)
    unknownNumbers: List[str] = Field(default_factory=list)

class EngagementMetrics(BaseModel):
    """Metrics for the engagement."""
    engagementDurationSeconds: int
    totalMessagesExchanged: int

class HoneypotResponse(BaseModel):
    """The outgoing response body from the honeypot API."""
    status: str = Field(..., description="success")
    scamDetected: bool
    engagementMetrics: EngagementMetrics
    extractedIntelligence: ExtractedIntelligence
    agentNotes: str

# --- 2. LangGraph State Model ---

class AgentState(TypedDict):
    """The state object for the LangGraph state machine."""
    sessionId: str
    conversationHistory: Annotated[List[Message], add]
    scamDetected: bool
    extractedIntelligence: ExtractedIntelligence
    agentNotes: str
    totalMessagesExchanged: int
    should_continue_engagement: bool
    agent_response_text: str
    callbackSent: bool