From dfca371d1567dd264aca65c544603077f8be5eb8 Mon Sep 17 00:00:00 2001
From: Thiago Castro Ferreira <thiago@aixplain.com>
Date: Thu, 21 Nov 2024 16:16:33 -0300
Subject: [PATCH 1/2] Validate prompt variables

---
 aixplain/modules/agent/__init__.py      | 17 ++++++++++++++++-
 aixplain/modules/team_agent/__init__.py | 17 ++++++++++++++++-
 tests/unit/agent_test.py                | 10 ++++++++++
 3 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/aixplain/modules/agent/__init__.py b/aixplain/modules/agent/__init__.py
index b7aad7aa..98e25d96 100644
--- a/aixplain/modules/agent/__init__.py
+++ b/aixplain/modules/agent/__init__.py
@@ -238,9 +238,24 @@ def run_async(
 
         headers = {"x-api-key": self.api_key, "Content-Type": "application/json"}
 
+        # build query
+        variables = re.findall(r"(?<!{){([^}]+)}(?!})", self.description)
+        input_data = {"input": FileFactory.to_link(query)}
+        for variable in variables:
+            if isinstance(data, dict):
+                assert (
+                    variable in data or variable in parameters
+                ), f"Variable '{variable}' not found in data or parameters. This variable is required by the agent according to its description ('{self.description}')."
+                input_data[variable] = data.pop(variable) if variable in data else parameters.pop(variable)
+            else:
+                assert (
+                    variable in parameters
+                ), f"Variable '{variable}' not found in parameters. This variable is required by the agent according to its description ('{self.description}')."
+                input_data[variable] = parameters.pop(variable)
+
         payload = {
             "id": self.id,
-            "query": FileFactory.to_link(query),
+            "query": input_data,
             "sessionId": session_id,
             "history": history,
             "executionParams": {
diff --git a/aixplain/modules/team_agent/__init__.py b/aixplain/modules/team_agent/__init__.py
index 08d820f0..9d2a8525 100644
--- a/aixplain/modules/team_agent/__init__.py
+++ b/aixplain/modules/team_agent/__init__.py
@@ -222,9 +222,24 @@ def run_async(
 
         headers = {"x-api-key": self.api_key, "Content-Type": "application/json"}
 
+        # build query
+        variables = re.findall(r"(?<!{){([^}]+)}(?!})", self.description)
+        input_data = {"input": FileFactory.to_link(query)}
+        for variable in variables:
+            if isinstance(data, dict):
+                assert (
+                    variable in data or variable in parameters
+                ), f"Variable '{variable}' not found in data or parameters. This variable is required by the agent according to its description ('{self.description}')."
+                input_data[variable] = data.pop(variable) if variable in data else parameters.pop(variable)
+            else:
+                assert (
+                    variable in parameters
+                ), f"Variable '{variable}' not found in parameters. This variable is required by the agent according to its description ('{self.description}')."
+                input_data[variable] = parameters.pop(variable)
+
         payload = {
             "id": self.id,
-            "query": FileFactory.to_link(query),
+            "query": input_data,
             "sessionId": session_id,
             "history": history,
             "executionParams": {
diff --git a/tests/unit/agent_test.py b/tests/unit/agent_test.py
index 9e38937f..5dcbfdfd 100644
--- a/tests/unit/agent_test.py
+++ b/tests/unit/agent_test.py
@@ -244,3 +244,13 @@ def test_run_success():
         )
     assert response["status"] == "IN_PROGRESS"
     assert response["url"] == ref_response["data"]
+
+
+def test_run_variable_error():
+    agent = Agent("123", "Test Agent", "Translate the input data into {target_language}")
+    with pytest.raises(Exception) as exc_info:
+        agent.run_async(data={"query": "Hello, how are you?"}, output_format=OutputFormat.MARKDOWN)
+    assert (
+        str(exc_info.value)
+        == "Variable 'target_language' not found in data or parameters. This variable is required by the agent according to its description ('Translate the input data into {target_language}')."
+    )

From caefbde7055b7bc56380cc04df6411028151e715 Mon Sep 17 00:00:00 2001
From: Thiago Castro Ferreira <thiago@aixplain.com>
Date: Fri, 22 Nov 2024 16:48:57 -0300
Subject: [PATCH 2/2] Isolate process variable process in a method

---
 aixplain/modules/agent/__init__.py      | 15 ++-------------
 aixplain/modules/agent/utils.py         | 22 ++++++++++++++++++++++
 aixplain/modules/team_agent/__init__.py | 15 ++-------------
 tests/unit/agent_test.py                | 11 +++++++++++
 4 files changed, 37 insertions(+), 26 deletions(-)
 create mode 100644 aixplain/modules/agent/utils.py

diff --git a/aixplain/modules/agent/__init__.py b/aixplain/modules/agent/__init__.py
index 98e25d96..c436b84a 100644
--- a/aixplain/modules/agent/__init__.py
+++ b/aixplain/modules/agent/__init__.py
@@ -36,6 +36,7 @@
 from aixplain.modules.agent.tool import Tool
 from aixplain.modules.agent.tool.model_tool import ModelTool
 from aixplain.modules.agent.tool.pipeline_tool import PipelineTool
+from aixplain.modules.agent.utils import process_variables
 from typing import Dict, List, Text, Optional, Union
 from urllib.parse import urljoin
 
@@ -239,19 +240,7 @@ def run_async(
         headers = {"x-api-key": self.api_key, "Content-Type": "application/json"}
 
         # build query
-        variables = re.findall(r"(?<!{){([^}]+)}(?!})", self.description)
-        input_data = {"input": FileFactory.to_link(query)}
-        for variable in variables:
-            if isinstance(data, dict):
-                assert (
-                    variable in data or variable in parameters
-                ), f"Variable '{variable}' not found in data or parameters. This variable is required by the agent according to its description ('{self.description}')."
-                input_data[variable] = data.pop(variable) if variable in data else parameters.pop(variable)
-            else:
-                assert (
-                    variable in parameters
-                ), f"Variable '{variable}' not found in parameters. This variable is required by the agent according to its description ('{self.description}')."
-                input_data[variable] = parameters.pop(variable)
+        input_data = process_variables(query, data, parameters, self.description)
 
         payload = {
             "id": self.id,
diff --git a/aixplain/modules/agent/utils.py b/aixplain/modules/agent/utils.py
new file mode 100644
index 00000000..03de61d1
--- /dev/null
+++ b/aixplain/modules/agent/utils.py
@@ -0,0 +1,22 @@
+from typing import Dict, Text, Union
+import re
+
+
+def process_variables(query: Text, data: Union[Dict, Text], parameters: Dict, agent_description: Text) -> Text:
+    from aixplain.factories.file_factory import FileFactory
+
+    variables = re.findall(r"(?<!{){([^}]+)}(?!})", agent_description)
+    input_data = {"input": FileFactory.to_link(query)}
+    for variable in variables:
+        if isinstance(data, dict):
+            assert (
+                variable in data or variable in parameters
+            ), f"Variable '{variable}' not found in data or parameters. This variable is required by the agent according to its description ('{agent_description}')."
+            input_data[variable] = data.pop(variable) if variable in data else parameters.pop(variable)
+        else:
+            assert (
+                variable in parameters
+            ), f"Variable '{variable}' not found in parameters. This variable is required by the agent according to its description ('{agent_description}')."
+            input_data[variable] = parameters.pop(variable)
+
+    return input_data
diff --git a/aixplain/modules/team_agent/__init__.py b/aixplain/modules/team_agent/__init__.py
index 9d2a8525..f92b437d 100644
--- a/aixplain/modules/team_agent/__init__.py
+++ b/aixplain/modules/team_agent/__init__.py
@@ -34,6 +34,7 @@
 from aixplain.enums.storage_type import StorageType
 from aixplain.modules.model import Model
 from aixplain.modules.agent import Agent, OutputFormat
+from aixplain.modules.agent.utils import process_variables
 from typing import Dict, List, Text, Optional, Union
 from urllib.parse import urljoin
 
@@ -223,19 +224,7 @@ def run_async(
         headers = {"x-api-key": self.api_key, "Content-Type": "application/json"}
 
         # build query
-        variables = re.findall(r"(?<!{){([^}]+)}(?!})", self.description)
-        input_data = {"input": FileFactory.to_link(query)}
-        for variable in variables:
-            if isinstance(data, dict):
-                assert (
-                    variable in data or variable in parameters
-                ), f"Variable '{variable}' not found in data or parameters. This variable is required by the agent according to its description ('{self.description}')."
-                input_data[variable] = data.pop(variable) if variable in data else parameters.pop(variable)
-            else:
-                assert (
-                    variable in parameters
-                ), f"Variable '{variable}' not found in parameters. This variable is required by the agent according to its description ('{self.description}')."
-                input_data[variable] = parameters.pop(variable)
+        input_data = process_variables(query, data, parameters, self.description)
 
         payload = {
             "id": self.id,
diff --git a/tests/unit/agent_test.py b/tests/unit/agent_test.py
index 5dcbfdfd..cb60c620 100644
--- a/tests/unit/agent_test.py
+++ b/tests/unit/agent_test.py
@@ -6,6 +6,7 @@
 from aixplain.utils import config
 from aixplain.factories import AgentFactory
 from aixplain.modules.agent import PipelineTool, ModelTool
+from aixplain.modules.agent.utils import process_variables
 from urllib.parse import urljoin
 
 
@@ -254,3 +255,13 @@ def test_run_variable_error():
         str(exc_info.value)
         == "Variable 'target_language' not found in data or parameters. This variable is required by the agent according to its description ('Translate the input data into {target_language}')."
     )
+
+
+def test_process_variables():
+    query = "Hello, how are you?"
+    data = {"target_language": "English"}
+    agent_description = "Translate the input data into {target_language}"
+    assert process_variables(query=query, data=data, parameters={}, agent_description=agent_description) == {
+        "input": "Hello, how are you?",
+        "target_language": "English",
+    }