Update JsonAgent and ZeroShotAgent components

This commit is contained in:
anovazzi1 2024-01-25 19:14:58 -03:00
commit d43d9d3760
2 changed files with 6 additions and 6 deletions

View file

@ -3,7 +3,7 @@ from langchain.agents import AgentExecutor, create_json_agent
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.json.toolkit import JsonToolkit
class JsonAgentComponent(CustomComponent): class JsonAgentComponent(CustomComponent):
@ -19,6 +19,6 @@ class JsonAgentComponent(CustomComponent):
def build( def build(
self, self,
llm: BaseLanguageModel, llm: BaseLanguageModel,
toolkit: BaseToolkit, toolkit: JsonToolkit,
) -> AgentExecutor: ) -> AgentExecutor:
return create_json_agent(llm=llm, toolkit=toolkit) return create_json_agent(llm=llm, toolkit=toolkit)

View file

@ -1,6 +1,6 @@
from typing import List from typing import List, Optional
from langchain.agents import ZeroShotAgent from langchain.agents.mrkl.base import ZeroShotAgent
from langchain_core.tools import BaseTool from langchain_core.tools import BaseTool
from langflow import CustomComponent from langflow import CustomComponent
from langflow.components.chains.LLMChain import LLMChain from langflow.components.chains.LLMChain import LLMChain
@ -21,8 +21,8 @@ class ZeroShotAgentComponent(CustomComponent):
def build( def build(
self, self,
llm: LLMChain, llm: LLMChain,
tools: List[BaseTool], tools: Optional[List[BaseTool]] = None,
prefix: str = "Answer the following questions as best you can. You have access to the following tools:", prefix: str = "Answer the following questions as best you can. You have access to the following tools:",
suffix: str = "Begin!\n\nQuestion: {input}\nThought:{agent_scratchpad}", suffix: str = "Begin!\n\nQuestion: {input}\nThought:{agent_scratchpad}",
) -> ZeroShotAgent: ) -> ZeroShotAgent:
return ZeroShotAgent(llm_chain=llm, tools=tools, prefix=prefix, suffix=suffix) return ZeroShotAgent(llm_chain=llm, allowed_tools=tools, prefix=prefix, suffix=suffix)