Merge branch 'main' into anthropic-implement

.env

@@ -1,11 +1 @@
 INVARIANT_API_URL=https://explorer.invariantlabs.ai
-
-# For the testing stack only
-# TODO: Check if these can be removed.
-# Currently we rely on the `invariant explorer` command
-# to setup a local instance of explorer to test the proxy.
-# That requires these variables to be set in the .env file.
-POSTGRES_USER=postgres
-POSTGRES_PASSWORD=postgres
-POSTGRES_DB=invariantmonitor
-POSTGRES_HOST=database

.github/workflows/tests_ci.yml

@@ -0,0 +1,22 @@
+name: Invariant proxy testing CI
+
+on:
+  push:
+    branches: [ "main" ]
+  pull_request:
+    branches: [ "main" ]
+
+permissions:
+  contents: read
+
+jobs:
+  test:
+    name: Build & Test
+    runs-on: ubuntu-latest
+    steps:
+        - name: Checkout
+          uses: actions/checkout@v1
+        - name: Run tests
+          env:
+            OPENAI_API_KEY: ${{ secrets.INVARIANT_TESTING_OPENAI_KEY }}
+          run: ./run.sh tests -s -vv

.gitignore

@@ -3,6 +3,7 @@ __pycache__/
 .pytest_cache/
 .py[oc]
 data/
+tests/results/
 
 # Coverage and build artifacts
 .coverage

proxy/routes/open_ai.py

@@ -26,7 +26,7 @@ proxy = APIRouter()
 MISSING_INVARIANT_AUTH_HEADER = "Missing invariant-authorization header"
 MISSING_AUTH_HEADER = "Missing authorization header"
 NOT_SUPPORTED_ENDPOINT = "Not supported OpenAI endpoint"
