256 lines
11 KiB
Python
256 lines
11 KiB
Python
import importlib
|
||
from typing import List, Union, Dict, Any
|
||
from loguru import logger
|
||
import os
|
||
from langchain.embeddings.base import Embeddings
|
||
from langchain.agents import Tool
|
||
from langchain.llms.base import BaseLLM, LLM
|
||
|
||
from coagent.retrieval.base_retrieval import IMRertrieval
|
||
from coagent.llm_models.llm_config import EmbedConfig, LLMConfig
|
||
from coagent.connector.phase import BasePhase
|
||
from coagent.connector.agents import BaseAgent
|
||
from coagent.connector.chains import BaseChain
|
||
from coagent.connector.schema import Message, Role, PromptField, ChainConfig
|
||
from coagent.tools import toLangchainTools, TOOL_DICT, TOOL_SETS
|
||
|
||
|
||
class AgentFlow:
|
||
def __init__(
|
||
self,
|
||
role_name: str,
|
||
agent_type: str,
|
||
role_type: str = "assistant",
|
||
agent_index: int = 0,
|
||
role_prompt: str = "",
|
||
prompt_config: List[Dict[str, Any]] = [],
|
||
prompt_manager_type: str = "PromptManager",
|
||
chat_turn: int = 3,
|
||
focus_agents: List[str] = [],
|
||
focus_messages: List[str] = [],
|
||
embeddings: Embeddings = None,
|
||
llm: BaseLLM = None,
|
||
doc_retrieval: IMRertrieval = None,
|
||
code_retrieval: IMRertrieval = None,
|
||
search_retrieval: IMRertrieval = None,
|
||
**kwargs
|
||
):
|
||
self.role_type = role_type
|
||
self.role_name = role_name
|
||
self.agent_type = agent_type
|
||
self.role_prompt = role_prompt
|
||
self.agent_index = agent_index
|
||
|
||
self.prompt_config = prompt_config
|
||
self.prompt_manager_type = prompt_manager_type
|
||
|
||
self.chat_turn = chat_turn
|
||
self.focus_agents = focus_agents
|
||
self.focus_messages = focus_messages
|
||
|
||
self.embeddings = embeddings
|
||
self.llm = llm
|
||
self.doc_retrieval = doc_retrieval
|
||
self.code_retrieval = code_retrieval
|
||
self.search_retrieval = search_retrieval
|
||
# self.build_config()
|
||
# self.build_agent()
|
||
|
||
def build_config(self, embeddings: Embeddings = None, llm: BaseLLM = None):
|
||
self.llm_config = LLMConfig(model_name="test", llm=self.llm or llm)
|
||
self.embed_config = EmbedConfig(embed_model="test", langchain_embeddings=self.embeddings or embeddings)
|
||
|
||
def build_agent(self,
|
||
embeddings: Embeddings = None, llm: BaseLLM = None,
|
||
doc_retrieval: IMRertrieval = None,
|
||
code_retrieval: IMRertrieval = None,
|
||
search_retrieval: IMRertrieval = None,
|
||
):
|
||
# 可注册个性化的agent，仅通过start_action和end_action来注册
|
||
# class ExtraAgent(BaseAgent):
|
||
# def start_action_step(self, message: Message) -> Message:
|
||
# pass
|
||
|
||
# def end_action_step(self, message: Message) -> Message:
|
||
# pass
|
||
# agent_module = importlib.import_module("coagent.connector.agents")
|
||
# setattr(agent_module, 'extraAgent', ExtraAgent)
|
||
|
||
# 可注册个性化的prompt组装方式，
|
||
# class CodeRetrievalPM(PromptManager):
|
||
# def handle_code_packages(self, **kwargs) -> str:
|
||
# if 'previous_agent_message' not in kwargs:
|
||
# return ""
|
||
# previous_agent_message: Message = kwargs['previous_agent_message']
|
||
# # 由于两个agent共用了同一个manager，所以临时性处理
|
||
# vertices = previous_agent_message.customed_kargs.get("RelatedVerticesRetrivalRes", {}).get("vertices", [])
|
||
# return ", ".join([str(v) for v in vertices])
|
||
|
||
# prompt_manager_module = importlib.import_module("coagent.connector.prompt_manager")
|
||
# setattr(prompt_manager_module, 'CodeRetrievalPM', CodeRetrievalPM)
|
||
|
||
# agent实例化
|
||
agent_module = importlib.import_module("coagent.connector.agents")
|
||
baseAgent: BaseAgent = getattr(agent_module, self.agent_type)
|
||
role = Role(
|
||
role_type=self.agent_type, role_name=self.role_name,
|
||
agent_type=self.agent_type, role_prompt=self.role_prompt,
|
||
)
|
||
|
||
self.build_config(embeddings, llm)
|
||
self.agent = baseAgent(
|
||
role=role,
|
||
prompt_config = [PromptField(**config) for config in self.prompt_config],
|
||
prompt_manager_type=self.prompt_manager_type,
|
||
chat_turn=self.chat_turn,
|
||
focus_agents=self.focus_agents,
|
||
focus_message_keys=self.focus_messages,
|
||
llm_config=self.llm_config,
|
||
embed_config=self.embed_config,
|
||
doc_retrieval=doc_retrieval or self.doc_retrieval,
|
||
code_retrieval=code_retrieval or self.code_retrieval,
|
||
search_retrieval=search_retrieval or self.search_retrieval,
|
||
)
|
||
|
||
|
||
|
||
class ChainFlow:
|
||
def __init__(
|
||
self,
|
||
chain_name: str,
|
||
chain_index: int = 0,
|
||
agent_flows: List[AgentFlow] = [],
|
||
chat_turn: int = 5,
|
||
do_checker: bool = False,
|
||
embeddings: Embeddings = None,
|
||
llm: BaseLLM = None,
|
||
doc_retrieval: IMRertrieval = None,
|
||
code_retrieval: IMRertrieval = None,
|
||
search_retrieval: IMRertrieval = None,
|
||
# chain_type: str = "BaseChain",
|
||
**kwargs
|
||
):
|
||
self.agent_flows = sorted(agent_flows, key=lambda x:x.agent_index)
|
||
self.chat_turn = chat_turn
|
||
self.do_checker = do_checker
|
||
self.chain_name = chain_name
|
||
self.chain_index = chain_index
|
||
self.chain_type = "BaseChain"
|
||
|
||
self.embeddings = embeddings
|
||
self.llm = llm
|
||
|
||
self.doc_retrieval = doc_retrieval
|
||
self.code_retrieval = code_retrieval
|
||
self.search_retrieval = search_retrieval
|
||
# self.build_config()
|
||
# self.build_chain()
|
||
|
||
def build_config(self, embeddings: Embeddings = None, llm: BaseLLM = None):
|
||
self.llm_config = LLMConfig(model_name="test", llm=self.llm or llm)
|
||
self.embed_config = EmbedConfig(embed_model="test", langchain_embeddings=self.embeddings or embeddings)
|
||
|
||
def build_chain(self,
|
||
embeddings: Embeddings = None, llm: BaseLLM = None,
|
||
doc_retrieval: IMRertrieval = None,
|
||
code_retrieval: IMRertrieval = None,
|
||
search_retrieval: IMRertrieval = None,
|
||
):
|
||
# chain 实例化
|
||
chain_module = importlib.import_module("coagent.connector.chains")
|
||
baseChain: BaseChain = getattr(chain_module, self.chain_type)
|
||
|
||
agent_names = [agent_flow.role_name for agent_flow in self.agent_flows]
|
||
chain_config = ChainConfig(chain_name=self.chain_name, agents=agent_names, do_checker=self.do_checker, chat_turn=self.chat_turn)
|
||
|
||
# agent 实例化
|
||
self.build_config(embeddings, llm)
|
||
for agent_flow in self.agent_flows:
|
||
agent_flow.build_agent(embeddings, llm)
|
||
|
||
self.chain = baseChain(
|
||
chain_config,
|
||
[agent_flow.agent for agent_flow in self.agent_flows],
|
||
embed_config=self.embed_config,
|
||
llm_config=self.llm_config,
|
||
doc_retrieval=doc_retrieval or self.doc_retrieval,
|
||
code_retrieval=code_retrieval or self.code_retrieval,
|
||
search_retrieval=search_retrieval or self.search_retrieval,
|
||
)
|
||
|
||
class PhaseFlow:
|
||
def __init__(
|
||
self,
|
||
phase_name: str,
|
||
chain_flows: List[ChainFlow],
|
||
embeddings: Embeddings = None,
|
||
llm: BaseLLM = None,
|
||
tools: List[Tool] = [],
|
||
doc_retrieval: IMRertrieval = None,
|
||
code_retrieval: IMRertrieval = None,
|
||
search_retrieval: IMRertrieval = None,
|
||
**kwargs
|
||
):
|
||
self.phase_name = phase_name
|
||
self.chain_flows = sorted(chain_flows, key=lambda x:x.chain_index)
|
||
self.phase_type = "BasePhase"
|
||
self.tools = tools
|
||
|
||
self.embeddings = embeddings
|
||
self.llm = llm
|
||
|
||
self.doc_retrieval = doc_retrieval
|
||
self.code_retrieval = code_retrieval
|
||
self.search_retrieval = search_retrieval
|
||
# self.build_config()
|
||
self.build_phase()
|
||
|
||
def __call__(self, params: dict) -> str:
|
||
|
||
# tools = toLangchainTools([TOOL_DICT[i] for i in TOOL_SETS if i in TOOL_DICT])
|
||
# query_content = "帮我确认下127.0.0.1这个服务器的在10点是否存在异常，请帮我判断一下"
|
||
try:
|
||
logger.info(f"params: {params}")
|
||
query_content = params.get("query") or params.get("input")
|
||
search_type = params.get("search_type")
|
||
query = Message(
|
||
role_name="human", role_type="user", tools=self.tools,
|
||
role_content=query_content, input_query=query_content, origin_query=query_content,
|
||
cb_search_type=search_type,
|
||
)
|
||
# phase.pre_print(query)
|
||
output_message, output_memory = self.phase.step(query)
|
||
output_content = "\n\n".join((output_memory.to_str_messages(return_all=True, content_key="parsed_output_list").split("\n\n")[1:])) or output_message.role_content
|
||
return output_content
|
||
except Exception as e:
|
||
logger.exception(e)
|
||
return f"Error {e}"
|
||
|
||
def build_config(self, embeddings: Embeddings = None, llm: BaseLLM = None):
|
||
self.llm_config = LLMConfig(model_name="test", llm=self.llm or llm)
|
||
self.embed_config = EmbedConfig(embed_model="test", langchain_embeddings=self.embeddings or embeddings)
|
||
|
||
def build_phase(self, embeddings: Embeddings = None, llm: BaseLLM = None):
|
||
# phase 实例化
|
||
phase_module = importlib.import_module("coagent.connector.phase")
|
||
basePhase: BasePhase = getattr(phase_module, self.phase_type)
|
||
|
||
# chain 实例化
|
||
self.build_config(self.embeddings or embeddings, self.llm or llm)
|
||
os.environ["log_verbose"] = "2"
|
||
for chain_flow in self.chain_flows:
|
||
chain_flow.build_chain(
|
||
self.embeddings or embeddings, self.llm or llm,
|
||
self.doc_retrieval, self.code_retrieval, self.search_retrieval
|
||
)
|
||
|
||
self.phase: BasePhase = basePhase(
|
||
phase_name=self.phase_name,
|
||
chains=[chain_flow.chain for chain_flow in self.chain_flows],
|
||
embed_config=self.embed_config,
|
||
llm_config=self.llm_config,
|
||
doc_retrieval=self.doc_retrieval,
|
||
code_retrieval=self.code_retrieval,
|
||
search_retrieval=self.search_retrieval
|
||
)
|