fix use AsyncHTTPHandler for logging langsmith requests

litellm/integrations/langsmith.py

@@ -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,84 +57,102 @@ 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):
+    def _prepare_log_data(
-        # Method definition
+        self, kwargs, response_obj, start_time, end_time, print_verbose
-        # inspired by Langsmith http api here: https://github.com/langchain-ai/langsmith-cookbook/blob/main/tracing-examples/rest/rest.ipynb
+    ):
-        metadata = (
+        import datetime
-            kwargs.get("litellm_params", {}).get("metadata", {}) or {}
+        from datetime import timezone
-        )  # if metadata is None
+
+        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(
+            start_time = kwargs["start_time"].astimezone(timezone.utc).isoformat()
-                f"Langsmith Logging - Enters logging function for model {kwargs}"
+            end_time = kwargs["end_time"].astimezone(timezone.utc).isoformat()
-            )
+        except:
-            import datetime
+            start_time = datetime.datetime.utcnow().isoformat()
-            from datetime import timezone
+            end_time = datetime.datetime.utcnow().isoformat()
 
-            import requests
+        # filter out kwargs to not include any dicts, langsmith throws an erros when trying to log kwargs
+        logged_kwargs = LangsmithInputs(**kwargs)
+        kwargs = logged_kwargs.model_dump()
 
+        new_kwargs = {}
+        for key in kwargs:
+            value = kwargs[key]
+            if key == "start_time" or key == "end_time" or value is None:
+                pass
+            elif type(value) == datetime.datetime:
+                new_kwargs[key] = value.isoformat()
+            elif type(value) != dict and is_serializable(value=value):
+                new_kwargs[key] = value
+
+        if isinstance(response_obj, BaseModel):
             try:
-                start_time = kwargs["start_time"].astimezone(timezone.utc).isoformat()
+                response_obj = response_obj.model_dump()
-                end_time = kwargs["end_time"].astimezone(timezone.utc).isoformat()
             except:
-                start_time = datetime.datetime.utcnow().isoformat()
+                response_obj = response_obj.dict()  # type: ignore
-                end_time = datetime.datetime.utcnow().isoformat()
 
-            # filter out kwargs to not include any dicts, langsmith throws an erros when trying to log kwargs
+        data = {
-            logged_kwargs = LangsmithInputs(**kwargs)
+            "name": run_name,
-            kwargs = logged_kwargs.model_dump()
+            "run_type": "llm",  # this should always be llm, since litellm always logs llm calls. Langsmith allow us to log "chain"
+            "inputs": new_kwargs,
+            "outputs": response_obj,
+            "session_name": project_name,
+            "start_time": start_time,
+            "end_time": end_time,
+            "id": run_id,
+        }
 
-            new_kwargs = {}
+        return data
-            for key in kwargs:
-                value = kwargs[key]
-                if key == "start_time" or key == "end_time" or value is None:
-                    pass
-                elif type(value) == datetime.datetime:
-                    new_kwargs[key] = value.isoformat()
-                elif type(value) != dict and is_serializable(value=value):
-                    new_kwargs[key] = value
 
-            if isinstance(response_obj, BaseModel):
+    async def async_log_event(
-                try:
+        self, kwargs, response_obj, start_time, end_time, print_verbose
-                    response_obj = response_obj.model_dump()
+    ):
-                except:
+        try:
-                    response_obj = response_obj.dict()  # type: ignore
+            data = self._prepare_log_data(
-
+                kwargs, response_obj, start_time, end_time, print_verbose
-            data = {
+            )
-                "name": run_name,
+            url = f"{self.langsmith_base_url}/runs"
-                "run_type": "llm",  # this should always be llm, since litellm always logs llm calls. Langsmith allow us to log "chain"
-                "inputs": new_kwargs,
-                "outputs": response_obj,
-                "session_name": project_name,
-                "start_time": start_time,
-                "end_time": end_time,
-                "id": run_id,
-            }
-
-            url = f"{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):