from typing import TypedDict, List, Optional

from pydantic import BaseModel, Field

from langchain_core.messages import BaseMessage, AIMessageChunk


class WorkflowInitArgs(BaseModel):
    alayamem_url: str = Field(..., description="AlayaMem URL")


class Input(BaseModel):
    messages: List[BaseMessage] = Field(..., description="List of input messages")


class WorkflowContext(BaseModel):
    user_id: str = Field(..., description="User ID")
    session_id: str = Field(..., description="Session ID")
    chat_model_id: str = Field(..., description="Chat Model ID (AlayaJet is preferred to enable KV cache reuse)")
    save_kv_cache_path: Optional[str] = Field(None, description="Path to save the kvcache file")
    load_kv_cache_path: Optional[str] = Field(None, description="Path to load the kvcache file")


class Output(BaseModel):
    chat_response: dict = Field(..., description="Chat response")


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