Refactor CSVAgent build method to include handle_parse_errors parameter

This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-03-25 14:44:14 -03:00
commit 81e366ec58

View file

@ -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 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): class CSVAgentComponent(CustomComponent):
display_name = "CSVAgent" display_name = "CSVAgent"
@ -12,12 +13,9 @@ class CSVAgentComponent(CustomComponent):
return { return {
"llm": {"display_name": "LLM", "type": BaseLanguageModel}, "llm": {"display_name": "LLM", "type": BaseLanguageModel},
"path": {"display_name": "Path", "field_type": "file", "suffixes": [".csv"], "file_types": [".csv"]}, "path": {"display_name": "Path", "field_type": "file", "suffixes": [".csv"], "file_types": [".csv"]},
"handle_parse_errors": {"display_name": "Handle Parse Errors", "advanced": True},
} }
def build( def build(self, llm: BaseLanguageModel, path: str, handle_parse_errors: bool = True) -> AgentExecutor:
self,
llm: BaseLanguageModel,
path: str,
) -> AgentExecutor:
# Instantiate and return the CSV agent class with the provided llm and path # 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))