File size: 1,672 Bytes
256a159 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 |
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
|