test: cleanup dead tests
This commit is contained in:
Krrish Dholakia
parent
a92b31a636
commit
25f2baad71
@ -46,45 +46,51 @@ def mock_snowflake_streaming_response_chunks() -> List[str]:
|
||||
Mock streaming response chunks for Snowflake.
|
||||
"""
|
||||
return [
|
||||
json.dumps({
|
||||
"id": "chatcmpl-snowflake-stream-123",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1700000000,
|
||||
"model": "mistral-7b",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"delta": {"role": "assistant", "content": "The"},
|
||||
"finish_reason": None,
|
||||
}
|
||||
],
|
||||
}),
|
||||
json.dumps({
|
||||
"id": "chatcmpl-snowflake-stream-123",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1700000000,
|
||||
"model": "mistral-7b",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"delta": {"content": " sky"},
|
||||
"finish_reason": None,
|
||||
}
|
||||
],
|
||||
}),
|
||||
json.dumps({
|
||||
"id": "chatcmpl-snowflake-stream-123",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1700000000,
|
||||
"model": "mistral-7b",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"delta": {"content": " is blue"},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
}),
|
||||
json.dumps(
|
||||
{
|
||||
"id": "chatcmpl-snowflake-stream-123",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1700000000,
|
||||
"model": "mistral-7b",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"delta": {"role": "assistant", "content": "The"},
|
||||
"finish_reason": None,
|
||||
}
|
||||
],
|
||||
}
|
||||
),
|
||||
json.dumps(
|
||||
{
|
||||
"id": "chatcmpl-snowflake-stream-123",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1700000000,
|
||||
"model": "mistral-7b",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"delta": {"content": " sky"},
|
||||
"finish_reason": None,
|
||||
}
|
||||
],
|
||||
}
|
||||
),
|
||||
json.dumps(
|
||||
{
|
||||
"id": "chatcmpl-snowflake-stream-123",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1700000000,
|
||||
"model": "mistral-7b",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"delta": {"content": " is blue"},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
}
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@ -120,6 +126,7 @@ def test_chat_completion_snowflake(sync_mode):
|
||||
async_handler = AsyncHTTPHandler()
|
||||
with patch.object(AsyncHTTPHandler, "post", return_value=mock_response):
|
||||
import asyncio
|
||||
|
||||
response = asyncio.run(
|
||||
acompletion(
|
||||
model="snowflake/mistral-7b",
|
||||
@ -148,16 +155,16 @@ def test_chat_completion_snowflake_stream(sync_mode):
|
||||
if sync_mode:
|
||||
sync_handler = HTTPHandler()
|
||||
mock_chunks = mock_snowflake_streaming_response_chunks()
|
||||
|
||||
|
||||
def mock_iter_lines():
|
||||
for chunk in mock_chunks:
|
||||
for line in [f"data: {chunk}", "data: [DONE]"]:
|
||||
yield line
|
||||
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.iter_lines.side_effect = mock_iter_lines
|
||||
mock_response.status_code = 200
|
||||
|
||||
|
||||
with patch.object(HTTPHandler, "post", return_value=mock_response):
|
||||
response = completion(
|
||||
model="snowflake/mistral-7b",
|
||||
@ -167,28 +174,28 @@ def test_chat_completion_snowflake_stream(sync_mode):
|
||||
api_base="https://exampleopenaiendpoint-production.up.railway.app/v1/chat/completions",
|
||||
client=sync_handler,
|
||||
)
|
||||
|
||||
|
||||
chunks_received = []
|
||||
for chunk in response:
|
||||
chunks_received.append(chunk)
|
||||
|
||||
|
||||
assert len(chunks_received) > 0
|
||||
else:
|
||||
async_handler = AsyncHTTPHandler()
|
||||
mock_chunks = mock_snowflake_streaming_response_chunks()
|
||||
|
||||
|
||||
async def mock_iter_lines():
|
||||
for chunk in mock_chunks:
|
||||
for line in [f"data: {chunk}", "data: [DONE]"]:
|
||||
yield line
|
||||
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.iter_lines.side_effect = mock_iter_lines
|
||||
mock_response.status_code = 200
|
||||
|
||||
|
||||
with patch.object(AsyncHTTPHandler, "post", return_value=mock_response):
|
||||
import asyncio
|
||||
|
||||
|
||||
async def test_async_stream():
|
||||
response = await acompletion(
|
||||
model="snowflake/mistral-7b",
|
||||
@ -198,78 +205,11 @@ def test_chat_completion_snowflake_stream(sync_mode):
|
||||
api_base="https://exampleopenaiendpoint-production.up.railway.app/v1/chat/completions",
|
||||
client=async_handler,
|
||||
)
|
||||
|
||||
|
||||
chunks_received = []
|
||||
async for chunk in response:
|
||||
chunks_received.append(chunk)
|
||||
|
||||
|
||||
assert len(chunks_received) > 0
|
||||
|
||||
|
||||
asyncio.run(test_async_stream())
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="Requires Snowflake credentials - run manually when needed")
|
||||
def test_snowflake_tool_calling_responses_api():
|
||||
"""
|
||||
Test Snowflake tool calling with Responses API.
|
||||
Requires SNOWFLAKE_JWT and SNOWFLAKE_ACCOUNT_ID environment variables.
|
||||
"""
|
||||
import litellm
|
||||
|
||||
# Skip if credentials not available
|
||||
if not os.getenv("SNOWFLAKE_JWT") or not os.getenv("SNOWFLAKE_ACCOUNT_ID"):
|
||||
pytest.skip("Snowflake credentials not available")
|
||||
|
||||
litellm.drop_params = False # We now support tools!
|
||||
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"name": "get_weather",
|
||||
"description": "Get the current weather in a given location",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"location": {
|
||||
"type": "string",
|
||||
"description": "The city and state, e.g. San Francisco, CA",
|
||||
}
|
||||
},
|
||||
"required": ["location"],
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
try:
|
||||
# Test with tool_choice to force tool use
|
||||
response = responses(
|
||||
model="snowflake/claude-3-5-sonnet",
|
||||
input="What's the weather in Paris?",
|
||||
tools=tools,
|
||||
tool_choice={"type": "function", "function": {"name": "get_weather"}},
|
||||
max_output_tokens=200,
|
||||
)
|
||||
|
||||
assert response is not None
|
||||
assert hasattr(response, "output")
|
||||
assert len(response.output) > 0
|
||||
|
||||
# Verify tool call was made
|
||||
tool_call_found = False
|
||||
for item in response.output:
|
||||
if hasattr(item, "type") and item.type == "function_call":
|
||||
tool_call_found = True
|
||||
assert item.name == "get_weather"
|
||||
assert hasattr(item, "arguments")
|
||||
print(f"✅ Tool call detected: {item.name}({item.arguments})")
|
||||
break
|
||||
|
||||
assert tool_call_found, "Expected tool call but none was found"
|
||||
|
||||
except APIConnectionError as e:
|
||||
if "JWT token is invalid" in str(e):
|
||||
pytest.skip("Invalid Snowflake JWT token")
|
||||
elif "Application failed to respond" in str(e) or "502" in str(e):
|
||||
pytest.skip(f"Snowflake API unavailable: {e}")
|
||||
else:
|
||||
raise
|
||||
|
||||
@ -1,704 +0,0 @@
|
||||
# # this tests if the router is initialized correctly
|
||||
# import asyncio
|
||||
# import os
|
||||
# import sys
|
||||
# import time
|
||||
# import traceback
|
||||
|
||||
# import pytest
|
||||
|
||||
# sys.path.insert(
|
||||
# 0, os.path.abspath("../..")
|
||||
# ) # Adds the parent directory to the system path
|
||||
# from collections import defaultdict
|
||||
# from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
# from dotenv import load_dotenv
|
||||
|
||||
# import litellm
|
||||
# from litellm import Router
|
||||
|
||||
# load_dotenv()
|
||||
|
||||
# # every time we load the router we should have 4 clients:
|
||||
# # Async
|
||||
# # Sync
|
||||
# # Async + Stream
|
||||
# # Sync + Stream
|
||||
|
||||
|
||||
# def test_init_clients():
|
||||
# litellm.set_verbose = True
|
||||
# import logging
|
||||
|
||||
# from litellm._logging import verbose_router_logger
|
||||
|
||||
# verbose_router_logger.setLevel(logging.DEBUG)
|
||||
# try:
|
||||
# print("testing init 4 clients with diff timeouts")
|
||||
# model_list = [
|
||||
# {
|
||||
# "model_name": "gpt-3.5-turbo",
|
||||
# "litellm_params": {
|
||||
# "model": "azure/gpt-4.1-mini",
|
||||
# "api_key": os.getenv("AZURE_AI_API_KEY"),
|
||||
# "api_version": os.getenv("AZURE_API_VERSION"),
|
||||
# "api_base": os.getenv("AZURE_AI_API_BASE"),
|
||||
# "timeout": 0.01,
|
||||
# "stream_timeout": 0.000_001,
|
||||
# "max_retries": 7,
|
||||
# },
|
||||
# },
|
||||
# ]
|
||||
# router = Router(model_list=model_list, set_verbose=True)
|
||||
# for elem in router.model_list:
|
||||
# model_id = elem["model_info"]["id"]
|
||||
# assert router.cache.get_cache(f"{model_id}_client") is not None
|
||||
# assert router.cache.get_cache(f"{model_id}_async_client") is not None
|
||||
# assert router.cache.get_cache(f"{model_id}_stream_client") is not None
|
||||
# assert router.cache.get_cache(f"{model_id}_stream_async_client") is not None
|
||||
|
||||
# # check if timeout for stream/non stream clients is set correctly
|
||||
# async_client = router.cache.get_cache(f"{model_id}_async_client")
|
||||
# stream_async_client = router.cache.get_cache(
|
||||
# f"{model_id}_stream_async_client"
|
||||
# )
|
||||
|
||||
# assert async_client.timeout == 0.01
|
||||
# assert stream_async_client.timeout == 0.000_001
|
||||
# print(vars(async_client))
|
||||
# print()
|
||||
# print(async_client._base_url)
|
||||
# assert (
|
||||
# async_client._base_url
|
||||
# == "https://openai-gpt-4-test-v-1.openai.azure.com/openai/"
|
||||
# )
|
||||
# assert (
|
||||
# stream_async_client._base_url
|
||||
# == "https://openai-gpt-4-test-v-1.openai.azure.com/openai/"
|
||||
# )
|
||||
|
||||
# print("PASSED !")
|
||||
|
||||
# except Exception as e:
|
||||
# traceback.print_exc()
|
||||
# pytest.fail(f"Error occurred: {e}")
|
||||
|
||||
|
||||
# # test_init_clients()
|
||||
|
||||
|
||||
# def test_init_clients_basic():
|
||||
# litellm.set_verbose = True
|
||||
# try:
|
||||
# print("Test basic client init")
|
||||
# model_list = [
|
||||
# {
|
||||
# "model_name": "gpt-3.5-turbo",
|
||||
# "litellm_params": {
|
||||
# "model": "azure/gpt-4.1-mini",
|
||||
# "api_key": os.getenv("AZURE_AI_API_KEY"),
|
||||
# "api_version": os.getenv("AZURE_API_VERSION"),
|
||||
# "api_base": os.getenv("AZURE_AI_API_BASE"),
|
||||
# },
|
||||
# },
|
||||
# ]
|
||||
# router = Router(model_list=model_list)
|
||||
# for elem in router.model_list:
|
||||
# model_id = elem["model_info"]["id"]
|
||||
# assert router.cache.get_cache(f"{model_id}_client") is not None
|
||||
# assert router.cache.get_cache(f"{model_id}_async_client") is not None
|
||||
# assert router.cache.get_cache(f"{model_id}_stream_client") is not None
|
||||
# assert router.cache.get_cache(f"{model_id}_stream_async_client") is not None
|
||||
# print("PASSED !")
|
||||
|
||||
# # see if we can init clients without timeout or max retries set
|
||||
# except Exception as e:
|
||||
# traceback.print_exc()
|
||||
# pytest.fail(f"Error occurred: {e}")
|
||||
|
||||
|
||||
# # test_init_clients_basic()
|
||||
|
||||
|
||||
# def test_init_clients_basic_azure_cloudflare():
|
||||
# # init azure + cloudflare
|
||||
# # init OpenAI gpt-3.5
|
||||
# # init OpenAI text-embedding
|
||||
# # init OpenAI comptaible - Mistral/mistral-medium
|
||||
# # init OpenAI compatible - xinference/bge
|
||||
# litellm.set_verbose = True
|
||||
# try:
|
||||
# print("Test basic client init")
|
||||
# model_list = [
|
||||
# {
|
||||
# "model_name": "azure-cloudflare",
|
||||
# "litellm_params": {
|
||||
# "model": "azure/gpt-4.1-mini",
|
||||
# "api_key": os.getenv("AZURE_AI_API_KEY"),
|
||||
# "api_version": os.getenv("AZURE_API_VERSION"),
|
||||
# "api_base": "https://gateway.ai.cloudflare.com/v1/0399b10e77ac6668c80404a5ff49eb37/litellm-test/azure-openai/openai-gpt-4-test-v-1",
|
||||
# },
|
||||
# },
|
||||
# {
|
||||
# "model_name": "gpt-openai",
|
||||
# "litellm_params": {
|
||||
# "model": "gpt-3.5-turbo",
|
||||
# "api_key": os.getenv("OPENAI_API_KEY"),
|
||||
# },
|
||||
# },
|
||||
# {
|
||||
# "model_name": "text-embedding-ada-002",
|
||||
# "litellm_params": {
|
||||
# "model": "text-embedding-ada-002",
|
||||
# "api_key": os.getenv("OPENAI_API_KEY"),
|
||||
# },
|
||||
# },
|
||||
# {
|
||||
# "model_name": "mistral",
|
||||
# "litellm_params": {
|
||||
# "model": "mistral/mistral-tiny",
|
||||
# "api_key": os.getenv("MISTRAL_API_KEY"),
|
||||
# },
|
||||
# },
|
||||
# {
|
||||
# "model_name": "bge-base-en",
|
||||
# "litellm_params": {
|
||||
# "model": "xinference/bge-base-en",
|
||||
# "api_base": "http://127.0.0.1:9997/v1",
|
||||
# "api_key": os.getenv("OPENAI_API_KEY"),
|
||||
# },
|
||||
# },
|
||||
# ]
|
||||
# router = Router(model_list=model_list)
|
||||
# for elem in router.model_list:
|
||||
# model_id = elem["model_info"]["id"]
|
||||
# assert router.cache.get_cache(f"{model_id}_client") is not None
|
||||
# assert router.cache.get_cache(f"{model_id}_async_client") is not None
|
||||
# assert router.cache.get_cache(f"{model_id}_stream_client") is not None
|
||||
# assert router.cache.get_cache(f"{model_id}_stream_async_client") is not None
|
||||
# print("PASSED !")
|
||||
|
||||
# # see if we can init clients without timeout or max retries set
|
||||
# except Exception as e:
|
||||
# traceback.print_exc()
|
||||
# pytest.fail(f"Error occurred: {e}")
|
||||
|
||||
|
||||
# # test_init_clients_basic_azure_cloudflare()
|
||||
|
||||
|
||||
# def test_timeouts_router():
|
||||
# """
|
||||
# Test the timeouts of the router with multiple clients. This HASas to raise a timeout error
|
||||
# """
|
||||
# import openai
|
||||
|
||||
# litellm.set_verbose = True
|
||||
# try:
|
||||
# print("testing init 4 clients with diff timeouts")
|
||||
# model_list = [
|
||||
# {
|
||||
# "model_name": "gpt-3.5-turbo",
|
||||
# "litellm_params": {
|
||||
# "model": "azure/gpt-4.1-mini",
|
||||
# "api_key": os.getenv("AZURE_AI_API_KEY"),
|
||||
# "api_version": os.getenv("AZURE_API_VERSION"),
|
||||
# "api_base": os.getenv("AZURE_AI_API_BASE"),
|
||||
# "timeout": 0.000001,
|
||||
# "stream_timeout": 0.000_001,
|
||||
# },
|
||||
# },
|
||||
# ]
|
||||
# router = Router(model_list=model_list, num_retries=0)
|
||||
|
||||
# print("PASSED !")
|
||||
|
||||
# async def test():
|
||||
# try:
|
||||
# await router.acompletion(
|
||||
# model="gpt-3.5-turbo",
|
||||
# messages=[
|
||||
# {"role": "user", "content": "hello, write a 20 pg essay"}
|
||||
# ],
|
||||
# )
|
||||
# except Exception as e:
|
||||
# raise e
|
||||
|
||||
# asyncio.run(test())
|
||||
# except openai.APITimeoutError as e:
|
||||
# print(
|
||||
# "Passed: Raised correct exception. Got openai.APITimeoutError\nGood Job", e
|
||||
# )
|
||||
# print(type(e))
|
||||
# pass
|
||||
# except Exception as e:
|
||||
# pytest.fail(
|
||||
# f"Did not raise error `openai.APITimeoutError`. Instead raised error type: {type(e)}, Error: {e}"
|
||||
# )
|
||||
|
||||
|
||||
# # test_timeouts_router()
|
||||
|
||||
|
||||
# def test_stream_timeouts_router():
|
||||
# """
|
||||
# Test the stream timeouts router. See if it selected the correct client with stream timeout
|
||||
# """
|
||||
# import openai
|
||||
|
||||
# litellm.set_verbose = True
|
||||
# try:
|
||||
# print("testing init 4 clients with diff timeouts")
|
||||
# model_list = [
|
||||
# {
|
||||
# "model_name": "gpt-3.5-turbo",
|
||||
# "litellm_params": {
|
||||
# "model": "azure/gpt-4.1-mini",
|
||||
# "api_key": os.getenv("AZURE_AI_API_KEY"),
|
||||
# "api_version": os.getenv("AZURE_API_VERSION"),
|
||||
# "api_base": os.getenv("AZURE_AI_API_BASE"),
|
||||
# "timeout": 200, # regular calls will not timeout, stream calls will
|
||||
# "stream_timeout": 10,
|
||||
# },
|
||||
# },
|
||||
# ]
|
||||
# router = Router(model_list=model_list)
|
||||
|
||||
# print("PASSED !")
|
||||
# data = {
|
||||
# "model": "gpt-3.5-turbo",
|
||||
# "messages": [{"role": "user", "content": "hello, write a 20 pg essay"}],
|
||||
# "stream": True,
|
||||
# }
|
||||
# selected_client = router._get_client(
|
||||
# deployment=router.model_list[0],
|
||||
# kwargs=data,
|
||||
# client_type=None,
|
||||
# )
|
||||
# print("Select client timeout", selected_client.timeout)
|
||||
# assert selected_client.timeout == 10
|
||||
|
||||
# # make actual call
|
||||
# response = router.completion(**data)
|
||||
|
||||
# for chunk in response:
|
||||
# print(f"chunk: {chunk}")
|
||||
# except openai.APITimeoutError as e:
|
||||
# print(
|
||||
# "Passed: Raised correct exception. Got openai.APITimeoutError\nGood Job", e
|
||||
# )
|
||||
# print(type(e))
|
||||
# pass
|
||||
# except Exception as e:
|
||||
# pytest.fail(
|
||||
# f"Did not raise error `openai.APITimeoutError`. Instead raised error type: {type(e)}, Error: {e}"
|
||||
# )
|
||||
|
||||
|
||||
# # test_stream_timeouts_router()
|
||||
|
||||
|
||||
# def test_xinference_embedding():
|
||||
# # [Test Init Xinference] this tests if we init xinference on the router correctly
|
||||
# # [Test Exception Mapping] tests that xinference is an openai comptiable provider
|
||||
# print("Testing init xinference")
|
||||
# print(
|
||||
# "this tests if we create an OpenAI client for Xinference, with the correct API BASE"
|
||||
# )
|
||||
|
||||
# model_list = [
|
||||
# {
|
||||
# "model_name": "xinference",
|
||||
# "litellm_params": {
|
||||
# "model": "xinference/bge-base-en",
|
||||
# "api_base": "os.environ/XINFERENCE_API_BASE",
|
||||
# },
|
||||
# }
|
||||
# ]
|
||||
|
||||
# router = Router(model_list=model_list)
|
||||
|
||||
# print(router.model_list)
|
||||
# print(router.model_list[0])
|
||||
|
||||
# assert (
|
||||
# router.model_list[0]["litellm_params"]["api_base"] == "http://0.0.0.0:9997"
|
||||
# ) # set in env
|
||||
|
||||
# openai_client = router._get_client(
|
||||
# deployment=router.model_list[0],
|
||||
# kwargs={"input": ["hello"], "model": "xinference"},
|
||||
# )
|
||||
|
||||
# assert openai_client._base_url == "http://0.0.0.0:9997"
|
||||
# assert "xinference" in litellm.openai_compatible_providers
|
||||
# print("passed")
|
||||
|
||||
|
||||
# # test_xinference_embedding()
|
||||
|
||||
|
||||
# def test_router_init_gpt_4_vision_enhancements():
|
||||
# try:
|
||||
# # tests base_url set when any base_url with /openai/deployments passed to router
|
||||
# print("Testing Azure GPT_Vision enhancements")
|
||||
|
||||
# model_list = [
|
||||
# {
|
||||
# "model_name": "gpt-4-vision-enhancements",
|
||||
# "litellm_params": {
|
||||
# "model": "azure/gpt-4-vision",
|
||||
# "api_key": os.getenv("AZURE_AI_API_KEY"),
|
||||
# "base_url": "https://gpt-4-vision-resource.openai.azure.com/openai/deployments/gpt-4-vision/extensions/",
|
||||
# "dataSources": [
|
||||
# {
|
||||
# "type": "AzureComputerVision",
|
||||
# "parameters": {
|
||||
# "endpoint": "os.environ/AZURE_VISION_ENHANCE_ENDPOINT",
|
||||
# "key": "os.environ/AZURE_VISION_ENHANCE_KEY",
|
||||
# },
|
||||
# }
|
||||
# ],
|
||||
# },
|
||||
# }
|
||||
# ]
|
||||
|
||||
# router = Router(model_list=model_list)
|
||||
|
||||
# print(router.model_list)
|
||||
# print(router.model_list[0])
|
||||
|
||||
# assert (
|
||||
# router.model_list[0]["litellm_params"]["base_url"]
|
||||
# == "https://gpt-4-vision-resource.openai.azure.com/openai/deployments/gpt-4-vision/extensions/"
|
||||
# ) # set in env
|
||||
|
||||
# assert (
|
||||
# router.model_list[0]["litellm_params"]["dataSources"][0]["parameters"][
|
||||
# "endpoint"
|
||||
# ]
|
||||
# == os.environ["AZURE_VISION_ENHANCE_ENDPOINT"]
|
||||
# )
|
||||
|
||||
# assert (
|
||||
# router.model_list[0]["litellm_params"]["dataSources"][0]["parameters"][
|
||||
# "key"
|
||||
# ]
|
||||
# == os.environ["AZURE_VISION_ENHANCE_KEY"]
|
||||
# )
|
||||
|
||||
# azure_client = router._get_client(
|
||||
# deployment=router.model_list[0],
|
||||
# kwargs={"stream": True, "model": "gpt-4-vision-enhancements"},
|
||||
# client_type="async",
|
||||
# )
|
||||
|
||||
# assert (
|
||||
# azure_client._base_url
|
||||
# == "https://gpt-4-vision-resource.openai.azure.com/openai/deployments/gpt-4-vision/extensions/"
|
||||
# )
|
||||
# print("passed")
|
||||
# except Exception as e:
|
||||
# pytest.fail(f"Error occurred: {e}")
|
||||
|
||||
|
||||
# @pytest.mark.parametrize("sync_mode", [True, False])
|
||||
# @pytest.mark.asyncio
|
||||
# async def test_openai_with_organization(sync_mode):
|
||||
# try:
|
||||
# print("Testing OpenAI with organization")
|
||||
# model_list = [
|
||||
# {
|
||||
# "model_name": "openai-bad-org",
|
||||
# "litellm_params": {
|
||||
# "model": "gpt-3.5-turbo",
|
||||
# "organization": "org-ikDc4ex8NB",
|
||||
# },
|
||||
# },
|
||||
# {
|
||||
# "model_name": "openai-good-org",
|
||||
# "litellm_params": {"model": "gpt-3.5-turbo"},
|
||||
# },
|
||||
# ]
|
||||
|
||||
# router = Router(model_list=model_list)
|
||||
|
||||
# print(router.model_list)
|
||||
# print(router.model_list[0])
|
||||
|
||||
# if sync_mode:
|
||||
# openai_client = router._get_client(
|
||||
# deployment=router.model_list[0],
|
||||
# kwargs={"input": ["hello"], "model": "openai-bad-org"},
|
||||
# )
|
||||
# print(vars(openai_client))
|
||||
|
||||
# assert openai_client.organization == "org-ikDc4ex8NB"
|
||||
|
||||
# # bad org raises error
|
||||
|
||||
# try:
|
||||
# response = router.completion(
|
||||
# model="openai-bad-org",
|
||||
# messages=[{"role": "user", "content": "this is a test"}],
|
||||
# )
|
||||
# pytest.fail(
|
||||
# "Request should have failed - This organization does not exist"
|
||||
# )
|
||||
# except Exception as e:
|
||||
# print("Got exception: " + str(e))
|
||||
# assert "header should match organization for API key" in str(
|
||||
# e
|
||||
# ) or "No such organization" in str(e)
|
||||
|
||||
# # good org works
|
||||
# response = router.completion(
|
||||
# model="openai-good-org",
|
||||
# messages=[{"role": "user", "content": "this is a test"}],
|
||||
# max_tokens=5,
|
||||
# )
|
||||
# else:
|
||||
# openai_client = router._get_client(
|
||||
# deployment=router.model_list[0],
|
||||
# kwargs={"input": ["hello"], "model": "openai-bad-org"},
|
||||
# client_type="async",
|
||||
# )
|
||||
# print(vars(openai_client))
|
||||
|
||||
# assert openai_client.organization == "org-ikDc4ex8NB"
|
||||
|
||||
# # bad org raises error
|
||||
|
||||
# try:
|
||||
# response = await router.acompletion(
|
||||
# model="openai-bad-org",
|
||||
# messages=[{"role": "user", "content": "this is a test"}],
|
||||
# )
|
||||
# pytest.fail(
|
||||
# "Request should have failed - This organization does not exist"
|
||||
# )
|
||||
# except Exception as e:
|
||||
# print("Got exception: " + str(e))
|
||||
# assert "header should match organization for API key" in str(
|
||||
# e
|
||||
# ) or "No such organization" in str(e)
|
||||
|
||||
# # good org works
|
||||
# response = await router.acompletion(
|
||||
# model="openai-good-org",
|
||||
# messages=[{"role": "user", "content": "this is a test"}],
|
||||
# max_tokens=5,
|
||||
# )
|
||||
|
||||
# except Exception as e:
|
||||
# pytest.fail(f"Error occurred: {e}")
|
||||
|
||||
|
||||
# def test_init_clients_azure_command_r_plus():
|
||||
# # This tests that the router uses the OpenAI client for Azure/Command-R+
|
||||
# # For azure/command-r-plus we need to use openai.OpenAI because of how the Azure provider requires requests being sent
|
||||
# litellm.set_verbose = True
|
||||
# import logging
|
||||
|
||||
# from litellm._logging import verbose_router_logger
|
||||
|
||||
# verbose_router_logger.setLevel(logging.DEBUG)
|
||||
# try:
|
||||
# print("testing init 4 clients with diff timeouts")
|
||||
# model_list = [
|
||||
# {
|
||||
# "model_name": "gpt-3.5-turbo",
|
||||
# "litellm_params": {
|
||||
# "model": "azure/command-r-plus",
|
||||
# "api_key": os.getenv("AZURE_COHERE_API_KEY"),
|
||||
# "api_base": os.getenv("AZURE_COHERE_API_BASE"),
|
||||
# "timeout": 0.01,
|
||||
# "stream_timeout": 0.000_001,
|
||||
# "max_retries": 7,
|
||||
# },
|
||||
# },
|
||||
# ]
|
||||
# router = Router(model_list=model_list, set_verbose=True)
|
||||
# for elem in router.model_list:
|
||||
# model_id = elem["model_info"]["id"]
|
||||
# async_client = router.cache.get_cache(f"{model_id}_async_client")
|
||||
# stream_async_client = router.cache.get_cache(
|
||||
# f"{model_id}_stream_async_client"
|
||||
# )
|
||||
# # Assert the Async Clients used are OpenAI clients and not Azure
|
||||
# # For using Azure/Command-R-Plus and Azure/Mistral the clients NEED to be OpenAI clients used
|
||||
# # this is weirdness introduced on Azure's side
|
||||
|
||||
# assert "openai.AsyncOpenAI" in str(async_client)
|
||||
# assert "openai.AsyncOpenAI" in str(stream_async_client)
|
||||
# print("PASSED !")
|
||||
|
||||
# except Exception as e:
|
||||
# traceback.print_exc()
|
||||
# pytest.fail(f"Error occurred: {e}")
|
||||
|
||||
|
||||
# @pytest.mark.asyncio
|
||||
# async def test_aaaaatext_completion_with_organization():
|
||||
# try:
|
||||
# print("Testing Text OpenAI with organization")
|
||||
# model_list = [
|
||||
# {
|
||||
# "model_name": "openai-bad-org",
|
||||
# "litellm_params": {
|
||||
# "model": "text-completion-openai/gpt-3.5-turbo-instruct",
|
||||
# "api_key": os.getenv("OPENAI_API_KEY", None),
|
||||
# "organization": "org-ikDc4ex8NB",
|
||||
# },
|
||||
# },
|
||||
# {
|
||||
# "model_name": "openai-good-org",
|
||||
# "litellm_params": {
|
||||
# "model": "text-completion-openai/gpt-3.5-turbo-instruct",
|
||||
# "api_key": os.getenv("OPENAI_API_KEY", None),
|
||||
# "organization": os.getenv("OPENAI_ORGANIZATION", None),
|
||||
# },
|
||||
# },
|
||||
# ]
|
||||
|
||||
# router = Router(model_list=model_list)
|
||||
|
||||
# print(router.model_list)
|
||||
# print(router.model_list[0])
|
||||
|
||||
# openai_client = router._get_client(
|
||||
# deployment=router.model_list[0],
|
||||
# kwargs={"input": ["hello"], "model": "openai-bad-org"},
|
||||
# )
|
||||
# print(vars(openai_client))
|
||||
|
||||
# assert openai_client.organization == "org-ikDc4ex8NB"
|
||||
|
||||
# # bad org raises error
|
||||
|
||||
# try:
|
||||
# response = await router.atext_completion(
|
||||
# model="openai-bad-org",
|
||||
# prompt="this is a test",
|
||||
# )
|
||||
# pytest.fail("Request should have failed - This organization does not exist")
|
||||
# except Exception as e:
|
||||
# print("Got exception: " + str(e))
|
||||
# assert "header should match organization for API key" in str(
|
||||
# e
|
||||
# ) or "No such organization" in str(e)
|
||||
|
||||
# # good org works
|
||||
# response = await router.atext_completion(
|
||||
# model="openai-good-org",
|
||||
# prompt="this is a test",
|
||||
# max_tokens=5,
|
||||
# )
|
||||
# print("working response: ", response)
|
||||
|
||||
# except Exception as e:
|
||||
# pytest.fail(f"Error occurred: {e}")
|
||||
|
||||
|
||||
# def test_init_clients_async_mode():
|
||||
# litellm.set_verbose = True
|
||||
# import logging
|
||||
|
||||
# from litellm._logging import verbose_router_logger
|
||||
# from litellm.types.router import RouterGeneralSettings
|
||||
|
||||
# verbose_router_logger.setLevel(logging.DEBUG)
|
||||
# try:
|
||||
# print("testing init 4 clients with diff timeouts")
|
||||
# model_list = [
|
||||
# {
|
||||
# "model_name": "gpt-3.5-turbo",
|
||||
# "litellm_params": {
|
||||
# "model": "azure/gpt-4.1-mini",
|
||||
# "api_key": os.getenv("AZURE_AI_API_KEY"),
|
||||
# "api_version": os.getenv("AZURE_API_VERSION"),
|
||||
# "api_base": os.getenv("AZURE_AI_API_BASE"),
|
||||
# "timeout": 0.01,
|
||||
# "stream_timeout": 0.000_001,
|
||||
# "max_retries": 7,
|
||||
# },
|
||||
# },
|
||||
# ]
|
||||
# router = Router(
|
||||
# model_list=model_list,
|
||||
# set_verbose=True,
|
||||
# router_general_settings=RouterGeneralSettings(async_only_mode=True),
|
||||
# )
|
||||
# for elem in router.model_list:
|
||||
# model_id = elem["model_info"]["id"]
|
||||
|
||||
# # sync clients not initialized in async_only_mode=True
|
||||
# assert router.cache.get_cache(f"{model_id}_client") is None
|
||||
# assert router.cache.get_cache(f"{model_id}_stream_client") is None
|
||||
|
||||
# # only async clients initialized in async_only_mode=True
|
||||
# assert router.cache.get_cache(f"{model_id}_async_client") is not None
|
||||
# assert router.cache.get_cache(f"{model_id}_stream_async_client") is not None
|
||||
# except Exception as e:
|
||||
# pytest.fail(f"Error occurred: {e}")
|
||||
|
||||
|
||||
# @pytest.mark.parametrize(
|
||||
# "environment,expected_models",
|
||||
# [
|
||||
# ("development", ["gpt-3.5-turbo"]),
|
||||
# ("production", ["gpt-4", "gpt-3.5-turbo", "gpt-4o"]),
|
||||
# ],
|
||||
# )
|
||||
# def test_init_router_with_supported_environments(environment, expected_models):
|
||||
# """
|
||||
# Tests that the correct models are setup on router when LITELLM_ENVIRONMENT is set
|
||||
# """
|
||||
# os.environ["LITELLM_ENVIRONMENT"] = environment
|
||||
# model_list = [
|
||||
# {
|
||||
# "model_name": "gpt-3.5-turbo",
|
||||
# "litellm_params": {
|
||||
# "model": "azure/gpt-4.1-mini",
|
||||
# "api_key": os.getenv("AZURE_AI_API_KEY"),
|
||||
# "api_version": os.getenv("AZURE_API_VERSION"),
|
||||
# "api_base": os.getenv("AZURE_AI_API_BASE"),
|
||||
# "timeout": 0.01,
|
||||
# "stream_timeout": 0.000_001,
|
||||
# "max_retries": 7,
|
||||
# },
|
||||
# "model_info": {"supported_environments": ["development", "production"]},
|
||||
# },
|
||||
# {
|
||||
# "model_name": "gpt-4",
|
||||
# "litellm_params": {
|
||||
# "model": "openai/gpt-4",
|
||||
# "api_key": os.getenv("OPENAI_API_KEY"),
|
||||
# "timeout": 0.01,
|
||||
# "stream_timeout": 0.000_001,
|
||||
# "max_retries": 7,
|
||||
# },
|
||||
# "model_info": {"supported_environments": ["production"]},
|
||||
# },
|
||||
# {
|
||||
# "model_name": "gpt-4o",
|
||||
# "litellm_params": {
|
||||
# "model": "openai/gpt-4o",
|
||||
# "api_key": os.getenv("OPENAI_API_KEY"),
|
||||
# "timeout": 0.01,
|
||||
# "stream_timeout": 0.000_001,
|
||||
# "max_retries": 7,
|
||||
# },
|
||||
# "model_info": {"supported_environments": ["production"]},
|
||||
# },
|
||||
# ]
|
||||
# router = Router(model_list=model_list, set_verbose=True)
|
||||
# _model_list = router.get_model_names()
|
||||
|
||||
# print("model_list: ", _model_list)
|
||||
# print("expected_models: ", expected_models)
|
||||
|
||||
# assert set(_model_list) == set(expected_models)
|
||||
|
||||
# os.environ.pop("LITELLM_ENVIRONMENT")
|
||||
@ -74,15 +74,13 @@ async def test_basic_s3_logging(sync_mode, streaming):
|
||||
s3.delete_object(Bucket="load-testing-oct", Key=key)
|
||||
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"streaming", [(True)]
|
||||
)
|
||||
@pytest.mark.parametrize("streaming", [True])
|
||||
@pytest.mark.flaky(retries=3, delay=1)
|
||||
async def test_basic_s3_v2_logging(streaming):
|
||||
from blockbuster import BlockBuster
|
||||
from litellm.integrations.s3_v2 import S3Logger
|
||||
|
||||
s3_v2_logger = S3Logger(s3_flush_interval=1)
|
||||
litellm.callbacks = [s3_v2_logger]
|
||||
blockbuster = BlockBuster()
|
||||
@ -120,7 +118,7 @@ async def test_basic_s3_v2_logging(streaming):
|
||||
|
||||
print(f"all_s3_keys: {all_s3_keys}")
|
||||
|
||||
#assert that atlest one key has response.id in it
|
||||
# assert that atlest one key has response.id in it
|
||||
assert any(response_id in key for key in all_s3_keys)
|
||||
s3 = boto3.client("s3")
|
||||
# delete all objects
|
||||
@ -134,22 +132,22 @@ async def test_basic_s3_v2_logging_failure():
|
||||
"""Test that S3 v2 logger makes httpx PUT request when logging failures"""
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
from litellm.integrations.s3_v2 import S3Logger
|
||||
|
||||
|
||||
# Create S3 logger with short flush interval
|
||||
s3_v2_logger = S3Logger(s3_flush_interval=1)
|
||||
|
||||
|
||||
# Mock the httpx client to capture the PUT request
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
|
||||
|
||||
s3_v2_logger.async_httpx_client = AsyncMock()
|
||||
s3_v2_logger.async_httpx_client.put.return_value = mock_response
|
||||
|
||||
|
||||
# Track the upload method calls
|
||||
original_upload = s3_v2_logger.async_upload_data_to_s3
|
||||
upload_called = False
|
||||
|
||||
|
||||
async def mock_upload(batch_logging_element):
|
||||
nonlocal upload_called
|
||||
upload_called = True
|
||||
@ -157,12 +155,12 @@ async def test_basic_s3_v2_logging_failure():
|
||||
url = f"https://test-bucket.s3.us-west-2.amazonaws.com/{batch_logging_element.s3_object_key}"
|
||||
headers = {"Content-Type": "application/json"}
|
||||
data = '{"model": "gpt-4o-mini"}'
|
||||
|
||||
|
||||
# Make the actual httpx call we want to test
|
||||
await s3_v2_logger.async_httpx_client.put(url=url, headers=headers, data=data)
|
||||
|
||||
|
||||
s3_v2_logger.async_upload_data_to_s3 = mock_upload
|
||||
|
||||
|
||||
# Configure S3 callback params
|
||||
litellm.callbacks = [s3_v2_logger]
|
||||
litellm.s3_callback_params = {
|
||||
@ -172,7 +170,7 @@ async def test_basic_s3_v2_logging_failure():
|
||||
"s3_region_name": "us-west-2",
|
||||
}
|
||||
litellm.set_verbose = True
|
||||
|
||||
|
||||
# Trigger a failure by using invalid API key
|
||||
try:
|
||||
response = await litellm.acompletion(
|
||||
@ -182,33 +180,33 @@ async def test_basic_s3_v2_logging_failure():
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"Expected error: {e}")
|
||||
|
||||
|
||||
# Wait for logger to process the failure
|
||||
await asyncio.sleep(5)
|
||||
|
||||
|
||||
# Verify that our mock upload was called
|
||||
assert upload_called, "S3 upload method was not called"
|
||||
print("✓ S3 upload method was called")
|
||||
|
||||
|
||||
# Verify that httpx PUT was called
|
||||
s3_v2_logger.async_httpx_client.put.assert_called()
|
||||
|
||||
|
||||
# Get the call arguments to verify the S3 URL
|
||||
call_args = s3_v2_logger.async_httpx_client.put.call_args
|
||||
assert call_args is not None
|
||||
url = call_args[1]['url'] if 'url' in call_args[1] else call_args[0][0]
|
||||
|
||||
url = call_args[1]["url"] if "url" in call_args[1] else call_args[0][0]
|
||||
|
||||
# Verify the URL contains expected S3 endpoint
|
||||
assert "test-bucket.s3.us-west-2.amazonaws.com" in url
|
||||
print(f"✓ S3 PUT request made to: {url}")
|
||||
|
||||
|
||||
# Verify headers include expected content type
|
||||
headers = call_args[1]['headers']
|
||||
assert headers['Content-Type'] == 'application/json'
|
||||
headers = call_args[1]["headers"]
|
||||
assert headers["Content-Type"] == "application/json"
|
||||
print("✓ S3 request headers are correct")
|
||||
|
||||
|
||||
# Verify JSON data was included
|
||||
data = call_args[1]['data']
|
||||
data = call_args[1]["data"]
|
||||
assert data is not None
|
||||
assert '"model": "gpt-4o-mini"' in data
|
||||
print("✓ S3 request data contains expected log payload")
|
||||
@ -411,83 +409,19 @@ async def make_async_calls():
|
||||
return total_time
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="flaky test on ci/cd")
|
||||
def test_s3_logging_r2():
|
||||
# all s3 requests need to be in one test function
|
||||
# since we are modifying stdout, and pytests runs tests in parallel
|
||||
# on circle ci - we only test litellm.acompletion()
|
||||
try:
|
||||
# redirect stdout to log_file
|
||||
# litellm.cache = litellm.Cache(
|
||||
# type="s3", s3_bucket_name="litellm-r2-bucket", s3_region_name="us-west-2"
|
||||
# )
|
||||
litellm.set_verbose = True
|
||||
from litellm._logging import verbose_logger
|
||||
import logging
|
||||
|
||||
verbose_logger.setLevel(level=logging.DEBUG)
|
||||
|
||||
litellm.success_callback = ["s3"]
|
||||
litellm.s3_callback_params = {
|
||||
"s3_bucket_name": "litellm-r2-bucket",
|
||||
"s3_aws_secret_access_key": "os.environ/R2_S3_ACCESS_KEY",
|
||||
"s3_aws_access_key_id": "os.environ/R2_S3_ACCESS_ID",
|
||||
"s3_endpoint_url": "os.environ/R2_S3_URL",
|
||||
"s3_region_name": "os.environ/R2_S3_REGION_NAME",
|
||||
}
|
||||
print("Testing async s3 logging")
|
||||
|
||||
expected_keys = []
|
||||
|
||||
import time
|
||||
|
||||
curr_time = str(time.time())
|
||||
|
||||
async def _test():
|
||||
return await litellm.acompletion(
|
||||
model="gpt-3.5-turbo",
|
||||
messages=[{"role": "user", "content": f"This is a test {curr_time}"}],
|
||||
max_tokens=10,
|
||||
temperature=0.7,
|
||||
user="ishaan-2",
|
||||
)
|
||||
|
||||
response = asyncio.run(_test())
|
||||
print(f"response: {response}")
|
||||
expected_keys.append(response.id)
|
||||
|
||||
import boto3
|
||||
|
||||
s3 = boto3.client(
|
||||
"s3",
|
||||
endpoint_url=os.getenv("R2_S3_URL"),
|
||||
region_name=os.getenv("R2_S3_REGION_NAME"),
|
||||
aws_access_key_id=os.getenv("R2_S3_ACCESS_ID"),
|
||||
aws_secret_access_key=os.getenv("R2_S3_ACCESS_KEY"),
|
||||
)
|
||||
|
||||
bucket_name = "litellm-r2-bucket"
|
||||
# List objects in the bucket
|
||||
response = s3.list_objects(Bucket=bucket_name)
|
||||
|
||||
except Exception as e:
|
||||
pytest.fail(f"An exception occurred - {e}")
|
||||
finally:
|
||||
# post, close log file and verify
|
||||
# Reset stdout to the original value
|
||||
print("Passed! Testing async s3 logging")
|
||||
|
||||
from litellm.integrations.s3_v2 import S3Logger
|
||||
|
||||
|
||||
class TestS3Logger(S3Logger):
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.recorded_requests = {}
|
||||
self.logged_standard_logging_payload: Optional[StandardLoggingPayload] = None
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
|
||||
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||||
self.recorded_requests[response_obj["id"]] = start_time
|
||||
print("recorded request", self.recorded_requests)
|
||||
self.logged_standard_logging_payload = kwargs["standard_logging_object"]
|
||||
return await super().async_log_success_event(kwargs, response_obj, start_time, end_time)
|
||||
|
||||
return await super().async_log_success_event(
|
||||
kwargs, response_obj, start_time, end_time
|
||||
)
|
||||
|
||||
@ -59,137 +59,3 @@ async def chat_completion(session, key, model="azure-gpt-3.5", request_metadata=
|
||||
|
||||
if status != 200:
|
||||
raise Exception(f"Request did not return a 200 status code: {status}")
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="flaky test - covered by simpler unit testing.")
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.flaky(retries=12, delay=2)
|
||||
async def test_aaateam_logging():
|
||||
"""
|
||||
-> Team 1 logs to project 1
|
||||
-> Create Key
|
||||
-> Make chat/completions call
|
||||
-> Fetch logs from langfuse
|
||||
"""
|
||||
try:
|
||||
async with aiohttp.ClientSession() as session:
|
||||
|
||||
key = await generate_key(
|
||||
session, models=["fake-openai-endpoint"], team_id="team-1"
|
||||
) # team-1 logs to project 1
|
||||
|
||||
from litellm._uuid import uuid
|
||||
|
||||
_trace_id = f"trace-{uuid.uuid4()}"
|
||||
_request_metadata = {
|
||||
"trace_id": _trace_id,
|
||||
}
|
||||
|
||||
await chat_completion(
|
||||
session,
|
||||
key["key"],
|
||||
model="fake-openai-endpoint",
|
||||
request_metadata=_request_metadata,
|
||||
)
|
||||
|
||||
# Test - if the logs were sent to the correct team on langfuse
|
||||
import langfuse
|
||||
|
||||
print(f"langfuse_public_key: {os.getenv('LANGFUSE_PROJECT1_PUBLIC')}")
|
||||
print(f"langfuse_secret_key: {os.getenv('LANGFUSE_HOST')}")
|
||||
langfuse_client = langfuse.Langfuse(
|
||||
public_key=os.getenv("LANGFUSE_PROJECT1_PUBLIC"),
|
||||
secret_key=os.getenv("LANGFUSE_PROJECT1_SECRET"),
|
||||
host="https://us.cloud.langfuse.com",
|
||||
)
|
||||
|
||||
await asyncio.sleep(30)
|
||||
|
||||
print(f"searching for trace_id={_trace_id} on langfuse")
|
||||
|
||||
generations = langfuse_client.get_generations(trace_id=_trace_id).data
|
||||
print(generations)
|
||||
assert len(generations) == 1
|
||||
except Exception as e:
|
||||
pytest.fail(f"Unexpected error: {str(e)}")
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="todo fix langfuse credential error")
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_2logging():
|
||||
"""
|
||||
-> Team 1 logs to project 2
|
||||
-> Create Key
|
||||
-> Make chat/completions call
|
||||
-> Fetch logs from langfuse
|
||||
"""
|
||||
langfuse_public_key = os.getenv("LANGFUSE_PROJECT2_PUBLIC")
|
||||
|
||||
print(f"langfuse_public_key: {langfuse_public_key}")
|
||||
langfuse_secret_key = os.getenv("LANGFUSE_PROJECT2_SECRET")
|
||||
print(f"langfuse_secret_key: {langfuse_secret_key}")
|
||||
langfuse_host = "https://us.cloud.langfuse.com"
|
||||
|
||||
try:
|
||||
assert langfuse_public_key is not None
|
||||
assert langfuse_secret_key is not None
|
||||
except Exception as e:
|
||||
# skip test if langfuse credentials are not set
|
||||
return
|
||||
|
||||
try:
|
||||
async with aiohttp.ClientSession() as session:
|
||||
|
||||
key = await generate_key(
|
||||
session, models=["fake-openai-endpoint"], team_id="team-2"
|
||||
) # team-1 logs to project 1
|
||||
|
||||
from litellm._uuid import uuid
|
||||
|
||||
_trace_id = f"trace-{uuid.uuid4()}"
|
||||
_request_metadata = {
|
||||
"trace_id": _trace_id,
|
||||
}
|
||||
|
||||
await chat_completion(
|
||||
session,
|
||||
key["key"],
|
||||
model="fake-openai-endpoint",
|
||||
request_metadata=_request_metadata,
|
||||
)
|
||||
|
||||
# Test - if the logs were sent to the correct team on langfuse
|
||||
import langfuse
|
||||
|
||||
langfuse_client = langfuse.Langfuse(
|
||||
public_key=langfuse_public_key,
|
||||
secret_key=langfuse_secret_key,
|
||||
host=langfuse_host,
|
||||
)
|
||||
|
||||
await asyncio.sleep(30)
|
||||
|
||||
print(f"searching for trace_id={_trace_id} on langfuse")
|
||||
|
||||
generations = langfuse_client.get_generations(trace_id=_trace_id).data
|
||||
print("Team 2 generations", generations)
|
||||
|
||||
# team-2 should have 1 generation with this trace id
|
||||
assert len(generations) == 1
|
||||
|
||||
# team-1 should have 0 generations with this trace id
|
||||
langfuse_client_1 = langfuse.Langfuse(
|
||||
public_key=os.getenv("LANGFUSE_PROJECT1_PUBLIC"),
|
||||
secret_key=os.getenv("LANGFUSE_PROJECT1_SECRET"),
|
||||
host="https://us.cloud.langfuse.com",
|
||||
)
|
||||
|
||||
generations_team_1 = langfuse_client_1.get_generations(
|
||||
trace_id=_trace_id
|
||||
).data
|
||||
print("Team 1 generations", generations_team_1)
|
||||
|
||||
assert len(generations_team_1) == 0
|
||||
|
||||
except Exception as e:
|
||||
pytest.fail("Team 2 logging failed: " + str(e))
|
||||
|
||||
Loading…
Reference in New Issue
Block a user