diff --git a/src/backend/langflow/components/agents/CSVAgent.py b/src/backend/langflow/components/agents/CSVAgent.py index cdbd30805..92cfac0db 100644 --- a/src/backend/langflow/components/agents/CSVAgent.py +++ b/src/backend/langflow/components/agents/CSVAgent.py @@ -14,8 +14,19 @@ class CSVAgentComponent(CustomComponent): "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}, + "agent_type": { + "display_name": "Agent Type", + "options": ["zero-shot-react-description", "openai-functions", "openai-tools"], + }, } - def build(self, llm: BaseLanguageModel, path: str, handle_parse_errors: bool = True) -> AgentExecutor: + def build( + self, llm: BaseLanguageModel, path: str, handle_parse_errors: bool = True, agent_type: str = "openai-tools" + ) -> AgentExecutor: # Instantiate and return the CSV agent class with the provided llm and path - return create_csv_agent(llm=llm, path=path, agent_executor_kwargs=dict(handle_parse_errors=handle_parse_errors)) + return create_csv_agent( + llm=llm, + path=path, + agent_type=agent_type, + agent_executor_kwargs=dict(handle_parse_errors=handle_parse_errors), + )