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/.github/workflows/tests_ci.yml b/.github/workflows/tests_ci.yml new file mode 100644 index 0000000..fa97fb5 --- /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 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/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", [])] 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"]) diff --git a/run.sh b/run.sh index ce3dfc3..7b7a251 100755 --- a/run.sh +++ b/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 $@ } diff --git a/tests/.env.test b/tests/.env.test new file mode 100644 index 0000000..e9e47a0 --- /dev/null +++ b/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 \ No newline at end of file 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/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 "}, + ), + 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 diff --git a/tests/docker-compose.test.yml b/tests/docker-compose.test.yml new file mode 100644 index 0000000..43d6e4d --- /dev/null +++ b/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: 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..422e5ff --- /dev/null +++ b/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 " + }, # 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 " + }, # 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 diff --git a/tests/open_ai/test_chat_without_tool_calls.py b/tests/open_ai/test_chat_without_tool_calls.py index 5f6c7b1..130b597 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,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 " + }, # 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, + }, + ] diff --git a/tests/requirements.txt b/tests/requirements.txt index 9160ba9..b3c6589 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 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