File size: 766 Bytes
3b769cb
963aee4
1af4802
963aee4
 
21b7409
13686ad
963aee4
 
 
 
 
49a6658
5767260
963aee4
 
c1688dc
1af4802
963aee4
 
fbb5eac
963aee4
1af4802
 
 
 
13686ad
f9cb6a0
21b7409
1af4802
3b769cb
1af4802
0d9117e
d7566e3
 
3b769cb
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
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