From fee9555afad67853ecc0af905ee769f3f0ffd3ca Mon Sep 17 00:00:00 2001
From: Dixing Xu <i@dex.moe>
Date: Thu, 9 Jan 2025 18:07:03 +0800
Subject: [PATCH] :sparkles: Add anthropic tool use

---
 aide/backend/backend_anthropic.py | 52 ++++++++++++++++++++++++++-----
 aide/backend/utils.py             | 18 +++++++++++
 2 files changed, 62 insertions(+), 8 deletions(-)

diff --git a/aide/backend/backend_anthropic.py b/aide/backend/backend_anthropic.py
index c411300..bfc90f3 100644
--- a/aide/backend/backend_anthropic.py
+++ b/aide/backend/backend_anthropic.py
@@ -1,11 +1,15 @@
 """Backend for Anthropic API."""
 
+import json
+import logging
 import time
 
 from .utils import FunctionSpec, OutputType, opt_messages_to_list, backoff_create
 from funcy import notnone, once, select_values
 import anthropic
 
+logger = logging.getLogger("aide")
+
 _client: anthropic.Anthropic = None  # type: ignore
 
 ANTHROPIC_TIMEOUT_EXCEPTIONS = (
@@ -15,6 +19,10 @@
     anthropic.InternalServerError,
 )
 
+ANTHROPIC_MODEL_ALIASES = {
+    "claude-3.5-sonnet": "claude-3-sonnet-20241022",
+}
+
 
 @once
 def _setup_anthropic_client():
@@ -28,23 +36,32 @@ def query(
     func_spec: FunctionSpec | None = None,
     **model_kwargs,
 ) -> tuple[OutputType, float, int, int, dict]:
+    """
+    Query Anthropic's API, optionally with tool use (Anthropic's equivalent to function calling).
+    """
     _setup_anthropic_client()
 
     filtered_kwargs: dict = select_values(notnone, model_kwargs)  # type: ignore
     if "max_tokens" not in filtered_kwargs:
         filtered_kwargs["max_tokens"] = 4096  # default for Claude models
 
-    if func_spec is not None:
-        raise NotImplementedError(
-            "Anthropic does not support function calling for now."
-        )
+    model_name = filtered_kwargs.get("model", "")
+    logger.debug(f"Anthropic query called with model='{model_name}'")
+
+    if model_name in ANTHROPIC_MODEL_ALIASES:
+        model_name = ANTHROPIC_MODEL_ALIASES[model_name]
+
+    if func_spec is not None and func_spec.name == "submit_review":
+        filtered_kwargs["tools"] = [func_spec.as_anthropic_tool_dict]
+        # Force tool use
+        filtered_kwargs["tool_choice"] = func_spec.anthropic_tool_choice_dict
 
-    # Anthropic doesn't allow not having a user messages
+    # Anthropic doesn't allow not having user messages
     # if we only have system msg -> use it as user msg
     if system_message is not None and user_message is None:
         system_message, user_message = user_message, system_message
 
-    # Anthropic passes the system messages as a separate argument
+    # Anthropic passes system messages as a separate argument
     if system_message is not None:
         filtered_kwargs["system"] = system_message
 
@@ -59,14 +76,33 @@ def query(
     )
     req_time = time.time() - t0
 
-    assert len(message.content) == 1 and message.content[0].type == "text"
+    # Handle tool calls if present
+    if (
+        func_spec is not None
+        and "tools" in filtered_kwargs
+        and len(message.content) > 0
+        and message.content[0].type == "tool_use"
+    ):
+        block = message.content[0]  # This is a "ToolUseBlock"
+        # block has attributes: type, id, name, input
+        assert (
+            block.name == func_spec.name
+        ), f"Function name mismatch: expected {func_spec.name}, got {block.name}"
+        output = block.input  # Anthropic calls the parameters "input"
+    else:
+        # For non-tool responses, ensure we have text content
+        assert len(message.content) == 1, "Expected single content item"
+        assert (
+            message.content[0].type == "text"
+        ), f"Expected text response, got {message.content[0].type}"
+        output = message.content[0].text
 
-    output: str = message.content[0].text
     in_tokens = message.usage.input_tokens
     out_tokens = message.usage.output_tokens
 
     info = {
         "stop_reason": message.stop_reason,
+        "model": message.model,
     }
 
     return output, req_time, in_tokens, out_tokens, info
diff --git a/aide/backend/utils.py b/aide/backend/utils.py
index 01499d8..7e0defa 100644
--- a/aide/backend/utils.py
+++ b/aide/backend/utils.py
@@ -66,6 +66,7 @@ def __post_init__(self):
 
     @property
     def as_openai_tool_dict(self):
+        """Convert to OpenAI's function format."""
         return {
             "type": "function",
             "function": {
@@ -81,3 +82,20 @@ def openai_tool_choice_dict(self):
             "type": "function",
             "function": {"name": self.name},
         }
+
+    @property
+    def as_anthropic_tool_dict(self):
+        """Convert to Anthropic's tool format."""
+        return {
+            "name": self.name,
+            "description": self.description,
+            "input_schema": self.json_schema,  # Anthropic uses input_schema instead of parameters
+        }
+
+    @property
+    def anthropic_tool_choice_dict(self):
+        """Convert to Anthropic's tool choice format."""
+        return {
+            "type": "tool",  # Anthropic uses "tool" instead of "function"
+            "name": self.name,
+        }