上传文件至 /

state.py

@@ -0,0 +1,64 @@
+from typing import Annotated, List, TypedDict, Literal, Dict
+from pydantic import BaseModel, Field
+import operator
+
+
+class Section(BaseModel):
+    name: str = Field(
+        description="章节的标题，应简洁明了地概括本章节的主题或主旨。",
+    )
+    description: str = Field(
+        description="简要说明这一章节将围绕哪些内容进行展开，字数在50-100字。描述应突出章节的写作重点，明确要讨论的问题、角度或结构。",
+    )
+    research: bool = Field(
+        description="判断该章节的写作是否需要进行网络搜索。若该部分需要引用外部资料、数据、案例或是知识储备不足以独立完成，则应标记为需要研究。"
+    )
+    content: str = Field(
+        description="章节的主要写作内容。暂时留空，后续将用于填写具体文本内容"
+    )
+
+
+class Sections(BaseModel):
+    sections: List[Section] = Field(
+        description="包含报告的各个章节",
+    )
+
+
+class SearchQuery(BaseModel):
+    search_query: str = Field(None, description="网络搜索的查询")
+
+
+class Queries(BaseModel):
+    thought: str = Field(None, description="思考过程")
+    queries: List[SearchQuery] = Field(
+        description="包含网络搜索的查询列表",
+    )
+
+
+'''作为整张图的状态输入输出，输入是报告主题，输出是最终报告'''
+
+
+class ReportStateInput(TypedDict):
+    topic: str  # 报告主题
+
+
+class ReportStateOutput(TypedDict):
+    final_report: str  # 最终报告
+
+
+'''管理报告的状态，有报告的各个章节'''
+
+
+class ReportState(TypedDict):
+    topic: str
+    feedback_on_report_plan: str
+    sections: list[Section]
+    # 维护一个推理链便于后面模型思考怎么做
+    thinking_process: Annotated[list[str], operator.add]
+    completed_sections: Annotated[list, operator.add]
+    current_section_index: int  # 目前正在研究的章节序号
+    final_report: str  # Final report
+    search_queries: Annotated[Dict[str, List[SearchQuery]], operator.or_]
+    # 键是来源，值是title,url等为键的字典
+    sources: Annotated[Dict[str, Dict[str, str]], operator.or_]
+    source_str: str