From 7c692fa6df333ed1ca5eea16ac64e9803fcc6b77 Mon Sep 17 00:00:00 2001
From: Hemang <hemang@invariantlabs.ai>
Date: Thu, 6 Feb 2025 14:59:23 +0100
Subject: [PATCH 01/15] Add a health check endpoint for proxy.

---
 proxy/serve.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/proxy/serve.py b/proxy/serve.py
index a0216c3..b2d923e 100644
--- a/proxy/serve.py
+++ b/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"])

From 084ef4f0e85fe54efe7ea691b74263e25242c8d6 Mon Sep 17 00:00:00 2001
From: Hemang <hemang@invariantlabs.ai>
Date: Thu, 6 Feb 2025 15:52:50 +0100
Subject: [PATCH 02/15] Update the testing setup.

---
 .env                          |  10 ----
 run.sh                        |  45 ++++++++++----
 tests/.env.test               |   6 ++
 tests/docker-compose.test.yml | 110 ++++++++++++++++++++++++++++++++++
 4 files changed, 151 insertions(+), 20 deletions(-)
 create mode 100644 tests/.env.test
 create mode 100644 tests/docker-compose.test.yml

diff --git a/.env b/.env
index 6aaa6ca..419976f 100644
--- a/.env
+++ b/.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
\ No newline at end of file
diff --git a/run.sh b/run.sh
index ce3dfc3..6a06d77 100755
--- a/run.sh
+++ b/run.sh
@@ -22,20 +22,45 @@ 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-explorer-web-test >/dev/null 2>&1 || \
+    docker network create invariant-explorer-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..."
+
+  docker build -t 'explorer-proxy-test' -f ./tests/Dockerfile.test ./tests
+
+  docker run \
     --mount type=bind,source=./tests,target=/tests \
     --network host \
     explorer-proxy-test $@
diff --git a/tests/.env.test b/tests/.env.test
new file mode 100644
index 0000000..d0b7167
--- /dev/null
+++ b/tests/.env.test
@@ -0,0 +1,6 @@
+INVARIANT_API_URL=http://127.0.0.1
+
+POSTGRES_USER=postgres
+POSTGRES_PASSWORD=postgres
+POSTGRES_DB=invariantmonitor
+POSTGRES_HOST=database
\ No newline at end of file
diff --git a/tests/docker-compose.test.yml b/tests/docker-compose.test.yml
new file mode 100644
index 0000000..92f0377
--- /dev/null
+++ b/tests/docker-compose.test.yml
@@ -0,0 +1,110 @@
+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-explorer-web-test.address=0.0.0.0:80
+    networks:
+      - invariant-explorer-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-explorer-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-explorer-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-explorer-web-test"
+      - "traefik.http.services.explorer-proxy-api.loadbalancer.server.port=8000"
+      - "traefik.docker.network=invariant-explorer-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-explorer-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-explorer-web-test"
+      - "traefik.http.services.explorer-test-api.loadbalancer.server.port=8000"
+      - "traefik.docker.network=invariant-explorer-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
+    volumes:
+      - type: bind
+        source: /tmp/invariant-explorer-test/data/database
+        target: /var/lib/postgresql/data
+    healthcheck:
+      test: [ "CMD-SHELL", "pg_isready -U postgres" ]
+      interval: 5s
+      timeout: 5s
+      retries: 5
+
+networks:
+  invariant-explorer-web-test:
+    external: true
+  internal:

From c2db9ede35bdf668ad8b5673e6340fdf3fc60f51 Mon Sep 17 00:00:00 2001
From: Hemang <hemang@invariantlabs.ai>
Date: Thu, 6 Feb 2025 17:11:07 +0100
Subject: [PATCH 03/15] Change the network name and the .env.test file so that
 test explorer proxy can communicate with the test explorer backend app-api.

---
 run.sh                        |  4 ++--
 tests/.env.test               |  5 ++++-
 tests/docker-compose.test.yml | 20 ++++++++++----------
 3 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/run.sh b/run.sh
