From e44479fb1679415e7c5fd64a5b0c356dca6cafec Mon Sep 17 00:00:00 2001
From: Cristhian Zanforlin Lousa <cristhian.lousa@gmail.com>
Date: Thu, 12 Sep 2024 17:03:05 -0300
Subject: [PATCH] fix: update csvAgent to work as expected (#3790)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* ✨ (CSVAgent.py): Add support for MessageTextInput input type and Output for building agent response
📝 (CSVAgent.py): Update input descriptions and add info for better understanding of inputs
📝 (CSVAgent.py): Update build_agent_response method to handle input and return response as a Message object
📝 (CSVAgent.py): Update build_agent method to return Union[AgentExecutor, Agent] for flexibility and handle parsing errors

* [autofix.ci] apply automated fixes

* fix: remove Union from func return type

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: italojohnny <italojohnnydosanjos@gmail.com>
---
 .../langflow/components/agents/CSVAgent.py    | 60 +++++++++++++++++--
 1 file changed, 56 insertions(+), 4 deletions(-)

diff --git a/src/backend/base/langflow/components/agents/CSVAgent.py b/src/backend/base/langflow/components/agents/CSVAgent.py
index f28691ac3..444a04ae2 100644
--- a/src/backend/base/langflow/components/agents/CSVAgent.py
+++ b/src/backend/base/langflow/components/agents/CSVAgent.py
@@ -1,8 +1,11 @@
 from langchain_experimental.agents.agent_toolkits.csv.base import create_csv_agent
-
 from langflow.base.agents.agent import LCAgentComponent
 from langflow.field_typing import AgentExecutor
 from langflow.inputs import HandleInput, FileInput, DropdownInput
+from langflow.inputs.inputs import MessageTextInput
+from langflow.schema.message import Message
+
+from langflow.template.field.base import Output
 
 
 class CSVAgentComponent(LCAgentComponent):
@@ -12,8 +15,21 @@ class CSVAgentComponent(LCAgentComponent):
     name = "CSVAgent"
 
     inputs = LCAgentComponent._base_inputs + [
-        HandleInput(name="llm", display_name="Language Model", input_types=["LanguageModel"], required=True),
-        FileInput(name="path", display_name="File Path", file_types=["csv"], required=True),
+        HandleInput(
+            name="llm",
+            display_name="Language Model",
+            input_types=["LanguageModel"],
+            required=True,
+            info="An LLM Model Object (It can be found in any LLM Component).",
+        ),
+        FileInput(
+            name="path",
+            display_name="File Path",
+            file_types=["csv"],
+            input_types=["str", "Message"],
+            required=True,
+            info="A CSV File or File Path.",
+        ),
         DropdownInput(
             name="agent_type",
             display_name="Agent Type",
@@ -21,7 +37,43 @@ class CSVAgentComponent(LCAgentComponent):
             options=["zero-shot-react-description", "openai-functions", "openai-tools"],
             value="openai-tools",
         ),
+        MessageTextInput(
+            name="input_value",
+            display_name="Text",
+            info="Text to be passed as input and extract info from the CSV File.",
+        ),
     ]
 
+    outputs = [
+        Output(display_name="Response", name="response", method="build_agent_response"),
+        Output(display_name="Agent", name="agent", method="build_agent"),
+    ]
+
+    def build_agent_response(self) -> Message:
+        agent_kwargs = {
+            "verbose": self.verbose,
+            "allow_dangerous_code": True,
+        }
+
+        agent_csv = create_csv_agent(
+            llm=self.llm, path=self.path, agent_type=self.agent_type, handle_parsing_errors=True, **agent_kwargs
+        )
+
+        print(f"self.inputs = {agent_csv}")
+
+        result = agent_csv.invoke({"input": self.input_value})
+        return Message(text=str(result["output"]))
+
     def build_agent(self) -> AgentExecutor:
-        return create_csv_agent(llm=self.llm, path=self.path, agent_type=self.agent_type, **self.get_agent_kwargs())
+        agent_kwargs = {
+            "verbose": self.verbose,
+            "allow_dangerous_code": True,
+        }
+
+        agent_csv = create_csv_agent(
+            llm=self.llm, path=self.path, agent_type=self.agent_type, handle_parsing_errors=True, **agent_kwargs
+        )
+
+        self.status = Message(text=str(agent_csv))
+
+        return agent_csv