-FAILED_TO_PUSH_TRACE = "Failed to push trace to the dataset: "
+FINISH_REASON_TO_PUSH_TRACE = ["stop", "length", "content_filter"]
 
 
 def validate_headers(
@@ -79,12 +79,11 @@ async def openai_proxy(
             request_body_json,
             invariant_authorization,
         )
-    else:
-        async with client:
-            response = await client.send(open_ai_request)
-            return await handle_non_streaming_response(
-                response, dataset_name, request_body_json, invariant_authorization
-            )
+    async with client:
+        response = await client.send(open_ai_request)
+        return await handle_non_streaming_response(
+            response, dataset_name, request_body_json, invariant_authorization
+        )
 
 
 async def stream_response(
@@ -226,6 +225,8 @@ def update_merged_response(
 
         existing_choice = merged_response["choices"][choice_mapping_by_index[index]]
         delta = choice.get("delta", {})
+        if choice.get("finish_reason"):
+            existing_choice["finish_reason"] = choice["finish_reason"]
 
         update_existing_choice_with_delta(
             existing_choice, delta, tool_call_mapping_by_index, choice_index=index
@@ -297,6 +298,13 @@ async def push_to_explorer(
     invariant_authorization: str,
 ) -> None:
     """Pushes the full trace to the Invariant Explorer"""
+    # Only push the trace to explorer if the message is an end turn message
+    if (
+        merged_response.get("choices")
+        and merged_response["choices"][0].get("finish_reason")
+        not in FINISH_REASON_TO_PUSH_TRACE
+    ):
+        return
     # 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", [])]

proxy/serve.py

@@ -15,6 +15,13 @@ app.add_middleware(CompressMiddleware)
 
 router = fastapi.APIRouter(prefix="/api/v1")
 
+
+@router.get("/proxy/health", tags=["health_check"], include_in_schema=False)
+async def get_proxy_info():
+    """Health check"""
+    return {"message": "Hello v1/proxy"}
+
+
 router.include_router(open_ai_proxy, prefix="/proxy", tags=["open_ai_proxy"])
 
 router.include_router(anthropic_proxy, prefix="/proxy", tags=["anthropic_proxy"])

run.sh

@@ -18,26 +18,57 @@ build() {
 down() {
   # Bring down local services
   docker compose -f docker-compose.local.yml down
+  docker compose -f tests/docker-compose.test.yml down
 }
 
 
 tests() {
-    # Run tests
-    pip install invariant-ai
-    invariant explorer up -d --build
+  echo "Setting up test environment..."
 
-    until curl -X GET -I http://127.0.0.1/api/v1 --fail --silent --output /dev/null; do
-      echo "Backend API not available yet - checking health..."
-      sleep 2
-    done
+  # Ensure test network exists
+  docker network inspect invariant-proxy-web-test >/dev/null 2>&1 || \
+    docker network create invariant-proxy-web-test
 
-    echo "Backend API is available. Running tests..."
+  # Setup the explorer.test.yml file
+  CONFIG_DIR="/tmp/invariant-proxy-test/configs"
+  FILE="$CONFIG_DIR/explorer.test.yml"
+  mkdir -p "$CONFIG_DIR"
+  # Download the file
+  curl -L -o "$FILE" https://raw.githubusercontent.com/invariantlabs-ai/explorer/main/configs/explorer.test.yml
+  # Verify if the file exists
+  if [ ! -f "$FILE" ]; then
+    echo "Error: File $FILE not found. Issue with download."
+    exit 1
+  fi
+  echo "File successfully downloaded: $FILE"
 
-    docker build -t 'explorer-proxy-test' -f ./tests/Dockerfile.test ./tests
+  # Start containers
+  docker compose -f tests/docker-compose.test.yml down
+  docker compose -f tests/docker-compose.test.yml build
+  docker compose -f tests/docker-compose.test.yml up -d
 
-    docker run \
+  until [ "$(docker inspect -f '{{.State.Health.Status}}' explorer-proxy-test-app-api)" = "healthy" ]; do
+    echo "explorer-proxy-test-app-api container starting..."
+    sleep 2
+  done
+
+  until [ "$(docker inspect -f '{{.State.Health.Status}}' explorer-proxy-test)" = "healthy" ]; do
+    echo "explorer-proxy-test container starting..."
+    sleep 2
+  done
+
+  echo "app-api and proxy are available. Running tests..."
+
+  # Make call to signup endpoint
+  curl -k -X POST http://127.0.0.1/api/v1/user/signup
+
+  docker build -t 'explorer-proxy-test' -f ./tests/Dockerfile.test ./tests
+
+  docker run \
     --mount type=bind,source=./tests,target=/tests \
-    --network host \
+    --network invariant-proxy-web-test \
+    -e OPENAI_API_KEY="$OPENAI_API_KEY" \
+    --env-file ./tests/.env.test \
     explorer-proxy-test $@
 }
 

tests/.env.test

@@ -0,0 +1,10 @@
+# To push traces to the local instance of the explorer app-api
+# from the proxy. Both proxy and app-api are on the same network:
+# invariant-proxy-web-test
+INVARIANT_API_URL=http://explorer-proxy-test-app-api:8000
+INVARIANT_PROXY_API_URL=http://explorer-proxy-test:8000
+
+POSTGRES_USER=postgres
+POSTGRES_PASSWORD=postgres
+POSTGRES_DB=invariantmonitor
+POSTGRES_HOST=database

tests/Dockerfile.test

@@ -3,6 +3,7 @@ FROM mcr.microsoft.com/playwright/python:v1.50.0-noble
 RUN mkdir -p /tests
 COPY ./requirements.txt /tests/requirements.txt
 WORKDIR /tests
+RUN pip install --upgrade pip
 RUN pip install --no-cache-dir -r requirements.txt
 
-ENTRYPOINT ["pytest", "-s", "-vv"]
+ENTRYPOINT ["pytest",  "--capture=tee-sys", "--tracing", "off", "--junit-xml=/tests/results/test-results-all.xml", "-s", "-vv"]

tests/anthropic/test_claude_weather_agent.py

@@ -1,26 +1,25 @@
-import anthropic
-from typing import Dict
-import os
-from tavily import TavilyClient
-import anthropic
-from httpx import Client
-import os
-# from invariant import testing
-tavily_client = TavilyClient(api_key=os.getenv("TAVILY_API_KEY"))
 import datetime
+import os
+from typing import Dict
+
+import anthropic
+import pytest
+from httpx import Client
+from tavily import TavilyClient
+
 
 class WeatherAgent:
     def __init__(self, api_key: str):
-        dataset_name = "claude_weather_agent_test" + str(datetime.datetime.now().strftime("%Y%m%d%H%M%S"))
-        invariant_api_key = os.environ.get("INVARIANT_API_KEY")
+        self.tavily_client = TavilyClient(api_key=os.getenv("TAVILY_API_KEY"))
+        dataset_name = "claude_weather_agent_test" + str(
+            datetime.datetime.now().strftime("%Y%m%d%H%M%S")
+        )
         self.client = anthropic.Anthropic(
             http_client=Client(
-                headers={
-                    "Invariant-Authorization": f"Bearer {invariant_api_key}"
-                },
-        ),
-        base_url=f"http://localhost/api/v1/proxy/{dataset_name}/anthropic",
-    )
+                headers={"Invariant-Authorization": "Bearer <some-api-key>"},
+            ),
+            base_url=f"http://localhost/api/v1/proxy/{dataset_name}/anthropic",
+        )
         self.get_weather_function = {
             "name": "get_weather",
             "description": "Get the current weather in a given location",
@@ -29,39 +28,34 @@ class WeatherAgent:
                 "properties": {
                     "location": {
                         "type": "string",
-                        "description": "The city and state, e.g. San Francisco, CA"
+                        "description": "The city and state, e.g. San Francisco, CA",
                     },
                     "unit": {
                         "type": "string",
                         "enum": ["celsius", "fahrenheit"],
-                        "description": "The unit of temperature, either \"celsius\" or \"fahrenheit\""
-                    }
+                        "description": 'The unit of temperature, either "celsius" or "fahrenheit"',
+                    },
                 },
-                "required": ["location"]
-            }
+                "required": ["location"],
+            },
         }
 
