Merge branch 'main' into anthropic-add

proxy/routes/open_ai.py

@@ -5,7 +5,7 @@ from typing import Any
 
 import httpx
 from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response
-from starlette.responses import StreamingResponse
+from fastapi.responses import StreamingResponse
 from utils.constants import CLIENT_TIMEOUT, IGNORED_HEADERS
 from utils.explorer import push_trace
 
@@ -101,7 +101,7 @@ async def stream_response(
     dataset_name: str,
     request_body_json: dict[str, Any],
     invariant_authorization: str,
-) -> StreamingResponse:
+) -> Response:
     """
     Handles streaming the OpenAI response to the client while building a merged_response
     The chunks are returned to the caller immediately
@@ -109,6 +109,16 @@ async def stream_response(
     It is sent to the Invariant Explorer at the end of the stream
     """
 
+    response = await client.send(open_ai_request, stream=True)
+    if response.status_code != 200:
+        error_content = await response.aread()
+        try:
+            error_json = json.loads(error_content.decode("utf-8"))
+            error_detail = error_json.get("error", "Unknown error from OpenAI API")
+        except json.JSONDecodeError:
+            error_detail = {"error": "Failed to parse OpenAI error response"}
+        raise HTTPException(status_code=response.status_code, detail=error_detail)
+
     async def event_generator() -> Any:
         # merged_response will be updated with the data from the chunks in the stream
         # At the end of the stream, this will be sent to the explorer
@@ -128,43 +138,31 @@ async def stream_response(
         # Combines the choice index and tool call index to uniquely identify a tool call
         tool_call_mapping_by_index = {}
 
-        async with client.stream(
-            "POST",
-            open_ai_request.url,
-            headers=open_ai_request.headers,
-            content=open_ai_request.content,
-        ) as response:
-            if response.status_code != 200:
-                yield json.dumps(
-                    {"error": f"Failed to fetch response: {response.status_code}"}
-                ).encode()
-                return
+        async for chunk in response.aiter_bytes():
+            chunk_text = chunk.decode().strip()
+            if not chunk_text:
+                continue
 
-            async for chunk in response.aiter_bytes():
-                chunk_text = chunk.decode().strip()
-                if not chunk_text:
-                    continue
+            # Yield chunk immediately to the client (proxy behavior)
+            yield chunk
 
-                # Yield chunk immediately to the client (proxy behavior)
-                yield chunk
-
-                # Process the chunk
-                # This will update merged_response with the data from the chunk
-                process_chunk_text(
-                    chunk_text,
-                    merged_response,
-                    choice_mapping_by_index,
-                    tool_call_mapping_by_index,
-                )
-
-            # Send full merged response to the explorer
-            await push_to_explorer(
-                dataset_name,
+            # Process the chunk
+            # This will update merged_response with the data from the chunk
+            process_chunk_text(
+                chunk_text,
                 merged_response,
-                request_body_json,
-                invariant_authorization,
+                choice_mapping_by_index,
+                tool_call_mapping_by_index,
             )
 
+        # Send full merged response to the explorer
+        await push_to_explorer(
+            dataset_name,
+            merged_response,
+            request_body_json,
+            invariant_authorization,
+        )
+
     return StreamingResponse(event_generator(), media_type="text/event-stream")
 
 
@@ -331,7 +329,18 @@ async def handle_non_streaming_response(
     invariant_authorization: str,
 ):
     """Handles non-streaming OpenAI responses"""
-    json_response = response.json()
+    try:
+        json_response = response.json()
+    except json.JSONDecodeError as e:
+        raise HTTPException(
+            status_code=response.status_code,
+            detail="Invalid JSON response received from OpenAI API",
+        ) from e
+    if response.status_code != 200:
+        raise HTTPException(
+            status_code=response.status_code,
+            detail=json_response.get("error", "Unknown error from OpenAI API"),
+        )
     await push_to_explorer(
         dataset_name, json_response, request_body_json, invariant_authorization
     )

tests/open_ai/test_chat_without_tool_calls.py

@@ -1,11 +1,14 @@
 """Test the chat completions proxy calls without tool calling."""
 
+import base64
 import os
 import sys
 import uuid
+from pathlib import Path
 
 import pytest
 from httpx import Client
+
 # add tests folder (parent) to sys.path
 from openai import OpenAI
 
@@ -74,3 +77,77 @@ async def test_chat_completion(context, explorer_api_url, proxy_url, do_stream):
             "content": expected_assistant_message,
         },
     ]
+
+
+@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="No OPENAI_API_KEY set")
+async def test_chat_completion_with_image(context, explorer_api_url, proxy_url):
+    """Test the chat completions proxy works with image."""
+    dataset_name = "test-dataset-open-ai-" + str(uuid.uuid4())
+
+    client = OpenAI(
+        http_client=Client(
+            headers={
+                "Invariant-Authorization": "Bearer <some-key>"
+            },  # This key is not used for local tests
+        ),
+        base_url=f"{proxy_url}/api/v1/proxy/{dataset_name}/openai",
+    )
+    image_path = Path(__file__).parent.parent / "images" / "two-cats.png"
+    with image_path.open("rb") as image_file:
+        base64_image = base64.b64encode(image_file.read()).decode("utf-8")
+
+        chat_response = client.chat.completions.create(
+            model="gpt-4o",
+            messages=[
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": "How many cats are there in this image?",
+                        },
+                        {
+                            "type": "image_url",
+                            "image_url": {
+                                "url": f"data:image/png;base64,{base64_image}"
+                            },
+                        },
+                    ],
+                }
+            ],
+            max_tokens=100,
+        )
+
+        assert "TWO" in chat_response.choices[0].message.content.upper()
+
+        # Fetch the trace ids for the dataset
+        traces_response = await context.request.get(
+            f"{explorer_api_url}/api/v1/dataset/byuser/developer/{dataset_name}/traces"
+        )
+        traces = await traces_response.json()
+        assert len(traces) == 1
+        trace_id = traces[0]["id"]
+
+        # Fetch the trace
+        trace_response = await context.request.get(
+            f"{explorer_api_url}/api/v1/trace/{trace_id}"
+        )
+        trace = await trace_response.json()
+
+        # Verify the trace messages
+        assert trace["messages"] == [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": "How many cats are there in this image?"},
+                    {
+                        "type": "image_url",
+                        "image_url": {"url": "data:image/png;base64," + base64_image},
+                    },
+                ],
+            },
+            {
+                "role": "assistant",
+                "content": chat_response.choices[0].message.content,
+            },
+        ]