Refactor JsonAgentComponent to use create_json_agent function

This commit is contained in:
anovazzi1 2024-01-15 17:52:10 -03:00
commit f322c77c5c

View file

@ -1,12 +1,10 @@
from langflow import CustomComponent from langflow import CustomComponent
from langchain.agents import AgentExecutor from langchain.agents import AgentExecutor,create_json_agent
from typing import Callable
from langflow.field_typing import ( from langflow.field_typing import (
BaseLanguageModel, BaseLanguageModel,
) )
from langchain_community.agent_toolkits.base import BaseToolkit from langchain_community.agent_toolkits.base import BaseToolkit
class JsonAgentComponent(CustomComponent): class JsonAgentComponent(CustomComponent):
display_name = "JsonAgent" display_name = "JsonAgent"
description = "Construct a json agent from an LLM and tools." description = "Construct a json agent from an LLM and tools."
@ -21,5 +19,5 @@ class JsonAgentComponent(CustomComponent):
self, self,
llm: BaseLanguageModel, llm: BaseLanguageModel,
toolkit: BaseToolkit, toolkit: BaseToolkit,
) -> Callable: ) -> AgentExecutor:
return AgentExecutor(llm=llm, toolkit=toolkit) return create_json_agent(llm=llm, toolkit=toolkit[0])