-        # self.system_prompt = """You are an assistant that can perform weather searches using function calls. 
-        #     When a user asks for weather information, respond with a JSON object specifying 
+        # self.system_prompt = """You are an assistant that can perform weather searches using function calls.
+        #     When a user asks for weather information, respond with a JSON object specifying
         #     the function name `get_weather` and the arguments latitude and longitude are needed."""
 
     def get_response(self, user_query: str) -> Dict:
         """
         Get the response from the agent for a given user query for weather.
         """
-        messages = [
-            {
-                "role": "user",
-                "content": user_query
-            }
-        ]
+        messages = [{"role": "user", "content": user_query}]
         while True:
             response = self.client.messages.create(
                 # system=self.system_prompt,
-                tools = [self.get_weather_function],
+                tools=[self.get_weather_function],
                 model="claude-3-5-sonnet-20241022",
                 max_tokens=1024,
-                messages=messages
+                messages=messages,
             )
 
             # If there's tool call, Extract the tool call parameters from the response
@@ -69,44 +63,49 @@ class WeatherAgent:
                 tool_call_params = response.content[1].input
                 tool_call_result = self.get_weather(tool_call_params["location"])
                 tool_call_id = response.content[1].id
-                messages.append({
-                    "role": response.role,
-                    "content": response.content
+                messages.append({"role": response.role, "content": response.content})
+                messages.append(
+                    {
+                        "role": "user",
+                        "content": [
+                            {
+                                "type": "tool_result",
+                                "tool_use_id": tool_call_id,
+                                "content": tool_call_result,
+                            }
+                        ],
                     }
                 )
-                messages.append({
-                    "role": "user",
-                    "content": [{
-                        "type": "tool_result",
-                        "tool_use_id": tool_call_id,
-                        "content": tool_call_result
-                    }]
-                })
             else:
                 return response.content[0].text
 
     def get_weather(self, location: str):
         """Get the current weather in a given location using latitude and longitude."""
         query = f"What is the weather in {location}?"
-        response = tavily_client.search(query)
+        response = self.tavily_client.search(query)
         response_content = response["results"][0]["content"]
         return response["results"][0]["title"] + ":\n" + response_content
 
 
-# Initialize agent with your Anthropic API key
-anthropic_api_key = os.getenv("ANTHROPIC_API_KEY")
-weather_agent = WeatherAgent(anthropic_api_key)
-    
-def test_proxy_response():    
+@pytest.mark.skipif(
+    not os.getenv("ANTHROPIC_API_KEY") or not os.getenv("TAVILY_API_KEY"),
+    reason="API keys not set",
+)
+def test_proxy_response():
+    """Test the proxy response for the weather agent."""
+    # Initialize agent with Anthropic API key
+    anthropic_api_key = os.getenv("ANTHROPIC_API_KEY")
+    weather_agent = WeatherAgent(anthropic_api_key)
+
     # Example queries
     queries = [
         "What's the weather like in Zurich city?",
         "Tell me the forecast for New York",
-        "How's the weather in London next week?"
+        "How's the weather in London next week?",
     ]
     cities = ["Zurich", "New York", "London"]
     # Process each query
-    for index,query in enumerate(queries):
+    for index, query in enumerate(queries):
         response = weather_agent.get_response(query)
         assert response is not None
-        assert cities[index] in response
+        assert cities[index] in response

tests/docker-compose.test.yml

@@ -0,0 +1,106 @@
+name: explorer-proxy-test
+services:
+  traefik:
+    image: traefik:v2.0
+    container_name: "explorer-proxy-test-traefik"
+    command:
+      - --providers.docker=true
+      # Enable the API handler in insecure mode,
+      # which means that the Traefik API will be available directly
+      # on the entry point named traefik.
+      - --api.insecure=true
+      # Define Traefik entry points to port [80] for http and port [443] for https.
+      - --entrypoints.invariant-proxy-web-test.address=0.0.0.0:80
+    networks:
+      - invariant-proxy-web-test
+    ports:
+      - '${PORT_HTTP:-80}:80'
+    volumes:
+      - /var/run/docker.sock:/var/run/docker.sock
+    labels:
+      - "traefik.enable=true"
+      - "traefik.http.routers.traefik-http.entrypoints=invariant-proxy-web-test"
+
+  explorer-proxy:
+    container_name: explorer-proxy-test
+    build:
+      context: ../proxy
+      dockerfile: ../proxy/Dockerfile.proxy
+    depends_on:
+      app-api:
+        condition: service_healthy
+    working_dir: /srv/proxy
+    env_file:
+      - .env.test
+    environment:
+      - DEV_MODE=true
+    volumes:
+      - type: bind
+        source: ../proxy
+        target: /srv/proxy
+    networks:
+      - invariant-proxy-web-test
+    ports: []
+    labels:
+      - "traefik.enable=true"
+      - "traefik.http.routers.explorer-proxy-api.rule=(Host(`localhost`) && PathPrefix(`/api/v1/proxy/`)) || (Host(`127.0.0.1`) && PathPrefix(`/api/v1/proxy/`))"
+      - "traefik.http.routers.explorer-proxy-api.entrypoints=invariant-proxy-web-test"
+      - "traefik.http.services.explorer-proxy-api.loadbalancer.server.port=8000"
+      - "traefik.docker.network=invariant-proxy-web-test"
+    healthcheck:
+      test: curl -X GET -I http://localhost:8000/api/v1/proxy/health --fail
+      interval: 1s
+      timeout: 5s
+
+  app-api:
+    container_name: explorer-proxy-test-app-api
+    image: ghcr.io/invariantlabs-ai/explorer/app-api:latest
+    platform: linux/amd64
+    depends_on:
+      database:
+        condition: service_healthy
+    working_dir: /srv/app
+    env_file:
+      - .env.test
+    environment:
+      - PROJECTS_DIR=/srv/projects
+      - KEYCLOAK_CLIENT_ID_SECRET=local-does-not-use-keycloak
+      - TZ=Europe/Berlin
+      - DEV_MODE=true
+      - APP_NAME=explorer-test
+      - CONFIG_FILE=/config/explorer.config.yml
+      - PORT_HTTP=8000
+      - PORT_API=80
+    networks:
+      - internal
+      - invariant-proxy-web-test
+    volumes:
+      - /tmp/invariant-proxy-test/configs/explorer.test.yml:/config/explorer.config.yml
+    labels:
+      - "traefik.enable=true"
+      - "traefik.http.routers.explorer-test-api.rule=(Host(`localhost`) && PathPrefix(`/api/`)) || (Host(`127.0.0.1`) && PathPrefix(`/api/`))"
+      - "traefik.http.routers.explorer-test-api.entrypoints=invariant-proxy-web-test"
+      - "traefik.http.services.explorer-test-api.loadbalancer.server.port=8000"
+      - "traefik.docker.network=invariant-proxy-web-test"
+    healthcheck:
+      test: curl -X GET -I http://localhost:8000/api/v1/ --fail
+      interval: 1s
+      timeout: 5s
+
+  database:
+    container_name: explorer-proxy-test-database
+    image: postgres:16
+    env_file:
+      - .env.test
+    networks:
+      - internal
+    healthcheck:
+      test: [ "CMD-SHELL", "pg_isready -U postgres" ]
+      interval: 5s
+      timeout: 5s
+      retries: 5
+
+networks:
+  invariant-proxy-web-test:
+    external: true
+  internal:

tests/open_ai/test_chat_with_tool_call.py

@@ -0,0 +1,231 @@
+"""Test the chat completions proxy calls with tool calling and processing response."""
+
+import json
+import os
+import sys
+import uuid
+
+import pytest
+from httpx import Client
+
+# add tests folder (parent) to sys.path
+from openai import OpenAI
+
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from util import *  # needed for pytest fixtures
+
+pytest_plugins = ("pytest_asyncio",)
+
+
+@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="No OPENAI_API_KEY set")
+async def test_chat_completion_with_tool_call_without_streaming(
+    context, explorer_api_url, proxy_url
+):
+    """
+    Test the chat completions proxy calls with tool calling and response processing
+    without streaming.
+    """
+    dataset_name = "test-dataset-open-ai-tool-call-" + 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",
+    )
+
+    chat_response = client.chat.completions.create(
+        model="gpt-4o",
+        messages=[{"role": "user", "content": "What is the weather in New York?"}],
+        tools=[
+            {
+                "type": "function",
+                "function": {
+                    "name": "get_weather",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {"location": {"type": "string"}},
+                        "required": ["location"],
+                    },
+                },
+            }
+        ],
+    )
+
+    assert chat_response.choices[0].message.role == "assistant"
+    # Extract tool call
+    assert len(chat_response.choices[0].message.tool_calls) == 1
+    assert chat_response.choices[0].message.tool_calls[0].function.name == "get_weather"
+    assert (
+        chat_response.choices[0].message.tool_calls[0].function.arguments
+        == '{"location":"New York"}'
+    )
+    tool_call = chat_response.choices[0].message.tool_calls[0]
+
+    # Mock response of tool call
+    tool_result = "The temperature in New York is 15°C and it is raining."
+    history = [
+        {"role": "user", "content": "What is the weather in New York?"},
+        {
+            "role": "assistant",
+            "tool_calls": [
+                {
+                    "function": {
+                        "arguments": tool_call.function.arguments,
+                        "name": tool_call.function.name,
+                    },
+                    "id": tool_call.id,
+                    "type": tool_call.type,
+                }
+            ],
+        },
+        {
+            "role": "tool",
+            "tool_call_id": tool_call.id,
+            "tool_name": "get_weather",
+            "content": tool_result,
+        },
+    ]
+
+    # Send mock response back to OpenAI with history of chat
+    chat_response_final = client.chat.completions.create(
+        model="gpt-4o",
+        messages=history,
+    )
+    assert "15°C" in chat_response_final.choices[0].message.content
+
+    # 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
+    expected_messages = history + [
+        {
+            "role": "assistant",
+            "content": chat_response_final.choices[0].message.content,
+        }
+    ]
+    expected_messages[1]["tool_calls"][0]["function"]["arguments"] = json.loads(
+        expected_messages[1]["tool_calls"][0]["function"]["arguments"]
+    )
+    assert trace["messages"] == expected_messages
+
+
+@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="No OPENAI_API_KEY set")
+async def test_chat_completion_with_tool_call_with_streaming(
+    context, explorer_api_url, proxy_url
+):
+    """
+    Test the chat completions proxy calls with tool calling and response processing
+    while streaming.
+    """
+    dataset_name = "test-dataset-open-ai-tool-call-" + 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",
+    )
+
+    chat_response = client.chat.completions.create(
+        model="gpt-4o",
+        messages=[{"role": "user", "content": "What is the weather in New York?"}],
+        tools=[
+            {
+                "type": "function",
+                "function": {
+                    "name": "get_weather",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {"location": {"type": "string"}},
+                        "required": ["location"],
+                    },
+                },
+            }
+        ],
+        stream=True,
+    )
+
+    tool_call = {"function": {}}
+    for chunk in chat_response:
+        if chunk.choices and chunk.choices[0].delta.tool_calls:
+            partial_tool_call = chunk.choices[0].delta.tool_calls[0]
+            tool_call.setdefault("id", partial_tool_call.id)
+            tool_call.setdefault("type", partial_tool_call.type)
+            tool_call["function"].setdefault("name", partial_tool_call.function.name)
+            tool_call["function"].setdefault("arguments", "")
+            tool_call["function"]["arguments"] += partial_tool_call.function.arguments
+
+    assert tool_call["function"]["name"] == "get_weather"
+    assert tool_call["function"]["arguments"] == '{"location":"New York"}'
+
+    # Mock response of tool call
+    tool_result = "The temperature in New York is 15°C and it is raining."
+    history = [
+        {"role": "user", "content": "What is the weather in New York?"},
+        {
+            "role": "assistant",
+            "tool_calls": [
+                {
+                    "function": {
+                        "arguments": tool_call["function"]["arguments"],
+                        "name": tool_call["function"]["name"],
+                    },
+                    "id": tool_call["id"],
+                    "type": tool_call["type"],
+                }
+            ],
+        },
+        {
+            "role": "tool",
+            "tool_call_id": tool_call["id"],
+            "tool_name": "get_weather",
+            "content": tool_result,
+        },
+    ]
+
+    # Send mock response back to OpenAI with history of chat
+    chat_response_final = client.chat.completions.create(
+        model="gpt-4o", messages=history, stream=True
+    )
+    final_response = {"role": "assistant", "content": ""}
+    for chunk in chat_response_final:
+        if chunk.choices and chunk.choices[0].delta.content:
+            final_response["content"] += chunk.choices[0].delta.content
+
+    # 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
+    expected_messages = history + [final_response]
+    expected_messages[1]["tool_calls"][0]["function"]["arguments"] = json.loads(
+        expected_messages[1]["tool_calls"][0]["function"]["arguments"]
+    )
+    assert trace["messages"] == expected_messages

