Merge pull request #467 from julep-ai/x/sample-test-fix

fix(agents-api): Make the sample work

agents-api/agents_api/activities/task_steps/base_evaluate.py

@@ -1,6 +1,8 @@
 from typing import Any
 
 from beartype import beartype
+from box import Box
+from openai import BaseModel
 from temporalio import activity
 
 from ...env import testing
@@ -15,6 +17,15 @@ async def base_evaluate(
     input_len = 1 if isinstance(exprs, str) else len(exprs)
     assert input_len > 0, "exprs must be a non-empty string, list or dict"
 
+    # Turn the nested dict values from pydantic to dicts where possible
+    values = {
+        k: v.model_dump() if isinstance(v, BaseModel) else v for k, v in values.items()
+    }
+
+    # TODO: We should make this frozen_box=True, but we need to make sure that
+    # we don't break anything
+    values = Box(values, frozen_box=False, conversion_box=False)
+
     evaluator = get_evaluator(names=values)
 
     try:

agents-api/agents_api/activities/utils.py

@@ -8,6 +8,7 @@
 from yaml import CSafeLoader
 
 ALLOWED_FUNCTIONS = {
+    "zip": zip,
     "len": len,
     "load_yaml": lambda string: yaml.load(string, Loader=CSafeLoader),
     "match_regex": lambda pattern, string: bool(re2.fullmatch(pattern, string)),

agents-api/agents_api/workflows/task_execution.py

@@ -294,15 +294,13 @@ async def transition(**kwargs) -> None:
                         previous_inputs + [item],
                     ]
 
+                    # TODO: We should parallelize this
                     # Execute the chosen branch and come back here
                     output = await workflow.execute_child_workflow(
                         TaskExecutionWorkflow.run,
                         args=map_reduce_args,
                     )
 
-                    if hasattr(output, "model_dump"):
-                        output = output.model_dump()
-
                     initial = await execute_activity(
                         task_steps.base_evaluate,
                         args=[
@@ -365,7 +363,7 @@ async def transition(**kwargs) -> None:
                 )
 
             case PromptStep(), StepOutcome(output=response):
-                state.output = response.get("choices", [{}])[0].get("message")
+                state.output = response
 
             case _:
                 raise ApplicationError("Not implemented")

agents-api/pyproject.toml

@@ -34,6 +34,7 @@ lz4 = "^4.3.3"
 
 pyyaml = "^6.0.2"
 google-re2 = "^1.1.20240702"
+python-box = "^7.2.0"
 [tool.poetry.group.dev.dependencies]
 ipython = "^8.26.0"
 ruff = "^0.5.5"

agents-api/tests/sample_tasks/find_selector.yaml

@@ -63,28 +63,24 @@ main:
               image_url:
                 url: "{{inputs[0].screenshot_base64}}"
 
-    over: _.parameters
+    over: _["parameters"]
     reduce: >-
       results
       + [
-        yaml.safe_load(_["choices"][0]["message"]["content"].trim())
+        load_yaml(_["choices"][0]["message"].content.strip())
       ]
 
   - evaluate:
       result: >-
         [
           {"value": result["value"], "network_request": request}
           for request in inputs[0]["network_requests"]
-          if result["value"] in nr["response"]["body"]
           for result in _
-          if result["found"]
+          if result["found"] and result["value"] in request["response"]["body"]
         ]
 
   - if: len(_["result"]) > 0
     then:
-      workflow: find_selectors
-      arguments:
-        results: list(zip(_, execution.input.network_requests))
-        parameters: execution.input.parameters
+      log: list(zip(_, inputs[0]["network_requests"]))
     else:
-      error: "Could not find the selector in any of the network requests"
+      error: "Could not find the selector in any of the network requests"

agents-api/tests/sample_tasks/test_find_selector.py

@@ -85,9 +85,9 @@ async def _(
     agent_id = str(agent.id)
     task_id = str(uuid4())
 
-    with patch_embed_acompletion(), open(
-        f"{this_dir}/find_selector.yaml", "r"
-    ) as sample_file:
+    with patch_embed_acompletion(
+        output={"role": "assistant", "content": "found: true\nvalue: 'Gaga'"}
+    ), open(f"{this_dir}/find_selector.yaml", "r") as sample_file:
         task_def = sample_file.read()
 
         async with patch_http_client_with_temporal(