from typing import TypedDict, Annotated
from operator import add

class WorkflowState(TypedDict):
    input: str
    output: str
    agent_responses: Annotated[list[dict], add]  # Accumulates
    metadata: dict

from langgraph.graph import MessagesState
from langgraph.graph.message import add_messages
from typing import Annotated

# Option 1: Use built-in MessagesState (recommended)
class AgentState(MessagesState):
    """Extends MessagesState with custom fields."""
    user_id: str
    context: dict

# Option 2: Define messages manually with add_messages reducer
class CustomState(TypedDict):
    messages: Annotated[list, add_messages]  # Smart append/update by ID
    metadata: dict

from pydantic import BaseModel, Field

class WorkflowState(BaseModel):
    input: str = Field(description="User input")
    output: str = ""
    agent_responses: list[dict] = Field(default_factory=list)

    def add_response(self, agent: str, result: str):
        self.agent_responses.append({"agent": agent, "result": result})

from typing import Annotated
from operator import add

class AnalysisState(TypedDict):
    url: str
    raw_content: str

    # Accumulate agent outputs
    findings: Annotated[list[Finding], add]
    embeddings: Annotated[list[Embedding], add]

    # Control flow
    current_agent: str
    agents_completed: list[str]
    quality_passed: bool

from typing import Annotated

def merge_dicts(a: dict, b: dict) -> dict:
    """Custom reducer that merges dictionaries."""
    return {**a, **b}

class State(TypedDict):
    config: Annotated[dict, merge_dicts]  # Merges updates

def last_value(a, b):
    """Keep only the latest value."""
    return b

class State(TypedDict):
    status: Annotated[str, last_value]  # Overwrites

def node(state: WorkflowState) -> WorkflowState:
    """Return new state, don't mutate in place."""
    # Wrong: state["output"] = "result"
    # Right:
    return {
        **state,
        "output": "result"
    }

from dataclasses import dataclass
from langgraph.graph import StateGraph

@dataclass
class ContextSchema:
    """Runtime configuration, not persisted in state."""
    llm_provider: str = "anthropic"
    temperature: float = 0.7
    max_retries: int = 3
    debug_mode: bool = False

# Create graph with context schema
graph = StateGraph(WorkflowState, context_schema=ContextSchema)

# Access context in nodes
def my_node(state: WorkflowState, context: ContextSchema):
    if context.llm_provider == "anthropic":
        response = call_claude(state["input"], context.temperature)
    else:
        response = call_openai(state["input"], context.temperature)

    if context.debug_mode:
        logger.debug(f"Response: {response}")

    return {"output": response}

# Invoke with context
graph.invoke(
    {"input": "Hello"},
    context={"llm_provider": "openai", "temperature": 0.5}
)

from langgraph.cache.memory import InMemoryCache
from langgraph.types import CachePolicy

# Add node with cache policy
builder.add_node(
    "embed_content",
    embed_content_node,
    cache_policy=CachePolicy(ttl=300)  # Cache for 5 minutes
)

builder.add_node(
    "llm_call",
    llm_node,
    cache_policy=CachePolicy(ttl=60)  # Cache for 1 minute
)

# Compile with cache
graph = builder.compile(cache=InMemoryCache())

from langgraph.types import RemainingSteps

def agent_node(state: WorkflowState, remaining: RemainingSteps):
    """Proactively handle recursion limit."""
    if remaining.steps < 5:
        # Running low on steps, wrap up
        return {
            "action": "summarize_and_exit",
            "reason": f"Only {remaining.steps} steps remaining"
        }

    # Continue normal processing
    return {"action": "continue"}