tests/open_ai/test_chat_without_tool_calls.py

@@ -1,9 +1,14 @@
 """Test the chat completions proxy calls without tool calling."""
 
 import os
+import sys
+import uuid
+
+import pytest
+from httpx import Client
 
 # add tests folder (parent) to sys.path
-import sys
+from openai import OpenAI
 
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
@@ -12,11 +17,61 @@ from util import *  # needed for pytest fixtures
 pytest_plugins = ("pytest_asyncio",)
 
 
-async def test_hello_world(context, url):
-    """Demo test"""
-    response = await context.request.get(
-        f"{url}/api/v1/dataset/byuser/developer/Welcome-to-Explorer"
+@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="No OPENAI_API_KEY set")
+@pytest.mark.parametrize("do_stream", [True, False])
+async def test_chat_completion(context, explorer_api_url, proxy_url, do_stream):
+    """Test the chat completions proxy calls without tool calling."""
+    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",
     )
-    dataset = await response.json()
-    assert dataset["name"] == "Welcome-to-Explorer"
-    assert dataset["user"]["username"] == "developer"
+
+    chat_response = client.chat.completions.create(
+        model="gpt-4o",
+        messages=[{"role": "user", "content": "What is the capital of France?"}],
+        stream=do_stream,
+    )
+
+    # Verify the chat response
+    if not do_stream:
+        assert "PARIS" in chat_response.choices[0].message.content.upper()
+        expected_assistant_message = chat_response.choices[0].message.content
+    else:
+        full_response = ""
+        for chunk in chat_response:
+            if chunk.choices and chunk.choices[0].delta.content:
+                full_response += chunk.choices[0].delta.content
+        assert "PARIS" in full_response.upper()
+        expected_assistant_message = full_response
+
+    # 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": "What is the capital of France?",
+        },
+        {
+            "role": "assistant",
+            "content": expected_assistant_message,
+        },
+    ]

tests/requirements.txt

@@ -1,4 +1,6 @@
+anthropic
 openai
 pytest
 pytest-asyncio
-pytest-playwright
+pytest-playwright
+tavily-python

tests/util.py

@@ -1,12 +1,23 @@
 """Util functions for tests"""
 
+import os
+
 import pytest
 from playwright.async_api import async_playwright
 
 
 @pytest.fixture
-def url():
-    return "http://127.0.0.1"
+def proxy_url():
+    if "INVARIANT_PROXY_API_URL" in os.environ:
+        return os.environ["INVARIANT_PROXY_API_URL"]
+    raise ValueError("Please set the INVARIANT_PROXY_API_URL environment variable")
+
+
+@pytest.fixture
+def explorer_api_url():
+    if "INVARIANT_API_URL" in os.environ:
+        return os.environ["INVARIANT_API_URL"]
+    raise ValueError("Please set the INVARIANT_API_URL environment variable")
 
 
 @pytest.fixture