litellm

1
# What this tests ?
2
## Tests /config/update + Test /chat/completions -> assert logs are sent to Langfuse
3

4
import pytest
5
import asyncio
6
import aiohttp
7
import os
8
import dotenv
9
from dotenv import load_dotenv
10
import pytest
11

12
load_dotenv()
13

14

15
async def config_update(session):
16
    url = "http://0.0.0.0:4000/config/update"
17
    headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
18
    data = {
19
        "litellm_settings": {
20
            "success_callback": ["langfuse"],
21
        },
22
        "environment_variables": {
23
            "LANGFUSE_PUBLIC_KEY": os.environ["LANGFUSE_PUBLIC_KEY"],
24
            "LANGFUSE_SECRET_KEY": os.environ["LANGFUSE_SECRET_KEY"],
25
        },
26
    }
27

28
    async with session.post(url, headers=headers, json=data) as response:
29
        status = response.status
30
        response_text = await response.text()
31

32
        print(response_text)
33
        print()
34

35
        if status != 200:
36
            raise Exception(f"Request did not return a 200 status code: {status}")
37
        return await response.json()
38

39

40
async def chat_completion(session, key, model="azure-gpt-3.5", request_metadata=None):
41
    url = "http://0.0.0.0:4000/chat/completions"
42
    headers = {
43
        "Authorization": f"Bearer {key}",
44
        "Content-Type": "application/json",
45
    }
46
    data = {
47
        "model": model,
48
        "messages": [
49
            {"role": "system", "content": "You are a helpful assistant."},
50
            {"role": "user", "content": "Hello!"},
51
        ],
52
        "metadata": request_metadata,
53
    }
54

55
    print("data sent in test=", data)
56

57
    async with session.post(url, headers=headers, json=data) as response:
58
        status = response.status
59
        response_text = await response.text()
60

61
        print(response_text)
62
        print()
63

64
        if status != 200:
65
            raise Exception(f"Request did not return a 200 status code: {status}")
66

67

68
@pytest.mark.asyncio
69
async def test_team_logging():
70
    """
71
    1. Add Langfuse as a callback with /config/update
72
    2. Call /chat/completions
73
    3. Assert the logs are sent to Langfuse
74
    """
75
    try:
76
        async with aiohttp.ClientSession() as session:
77

78
            # Add Langfuse as a callback with /config/update
79
            await config_update(session)
80

81
            # 2. Call /chat/completions with a specific trace id
82
            import uuid
83

84
            _trace_id = f"trace-{uuid.uuid4()}"
85
            _request_metadata = {
86
                "trace_id": _trace_id,
87
            }
88

89
            await chat_completion(
90
                session,
91
                key="sk-1234",
92
                model="fake-openai-endpoint",
93
                request_metadata=_request_metadata,
94
            )
95

96
            # Test - if the logs were sent to the correct team on langfuse
97
            import langfuse
98

99
            langfuse_client = langfuse.Langfuse(
100
                public_key=os.getenv("LANGFUSE_PUBLIC_KEY"),
101
                secret_key=os.getenv("LANGFUSE_SECRET_KEY"),
102
            )
103

104
            await asyncio.sleep(10)
105

106
            print(f"searching for trace_id={_trace_id} on langfuse")
107

108
            generations = langfuse_client.get_generations(trace_id=_trace_id).data
109

110
            # 1 generation with this trace id
111
            assert len(generations) == 1
112

113
    except Exception as e:
114
        pytest.fail("Team 2 logging failed: " + str(e))
115