from typing import TypedDict, List, Optional, ClassVar

from langchain_core.messages import BaseMessage, AIMessageChunk


class WorkflowState(TypedDict):
    user_id: str
    session_id: str
    messages: List[BaseMessage]
    memory_initialized: Optional[bool]
    stream_chunks: List[AIMessageChunk]
    chat_response: Optional[dict]
    retrieved_docs: Optional[List[str]]
    context: Optional[str]


def get_state_schema() -> ClassVar[WorkflowState]:
    return WorkflowState
