Merge pull request #1076 from julep-ai/x/evaluate-step

Ahmad-mtos · web-flow · commit 9d3633ae5863 · 2025-01-22T15:36:36.000+03:00
fix(agents-api): allow nested dictionaries in ``EvaluateStep`` and ``SetStep``
diff --git a/agents-api/agents_api/activities/utils.py b/agents-api/agents_api/activities/utils.py
@@ -422,7 +422,9 @@ def get_evaluator(
 
 
 @beartype
-def simple_eval_dict(exprs: dict[str, str], values: dict[str, Any]) -> dict[str, Any]:
+def simple_eval_dict(
+    exprs: dict[str, str | dict[str, Any]], values: dict[str, Any]
+) -> dict[str, Any]:
     if len(exprs) > MAX_COLLECTION_SIZE:
         msg = f"Too many expressions (max {MAX_COLLECTION_SIZE})"
         raise ValueError(msg)
@@ -433,7 +435,10 @@ def simple_eval_dict(exprs: dict[str, str], values: dict[str, Any]) -> dict[str,
             raise ValueError(msg)
 
     evaluator = get_evaluator(names=values)
-    return {k: evaluator.eval(v) for k, v in exprs.items()}
+    return {
+        k: evaluator.eval(v) if isinstance(v, str) else simple_eval_dict(v, values)
+        for k, v in exprs.items()
+    }
 
 
 def get_handler_with_filtered_params(system: SystemDef) -> Callable:
diff --git a/agents-api/agents_api/autogen/Tasks.py b/agents-api/agents_api/autogen/Tasks.py
@@ -247,7 +247,7 @@ class EvaluateStep(BaseModel):
     """
     The label of this step for referencing it from other steps
     """
-    evaluate: dict[str, list[str] | dict[str, str] | list[dict[str, str]] | str]
+    evaluate: dict[str, dict[str, Any] | str]
     """
     The expression to evaluate
     """
@@ -861,7 +861,7 @@ class SetStep(BaseModel):
     """
     The label of this step for referencing it from other steps
     """
-    set: dict[str, str]
+    set: dict[str, dict[str, Any] | str]
     """
     The value to set
     """
diff --git a/agents-api/agents_api/clients/litellm.py b/agents-api/agents_api/clients/litellm.py
@@ -121,15 +121,14 @@ async def get_model_list(*, custom_api_key: str | None = None) -> list[dict]:
         list[dict]: A list of model information dictionaries
     """
 
-    headers = {
-        "accept": "application/json",
-        "x-api-key": custom_api_key or litellm_master_key
-    }
-
-    async with aiohttp.ClientSession() as session, session.get(
-        url=f"{litellm_url}/models" if not custom_api_key else "/models",
-        headers=headers
-    ) as response:
+    headers = {"accept": "application/json", "x-api-key": custom_api_key or litellm_master_key}
+
+    async with (
+        aiohttp.ClientSession() as session,
+        session.get(
+            url=f"{litellm_url}/models" if not custom_api_key else "/models", headers=headers
+        ) as response,
+    ):
         response.raise_for_status()
         data = await response.json()
         return data["data"]
diff --git a/agents-api/agents_api/routers/agents/create_agent.py b/agents-api/agents_api/routers/agents/create_agent.py
@@ -19,7 +19,6 @@ async def create_agent(
     x_developer_id: Annotated[UUID, Depends(get_developer_id)],
     data: CreateAgentRequest,
 ) -> ResourceCreatedResponse:
-
     if data.model:
         await validate_model(data.model)
 
diff --git a/agents-api/agents_api/routers/agents/create_or_update_agent.py b/agents-api/agents_api/routers/agents/create_or_update_agent.py
@@ -22,7 +22,6 @@ async def create_or_update_agent(
     data: CreateOrUpdateAgentRequest,
     x_developer_id: Annotated[UUID, Depends(get_developer_id)],
 ) -> ResourceCreatedResponse:
-
     if data.model:
         await validate_model(data.model)
 
diff --git a/agents-api/agents_api/routers/agents/patch_agent.py b/agents-api/agents_api/routers/agents/patch_agent.py
@@ -22,7 +22,6 @@ async def patch_agent(
     agent_id: UUID,
     data: PatchAgentRequest,
 ) -> ResourceUpdatedResponse:
-
     if data.model:
         await validate_model(data.model)
 
diff --git a/agents-api/agents_api/routers/agents/update_agent.py b/agents-api/agents_api/routers/agents/update_agent.py
@@ -21,8 +21,7 @@ async def update_agent(
     x_developer_id: Annotated[UUID, Depends(get_developer_id)],
     agent_id: UUID,
     data: UpdateAgentRequest,
-    ) -> ResourceUpdatedResponse:
-
+) -> ResourceUpdatedResponse:
     if data.model:
         await validate_model(data.model)
 
diff --git a/agents-api/agents_api/routers/utils/model_validation.py b/agents-api/agents_api/routers/utils/model_validation.py
@@ -15,5 +15,5 @@ async def validate_model(model_name: str) -> None:
     if model_name not in available_models:
         raise HTTPException(
             status_code=HTTP_400_BAD_REQUEST,
-            detail=f"Model {model_name} not available. Available models: {available_models}"
+            detail=f"Model {model_name} not available. Available models: {available_models}",
         )
diff --git a/agents-api/tests/fixtures.py b/agents-api/tests/fixtures.py
@@ -450,9 +450,14 @@ async def test_tool(
 
 @fixture(scope="global")
 def client(_dsn=pg_dsn):
-    with TestClient(app=app) as client:
-        with patch("agents_api.routers.utils.model_validation.get_model_list", return_value=SAMPLE_MODELS):
-            yield client
+    with (
+        TestClient(app=app) as client,
+        patch(
+            "agents_api.routers.utils.model_validation.get_model_list",
+            return_value=SAMPLE_MODELS,
+        ),
+    ):
+        yield client
 
 
 @fixture(scope="global")
diff --git a/agents-api/tests/test_simple_eval_dict.py b/agents-api/tests/test_simple_eval_dict.py
@@ -0,0 +1,26 @@
+from ward import test, raises
+from agents_api.activities.utils import simple_eval_dict, MAX_STRING_LENGTH, MAX_COLLECTION_SIZE
+from simpleeval import NameNotDefined
+
+@test("utility: simple_eval_dict - string length overflow")
+async def _():
+    with raises(ValueError):
+        simple_eval_dict({"a": "b" * (MAX_STRING_LENGTH + 1)}, {})
+
+@test("utility: simple_eval_dict - collection size overflow")
+async def _():
+    with raises(ValueError):
+        simple_eval_dict({str(i): "b" for i in range(MAX_COLLECTION_SIZE + 1)}, {})
+
+@test("utility: simple_eval_dict - value undefined")
+async def _():
+    with raises(NameNotDefined):
+        simple_eval_dict({"a": "b"}, {})
+
+@test("utility: simple_eval_dict")
+async def _():
+    exprs = {"a": {"b": "x + 5", "c": "x + 6"}}
+    values = {"x": 5}
+    result = simple_eval_dict(exprs, values)
+    assert result == {"a": {"b": 10, "c": 11}}
+
diff --git a/agents-api/tests/test_workflow_routes.py b/agents-api/tests/test_workflow_routes.py
@@ -80,6 +80,51 @@ async def _(
         ).raise_for_status()
 
 
+@test("workflow route: evaluate step single with yaml - nested")
+async def _(
+    make_request=make_request,
+    agent=test_agent,
+):
+    agent_id = str(agent.id)
+
+    async with patch_testing_temporal():
+        task_data = """
+name: test task
+description: test task about
+input_schema:
+  type: object
+  additionalProperties: true
+
+main:
+  - evaluate:
+      hello: '"world"'
+      hello2:
+        hello3:
+          hello4: inputs[0]['test']
+"""
+
+        result = (
+            make_request(
+                method="POST",
+                url=f"/agents/{agent_id}/tasks",
+                content=task_data.encode("utf-8"),
+                headers={"Content-Type": "text/yaml"},
+            )
+            .raise_for_status()
+            .json()
+        )
+
+        task_id = result["id"]
+
+        execution_data = {"input": {"test": "input"}}
+
+        make_request(
+            method="POST",
+            url=f"/tasks/{task_id}/executions",
+            json=execution_data,
+        ).raise_for_status()
+
+
 @test("workflow route: create or update: evaluate step single with yaml")
 async def _(
     make_request=make_request,
diff --git a/typespec/tasks/steps.tsp b/typespec/tasks/steps.tsp
@@ -134,7 +134,7 @@ model EvaluateStep extends BaseWorkflowStep<"evaluate"> {
 
 model EvaluateStepDef {
     /** The expression to evaluate */
-    evaluate: ExpressionObject<unknown>;
+    evaluate: Record<TypedExpression<unknown> | Record<unknown>>;
 }
 
 model WaitForInputStep extends BaseWorkflowStep<"wait_for_input"> {
@@ -191,7 +191,7 @@ model SetStep extends BaseWorkflowStep<"set"> {
 
 model SetStepDef {
     /** The value to set */
-    set: Record<TypedExpression<unknown>>;
+    set: Record<TypedExpression<unknown> | Record<unknown>>;
 }
 
 ///////////////////////
diff --git a/typespec/tsp-output/@typespec/openapi3/openapi-1.0.0.yaml b/typespec/tsp-output/@typespec/openapi3/openapi-1.0.0.yaml
@@ -4887,17 +4887,8 @@ components:
           additionalProperties:
             anyOf:
               - $ref: '#/components/schemas/Common.PyExpression'
-              - type: array
-                items:
-                  $ref: '#/components/schemas/Common.PyExpression'
               - type: object
-                additionalProperties:
-                  $ref: '#/components/schemas/Common.PyExpression'
-              - type: array
-                items:
-                  type: object
-                  additionalProperties:
-                    $ref: '#/components/schemas/Common.PyExpression'
+                additionalProperties: {}
           description: The expression to evaluate
       allOf:
         - type: object
@@ -5919,7 +5910,10 @@ components:
         set:
           type: object
           additionalProperties:
-            $ref: '#/components/schemas/Common.PyExpression'
+            anyOf:
+              - $ref: '#/components/schemas/Common.PyExpression'
+              - type: object
+                additionalProperties: {}
           description: The value to set
       allOf:
         - type: object

Original file line number	Diff line number	Diff line change
`@@ -15,5 +15,5 @@ async def validate_model(model_name: str) -> None:`
`15`	`15`	`if model_name not in available_models:`
`16`	`16`	`raise HTTPException(`
`17`	`17`	`status_code=HTTP_400_BAD_REQUEST,`
`18`		`- detail=f"Model {model_name} not available. Available models: {available_models}"`
	`18`	`+ detail=f"Model {model_name} not available. Available models: {available_models}",`
`19`	`19`	`)`
Original file line number	Diff line number	Diff line change
`@@ -134,7 +134,7 @@ model EvaluateStep extends BaseWorkflowStep<"evaluate"> {`
`134`	`134`
`135`	`135`	`model EvaluateStepDef {`
`136`	`136`	`/** The expression to evaluate */`
`137`		`- evaluate: ExpressionObject<unknown>;`
	`137`	`+ evaluate: Record<TypedExpression<unknown> \| Record<unknown>>;`
`138`	`138`	`}`
`139`	`139`
`140`	`140`	`model WaitForInputStep extends BaseWorkflowStep<"wait_for_input"> {`
`@@ -191,7 +191,7 @@ model SetStep extends BaseWorkflowStep<"set"> {`
`191`	`191`
`192`	`192`	`model SetStepDef {`
`193`	`193`	`/** The value to set */`
`194`		`- set: Record<TypedExpression<unknown>>;`
	`194`	`+ set: Record<TypedExpression<unknown> \| Record<unknown>>;`
`195`	`195`	`}`
`196`	`196`
`197`	`197`	`///////////////////////`