SurfSense/surfsense_backend/app/agents/researcher/state.py

"""Define the state structures for the agent."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

from sqlalchemy.ext.asyncio import AsyncSession

from app.services.streaming_service import StreamingService


@dataclass
class State:
    """Defines the dynamic state for the agent during execution.

    This state tracks the database session and the outputs generated by the agent's nodes.
    See: https://langchain-ai.github.io/langgraph/concepts/low_level/#state
    for more information.
    """

    # Runtime context (not part of actual graph state)
    db_session: AsyncSession

    # Streaming service
    streaming_service: StreamingService

    chat_history: list[Any] | None = field(default_factory=list)

    reformulated_query: str | None = field(default=None)
    # Using field to explicitly mark as part of state
    answer_outline: Any | None = field(default=None)
    further_questions: Any | None = field(default=None)

    # Temporary field to hold reranked documents from sub-agents for further question generation
    reranked_documents: list[Any] | None = field(default=None)

    # OUTPUT: Populated by agent nodes
    # Using field to explicitly mark as part of state
    final_written_report: str | None = field(default=None)