planning_ai/states.py

from typing import Annotated, TypedDict

import polars as pl
from langchain_core.documents import Document
from pydantic import BaseModel

from planning_ai.chains.hallucination_chain import HallucinationChecker
from planning_ai.common.utils import filename_reducer


class DocumentState(TypedDict):
    document: Document
    filename: int
    doc_type: str

    entities: list[dict]
    themes: list[dict]

    summary: BaseModel
    hallucination: HallucinationChecker
    chapters: list[str]

    is_hallucinated: bool
    refinement_attempts: int
    failed: bool
    processed: bool


class OverallState(TypedDict):
    documents: Annotated[list, filename_reducer]
    executive: str
    policies: pl.DataFrame

    unused_documents: list[int]

    n_docs: int