From 70ffff186aa9dfe80473dc4fe98a2bfbee62d92b Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Mon, 24 Jun 2024 08:32:53 -0300 Subject: [PATCH] refactor: Update ToolCallingAgent.py imports and type annotations --- .../langflow/components/agents/ToolCallingAgent.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/src/backend/base/langflow/components/agents/ToolCallingAgent.py b/src/backend/base/langflow/components/agents/ToolCallingAgent.py index 318c5c622..e2b643f1c 100644 --- a/src/backend/base/langflow/components/agents/ToolCallingAgent.py +++ b/src/backend/base/langflow/components/agents/ToolCallingAgent.py @@ -1,12 +1,14 @@ -from typing import List, Dict +from typing import Dict, List, cast + +from langchain.agents import AgentExecutor, BaseSingleActionAgent from langchain.agents.tool_calling_agent.base import create_tool_calling_agent -from langchain_core.prompts import ChatPromptTemplate -from langchain.agents import AgentExecutor from langchain_core.messages import BaseMessage -from langflow.schema.message import Message +from langchain_core.prompts import ChatPromptTemplate + from langflow.custom import Component -from langflow.io import HandleInput, BoolInput, Output, MessageTextInput +from langflow.io import BoolInput, HandleInput, MessageTextInput, Output from langflow.schema import Data +from langflow.schema.message import Message class ToolCallingAgentComponent(Component): @@ -80,7 +82,7 @@ class ToolCallingAgentComponent(Component): agent = create_tool_calling_agent(self.llm, self.tools, prompt) runnable = AgentExecutor.from_agent_and_tools( - agent=agent, + agent=cast(BaseSingleActionAgent, agent), tools=self.tools, verbose=True, handle_parsing_errors=self.handle_parsing_errors,