fix langsmith logging for streaming

This commit is contained in:
Ishaan Jaff 2024-07-17 16:04:45 -07:00
parent d3ee7a947c
commit 1abd66db1b
2 changed files with 21 additions and 2 deletions

View file

@ -5,7 +5,7 @@ import os
import traceback import traceback
import types import types
from datetime import datetime from datetime import datetime
from typing import Any, List, Optional from typing import Any, List, Optional, Union
import dotenv # type: ignore import dotenv # type: ignore
import requests # type: ignore import requests # type: ignore
@ -29,7 +29,7 @@ class LangsmithInputs(BaseModel):
custom_llm_provider: Optional[str] = None custom_llm_provider: Optional[str] = None
input: Optional[List[Any]] = None input: Optional[List[Any]] = None
log_event_type: Optional[str] = None log_event_type: Optional[str] = None
original_response: Optional[str] = None original_response: Optional[Any] = None
response_cost: Optional[float] = None response_cost: Optional[float] = None
# LiteLLM Virtual Key specific fields # LiteLLM Virtual Key specific fields
@ -96,10 +96,14 @@ class LangsmithLogger(CustomLogger):
value = kwargs[key] value = kwargs[key]
if key == "start_time" or key == "end_time" or value is None: if key == "start_time" or key == "end_time" or value is None:
pass pass
elif key == "original_response" and not isinstance(value, str):
new_kwargs[key] = str(value)
elif type(value) == datetime.datetime: elif type(value) == datetime.datetime:
new_kwargs[key] = value.isoformat() new_kwargs[key] = value.isoformat()
elif type(value) != dict and is_serializable(value=value): elif type(value) != dict and is_serializable(value=value):
new_kwargs[key] = value new_kwargs[key] = value
elif not is_serializable(value=value):
continue
if isinstance(response_obj, BaseModel): if isinstance(response_obj, BaseModel):
try: try:
@ -122,6 +126,11 @@ class LangsmithLogger(CustomLogger):
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
try: try:
verbose_logger.debug(
"Langsmith Async Layer Logging - kwargs: %s, response_obj: %s",
kwargs,
response_obj,
)
data = self._prepare_log_data(kwargs, response_obj, start_time, end_time) data = self._prepare_log_data(kwargs, response_obj, start_time, end_time)
url = f"{self.langsmith_base_url}/runs" url = f"{self.langsmith_base_url}/runs"
verbose_logger.debug(f"Langsmith Logging - About to send data to {url} ...") verbose_logger.debug(f"Langsmith Logging - About to send data to {url} ...")
@ -147,6 +156,11 @@ class LangsmithLogger(CustomLogger):
def log_success_event(self, kwargs, response_obj, start_time, end_time): def log_success_event(self, kwargs, response_obj, start_time, end_time):
try: try:
verbose_logger.debug(
"Langsmith Sync Layer Logging - kwargs: %s, response_obj: %s",
kwargs,
response_obj,
)
data = self._prepare_log_data(kwargs, response_obj, start_time, end_time) data = self._prepare_log_data(kwargs, response_obj, start_time, end_time)
url = f"{self.langsmith_base_url}/runs" url = f"{self.langsmith_base_url}/runs"
verbose_logger.debug(f"Langsmith Logging - About to send data to {url} ...") verbose_logger.debug(f"Langsmith Logging - About to send data to {url} ...")

View file

@ -427,6 +427,11 @@ def function_setup(
isinstance(cb, type(callback_class)) for cb in litellm.callbacks isinstance(cb, type(callback_class)) for cb in litellm.callbacks
): ):
litellm.callbacks.append(callback_class) # type: ignore litellm.callbacks.append(callback_class) # type: ignore
litellm.input_callback.append(callback_class) # type: ignore
litellm.success_callback.append(callback_class) # type: ignore
litellm.failure_callback.append(callback_class) # type: ignore
litellm._async_success_callback.append(callback_class) # type: ignore
litellm._async_failure_callback.append(callback_class) # type: ignore
# Pop the async items from success_callback in reverse order to avoid index issues # Pop the async items from success_callback in reverse order to avoid index issues
for index in reversed(removed_async_items): for index in reversed(removed_async_items):