litellm-mirror/tests/local_testing/test_braintrust.py
Krish Dholakia a187cee538
Litellm dev 01 07 2025 p3 (#7635)
* fix(__init__.py): fix mistral large tool calling

map bedrock mistral large to converse endpoint

Fixes https://github.com/BerriAI/litellm/issues/7521

* braintrust logging: respect project_id, add more metrics + more (#7613)

* braintrust logging: respect project_id, add more metrics

* braintrust logger: improve json formatting

* braintrust logger: add test for passing specific project_id

* rm unneeded import

* braintrust logging: rm unneeded var in tets

* add project_name

* update docs

---------

Co-authored-by: H <no@email.com>

---------

Co-authored-by: hi019 <65871571+hi019@users.noreply.github.com>
Co-authored-by: H <no@email.com>
2025-01-08 11:46:24 -08:00

77 lines
2 KiB
Python

# What is this?
## This tests the braintrust integration
import asyncio
import os
import random
import sys
import time
import traceback
from datetime import datetime
from dotenv import load_dotenv
from fastapi import Request
load_dotenv()
import os
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
import asyncio
import logging
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
import litellm
from litellm.llms.custom_httpx.http_handler import HTTPHandler
def test_braintrust_logging():
import litellm
litellm.set_verbose = True
http_client = HTTPHandler()
with patch.object(
litellm.integrations.braintrust_logging.global_braintrust_sync_http_handler,
"post",
new=MagicMock(),
) as mock_client:
# set braintrust as a callback, litellm will send the data to braintrust
litellm.callbacks = ["braintrust"]
# openai call
response = litellm.completion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "Hi 👋 - i'm openai"}],
)
time.sleep(2)
mock_client.assert_called()
def test_braintrust_logging_specific_project_id():
import litellm
litellm.set_verbose = True
with patch.object(
litellm.integrations.braintrust_logging.global_braintrust_sync_http_handler,
"post",
new=MagicMock(),
) as mock_client:
# set braintrust as a callback, litellm will send the data to braintrust
litellm.callbacks = ["braintrust"]
response = litellm.completion(model="openai/gpt-4o", messages=[{ "content": "Hello, how are you?","role": "user"}], metadata={"project_id": "123"})
time.sleep(2)
# Check that the log was inserted into the correct project
mock_client.assert_called()
_, kwargs = mock_client.call_args
assert 'url' in kwargs
assert kwargs['url'] == "https://api.braintrustdata.com/v1/project_logs/123/insert"