65 lines
2.1 KiB
Python
65 lines
2.1 KiB
Python
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
|