fix use AsyncHTTPHandler for logging langsmith requests

This commit is contained in:
Ishaan Jaff 2024-07-17 11:49:48 -07:00
parent 1cfbddeb43
commit 0e170c75de

View file

@ -11,6 +11,9 @@ import dotenv # type: ignore
import requests # type: ignore
from pydantic import BaseModel # type: ignore
from litellm.integrations.custom_logger import CustomLogger
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
class LangsmithInputs(BaseModel):
model: Optional[str] = None
@ -43,7 +46,7 @@ def is_serializable(value):
return not isinstance(value, non_serializable_types)
class LangsmithLogger:
class LangsmithLogger(CustomLogger):
# Class variables or attributes
def __init__(self):
self.langsmith_api_key = os.getenv("LANGSMITH_API_KEY")
@ -54,43 +57,28 @@ class LangsmithLogger:
self.langsmith_base_url = os.getenv(
"LANGSMITH_BASE_URL", "https://api.smith.langchain.com"
)
self.async_httpx_client = AsyncHTTPHandler()
def log_event(self, kwargs, response_obj, start_time, end_time, print_verbose):
# Method definition
# inspired by Langsmith http api here: https://github.com/langchain-ai/langsmith-cookbook/blob/main/tracing-examples/rest/rest.ipynb
metadata = (
kwargs.get("litellm_params", {}).get("metadata", {}) or {}
) # if metadata is None
def _prepare_log_data(
self, kwargs, response_obj, start_time, end_time, print_verbose
):
import datetime
from datetime import timezone
metadata = kwargs.get("litellm_params", {}).get("metadata", {}) or {}
# set user_api_key, user_team_id, user_api_key_user_id
kwargs["user_api_key"] = metadata.get("user_api_key", None)
kwargs["user_api_key_user_id"] = metadata.get("user_api_key_user_id", None)
kwargs["user_api_key_team_alias"] = metadata.get(
"user_api_key_team_alias", None
)
# set project name and run_name for langsmith logging
# users can pass project_name and run name to litellm.completion()
# Example: litellm.completion(model, messages, metadata={"project_name": "my-litellm-project", "run_name": "my-langsmith-run"})
# if not set litellm will fallback to the environment variable LANGSMITH_PROJECT, then to the default project_name = litellm-completion, run_name = LLMRun
project_name = metadata.get("project_name", self.langsmith_project)
run_name = metadata.get("run_name", self.langsmith_default_run_name)
run_id = metadata.get("id", None)
print_verbose(
f"Langsmith Logging - project_name: {project_name}, run_name {run_name}"
)
langsmith_base_url = os.getenv(
"LANGSMITH_BASE_URL", "https://api.smith.langchain.com"
)
try:
print_verbose(
f"Langsmith Logging - Enters logging function for model {kwargs}"
)
import datetime
from datetime import timezone
import requests
try:
start_time = kwargs["start_time"].astimezone(timezone.utc).isoformat()
@ -130,8 +118,41 @@ class LangsmithLogger:
"id": run_id,
}
url = f"{langsmith_base_url}/runs"
return data
async def async_log_event(
self, kwargs, response_obj, start_time, end_time, print_verbose
):
try:
data = self._prepare_log_data(
kwargs, response_obj, start_time, end_time, print_verbose
)
url = f"{self.langsmith_base_url}/runs"
print_verbose(f"Langsmith Logging - About to send data to {url} ...")
headers = {"x-api-key": self.langsmith_api_key}
response = await self.async_httpx_client.post(
url=url, json=data, headers=headers
)
if response.status_code >= 300:
print_verbose(f"Error: {response.status_code}")
else:
print_verbose("Run successfully created")
print_verbose(
f"Langsmith Layer Logging - final response object: {response_obj}. Response text from langsmith={response.text}"
)
except:
print_verbose(f"Langsmith Layer Error - {traceback.format_exc()}")
def log_event(self, kwargs, response_obj, start_time, end_time, print_verbose):
try:
data = self._prepare_log_data(
kwargs, response_obj, start_time, end_time, print_verbose
)
url = f"{self.langsmith_base_url}/runs"
print_verbose(f"Langsmith Logging - About to send data to {url} ...")
response = requests.post(
url=url,
json=data,
@ -145,10 +166,8 @@ class LangsmithLogger:
print_verbose(
f"Langsmith Layer Logging - final response object: {response_obj}. Response text from langsmith={response.text}"
)
return
except:
print_verbose(f"Langsmith Layer Error - {traceback.format_exc()}")
pass
def get_run_by_id(self, run_id):