Add agent_type parameter to CSVAgent build method

This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-03-25 14:45:31 -03:00
commit 611a2f8d38

View file

@ -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),
)