"""State management for the APM graph.

This module defines the state structures used in the APM graph. It includes
definitions for agent state, input state, and router classification schema.
"""

from dataclasses import dataclass, field
from typing import Optional, Literal, List, Tuple
from typing_extensions import TypedDict

class Router(TypedDict):
    """Classify a user query."""
    logic: str
    datasource: Optional[Literal["vectorstore", "websearch"]]

# Optional, the InputState is a restricted version of the State that is used to
# define a narrower interface to the outside world vs. what is maintained
# internally.
@dataclass(kw_only=True)
class InputState:
    """Represents the input state for the agent.

    This class defines the structure of the input state, which includes
    the messages exchanged between the user and the agent. It serves as
    a restricted version of the full State, providing a narrower interface
    to the outside world compared to what is maintained internally.
    """

    """Attributes:
        question: user question
    """
    question: str

@dataclass(kw_only=True)
class OutputState:
    """Represents the output schema for the APM agent."""
    question: str
    messages: Optional[List[str]] = None
    generation: Optional[str] = None
    source: Optional[str] = None
    """Answer to user's Architecture IT Landscape question about ."""

@dataclass(kw_only=True)
class OverallState(InputState, OutputState):
    """State of the APM graph / agent."""
    
    """
    safety_status: user question's safeguarding status, justification, rephrased question
    router: classification of the user's query
    source: RAG or websearch
    retrieved: list of documents retrieved by the retriever
    rag: last RAG approach used
    chat_memory: user chat memory
    """
    safety_status: Optional[Tuple[str, str, str]] = None
    router: Optional[Router] = None
    rag: Optional[str] = None
    chat_memory: Optional[str] = None
    retrieved: Optional[List[str]] = None