fix(agents-api): added anthropic call in chat endpoint + typespec changed to include imageurl in ChatInput

agents-api/agents_api/activities/embed_docs.py

@@ -13,7 +13,7 @@
 from .types import EmbedDocsPayload
 
 
-@auto_blob_store
+@auto_blob_store(deep=True)
 @beartype
 async def embed_docs(
     payload: EmbedDocsPayload, cozo_client=None, max_batch_size: int = 100

agents-api/agents_api/activities/excecute_api_call.py

@@ -20,7 +20,7 @@ class RequestArgs(TypedDict):
     headers: Optional[dict[str, str]]
 
 
-@auto_blob_store
+@auto_blob_store(deep=True)
 @beartype
 async def execute_api_call(
     api_call: ApiCallDef,

agents-api/agents_api/activities/execute_system.py

@@ -7,10 +7,12 @@
 from fastapi.background import BackgroundTasks
 from temporalio import activity
 
-from ..autogen.Chat import ChatInput
-from ..autogen.Docs import (
+from ..autogen.openapi_model import (
+    ChatInput,
     CreateDocRequest,
+    CreateSessionRequest,
     HybridDocSearchRequest,
+    SystemDef,
     TextOnlyDocSearchRequest,
     VectorDocSearchRequest,
 )
@@ -23,7 +25,7 @@
 from .utils import get_handler
 
 
-@auto_blob_store
+@auto_blob_store(deep=True)
 @beartype
 async def execute_system(
     context: StepContext,
@@ -82,9 +84,9 @@ async def execute_system(
 
         # Handle chat operations
         if system.operation == "chat" and system.resource == "session":
-            developer = get_developer(developer_id=arguments.pop("developer_id"))
-            session_id = arguments.pop("session_id")
-            x_custom_api_key = arguments.pop("x_custom_api_key", None)
+            developer = get_developer(developer_id=arguments.get("developer_id"))
+            session_id = arguments.get("session_id")
+            x_custom_api_key = arguments.get("x_custom_api_key", None)
             chat_input = ChatInput(**arguments)
             bg_runner = BackgroundTasks()
             res = await handler(

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

@@ -8,7 +8,7 @@
 from ...env import testing
 
 
-@auto_blob_store
+@auto_blob_store(deep=True)
 @beartype
 async def cozo_query_step(
     query_name: str,

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

@@ -9,7 +9,7 @@
 from ...env import testing
 
 
-@auto_blob_store
+@auto_blob_store(deep=True)
 @beartype
 async def evaluate_step(
     context: StepContext,

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

@@ -11,7 +11,7 @@
 from .base_evaluate import base_evaluate
 
 
-@auto_blob_store
+@auto_blob_store(deep=True)
 @beartype
 async def for_each_step(context: StepContext) -> StepOutcome:
     try:

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

@@ -8,7 +8,7 @@
 
 # TODO: We should use this step to query the parent workflow and get the value from the workflow context
 # SCRUM-1
-@auto_blob_store
+@auto_blob_store(deep=True)
 @beartype
 async def get_value_step(
     context: StepContext,

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

@@ -11,7 +11,7 @@
 from .base_evaluate import base_evaluate
 
 
-@auto_blob_store
+@auto_blob_store(deep=True)
 @beartype
 async def if_else_step(context: StepContext) -> StepOutcome:
     # NOTE: This activity is only for logging, so we just evaluate the expression

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

@@ -11,7 +11,7 @@
 from ...env import testing
 
 
-@auto_blob_store
+@auto_blob_store(deep=True)
 @beartype
 async def log_step(context: StepContext) -> StepOutcome:
     # NOTE: This activity is only for logging, so we just evaluate the expression

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

@@ -13,7 +13,7 @@
 from .base_evaluate import base_evaluate
 
 
-@auto_blob_store
+@auto_blob_store(deep=True)
 @beartype
 async def map_reduce_step(context: StepContext) -> StepOutcome:
     try:

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

@@ -1,46 +1,28 @@
-from datetime import datetime
 from typing import Callable
 
-from anthropic import AsyncAnthropic  # Import AsyncAnthropic client
 from anthropic.types.beta.beta_message import BetaMessage
 from beartype import beartype
 from langchain_core.tools import BaseTool
 from langchain_core.tools.convert import tool as tool_decorator
-from litellm import ChatCompletionMessageToolCall, Function, Message
-from litellm.types.utils import Choices, ModelResponse
+from litellm.types.utils import ModelResponse
 from temporalio import activity
 from temporalio.exceptions import ApplicationError
 
-from ...autogen.Tools import Tool
+from ...autogen.openapi_model import Tool
 from ...clients import (
     litellm,  # We dont directly import `acompletion` so we can mock it
 )
 from ...common.protocol.tasks import StepContext, StepOutcome
 from ...common.storage_handler import auto_blob_store
 from ...common.utils.template import render_template
-from ...env import anthropic_api_key, debug
 from ..utils import get_handler_with_filtered_params
 from .base_evaluate import base_evaluate
+from ...env import debug
 
 COMPUTER_USE_BETA_FLAG = "computer-use-2024-10-22"
 
 
 def format_tool(tool: Tool) -> dict:
-    if tool.type == "computer_20241022":
-        return {
-            "type": tool.type,
-            "name": tool.name,
-            "display_width_px": tool.computer_20241022
-            and tool.computer_20241022.display_width_px,
-            "display_height_px": tool.computer_20241022
-            and tool.computer_20241022.display_height_px,
-            "display_number": tool.computer_20241022
-            and tool.computer_20241022.display_number,
-        }
-
-    if tool.type in ["bash_20241022", "text_editor_20241022"]:
-        return tool.model_dump(include={"type", "name"})
-
     if tool.type == "function":
         return {
             "type": "function",
@@ -85,7 +67,7 @@ def format_tool(tool: Tool) -> dict:
 
 
 @activity.defn
-@auto_blob_store
+@auto_blob_store(deep=True)
 @beartype
 async def prompt_step(context: StepContext) -> StepOutcome:
     # Get context data
@@ -162,114 +144,50 @@ async def prompt_step(context: StepContext) -> StepOutcome:
         for fmt_tool, orig_tool in zip(formatted_tools, context.tools)
     }
 
-    # Check if the model is Anthropic
-    if agent_model.lower().startswith("claude-3.5") and any(
-        tool["type"] in ["computer_20241022", "bash_20241022", "text_editor_20241022"]
-        for tool in formatted_tools
-    ):
-        # Retrieve the API key from the environment variable
-        betas = [COMPUTER_USE_BETA_FLAG]
-        # Use Anthropic API directly
-        client = AsyncAnthropic(api_key=anthropic_api_key)
-
-        # Reformat the prompt for Anthropic
-        # Anthropic expects a list of messages with role and content (and no name etc)
-        prompt = [{"role": "user", "content": message["content"]} for message in prompt]
-
-        # Filter tools for specific types
-        filtered_tools = [
-            tool
-            for tool in formatted_tools
-            if tool["type"]
-            in ["computer_20241022", "bash_20241022", "text_editor_20241022"]
-        ]
-
-        # Claude Response
-        claude_response: BetaMessage = await client.beta.messages.create(
-            model="claude-3-5-sonnet-20241022",
-            messages=prompt,
-            tools=filtered_tools,
-            max_tokens=1024,
-            betas=betas,
-        )
-
-        # Claude returns [ToolUse | TextBlock]
-        # We need to convert tool_use to tool_calls
-        # And set content = TextBlock.text
-        # But we need to ensure no more than one text block is returned
-        if (
-            len([block for block in claude_response.content if block.type == "text"])
-            > 1
-        ):
-            raise ApplicationError("Claude should only return one message")
-
-        text_block = next(
-            (block for block in claude_response.content if block.type == "text"),
-            None,
-        )
-
-        stop_reason = claude_response.stop_reason
-
-        if stop_reason == "tool_use":
-            choice = Choices(
-                message=Message(
-                    role="assistant",
-                    content=text_block.text if text_block else None,
-                    tool_calls=[
-                        ChatCompletionMessageToolCall(
-                            type="function",
-                            function=Function(
-                                name=block.name,
-                                arguments=block.input,
-                            ),
-                        )
-                        for block in claude_response.content
-                        if block.type == "tool_use"
-                    ],
-                ),
-                finish_reason="tool_calls",
-            )
-        else:
-            assert (
-                text_block
-            ), "Claude should always return a text block for stop_reason=stop"
-
-            choice = Choices(
-                message=Message(
-                    role="assistant",
-                    content=text_block.text,
-                ),
-                finish_reason="stop",
-            )
-
-        response: ModelResponse = ModelResponse(
-            id=claude_response.id,
-            choices=[choice],
-            created=int(datetime.now().timestamp()),
-            model=claude_response.model,
-            object="text_completion",
-        )
-
-    else:
-        # FIXME: hardcoded tool to a None value as the tool calls are not implemented yet
+    # Check if using Claude model and has specific tool types
+    is_claude_model = agent_model.lower().startswith("claude-3.5")
+
+    # FIXME: Hack to make the computer use tools compatible with litellm
+    # Issue was: litellm expects type to be `computer_20241022` and spec to be
+    # `function` (see: https://docs.litellm.ai/docs/providers/anthropic#computer-tools)
+    # but we don't allow that (spec should match type).
+    formatted_tools = []
+    for i, tool in enumerate(context.tools):
+        if tool.type == "computer_20241022" and tool.computer_20241022:
+            function = tool.computer_20241022
+            tool = {
+                "type": tool.type,
+                "function": {
+                    "name": tool.name,
+                    "parameters": {
+                        k: v
+                        for k, v in function.model_dump().items()
+                        if k not in ["name", "type"]
+                    },
+                },
+            }
+            formatted_tools.append(tool)
+
+    if not is_claude_model:
         formatted_tools = None
-        # Use litellm for other models
-        completion_data: dict = {
-            "model": agent_model,
-            "tools": formatted_tools or None,
-            "messages": prompt,
-            **agent_default_settings,
-            **passed_settings,
-        }
 
-        extra_body = {
-            "cache": {"no-cache": debug or context.current_step.disable_cache},
-        }
+    # Use litellm for other models
+    completion_data: dict = {
+        "model": agent_model,
+        "tools": formatted_tools or None,
+        "messages": prompt,
+        **agent_default_settings,
+        **passed_settings,
+    }
 
-        response: ModelResponse = await litellm.acompletion(
-            **completion_data,
-            extra_body=extra_body,
-        )
+    extra_body = {
+        "cache": {"no-cache": debug or context.current_step.disable_cache},
+    }
+
+    response: ModelResponse = await litellm.acompletion(
+        **completion_data,
+        extra_body=extra_body,
+    )
 
     if context.current_step.unwrap:
         if len(response.choices) != 1:

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

@@ -11,7 +11,7 @@
 
 
 @activity.defn
-@auto_blob_store
+@auto_blob_store(deep=True)
 @beartype
 async def raise_complete_async(context: StepContext, output: Any) -> None:
     activity_info = activity.info()

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

@@ -11,7 +11,7 @@
 from .base_evaluate import base_evaluate
 
 
-@auto_blob_store
+@auto_blob_store(deep=True)
 @beartype
 async def return_step(context: StepContext) -> StepOutcome:
     try:

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

@@ -11,7 +11,7 @@
 
 # TODO: We should use this step to signal to the parent workflow and set the value on the workflow context
 # SCRUM-2
-@auto_blob_store
+@auto_blob_store(deep=True)
 @beartype
 async def set_value_step(
     context: StepContext,

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

@@ -11,7 +11,7 @@
 from ..utils import get_evaluator
 
 
-@auto_blob_store
+@auto_blob_store(deep=True)
 @beartype
 async def switch_step(context: StepContext) -> StepOutcome:
     try:

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

@@ -47,7 +47,7 @@ def construct_tool_call(
 
 
 @activity.defn
-@auto_blob_store
+@auto_blob_store(deep=True)
 @beartype
 async def tool_call_step(context: StepContext) -> StepOutcome:
     assert isinstance(context.current_step, ToolCallStep)

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

@@ -8,7 +8,7 @@
 from .base_evaluate import base_evaluate
 
 
-@auto_blob_store
+@auto_blob_store(deep=True)
 @beartype
 async def wait_for_input_step(context: StepContext) -> StepOutcome:
     try:

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

@@ -10,7 +10,7 @@
 from .base_evaluate import base_evaluate
 
 
-@auto_blob_store
+@auto_blob_store(deep=True)
 @beartype
 async def yield_step(context: StepContext) -> StepOutcome:
     try: