Skip to content

Commit c138716

Browse files
shah-sidddgustavocidornelas
shah-siddd
authored and
gustavocidornelas
committed
style: lint fixes
1 parent 49c1cb1 commit c138716

File tree

3 files changed

+327
-22
lines changed

3 files changed

+327
-22
lines changed

‎examples/tracing/litellm/litellm_tracing.ipynb‎

Lines changed: 8 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -94,11 +94,7 @@
9494
" temperature=0.7,\n",
9595
" max_tokens=100,\n",
9696
" inference_id=\"litellm-openai-example-1\" # Optional: custom inference ID\n",
97-
")\n",
98-
"\n",
99-
"print(f\"Response: {response.choices[0].message.content}\")\n",
100-
"print(f\"Model used: {response.model}\")\n",
101-
"print(f\"Tokens used: {response.usage.total_tokens}\")\n"
97+
")\n"
10298
]
10399
},
104100
{
@@ -127,22 +123,13 @@
127123
"]\n",
128124
"\n",
129125
"for model in models_to_test:\n",
130-
" try:\n",
131-
" print(f\"\\n--- Testing {model} ---\")\n",
132-
" response = litellm.completion(\n",
133-
" model=model,\n",
134-
" messages=messages,\n",
135-
" temperature=0.5,\n",
136-
" max_tokens=150,\n",
137-
" inference_id=f\"multi-provider-{model.replace('/', '-')}\"\n",
138-
" )\n",
139-
" \n",
140-
" print(f\"Model: {response.model}\")\n",
141-
" print(f\"Response: {response.choices[0].message.content[:200]}...\")\n",
142-
" print(f\"Tokens: {response.usage.total_tokens}\")\n",
143-
" \n",
144-
" except Exception as e:\n",
145-
" print(f\"Failed to test {model}: {e}\")\n"
126+
" response = litellm.completion(\n",
127+
" model=model,\n",
128+
" messages=messages,\n",
129+
" temperature=0.5,\n",
130+
" max_tokens=150,\n",
131+
" inference_id=f\"multi-provider-{model.replace('/', '-')}\"\n",
132+
" )\n"
146133
]
147134
},
148135
{

‎examples/tracing/litellm/test_local_litellm.py‎

Lines changed: 317 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,317 @@
1+
#!/usr/bin/env python3
2+
"""
3+
Test script for LiteLLM tracing with local LiteLLM server.
4+
5+
This script demonstrates how to test the LiteLLM integration with:
6+
1. Local LiteLLM proxy server
7+
2. Custom API base URLs
8+
3. Various providers and models
9+
10+
Prerequisites:
11+
- LiteLLM server running locally (e.g., litellm --port 4000)
12+
- API keys configured in environment or LiteLLM config
13+
"""
14+
15+
import os
16+
import sys
17+
import time
18+
from typing import Dict, Any
19+
20+
# Add the src directory to the path for local testing
21+
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../../../src'))
22+
23+
try:
24+
import litellm
25+
from openlayer.lib import trace_litellm
26+
from openlayer.lib.tracing import tracer
27+
from openlayer.lib.tracing.tracer import configure
28+
except ImportError as e:
29+
print(f"Import error: {e}")
30+
print("Make sure to install required dependencies:")
31+
print("pip install litellm openlayer")
32+
sys.exit(1)
33+
34+
35+
class LiteLLMTester:
36+
"""Test LiteLLM tracing with various configurations."""
37+
38+
def __init__(self, base_url: str = None, api_key: str = None, openlayer_base_url: str = None):
39+
"""Initialize the tester with optional custom base URL and API key."""
40+
self.base_url = base_url or "http://localhost:4000"
41+
self.api_key = api_key or os.getenv("LITELLM_API_KEY", "sk-1234")
42+
self.openlayer_base_url = openlayer_base_url or "http://localhost:8080/v1"
43+
44+
# Configure OpenLayer base URL programmatically
45+
configure(base_url=self.openlayer_base_url)
46+
print(f"🔧 OpenLayer configured for: {self.openlayer_base_url}")
47+
48+
# Configure LiteLLM for local testing
49+
if base_url:
50+
# Set custom API base for testing with local LiteLLM server
51+
os.environ["LITELLM_BASE_URL"] = self.base_url
52+
53+
# Enable tracing
54+
trace_litellm()
55+
print(f"✅ LiteLLM tracing enabled")
56+
print(f"🔗 LiteLLM Base URL: {self.base_url}")
57+
print(f"🏠 OpenLayer Base URL: {self.openlayer_base_url}")
58+
59+
def test_basic_completion(self, model: str = "gpt-3.5-turbo") -> Dict[str, Any]:
60+
"""Test basic completion with tracing."""
61+
print(f"\n📝 Testing basic completion with {model}")
62+
63+
try:
64+
response = litellm.completion(
65+
model=model,
66+
messages=[
67+
{"role": "system", "content": "You are a helpful assistant."},
68+
{"role": "user", "content": "What is 2 + 2?"}
69+
],
70+
temperature=0.5,
71+
max_tokens=50,
72+
api_base=self.base_url,
73+
api_key=self.api_key,
74+
inference_id=f"test-basic-{int(time.time())}"
75+
)
76+
77+
result = {
78+
"status": "success",
79+
"model": response.model,
80+
"content": response.choices[0].message.content,
81+
"usage": response.usage.model_dump() if response.usage else None,
82+
"provider": getattr(response, '_hidden_params', {}).get('custom_llm_provider', 'unknown')
83+
}
84+
85+
print(f"✅ Success: {result['content'][:100]}...")
86+
print(f"📊 Usage: {result['usage']}")
87+
print(f"🏢 Provider: {result['provider']}")
88+
89+
return result
90+
91+
except Exception as e:
92+
print(f"❌ Error: {e}")
93+
return {"status": "error", "error": str(e)}
94+
95+
def test_streaming_completion(self, model: str = "gpt-3.5-turbo") -> Dict[str, Any]:
96+
"""Test streaming completion with tracing."""
97+
print(f"\n🌊 Testing streaming completion with {model}")
98+
99+
try:
100+
stream = litellm.completion(
101+
model=model,
102+
messages=[
103+
{"role": "user", "content": "Count from 1 to 5, one number per line."}
104+
],
105+
stream=True,
106+
temperature=0.3,
107+
max_tokens=50,
108+
api_base=self.base_url,
109+
api_key=self.api_key,
110+
inference_id=f"test-stream-{int(time.time())}"
111+
)
112+
113+
collected_content = []
114+
chunk_count = 0
115+
116+
for chunk in stream:
117+
chunk_count += 1
118+
if hasattr(chunk.choices[0].delta, 'content') and chunk.choices[0].delta.content:
119+
content = chunk.choices[0].delta.content
120+
collected_content.append(content)
121+
print(content, end="", flush=True)
122+
123+
full_content = "".join(collected_content)
124+
125+
result = {
126+
"status": "success",
127+
"model": model,
128+
"content": full_content,
129+
"chunks": chunk_count,
130+
"provider": "streamed" # Provider detection in streaming is complex
131+
}
132+
133+
print(f"\n✅ Streaming complete: {chunk_count} chunks")
134+
print(f"📝 Content: {full_content}")
135+
136+
return result
137+
138+
except Exception as e:
139+
print(f"❌ Streaming error: {e}")
140+
return {"status": "error", "error": str(e)}
141+
142+
def test_multiple_providers(self, models: list = None) -> Dict[str, Any]:
143+
"""Test multiple providers/models with tracing."""
144+
if models is None:
145+
models = [
146+
"gpt-3.5-turbo",
147+
"claude-3-haiku-20240307",
148+
"gemini-pro",
149+
"llama-2-7b-chat"
150+
]
151+
152+
print(f"\n🔄 Testing multiple providers: {models}")
153+
154+
results = {}
155+
prompt = "What is the capital of Japan?"
156+
157+
with tracer.create_step(
158+
name="Multi-Provider Test",
159+
metadata={"test_type": "provider_comparison", "models": models}
160+
) as step:
161+
162+
for model in models:
163+
try:
164+
print(f"\n🧪 Testing {model}...")
165+
166+
response = litellm.completion(
167+
model=model,
168+
messages=[{"role": "user", "content": prompt}],
169+
temperature=0.5,
170+
max_tokens=30,
171+
api_base=self.base_url,
172+
api_key=self.api_key,
173+
inference_id=f"multi-test-{model.replace('/', '-')}-{int(time.time())}"
174+
)
175+
176+
results[model] = {
177+
"status": "success",
178+
"content": response.choices[0].message.content,
179+
"usage": response.usage.model_dump() if response.usage else None,
180+
"provider": getattr(response, '_hidden_params', {}).get('custom_llm_provider', 'unknown')
181+
}
182+
183+
print(f"✅ {model}: {results[model]['content'][:50]}...")
184+
185+
except Exception as e:
186+
results[model] = {"status": "error", "error": str(e)}
187+
print(f"❌ {model}: {e}")
188+
189+
step.log(results=results)
190+
191+
return results
192+
193+
def test_function_calling(self, model: str = "gpt-3.5-turbo") -> Dict[str, Any]:
194+
"""Test function calling with tracing."""
195+
print(f"\n🔧 Testing function calling with {model}")
196+
197+
functions = [
198+
{
199+
"name": "get_current_weather",
200+
"description": "Get the current weather in a given location",
201+
"parameters": {
202+
"type": "object",
203+
"properties": {
204+
"location": {
205+
"type": "string",
206+
"description": "The city and state, e.g. San Francisco, CA"
207+
},
208+
"unit": {
209+
"type": "string",
210+
"enum": ["celsius", "fahrenheit"]
211+
}
212+
},
213+
"required": ["location"]
214+
}
215+
}
216+
]
217+
218+
try:
219+
response = litellm.completion(
220+
model=model,
221+
messages=[
222+
{"role": "user", "content": "What's the weather like in Tokyo?"}
223+
],
224+
functions=functions,
225+
function_call="auto",
226+
api_base=self.base_url,
227+
api_key=self.api_key,
228+
inference_id=f"test-func-{int(time.time())}"
229+
)
230+
231+
message = response.choices[0].message
232+
233+
if message.function_call:
234+
result = {
235+
"status": "success",
236+
"function_name": message.function_call.name,
237+
"arguments": message.function_call.arguments,
238+
"usage": response.usage.model_dump() if response.usage else None
239+
}
240+
print(f"✅ Function called: {result['function_name']}")
241+
print(f"📋 Arguments: {result['arguments']}")
242+
else:
243+
result = {
244+
"status": "success",
245+
"content": message.content,
246+
"note": "No function call triggered",
247+
"usage": response.usage.model_dump() if response.usage else None
248+
}
249+
print(f"✅ Regular response: {result['content']}")
250+
251+
return result
252+
253+
except Exception as e:
254+
print(f"❌ Function calling error: {e}")
255+
return {"status": "error", "error": str(e)}
256+
257+
def run_all_tests(self):
258+
"""Run all test scenarios."""
259+
print("🚀 Starting comprehensive LiteLLM tracing tests")
260+
print("=" * 60)
261+
262+
results = {
263+
"basic": self.test_basic_completion(),
264+
"streaming": self.test_streaming_completion(),
265+
"multi_provider": self.test_multiple_providers(),
266+
"function_calling": self.test_function_calling(),
267+
}
268+
269+
print("\n" + "=" * 60)
270+
print("📊 Test Summary:")
271+
272+
for test_name, result in results.items():
273+
status = result.get("status", "unknown")
274+
emoji = "✅" if status == "success" else "❌"
275+
print(f"{emoji} {test_name}: {status}")
276+
277+
return results
278+
279+
280+
def main():
281+
"""Main test function."""
282+
print("🧪 LiteLLM Tracing Test Suite")
283+
print("=" * 40)
284+
285+
# Configuration
286+
base_url = os.getenv("LITELLM_BASE_URL", "http://localhost:4000")
287+
api_key = os.getenv("LITELLM_API_KEY", "sk-1234")
288+
openlayer_base_url = os.getenv("OPENLAYER_BASE_URL", "http://localhost:8080/v1")
289+
290+
# You can also set OpenLayer configuration
291+
os.environ.setdefault("OPENLAYER_API_KEY", "sk-ol-vMcEc8O_Tw52HDIF8ihNsiIlzmHLnXxC")
292+
os.environ.setdefault("OPENLAYER_INFERENCE_PIPELINE_ID", "efefdd4f-12ab-4343-a164-7c10d2d48d61")
293+
294+
print(f"🔗 LiteLLM Base URL: {base_url}")
295+
print(f"🏠 OpenLayer Base URL: {openlayer_base_url}")
296+
print(f"🔑 API Key: {api_key[:8]}...")
297+
298+
# Initialize tester
299+
tester = LiteLLMTester(base_url=base_url, api_key=api_key, openlayer_base_url=openlayer_base_url)
300+
301+
# Run tests
302+
try:
303+
results = tester.run_all_tests()
304+
305+
print("\n🎯 All tests completed!")
306+
print("Check your OpenLayer dashboard for detailed traces.")
307+
308+
except KeyboardInterrupt:
309+
print("\n⏹️ Tests interrupted by user")
310+
except Exception as e:
311+
print(f"\n💥 Unexpected error: {e}")
312+
import traceback
313+
traceback.print_exc()
314+
315+
316+
if __name__ == "__main__":
317+
main()

‎tests/test_litellm_integration.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,9 @@
11
"""Test LiteLLM integration."""
22

33
import builtins
4+
from unittest.mock import Mock, patch
5+
46
import pytest
5-
from unittest.mock import Mock, patch, MagicMock
67

78

89
class TestLiteLLMIntegration:

0 commit comments

Comments
 (0)