mirror of
https://github.com/open-compass/opencompass.git
synced 2025-05-30 16:03:24 +08:00
54 lines
1.6 KiB
Python
54 lines
1.6 KiB
Python
|
from typing import List, Tuple
|
||
|
|
|
||
|
|
from mmengine.registry import Registry
|
||
|
|
|
||
|
|
REGISTRY = Registry('helper')
|
||
|
|
|
||
|
|
|
||
|
|
class LangchainAgent:
|
||
|
|
"""Agent wrapper for Langchain.
|
||
|
|
|
||
|
|
https://github.com/langchain-ai/langchain.
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, agent_type, llm, tools) -> None:
|
||
|
|
from langchain.agents import initialize_agent, load_tools
|
||
|
|
|
||
|
|
llm = REGISTRY.build(llm)
|
||
|
|
tools = load_tools(tools, llm=llm)
|
||
|
|
self.agent = initialize_agent(tools,
|
||
|
|
llm,
|
||
|
|
agent=agent_type,
|
||
|
|
return_intermediate_steps=True)
|
||
|
|
|
||
|
|
def chat(self, user_input, ice=None) -> Tuple[str, List[dict]]:
|
||
|
|
from langchain.schema import AgentAction
|
||
|
|
try:
|
||
|
|
generation = self.agent(user_input)
|
||
|
|
answer = generation['output']
|
||
|
|
steps = []
|
||
|
|
for step in generation['intermediate_steps']:
|
||
|
|
action: AgentAction = step[0]
|
||
|
|
steps.append(
|
||
|
|
dict(
|
||
|
|
type=action.tool,
|
||
|
|
args=action.tool_input,
|
||
|
|
result=step[1],
|
||
|
|
thought=action.log,
|
||
|
|
state=0,
|
||
|
|
errmsg=None,
|
||
|
|
))
|
||
|
|
except Exception as e:
|
||
|
|
answer = None
|
||
|
|
steps = [
|
||
|
|
dict(
|
||
|
|
type='InvalidAction',
|
||
|
|
args={},
|
||
|
|
result=None,
|
||
|
|
thought=None,
|
||
|
|
state=-1002,
|
||
|
|
errmsg=str(e),
|
||
|
|
)
|
||
|
|
]
|
||
|
|
return answer, steps
|