Spaces:
Configuration error
Configuration error
| # What this tests ? | |
| ## Tests /chat/completions by generating a key and then making a chat completions request | |
| import pytest | |
| import asyncio | |
| import aiohttp, openai | |
| from openai import OpenAI, AsyncOpenAI | |
| from typing import Optional, List, Union | |
| import uuid | |
| async def generate_key( | |
| session, | |
| models=[ | |
| "gpt-4", | |
| "text-embedding-ada-002", | |
| "dall-e-2", | |
| "fake-openai-endpoint", | |
| "mistral-embed", | |
| ], | |
| ): | |
| url = "http://0.0.0.0:4000/key/generate" | |
| headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"} | |
| data = { | |
| "models": models, | |
| "duration": None, | |
| } | |
| async with session.post(url, headers=headers, json=data) as response: | |
| status = response.status | |
| response_text = await response.text() | |
| print(response_text) | |
| print() | |
| if status != 200: | |
| raise Exception(f"Request did not return a 200 status code: {status}") | |
| return await response.json() | |
| async def chat_completion(session, key, model: Union[str, List] = "gpt-4"): | |
| url = "http://0.0.0.0:4000/chat/completions" | |
| headers = { | |
| "Authorization": f"Bearer {key}", | |
| "Content-Type": "application/json", | |
| } | |
| data = { | |
| "model": model, | |
| "messages": [ | |
| {"role": "user", "content": f"Hello! {str(uuid.uuid4())}"}, | |
| ], | |
| } | |
| async with session.post(url, headers=headers, json=data) as response: | |
| status = response.status | |
| response_text = await response.text() | |
| print(response_text) | |
| print() | |
| if status != 200: | |
| raise Exception(f"Request did not return a 200 status code: {status}") | |
| return await response.json() | |
| async def get_otel_spans(session, key): | |
| url = "http://0.0.0.0:4000/otel-spans" | |
| headers = { | |
| "Authorization": f"Bearer {key}", | |
| "Content-Type": "application/json", | |
| } | |
| async with session.get(url, headers=headers) as response: | |
| status = response.status | |
| response_text = await response.text() | |
| print(response_text) | |
| print() | |
| if status != 200: | |
| raise Exception(f"Request did not return a 200 status code: {status}") | |
| return await response.json() | |
| async def test_chat_completion_check_otel_spans(): | |
| """ | |
| - Create key | |
| Make chat completion call | |
| - Create user | |
| make chat completion call | |
| """ | |
| async with aiohttp.ClientSession() as session: | |
| key_gen = await generate_key(session=session) | |
| key = key_gen["key"] | |
| await chat_completion(session=session, key=key, model="fake-openai-endpoint") | |
| await asyncio.sleep(3) | |
| otel_spans = await get_otel_spans(session=session, key=key) | |
| print("otel_spans: ", otel_spans) | |
| all_otel_spans = otel_spans["otel_spans"] | |
| most_recent_parent = str(otel_spans["most_recent_parent"]) | |
| print("Most recent OTEL parent: ", most_recent_parent) | |
| print("\n spans grouped by parent: ", otel_spans["spans_grouped_by_parent"]) | |
| parent_trace_spans = otel_spans["spans_grouped_by_parent"][most_recent_parent] | |
| print("Parent trace spans: ", parent_trace_spans) | |
| # either 5 or 6 traces depending on how many redis calls were made | |
| assert len(parent_trace_spans) >= 5 | |
| # 'postgres', 'redis', 'raw_gen_ai_request', 'litellm_request', 'Received Proxy Server Request' in the span | |
| assert "postgres" in parent_trace_spans | |
| assert "redis" in parent_trace_spans | |
| assert "raw_gen_ai_request" in parent_trace_spans | |
| assert "litellm_request" in parent_trace_spans | |
| assert "batch_write_to_db" in parent_trace_spans | |