Spaces:
Configuration error
Configuration error
File size: 3,714 Bytes
447ebeb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 |
# 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()
@pytest.mark.asyncio
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
|