Implement Anthropic proxy

Anthropic implement

proxy/requirements.txt

@@ -2,4 +2,6 @@ fastapi==0.115.7
 httpx==0.28.1
 uvicorn==0.34.0
 invariant-sdk
-starlette-compress==1.4.0
+starlette-compress==1.4.0
+tavily-python
+anthropic

proxy/routes/anthropic.py

@@ -1,5 +1,186 @@
 """Proxy service to forward requests to the Anthropic APIs"""
 
-from fastapi import APIRouter
+from fastapi import APIRouter, Header, HTTPException, Depends, Request
+import json
+import httpx
+from typing import Any
+from utils.explorer import push_trace
+# from .open_ai import push_to_explorer
 
 proxy = APIRouter()
+
+ALLOWED_ANTHROPIC_ENDPOINTS = {"v1/messages"}
+IGNORED_HEADERS = [
+    "accept-encoding",
+    "host",
+    "invariant-authorization",
+    "x-forwarded-for",
+    "x-forwarded-host",
+    "x-forwarded-port",
+    "x-forwarded-proto",
+    "x-forwarded-server",
+    "x-real-ip",
+]
+
+MISSING_INVARIANT_AUTH_HEADER = "Missing invariant-authorization header"
+MISSING_ANTHROPIC_AUTH_HEADER = "Missing athropic authorization header"
+NOT_SUPPORTED_ENDPOINT = "Not supported OpenAI endpoint"
+FAILED_TO_PUSH_TRACE = "Failed to push trace to the dataset: "
+END_REASONS = [
+    "end_turn",
+    "max_tokens",
+    "stop_sequence"
+]
+
+def validate_headers(
+        invariant_authorization: str = Header(None), x_api_key: str = Header(None)
+):
+    """Require the invariant-authorization and authorization headers to be present"""
+    if invariant_authorization is None:
+        raise HTTPException(status_code=400, detail=MISSING_INVARIANT_AUTH_HEADER)
+    if x_api_key is None:
+        raise HTTPException(status_code=400, detail=MISSING_ANTHROPIC_AUTH_HEADER)
+    
+@proxy.post(
+    "/{dataset_name}/anthropic/{endpoint:path}",
+    dependencies=[Depends(validate_headers)],
+) 
+async def anthropic_proxy(
+    dataset_name: str,
+    endpoint: str,
+    request: Request,
+):
+    """Proxy calls to the Anthropic APIs"""
+    if endpoint not in ALLOWED_ANTHROPIC_ENDPOINTS:
+        raise HTTPException(status_code=404, detail=NOT_SUPPORTED_ENDPOINT)
+    headers = {
+        k: v for k, v in request.headers.items() if k.lower() not in IGNORED_HEADERS
+    }
+
+    request_body = await request.body()
+
+    request_body_json = json.loads(request_body)
+
+    anthropic_url = f"https://api.anthropic.com/{endpoint}"
+    client = httpx.AsyncClient()
+
+    anthropic_request = client.build_request(
+        "POST", 
+        anthropic_url, 
+        headers=headers, 
+        data=request_body
+    )
+
+    invariant_authorization = request.headers.get("invariant-authorization")
+
+    async with client:
+        response = await client.send(anthropic_request)
+        await handle_non_streaming_response(
+            response, dataset_name, request_body_json, invariant_authorization
+        )
+        return response.json()
+
+async def push_to_explorer(
+    dataset_name: str,
+    merged_response: dict[str, Any],
+    request_body: dict[str, Any],
+    invariant_authorization: str,
+) -> None:
+    """Pushes the full trace to the Invariant Explorer"""
+    # Combine the messages from the request body and Anthropic response
+    messages = request_body.get("messages", [])
+    messages += [merged_response]
+
+    messages = anthropic_to_invariant_messages(messages)
+    _ = await push_trace(
+        dataset_name=dataset_name,
+        messages=[messages],
+        invariant_authorization=invariant_authorization,
+    )
+
+async def handle_non_streaming_response(
+    response: httpx.Response,
+    dataset_name: str,
+    request_body_json: dict[str, Any],
+    invariant_authorization: str,
+):
+    """Handles non-streaming Anthropic responses"""
+    json_response = response.json()
+    # Only push the trace to explorer if the last message is an end turn message
+    if json_response.get("stop_reason") in END_REASONS:
+        await push_to_explorer(
+            dataset_name,
+            json_response,
+            request_body_json,
+            invariant_authorization,
+        )
+
+def anthropic_to_invariant_messages(
+    messages: list[dict], keep_empty_tool_response: bool = False
+) -> list[dict]:
+    """Converts a list of messages from the Anthropic API to the Invariant API format."""
+    output = []
+    role_mapping = {
+        "system": lambda msg: {"role": "system", "content": msg["content"]},
+        "user": lambda msg: handle_user_message(msg, keep_empty_tool_response),
+        "assistant": lambda msg: handle_assistant_message(msg),
+    }
+
+    for message in messages:
+        handler = role_mapping.get(message["role"])
+        if handler:
+            output.extend(handler(message))
+
+    return output
+
+def handle_user_message(message, keep_empty_tool_response):
+    output = []
+    content = message["content"]
+    if isinstance(content, list):
+        for sub_message in content:
+            if sub_message["type"] == "tool_result":
+                if sub_message["content"]:
+                    output.append(
+                        {
+                            "role": "tool",
+                            "content": sub_message["content"],
+                            "tool_id": sub_message["tool_use_id"],
+                        }
+                    )
+                elif keep_empty_tool_response and any(sub_message.values()):
+                    output.append(
+                        {
+                            "role": "tool",
+                            "content": {"is_error": True} if sub_message["is_error"] else {},
+                            "tool_id": sub_message["tool_use_id"],
+                        }
+                    )
+            elif sub_message["type"] == "text":
+                output.append({"role": "user", "content": sub_message["text"]})
+    else:
+        output.append({"role": "user", "content": content})
+    return output
+
+def handle_assistant_message(message):
+    output = []
+    for sub_message in message["content"]:
+        if sub_message["type"] == "text":
+            output.append({"role": "assistant", "content": sub_message.get("text")})
+        elif sub_message["type"] == "tool_use":
+            output.append(
+                {
+                    "role": "assistant",
+                    "content": None,
+                    "tool_calls": [
+                        {
+                            "tool_id": sub_message.get("id"),
+                            "type": "function",
+                            "function": {
+                                "name": sub_message.get("name"),
+                                "arguments": sub_message.get("input"),
+                            },
+                        }
+                    ],
+                }
+            )
+    return output

proxy/routes/open_ai.py

@@ -300,7 +300,6 @@ async def push_to_explorer(
     # Combine the messages from the request body and the choices from the OpenAI response
     messages = request_body.get("messages", [])
     messages += [choice["message"] for choice in merged_response.get("choices", [])]
-
     _ = await push_trace(
         dataset_name=dataset_name,
         messages=[messages],