index 6a06d77..2a5211d 100755
--- a/run.sh
+++ b/run.sh
@@ -25,8 +25,8 @@ tests() {
   echo "Setting up test environment..."
 
   # Ensure test network exists
-  docker network inspect invariant-explorer-web-test >/dev/null 2>&1 || \
-    docker network create invariant-explorer-web-test
+  docker network inspect invariant-proxy-web-test >/dev/null 2>&1 || \
+    docker network create invariant-proxy-web-test
 
   # Setup the explorer.test.yml file
   CONFIG_DIR="/tmp/invariant-proxy-test/configs"
diff --git a/tests/.env.test b/tests/.env.test
index d0b7167..f52d757 100644
--- a/tests/.env.test
+++ b/tests/.env.test
@@ -1,4 +1,7 @@
-INVARIANT_API_URL=http://127.0.0.1
+# 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
 
 POSTGRES_USER=postgres
 POSTGRES_PASSWORD=postgres
diff --git a/tests/docker-compose.test.yml b/tests/docker-compose.test.yml
index 92f0377..dd46370 100644
--- a/tests/docker-compose.test.yml
+++ b/tests/docker-compose.test.yml
@@ -10,16 +10,16 @@ services:
       # 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-explorer-web-test.address=0.0.0.0:80
+      - --entrypoints.invariant-proxy-web-test.address=0.0.0.0:80
     networks:
-      - invariant-explorer-web-test
+      - 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-explorer-web-test"
+      - "traefik.http.routers.traefik-http.entrypoints=invariant-proxy-web-test"
 
   explorer-proxy:
     container_name: explorer-proxy-test
@@ -39,14 +39,14 @@ services:
         source: ../proxy
         target: /srv/proxy
     networks:
-      - invariant-explorer-web-test
+      - 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-explorer-web-test"
+      - "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-explorer-web-test"
+      - "traefik.docker.network=invariant-proxy-web-test"
     healthcheck:
       test: curl -X GET -I http://localhost:8000/api/v1/proxy/health --fail
       interval: 1s
@@ -73,15 +73,15 @@ services:
       - PORT_API=80
     networks:
       - internal
-      - invariant-explorer-web-test
+      - 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-explorer-web-test"
+      - "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-explorer-web-test"
+      - "traefik.docker.network=invariant-proxy-web-test"
     healthcheck:
       test: curl -X GET -I http://localhost:8000/api/v1 --fail
       interval: 1s
@@ -105,6 +105,6 @@ services:
       retries: 5
 
 networks:
-  invariant-explorer-web-test:
+  invariant-proxy-web-test:
     external: true
   internal:

From d98b73253b2fd0d99d07b8f23eb978da3d639842 Mon Sep 17 00:00:00 2001
From: Hemang <hemang@invariantlabs.ai>
Date: Mon, 10 Feb 2025 05:43:06 +0100
Subject: [PATCH 04/15] Skip Anthropic test if API key not set.

---
 tests/anthropic/test_claude_weather_agent.py | 103 +++++++++----------
 1 file changed, 51 insertions(+), 52 deletions(-)

diff --git a/tests/anthropic/test_claude_weather_agent.py b/tests/anthropic/test_claude_weather_agent.py
index bb61c3c..9315225 100644
--- a/tests/anthropic/test_claude_weather_agent.py
+++ b/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
\ No newline at end of file
+        assert cities[index] in response

From 094970f6b8eec3040015203e097b7567479cc7d8 Mon Sep 17 00:00:00 2001
From: Hemang <hemang@invariantlabs.ai>
Date: Mon, 10 Feb 2025 05:43:27 +0100
Subject: [PATCH 05/15] Add openai test for chat completions without streaming
 and without tool calls.

---
 .gitignore                                    |  1 +
 run.sh                                        |  7 ++-
 tests/.env.test                               |  1 +
 tests/Dockerfile.test                         |  3 +-
 tests/docker-compose.test.yml                 |  2 +-
 tests/open_ai/test_chat_without_tool_calls.py | 57 ++++++++++++++++---
 tests/requirements.txt                        |  4 +-
 tests/util.py                                 | 15 ++++-
 8 files changed, 76 insertions(+), 14 deletions(-)

diff --git a/.gitignore b/.gitignore
index 3f5d309..c62d361 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,6 +3,7 @@ __pycache__/
 .pytest_cache/
 .py[oc]
 data/
+tests/results/
 
 # Coverage and build artifacts
 .coverage
diff --git a/run.sh b/run.sh
index 2a5211d..86fc5d1 100755
--- a/run.sh
+++ b/run.sh
@@ -58,11 +58,16 @@ tests() {
 
   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 $@
 }
 
diff --git a/tests/.env.test b/tests/.env.test
index f52d757..e9e47a0 100644
--- a/tests/.env.test
+++ b/tests/.env.test
@@ -2,6 +2,7 @@
 # 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
diff --git a/tests/Dockerfile.test b/tests/Dockerfile.test
index 54476b4..54514e0 100644
--- a/tests/Dockerfile.test
+++ b/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"]
\ No newline at end of file
+ENTRYPOINT ["pytest",  "--capture=tee-sys", "--tracing", "off", "--junit-xml=/tests/results/test-results-all.xml", "-s", "-vv"]
\ No newline at end of file
diff --git a/tests/docker-compose.test.yml b/tests/docker-compose.test.yml
index dd46370..acbed5b 100644
--- a/tests/docker-compose.test.yml
+++ b/tests/docker-compose.test.yml
@@ -83,7 +83,7 @@ services:
       - "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
+      test: curl -X GET -I http://localhost:8000/api/v1/ --fail
       interval: 1s
       timeout: 5s
 
diff --git a/tests/open_ai/test_chat_without_tool_calls.py b/tests/open_ai/test_chat_without_tool_calls.py
index 5f6c7b1..5ad439a 100644
--- a/tests/open_ai/test_chat_without_tool_calls.py
+++ b/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,47 @@ 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")
+async def test_chat_completion_without_streaming(context, explorer_api_url, proxy_url):
+    """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?"}],
+    )
+
+    # Verify the chat response
+    assert "PARIS" 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 len(trace["messages"]) == 2
+    assert trace["messages"][0] == {
+        "role": "user",
+        "content": "What is the capital of France?",
+    }
+    assert trace["messages"][1]["role"] == "assistant"
+    assert "PARIS" in trace["messages"][1]["content"].upper()
diff --git a/tests/requirements.txt b/tests/requirements.txt
index 9160ba9..9f2ec8a 100644
--- a/tests/requirements.txt
+++ b/tests/requirements.txt
@@ -1,4 +1,6 @@
+anthropic
 openai
 pytest
 pytest-asyncio
-pytest-playwright
\ No newline at end of file
+pytest-playwright
+tavily-python
\ No newline at end of file
diff --git a/tests/util.py b/tests/util.py
index d8bfb23..30826aa 100644
--- a/tests/util.py
+++ b/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

From 5fc05ee39e285d449416fd312e8e55c497c000aa Mon Sep 17 00:00:00 2001
From: Hemang <hemang@invariantlabs.ai>
Date: Mon, 10 Feb 2025 09:57:04 +0100
Subject: [PATCH 06/15] Use the chat completions API response to verify the
 trace in explorer.

---
 tests/open_ai/test_chat_without_tool_calls.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/tests/open_ai/test_chat_without_tool_calls.py b/tests/open_ai/test_chat_without_tool_calls.py
index 5ad439a..c669b61 100644
--- a/tests/open_ai/test_chat_without_tool_calls.py
+++ b/tests/open_ai/test_chat_without_tool_calls.py
@@ -38,6 +38,10 @@ async def test_chat_completion_without_streaming(context, explorer_api_url, prox
 
     # Verify the chat response
     assert "PARIS" in chat_response.choices[0].message.content.upper()
+    expect_assistant_message = {
+        "role": chat_response.choices[0].message.role,
+        "content": chat_response.choices[0].message.content,
+    }
 
     # Fetch the trace ids for the dataset
     traces_response = await context.request.get(
@@ -59,5 +63,4 @@ async def test_chat_completion_without_streaming(context, explorer_api_url, prox
         "role": "user",
         "content": "What is the capital of France?",
     }
-    assert trace["messages"][1]["role"] == "assistant"
-    assert "PARIS" in trace["messages"][1]["content"].upper()
+    assert trace["messages"][1] == expect_assistant_message

From 8ad4a19c21f5b8ffdbd9ca7502c1a76f6a473b13 Mon Sep 17 00:00:00 2001
From: Hemang <hemang@invariantlabs.ai>
Date: Mon, 10 Feb 2025 10:09:42 +0100
Subject: [PATCH 07/15] Modify the test to verify streaming case.

---
 tests/open_ai/test_chat_without_tool_calls.py | 35 ++++++++++++-------
 1 file changed, 23 insertions(+), 12 deletions(-)

diff --git a/tests/open_ai/test_chat_without_tool_calls.py b/tests/open_ai/test_chat_without_tool_calls.py
index c669b61..130b597 100644
--- a/tests/open_ai/test_chat_without_tool_calls.py
+++ b/tests/open_ai/test_chat_without_tool_calls.py
@@ -18,7 +18,8 @@ pytest_plugins = ("pytest_asyncio",)
 
 
 @pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="No OPENAI_API_KEY set")
-async def test_chat_completion_without_streaming(context, explorer_api_url, proxy_url):
+@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())
 
@@ -34,14 +35,20 @@ async def test_chat_completion_without_streaming(context, explorer_api_url, prox
     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
-    assert "PARIS" in chat_response.choices[0].message.content.upper()
-    expect_assistant_message = {
-        "role": chat_response.choices[0].message.role,
-        "content": chat_response.choices[0].message.content,
-    }
+    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(
@@ -58,9 +65,13 @@ async def test_chat_completion_without_streaming(context, explorer_api_url, prox
     trace = await trace_response.json()
 
     # Verify the trace messages
-    assert len(trace["messages"]) == 2
-    assert trace["messages"][0] == {
-        "role": "user",
-        "content": "What is the capital of France?",
-    }
-    assert trace["messages"][1] == expect_assistant_message
+    assert trace["messages"] == [
+        {
+            "role": "user",
+            "content": "What is the capital of France?",
+        },
+        {
+            "role": "assistant",
+            "content": expected_assistant_message,
+        },
+    ]

From 671b7a230a0b9131655d63b75520f241a9332faf Mon Sep 17 00:00:00 2001
From: Hemang <hemang@invariantlabs.ai>
Date: Mon, 10 Feb 2025 13:15:45 +0100
Subject: [PATCH 08/15] Modify run.sh down() to include test infra.

---
 run.sh | 1 +
 1 file changed, 1 insertion(+)

diff --git a/run.sh b/run.sh
index 86fc5d1..7b7a251 100755
--- a/run.sh
+++ b/run.sh
@@ -18,6 +18,7 @@ build() {
 down() {
   # Bring down local services
   docker compose -f docker-compose.local.yml down
+  docker compose -f tests/docker-compose.test.yml down
 }
 
 

From ba26d7d0c711e7447663b4714a6b2719ae331d65 Mon Sep 17 00:00:00 2001
From: Hemang <hemang@invariantlabs.ai>
Date: Mon, 10 Feb 2025 13:45:34 +0100
Subject: [PATCH 09/15] Use finish_turn to decide when to push trace to the
 explorer dataset for openai.

---
 proxy/routes/open_ai.py | 22 +++++++++++++++-------
 1 file changed, 15 insertions(+), 7 deletions(-)

diff --git a/proxy/routes/open_ai.py b/proxy/routes/open_ai.py
index 74fbb44..37807f4 100644
--- a/proxy/routes/open_ai.py
+++ b/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", [])]

From 64ddd9f96d7d870ca760fb860ff75f9401c68399 Mon Sep 17 00:00:00 2001
From: Hemang <hemang@invariantlabs.ai>
Date: Mon, 10 Feb 2025 15:21:46 +0100
Subject: [PATCH 10/15] Add test for openai chat completions API via proxy with
 tool call.

---
 tests/open_ai/test_chat_with_tool_call.py | 140 ++++++++++++++++++++++
 1 file changed, 140 insertions(+)
 create mode 100644 tests/open_ai/test_chat_with_tool_call.py

diff --git a/tests/open_ai/test_chat_with_tool_call.py b/tests/open_ai/test_chat_with_tool_call.py
new file mode 100644
index 0000000..fbd15db
--- /dev/null
+++ b/tests/open_ai/test_chat_with_tool_call.py
@@ -0,0 +1,140 @@
+"""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."""
+    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
+    assert trace["messages"] == [
+        {
+            "role": "user",
+            "content": "What is the weather in New York?",
+        },
+        {
+            "role": "assistant",
+            "tool_calls": [
+                {
+                    "function": {
+                        "arguments": json.loads(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,
+        },
+        {
+            "role": "assistant",
+            "content": chat_response_final.choices[0].message.content,
+        },
+    ]

From 6f0a34c070ddf97436c9af588990b7ecd6cd9cf5 Mon Sep 17 00:00:00 2001
From: Hemang <hemang@invariantlabs.ai>
Date: Mon, 10 Feb 2025 16:06:20 +0100
Subject: [PATCH 11/15] Add test for openai chat completions via proxy with
 tool call and with streaming.

---
 tests/open_ai/test_chat_with_tool_call.py | 119 +++++++++++++++++++---
 1 file changed, 105 insertions(+), 14 deletions(-)

diff --git a/tests/open_ai/test_chat_with_tool_call.py b/tests/open_ai/test_chat_with_tool_call.py
index fbd15db..422e5ff 100644
--- a/tests/open_ai/test_chat_with_tool_call.py
+++ b/tests/open_ai/test_chat_with_tool_call.py
@@ -22,7 +22,10 @@ pytest_plugins = ("pytest_asyncio",)
 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."""
+    """
+    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(
@@ -109,32 +112,120 @@ async def test_chat_completion_with_tool_call_without_streaming(
     trace = await trace_response.json()
 
     # Verify the trace messages
-    assert trace["messages"] == [
+    expected_messages = history + [
         {
-            "role": "user",
-            "content": "What is the weather in New York?",
-        },
+            "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": json.loads(tool_call.function.arguments),
-                        "name": tool_call.function.name,
+                        "arguments": tool_call["function"]["arguments"],
+                        "name": tool_call["function"]["name"],
                     },
-                    "id": tool_call.id,
-                    "type": tool_call.type,
+                    "id": tool_call["id"],
+                    "type": tool_call["type"],
                 }
             ],
         },
         {
             "role": "tool",
-            "tool_call_id": tool_call.id,
+            "tool_call_id": tool_call["id"],
             "tool_name": "get_weather",
             "content": tool_result,
         },
-        {
-            "role": "assistant",
-            "content": chat_response_final.choices[0].message.content,
-        },
     ]
+
+    # 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

From 46417a804f242d64c035a91bad2120beb4a2083f Mon Sep 17 00:00:00 2001
From: Hemang Sarkar <hemang@invariantlabs.ai>
Date: Tue, 11 Feb 2025 10:16:17 +0100
Subject: [PATCH 12/15] Create tests_ci.yml

---
 .github/workflows/tests_ci.yml | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)
 create mode 100644 .github/workflows/tests_ci.yml

diff --git a/.github/workflows/tests_ci.yml b/.github/workflows/tests_ci.yml
new file mode 100644
index 0000000..e000800
--- /dev/null
+++ b/.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 --cov=proxy --cov-report=term

From a26ed0eebb1e3d0f86d2e46b7b07f39111678f1e Mon Sep 17 00:00:00 2001
From: Hemang <hemang@invariantlabs.ai>
Date: Tue, 11 Feb 2025 10:21:03 +0100
Subject: [PATCH 13/15] Remove volume bind from test docker compose.

---
 tests/docker-compose.test.yml | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/tests/docker-compose.test.yml b/tests/docker-compose.test.yml
index acbed5b..43d6e4d 100644
--- a/tests/docker-compose.test.yml
+++ b/tests/docker-compose.test.yml
@@ -94,10 +94,6 @@ services:
       - .env.test
     networks:
       - internal
-    volumes:
-      - type: bind
-        source: /tmp/invariant-explorer-test/data/database
-        target: /var/lib/postgresql/data
     healthcheck:
       test: [ "CMD-SHELL", "pg_isready -U postgres" ]
       interval: 5s

From c0abf96879b14fe33b671d80fafc3825e1fdf7a4 Mon Sep 17 00:00:00 2001
From: Hemang Sarkar <hemang@invariantlabs.ai>
Date: Tue, 11 Feb 2025 10:24:41 +0100
Subject: [PATCH 14/15] Update requirements.txt

---
 tests/requirements.txt | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tests/requirements.txt b/tests/requirements.txt
index 9f2ec8a..08b482c 100644
--- a/tests/requirements.txt
+++ b/tests/requirements.txt
@@ -2,5 +2,6 @@ anthropic
 openai
 pytest
 pytest-asyncio
+pytest-cov
 pytest-playwright
-tavily-python
\ No newline at end of file
+tavily-python

From 472136a2adbc42783855fa703b6cbc4862b5044e Mon Sep 17 00:00:00 2001
From: Hemang <hemang@invariantlabs.ai>
Date: Tue, 11 Feb 2025 10:29:04 +0100
Subject: [PATCH 15/15] Remove coverage for now.

---
 .github/workflows/tests_ci.yml | 2 +-
 tests/requirements.txt         | 1 -
 2 files changed, 1 insertion(+), 2 deletions(-)

diff --git a/.github/workflows/tests_ci.yml b/.github/workflows/tests_ci.yml
index e000800..fa97fb5 100644
--- a/.github/workflows/tests_ci.yml
+++ b/.github/workflows/tests_ci.yml
@@ -19,4 +19,4 @@ jobs:
         - name: Run tests
           env:
             OPENAI_API_KEY: ${{ secrets.INVARIANT_TESTING_OPENAI_KEY }}
-          run: ./run.sh tests -s -vv --cov=proxy --cov-report=term
+          run: ./run.sh tests -s -vv
diff --git a/tests/requirements.txt b/tests/requirements.txt
index 08b482c..b3c6589 100644
--- a/tests/requirements.txt
+++ b/tests/requirements.txt
@@ -2,6 +2,5 @@ anthropic
 openai
 pytest
 pytest-asyncio
-pytest-cov
 pytest-playwright
 tavily-python