ForcePilot/src/agents/registry.py

54 lines
1.7 KiB
Python
Raw Normal View History

2025-03-28 11:40:46 +08:00
from __future__ import annotations
2025-03-24 19:07:51 +08:00
from typing import Type, Annotated, Optional, TypedDict
from abc import abstractmethod
2025-03-28 11:40:46 +08:00
from dataclasses import dataclass, fields, field
2025-03-24 19:07:51 +08:00
2025-03-28 11:40:46 +08:00
from langchain_core.runnables import RunnableConfig
2025-03-25 05:40:07 +08:00
from langchain_core.messages import BaseMessage
2025-03-24 19:07:51 +08:00
from langgraph.graph.state import CompiledStateGraph
2025-03-25 05:40:07 +08:00
from langgraph.graph.message import add_messages
2025-03-24 19:07:51 +08:00
2025-03-29 17:33:09 +08:00
from src.config import SimpleConfig
2025-03-24 19:07:51 +08:00
class State(TypedDict):
"""
定义一个基础 State 各类 graph 继承, 其中:
1. messages 为所有 graph 的核心信息队列, 所有聊天工作流均应该将关键信息补充到此队列中;
2. history 为所有工作流单次启动时获取 history_len messages 所用(节约成本, 及防止单轮对话 tokens 占用长度达到 llm 支持上限),
history 中的信息理应是可以被丢弃的.
"""
messages: Annotated[list[BaseMessage], add_messages]
history: Optional[list[BaseMessage]]
@dataclass(kw_only=True)
2025-03-29 17:33:09 +08:00
class Configuration(SimpleConfig):
2025-03-24 19:07:51 +08:00
"""
定义一个基础 Configuration 各类 graph 继承
"""
2025-03-28 11:40:46 +08:00
@classmethod
def from_runnable_config(
cls, config: Optional[RunnableConfig] = None
) -> Configuration:
"""Create a Configuration instance from a RunnableConfig object."""
configurable = (config.get("configurable") or {}) if config else {}
_fields = {f.name for f in fields(cls) if f.init}
return cls(**{k: v for k, v in configurable.items() if k in _fields})
@classmethod
def to_dict(cls):
2025-03-29 17:33:09 +08:00
return {f.name: getattr(cls, f.name) for f in fields(cls) if f.init}
2025-03-25 05:40:07 +08:00
2025-03-24 19:07:51 +08:00
class BaseAgent():
2025-03-28 11:40:46 +08:00
def __init__(self, **kwargs):
pass
2025-03-24 19:07:51 +08:00
@abstractmethod
def get_graph(self) -> CompiledStateGraph:
pass