feat: Add test for evaluate step

agents-api/agents_api/activities/demo.py

@@ -1,3 +1,5 @@
+from typing import Callable
+
 from temporalio import activity
 
 from ..env import testing
@@ -12,6 +14,6 @@ async def mock_demo_activity(a: int, b: int) -> int:
     return a + b
 
 
-demo_activity = activity.defn(name="demo_activity")(
+demo_activity: Callable[[int, int], int] = activity.defn(name="demo_activity")(
     demo_activity if not testing else mock_demo_activity
 )

agents-api/agents_api/activities/embed_docs.py

@@ -1,8 +1,8 @@
 from beartype import beartype
 from temporalio import activity
 
+from ..clients import cozo
 from ..clients import embed as embedder
-from ..clients.cozo import get_cozo_client
 from ..env import testing
 from ..models.docs.embed_snippets import embed_snippets as embed_snippets_query
 from .types import EmbedDocsPayload
@@ -28,7 +28,7 @@ async def embed_docs(payload: EmbedDocsPayload, cozo_client=None) -> None:
         doc_id=payload.doc_id,
         snippet_indices=indices,
         embeddings=embeddings,
-        client=cozo_client or get_cozo_client(),
+        client=cozo_client or cozo.get_cozo_client(),
     )
 
 

agents-api/agents_api/activities/logger.py

@@ -1,7 +1,8 @@
 import logging
+from typing import TextIO
 
-logger = logging.getLogger(__name__)
-h = logging.StreamHandler()
-fmt = logging.Formatter("[%(asctime)s/%(levelname)s] - %(message)s")
+logger: logging.Logger = logging.getLogger(__name__)
+h: logging.StreamHandler[TextIO] = logging.StreamHandler()
+fmt: logging.Formatter = logging.Formatter("[%(asctime)s/%(levelname)s] - %(message)s")
 h.setFormatter(fmt)
 logger.addHandler(h)

agents-api/agents_api/activities/summarization.py

@@ -12,11 +12,11 @@
 
 
 # TODO: remove stubs
-def entries_summarization_query(*args, **kwargs):
+def entries_summarization_query(*args, **kwargs) -> pd.DataFrame:
     return pd.DataFrame()
 
 
-def get_toplevel_entries_query(*args, **kwargs):
+def get_toplevel_entries_query(*args, **kwargs) -> pd.DataFrame:
     return pd.DataFrame()
 
 

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

@@ -1,6 +1,3 @@
-from typing import Any
-
-from beartype import beartype
 from temporalio import activity
 
 from ...activities.task_steps.utils import simple_eval_dict
@@ -12,14 +9,16 @@
 from ...env import testing
 
 
-@beartype
 async def evaluate_step(
     context: StepContext[EvaluateStep],
-) -> StepOutcome[dict[str, Any]]:
-    exprs = context.definition.arguments
+) -> StepOutcome:
+    assert isinstance(context.current_step, EvaluateStep)
+
+    exprs = context.current_step.evaluate
     output = simple_eval_dict(exprs, values=context.model_dump())
 
-    return StepOutcome(output=output)
+    result = StepOutcome(output=output)
+    return result
 
 
 # Note: This is here just for clarity. We could have just imported evaluate_step directly

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

@@ -16,9 +16,9 @@ async def if_else_step(context: StepContext[IfElseWorkflowStep]) -> dict:
     # context_data: dict = context.model_dump()
 
     # next_workflow = (
-    #     context.definition.then
-    #     if simple_eval(context.definition.if_, names=context_data)
-    #     else context.definition.else_
+    #     context.current_step.then
+    #     if simple_eval(context.current_step.if_, names=context_data)
+    #     else context.current_step.else_
     # )
 
     # return {"goto_workflow": next_workflow}

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

@@ -25,9 +25,9 @@ async def prompt_step(context: StepContext[PromptStep]) -> StepOutcome:
 
     # Render template messages
     prompt = (
-        [InputChatMLMessage(content=context.definition.prompt)]
-        if isinstance(context.definition.prompt, str)
-        else context.definition.prompt
+        [InputChatMLMessage(content=context.current_step.prompt)]
+        if isinstance(context.current_step.prompt, str)
+        else context.current_step.prompt
     )
 
     template_messages: list[InputChatMLMessage] = prompt
@@ -47,7 +47,7 @@ async def prompt_step(context: StepContext[PromptStep]) -> StepOutcome:
         for m in messages
     ]
 
-    settings: dict = context.definition.settings.model_dump()
+    settings: dict = context.current_step.settings.model_dump()
     # Get settings and run llm
     response = await litellm.acompletion(
         messages=messages,

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

@@ -10,10 +10,10 @@
 @beartype
 async def tool_call_step(context: StepContext) -> dict:
     raise NotImplementedError()
-    # assert isinstance(context.definition, ToolCallStep)
+    # assert isinstance(context.current_step, ToolCallStep)
 
-    # context.definition.tool_id
-    # context.definition.arguments
+    # context.current_step.tool_id
+    # context.current_step.arguments
     # # get tool by id
     # # call tool
 

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

@@ -1,4 +1,4 @@
-from typing import Any
+from typing import Callable
 
 from beartype import beartype
 from temporalio import activity
@@ -12,7 +12,7 @@
 
 
 @beartype
-async def yield_step(context: StepContext[YieldStep]) -> StepOutcome[dict[str, Any]]:
+async def yield_step(context: StepContext[YieldStep]) -> StepOutcome:
     all_workflows = context.execution_input.task.workflows
     workflow = context.current_step.workflow
 
@@ -35,8 +35,8 @@ async def yield_step(context: StepContext[YieldStep]) -> StepOutcome[dict[str, A
 
 # Note: This is here just for clarity. We could have just imported yield_step directly
 # They do the same thing, so we dont need to mock the yield_step function
-mock_yield_step = yield_step
+mock_yield_step: Callable[[StepContext], StepOutcome] = yield_step
 
-yield_step = activity.defn(name="yield_step")(
+yield_step: Callable[[StepContext], StepOutcome] = activity.defn(name="yield_step")(
     yield_step if not testing else mock_yield_step
 )