fix test traceloop.py

This commit is contained in:
Ishaan Jaff 2024-08-01 19:49:53 -07:00
parent 540d5f318b
commit 85ced920a0

View file

@ -1,11 +1,13 @@
import sys
import os import os
import sys
import time import time
import pytest import pytest
import litellm
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
from traceloop.sdk import Traceloop from traceloop.sdk import Traceloop
import litellm
sys.path.insert(0, os.path.abspath("../..")) sys.path.insert(0, os.path.abspath("../.."))
@ -25,11 +27,11 @@ def exporter():
@pytest.mark.parametrize("model", ["claude-instant-1.2", "gpt-3.5-turbo"]) @pytest.mark.parametrize("model", ["claude-instant-1.2", "gpt-3.5-turbo"])
def test_traceloop_logging(exporter, model): def test_traceloop_logging(exporter, model):
litellm.completion( litellm.completion(
model=model, model=model,
messages=[{"role": "user", "content": "This is a test"}], messages=[{"role": "user", "content": "This is a test"}],
max_tokens=1000, max_tokens=1000,
temperature=0.7, temperature=0.7,
timeout=5, timeout=5,
mock_response="hi",
) )