diff --git a/proxy/routes/open_ai.py b/proxy/routes/open_ai.py index ba33d94..a18c084 100644 --- a/proxy/routes/open_ai.py +++ b/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 ) diff --git a/tests/images/two-cats.png b/tests/images/two-cats.png new file mode 100644 index 0000000..6820626 Binary files /dev/null and b/tests/images/two-cats.png differ diff --git a/tests/open_ai/test_chat_without_tool_calls.py b/tests/open_ai/test_chat_without_tool_calls.py index b96ae43..f6e43cb 100644 --- a/tests/open_ai/test_chat_without_tool_calls.py +++ b/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 " + }, # 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, + }, + ]