litellm/tests/local_testing/test_langsmith.py
Ishaan Jaff c3bc9e6b12
(Feat) Add langsmith key based logging (#6682)
* add langsmith_api_key to StandardCallbackDynamicParams

* create a file for langsmith types

* langsmith add key / team based logging

* add key based logging for langsmith

* fix langsmith key based logging

* fix linting langsmith

* remove NOQA violation

* add unit test coverage for all helpers in test langsmith

* test_langsmith_key_based_logging

* docs langsmith key based logging

* run langsmith tests in logging callback tests

* fix logging testing

* test_langsmith_key_based_logging

* test_add_callback_via_key_litellm_pre_call_utils_langsmith

* add debug statement langsmith key based logging

* test_langsmith_key_based_logging
2024-11-11 13:58:06 -08:00

115 lines
3.5 KiB
Python

import io
import os
import sys
sys.path.insert(0, os.path.abspath("../.."))
import asyncio
import logging
import uuid
import pytest
import litellm
from litellm import completion
from litellm._logging import verbose_logger
from litellm.integrations.langsmith import LangsmithLogger
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
verbose_logger.setLevel(logging.DEBUG)
litellm.set_verbose = True
import time
# test_langsmith_logging()
@pytest.mark.skip(reason="Flaky test. covered by unit tests on custom logger.")
def test_async_langsmith_logging_with_metadata():
try:
litellm.success_callback = ["langsmith"]
litellm.set_verbose = True
response = completion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "what llm are u"}],
max_tokens=10,
temperature=0.2,
)
print(response)
time.sleep(3)
for cb in litellm.callbacks:
if isinstance(cb, LangsmithLogger):
cb.async_httpx_client.close()
except Exception as e:
pytest.fail(f"Error occurred: {e}")
print(e)
@pytest.mark.skip(reason="Flaky test. covered by unit tests on custom logger.")
@pytest.mark.parametrize("sync_mode", [False, True])
@pytest.mark.asyncio
async def test_async_langsmith_logging_with_streaming_and_metadata(sync_mode):
try:
test_langsmith_logger = LangsmithLogger()
litellm.success_callback = ["langsmith"]
litellm.set_verbose = True
run_id = str(uuid.uuid4())
messages = [{"role": "user", "content": "what llm are u"}]
if sync_mode is True:
response = completion(
model="gpt-3.5-turbo",
messages=messages,
max_tokens=10,
temperature=0.2,
stream=True,
metadata={"id": run_id},
)
for cb in litellm.callbacks:
if isinstance(cb, LangsmithLogger):
cb.async_httpx_client = AsyncHTTPHandler()
for chunk in response:
continue
time.sleep(3)
else:
response = await litellm.acompletion(
model="gpt-3.5-turbo",
messages=messages,
max_tokens=10,
temperature=0.2,
mock_response="This is a mock request",
stream=True,
metadata={"id": run_id},
)
for cb in litellm.callbacks:
if isinstance(cb, LangsmithLogger):
cb.async_httpx_client = AsyncHTTPHandler()
async for chunk in response:
continue
await asyncio.sleep(3)
print("run_id", run_id)
logged_run_on_langsmith = test_langsmith_logger.get_run_by_id(run_id=run_id)
print("logged_run_on_langsmith", logged_run_on_langsmith)
print("fields in logged_run_on_langsmith", logged_run_on_langsmith.keys())
input_fields_on_langsmith = logged_run_on_langsmith.get("inputs")
extra_fields_on_langsmith = logged_run_on_langsmith.get("extra").get(
"invocation_params"
)
assert logged_run_on_langsmith.get("run_type") == "llm"
print("\nLogged INPUT ON LANGSMITH", input_fields_on_langsmith)
print("\nextra fields on langsmith", extra_fields_on_langsmith)
assert isinstance(input_fields_on_langsmith, dict)
except Exception as e:
pytest.fail(f"Error occurred: {e}")
print(e)