writer/graph_base.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2024/6/27 下午5:24
# @Author  : Tim-Saijun https://zair.top
# @File    : graph_base.py
# @Project : SAgent
from typing import List, Annotated, TypedDict

from langchain_core.pydantic_v1 import BaseModel, Field
from langgraph.graph.message import add_messages


class GraphState(TypedDict):
    abstract: str
    input: str
    messages: Annotated[list, add_messages]
    plan: List[str]
    time: int
    result: list[str]


class Plan(BaseModel):
    steps: List[str] = Field(
        description="Each element string contains the secondary title and theme, the tertiary title and theme, and the SEO word distribution strategy"
    )