src/rag/utilclasses.py

from dataclasses import dataclass, field
from typing import List, Literal, Optional

from pydantic import BaseModel, Field
from typing_extensions import TypedDict
from langchain.agents import AgentState
from langchain_core.messages import AnyMessage


@dataclass
class AgentContext:
    agent_name: str


@dataclass
class LeadAgentQueryResponse:
    response: str
    language: str
    processed_query: str = None
    confidence_fallback: bool = False
    max_turns_reached: bool = False
    should_cache: bool = False
    appointment_requested: bool = False
    show_booking_widget: bool = False
    relevant_programs: List[str] = field(default_factory=list)


class StructuredAgentResponse(BaseModel):
    response: str = Field(description="Main response to the query.")
    is_context_dependent: bool = Field(
        default=True,
        description=(
            "Set to False only if the question can be answered without using any user-specific "
            "information (e.g. name, age, preferences, extracted profile data) and without relying "
            "on prior conversation turns or conversation history. "
            "Must be True for responses involving eligibility, recommendations, comparisons after prior turns, "
            "or any answer influenced by user profile data or conversation context."
        )
    )
    appointment_requested: bool = Field(
        default=False,
        description="Set to True ONLY if the user explicitly asks to book, schedule, speak with admissions/an advisor, see appointment slots, or accepts a previous consultation offer. Routine pricing, comparisons, recommendations, and exploratory fit questions must be False."
    )
    show_booking_widget: bool = Field(
        default=False,
        description="Set to True ONLY when appointment_requested is True and the booking widget should be shown now. Never use this for soft contact mentions or routine informational answers."
    )
    relevant_programs: Optional[List[Literal["emba", "iemba", "emba_x"]]] = Field(
        default=None,
        description="If appointment_requested is True, list the programs relevant to the user. Options: 'emba', 'iemba', 'emba_x'. If the user is undecided or general, leave this list empty."
    )


class State(TypedDict):
    messages: list[AnyMessage]
    answer: str


class ConversationState(TypedDict):
    """Tracks user profile and conversation context"""
    user_id: str  # Unique session identifier
    user_language: str | None  # Locked after first message
    user_name: str | None  # User's name extracted from conversation
    experience_years: int | None  # Years of professional experience
    leadership_years: int | None  # Years of leadership experience
    field: str | None  # Professional field/industry
    interest: str | None  # Content interests
    qualification_level: str | None  # "bachelor", "master", "MBA", etc.
    program_interest: list[str]  # ["EMBA", "IEMBA", "EMBAX"]
    suggested_program: str | None  # Recommended program based on conversation
    handover_requested: bool | None  # True if appointment requested, False if declined, None if session active
    topics_discussed: list[str]  # Track what's been covered
    preferences_known: bool  # Whether we have enough context


class LeadInformationState(AgentState):
    lead_name: str
    lead_age: int
    lead_language_knowledge: list
    lead_work_experience: dict
    lead_motivation: list
    # Enhanced state tracking
    conversation_state: ConversationState