From d43d9d3760c85e64c3d1c50d7fb3d14b68e0a0a4 Mon Sep 17 00:00:00 2001
From: anovazzi1 <otavio2204@gmail.com>
Date: Thu, 25 Jan 2024 19:14:58 -0300
Subject: [PATCH] Update JsonAgent and ZeroShotAgent components

---
 src/backend/langflow/components/agents/JsonAgent.py     | 4 ++--
 src/backend/langflow/components/agents/ZeroShotAgent.py | 8 ++++----
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/src/backend/langflow/components/agents/JsonAgent.py b/src/backend/langflow/components/agents/JsonAgent.py
index d15366b77..0197b9210 100644
--- a/src/backend/langflow/components/agents/JsonAgent.py
+++ b/src/backend/langflow/components/agents/JsonAgent.py
@@ -3,7 +3,7 @@ from langchain.agents import AgentExecutor, create_json_agent
 from langflow.field_typing import (
     BaseLanguageModel,
 )
-from langchain_community.agent_toolkits.base import BaseToolkit
+from langchain_community.agent_toolkits.json.toolkit import JsonToolkit
 
 
 class JsonAgentComponent(CustomComponent):
@@ -19,6 +19,6 @@ class JsonAgentComponent(CustomComponent):
     def build(
         self,
         llm: BaseLanguageModel,
-        toolkit: BaseToolkit,
+        toolkit: JsonToolkit,
     ) -> AgentExecutor:
         return create_json_agent(llm=llm, toolkit=toolkit)
diff --git a/src/backend/langflow/components/agents/ZeroShotAgent.py b/src/backend/langflow/components/agents/ZeroShotAgent.py
index b65e14e41..538e693cd 100644
--- a/src/backend/langflow/components/agents/ZeroShotAgent.py
+++ b/src/backend/langflow/components/agents/ZeroShotAgent.py
@@ -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 langflow import CustomComponent
 from langflow.components.chains.LLMChain import LLMChain
@@ -21,8 +21,8 @@ class ZeroShotAgentComponent(CustomComponent):
     def build(
         self,
         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:",
         suffix: str = "Begin!\n\nQuestion: {input}\nThought:{agent_scratchpad}",
     ) -> ZeroShotAgent:
-        return ZeroShotAgent(llm_chain=llm, tools=tools, prefix=prefix, suffix=suffix)
+        return ZeroShotAgent(llm_chain=llm, allowed_tools=tools, prefix=prefix, suffix=suffix)