ea4all/src/ea4all_indexer/state.py

"""State management for the index graph."""

from dataclasses import dataclass
from typing import Annotated, Optional

from langchain_core.documents import Document

from ea4all.src.shared.state import reduce_docs

@dataclass(kw_only=True)
class InputState:
    """Represents the input state for the index graph.

    This class is used to pass the input documents to the index graph.
    It contains a single field, `path`, which is the source of documents.
    """

    path: Optional[str] = None
    """Document source path to be indexed by the graph."""


# The index state defines the simple IO for the single-node index graph
@dataclass(kw_only=True)
class OutputState:
    """Represents the state for document indexing and retrieval.

    This class defines the structure of the index state, which includes
    the documents to be indexed and the retriever used for searching
    these documents.
    """

    docs: Annotated[list[Document], reduce_docs]
    """A list of documents that the agent can index."""

@dataclass(kw_only=True)
class OverallState(InputState):
    """Represents the overall state of the index graph.

    This class combines the input and output states, allowing for
    both input documents and indexed documents to be managed within
    the same state.
    """

    pass