mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-25 18:54:30 +00:00
117 lines
3.5 KiB
Python
117 lines
3.5 KiB
Python
import sys
|
|
import os
|
|
import json
|
|
import time
|
|
|
|
import logfire
|
|
import litellm
|
|
import pytest
|
|
from logfire.testing import TestExporter, SimpleSpanProcessor
|
|
|
|
sys.path.insert(0, os.path.abspath("../.."))
|
|
|
|
# Testing scenarios for logfire logging:
|
|
# 1. Test logfire logging for completion
|
|
# 2. Test logfire logging for acompletion
|
|
# 3. Test logfire logging for completion while streaming is enabled
|
|
# 4. Test logfire logging for completion while streaming is enabled
|
|
|
|
|
|
@pytest.mark.parametrize("stream", [False, True])
|
|
def test_completion_logfire_logging(stream):
|
|
litellm.success_callback = ["logfire"]
|
|
litellm.set_verbose = True
|
|
|
|
exporter = TestExporter()
|
|
logfire.configure(
|
|
send_to_logfire=False,
|
|
console=False,
|
|
processors=[SimpleSpanProcessor(exporter)],
|
|
collect_system_metrics=False,
|
|
)
|
|
messages = [{"role": "user", "content": "what llm are u"}]
|
|
temperature = 0.3
|
|
max_tokens = 10
|
|
response = litellm.completion(
|
|
model="gpt-3.5-turbo",
|
|
messages=messages,
|
|
max_tokens=max_tokens,
|
|
temperature=temperature,
|
|
stream=stream,
|
|
)
|
|
print(response)
|
|
|
|
if stream:
|
|
for chunk in response:
|
|
print(chunk)
|
|
|
|
time.sleep(5)
|
|
exported_spans = exporter.exported_spans_as_dict()
|
|
|
|
assert len(exported_spans) == 1
|
|
assert (
|
|
exported_spans[0]["attributes"]["logfire.msg"]
|
|
== "Chat Completion with 'gpt-3.5-turbo'"
|
|
)
|
|
|
|
request_data = json.loads(exported_spans[0]["attributes"]["request_data"])
|
|
|
|
assert request_data["model"] == "gpt-3.5-turbo"
|
|
assert request_data["messages"] == messages
|
|
|
|
assert "completion_tokens" in request_data["usage"]
|
|
assert "prompt_tokens" in request_data["usage"]
|
|
assert "total_tokens" in request_data["usage"]
|
|
assert request_data["response"]["choices"][0]["message"]["content"]
|
|
assert request_data["modelParameters"]["max_tokens"] == max_tokens
|
|
assert request_data["modelParameters"]["temperature"] == temperature
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("stream", [False, True])
|
|
async def test_acompletion_logfire_logging(stream):
|
|
litellm.success_callback = ["logfire"]
|
|
litellm.set_verbose = True
|
|
|
|
exporter = TestExporter()
|
|
logfire.configure(
|
|
send_to_logfire=False,
|
|
console=False,
|
|
processors=[SimpleSpanProcessor(exporter)],
|
|
collect_system_metrics=False,
|
|
)
|
|
messages = [{"role": "user", "content": "what llm are u"}]
|
|
temperature = 0.3
|
|
max_tokens = 10
|
|
response = await litellm.acompletion(
|
|
model="gpt-3.5-turbo",
|
|
messages=messages,
|
|
max_tokens=max_tokens,
|
|
temperature=temperature,
|
|
)
|
|
print(response)
|
|
if stream:
|
|
for chunk in response:
|
|
print(chunk)
|
|
|
|
time.sleep(5)
|
|
exported_spans = exporter.exported_spans_as_dict()
|
|
print("exported_spans", exported_spans)
|
|
|
|
assert len(exported_spans) == 1
|
|
assert (
|
|
exported_spans[0]["attributes"]["logfire.msg"]
|
|
== "Chat Completion with 'gpt-3.5-turbo'"
|
|
)
|
|
|
|
request_data = json.loads(exported_spans[0]["attributes"]["request_data"])
|
|
|
|
assert request_data["model"] == "gpt-3.5-turbo"
|
|
assert request_data["messages"] == messages
|
|
|
|
assert "completion_tokens" in request_data["usage"]
|
|
assert "prompt_tokens" in request_data["usage"]
|
|
assert "total_tokens" in request_data["usage"]
|
|
assert request_data["response"]["choices"][0]["message"]["content"]
|
|
assert request_data["modelParameters"]["max_tokens"] == max_tokens
|
|
assert request_data["modelParameters"]["temperature"] == temperature
|