codefuse-chatbot/coagent/connector/schema/general_schema.py

309 lines
8.5 KiB
Python
Raw Normal View History

from pydantic import BaseModel
from typing import List, Dict, Optional, Union
from enum import Enum
import re
import json
from loguru import logger
from langchain.tools import BaseTool
class ActionStatus(Enum):
DEFAUILT = "default"
FINISHED = "finished"
STOPPED = "stopped"
CONTINUED = "continued"
TOOL_USING = "tool_using"
CODING = "coding"
CODE_EXECUTING = "code_executing"
CODING2FILE = "coding2file"
PLANNING = "planning"
UNCHANGED = "unchanged"
ADJUSTED = "adjusted"
CODE_RETRIEVAL = "code_retrieval"
def __eq__(self, other):
if isinstance(other, str):
return self.value.lower() == other.lower()
return super().__eq__(other)
class Action(BaseModel):
action_name: str
description: str
class FinishedAction(Action):
action_name: str = ActionStatus.FINISHED
description: str = "provide the final answer to the original query to break the chain answer"
class StoppedAction(Action):
action_name: str = ActionStatus.STOPPED
description: str = "provide the final answer to the original query to break the agent answer"
class ContinuedAction(Action):
action_name: str = ActionStatus.CONTINUED
description: str = "cant't provide the final answer to the original query"
class ToolUsingAction(Action):
action_name: str = ActionStatus.TOOL_USING
description: str = "proceed with using the specified tool."
class CodingdAction(Action):
action_name: str = ActionStatus.CODING
description: str = "provide the answer by writing code"
class Coding2FileAction(Action):
action_name: str = ActionStatus.CODING2FILE
description: str = "provide the answer by writing code and filename"
class CodeExecutingAction(Action):
action_name: str = ActionStatus.CODE_EXECUTING
description: str = "provide the answer by writing executable code"
class PlanningAction(Action):
action_name: str = ActionStatus.PLANNING
description: str = "provide a sequence of tasks"
class UnchangedAction(Action):
action_name: str = ActionStatus.UNCHANGED
description: str = "this PLAN has no problem, just set PLAN_STEP to CURRENT_STEP+1."
class AdjustedAction(Action):
action_name: str = ActionStatus.ADJUSTED
description: str = "the PLAN is to provide an optimized version of the original plan."
# extended action exmaple
class CodeRetrievalAction(Action):
action_name: str = ActionStatus.CODE_RETRIEVAL
description: str = "execute the code retrieval to acquire more code information"
class RoleTypeEnums(Enum):
SYSTEM = "system"
USER = "user"
ASSISTANT = "assistant"
FUNCTION = "function"
OBSERVATION = "observation"
SUMMARY = "summary"
def __eq__(self, other):
if isinstance(other, str):
return self.value == other
return super().__eq__(other)
class PromptKey(BaseModel):
key_name: str
description: str
class PromptKeyEnums(Enum):
# Origin Query is ui's user question
ORIGIN_QUERY = "origin_query"
# agent's input from last agent
CURRENT_QUESTION = "current_question"
# ui memory contaisn (user and assistants)
UI_MEMORY = "ui_memory"
# agent's memory
SELF_MEMORY = "self_memory"
# chain memory
CHAIN_MEMORY = "chain_memory"
# agent's memory
SELF_LOCAL_MEMORY = "self_local_memory"
# chain memory
CHAIN_LOCAL_MEMORY = "chain_local_memory"
# Doc Infomations contains (Doc\Code\Search)
DOC_INFOS = "doc_infos"
def __eq__(self, other):
if isinstance(other, str):
return self.value == other
return super().__eq__(other)
class Doc(BaseModel):
title: str
snippet: str
link: str
index: int
def get_title(self):
return self.title
def get_snippet(self, ):
return self.snippet
def get_link(self, ):
return self.link
def get_index(self, ):
return self.index
def to_json(self):
return vars(self)
def __str__(self,):
return f"""出处 [{self.index + 1}] 标题 [{self.title}]\n\n来源 ({self.link}) \n\n内容 {self.snippet}\n\n"""
class CodeDoc(BaseModel):
code: str
related_nodes: list
index: int
def get_code(self, ):
return self.code
def get_related_node(self, ):
return self.related_nodes
def get_index(self, ):
return self.index
def to_json(self):
return vars(self)
def __str__(self,):
return f"""出处 [{self.index + 1}] \n\n来源 ({self.related_nodes}) \n\n内容 {self.code}\n\n"""
class LogVerboseEnum(Enum):
Log0Level = "0" # don't print log
Log1Level = "1" # print level-1 log
Log2Level = "2" # print level-2 log
Log3Level = "3" # print level-3 log
def __eq__(self, other):
if isinstance(other, str):
return self.value.lower() == other.lower()
if isinstance(other, LogVerboseEnum):
return self.value == other.value
return False
def __ge__(self, other):
if isinstance(other, LogVerboseEnum):
return int(self.value) >= int(other.value)
if isinstance(other, str):
return int(self.value) >= int(other)
return NotImplemented
def __le__(self, other):
if isinstance(other, LogVerboseEnum):
return int(self.value) <= int(other.value)
if isinstance(other, str):
return int(self.value) <= int(other)
return NotImplemented
@classmethod
def ge(self, enum_value: 'LogVerboseEnum', other: Union[str, 'LogVerboseEnum']):
return enum_value <= other
class Task(BaseModel):
task_type: str
task_name: str
task_desc: str
task_prompt: str
class Env(BaseModel):
env_type: str
env_name: str
env_desc:str
class Role(BaseModel):
role_type: str
role_name: str
role_desc: str
agent_type: str = ""
role_prompt: str = ""
template_prompt: str = ""
class ChainConfig(BaseModel):
chain_name: str
chain_type: str
agents: List[str]
do_checker: bool = False
chat_turn: int = 1
class PromptField(BaseModel):
field_name: str # 假设这是一个函数类型,您可以根据需要更改
function_name: str
title: Optional[str] = None
description: Optional[str] = None
is_context: Optional[bool] = True
omit_if_empty: Optional[bool] = True
class AgentConfig(BaseModel):
role: Role
prompt_config: List[PromptField]
prompt_manager_type: str = "PromptManager"
chat_turn: int = 1
focus_agents: List = []
focus_message_keys: List = []
group_agents: List = []
stop: str = ""
class PhaseConfig(BaseModel):
phase_name: str
phase_type: str
chains: List[str]
do_summary: bool = False
do_search: bool = False
do_doc_retrieval: bool = False
do_code_retrieval: bool = False
do_tool_retrieval: bool = False
class CompleteChainConfig(BaseModel):
chain_name: str
chain_type: str
agents: Dict[str, AgentConfig]
do_checker: bool = False
chat_turn: int = 1
class CompletePhaseConfig(BaseModel):
phase_name: str
phase_type: str
chains: Dict[str, CompleteChainConfig]
do_summary: bool = False
do_search: bool = False
do_doc_retrieval: bool = False
do_code_retrieval: bool = False
do_tool_retrieval: bool = False
def load_role_configs(config) -> Dict[str, AgentConfig]:
if isinstance(config, str):
with open(config, 'r', encoding="utf8") as file:
configs = json.load(file)
else:
configs = config
# logger.debug(configs)
return {name: AgentConfig(**v) for name, v in configs.items()}
def load_chain_configs(config) -> Dict[str, ChainConfig]:
if isinstance(config, str):
with open(config, 'r', encoding="utf8") as file:
configs = json.load(file)
else:
configs = config
return {name: ChainConfig(**v) for name, v in configs.items()}
def load_phase_configs(config) -> Dict[str, PhaseConfig]:
if isinstance(config, str):
with open(config, 'r', encoding="utf8") as file:
configs = json.load(file)
else:
configs = config
return {name: PhaseConfig(**v) for name, v in configs.items()}
# AgentConfig.update_forward_refs()