From 81e366ec5818c205553e610c3c7aa7ad2742f24f Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Mon, 25 Mar 2024 14:44:14 -0300 Subject: [PATCH] Refactor CSVAgent build method to include handle_parse_errors parameter --- src/backend/langflow/components/agents/CSVAgent.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/src/backend/langflow/components/agents/CSVAgent.py b/src/backend/langflow/components/agents/CSVAgent.py index b54e5d90d..cdbd30805 100644 --- a/src/backend/langflow/components/agents/CSVAgent.py +++ b/src/backend/langflow/components/agents/CSVAgent.py @@ -1,7 +1,8 @@ -from langflow import CustomComponent -from langflow.field_typing import BaseLanguageModel, AgentExecutor from langchain_experimental.agents.agent_toolkits.csv.base import create_csv_agent +from langflow import CustomComponent +from langflow.field_typing import AgentExecutor, BaseLanguageModel + class CSVAgentComponent(CustomComponent): display_name = "CSVAgent" @@ -12,12 +13,9 @@ class CSVAgentComponent(CustomComponent): return { "llm": {"display_name": "LLM", "type": BaseLanguageModel}, "path": {"display_name": "Path", "field_type": "file", "suffixes": [".csv"], "file_types": [".csv"]}, + "handle_parse_errors": {"display_name": "Handle Parse Errors", "advanced": True}, } - def build( - self, - llm: BaseLanguageModel, - path: str, - ) -> AgentExecutor: + def build(self, llm: BaseLanguageModel, path: str, handle_parse_errors: bool = True) -> AgentExecutor: # Instantiate and return the CSV agent class with the provided llm and path - return create_csv_agent(llm=llm, path=path) + return create_csv_agent(llm=llm, path=path, agent_executor_kwargs=dict(handle_parse_errors=handle_parse_errors))