ENG-1094: Validate prompt variables

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
 
@@ -238,9 +239,12 @@ def run_async(
 
         headers = {"x-api-key": self.api_key, "Content-Type": "application/json"}
 
+        # build query
+        input_data = process_variables(query, data, parameters, self.description)
+
         payload = {
             "id": self.id,
-            "query": FileFactory.to_link(query),
+            "query": input_data,
             "sessionId": session_id,
             "history": history,
             "executionParams": {

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

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
 
@@ -222,9 +223,12 @@ def run_async(
 
         headers = {"x-api-key": self.api_key, "Content-Type": "application/json"}
 
+        # build query
+        input_data = process_variables(query, data, parameters, self.description)
+
         payload = {
             "id": self.id,
-            "query": FileFactory.to_link(query),
+            "query": input_data,
             "sessionId": session_id,
             "history": history,
             "executionParams": {

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
 
 
@@ -244,3 +245,23 @@ 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}')."
+    )
+
+
+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",
+    }