push

Author: hassiebp

langfuse/_client/client.py

@@ -1304,22 +1304,20 @@ def create_score(
         score_id = score_id or self._create_observation_id()
 
         try:
-            score_event = {
-                "id": score_id,
-                "session_id": session_id,
-                "dataset_run_id": dataset_run_id,
-                "trace_id": trace_id,
-                "observation_id": observation_id,
-                "name": name,
-                "value": value,
-                "data_type": data_type,
-                "comment": comment,
-                "config_id": config_id,
-                "environment": self._environment,
-                "metadata": metadata,
-            }
-
-            new_body = ScoreBody(**score_event)
+            new_body = ScoreBody(
+                id=score_id,
+                session_id=session_id,
+                dataset_run_id=dataset_run_id,
+                trace_id=trace_id,
+                observation_id=observation_id,
+                name=name,
+                value=value,
+                data_type=data_type,
+                comment=comment,
+                config_id=config_id,
+                environment=self._environment,
+                metadata=metadata,
+            )
 
             event = {
                 "id": self.create_trace_id(),
@@ -1960,7 +1958,7 @@ def get_prompt(
                         f"Returning fallback prompt for '{cache_key}' due to fetch error: {e}"
                     )
 
-                    fallback_client_args = {
+                    fallback_client_args: Dict[str, Any] = {
                         "name": name,
                         "prompt": fallback,
                         "type": type,

langfuse/_client/resource_manager.py

@@ -18,7 +18,7 @@
 import os
 import threading
 from queue import Full, Queue
-from typing import Dict, Optional, cast
+from typing import Dict, Optional, cast, Any
 
 import httpx
 from opentelemetry import trace as otel_trace_api
@@ -148,7 +148,7 @@ def _initialize_instance(
         )
         tracer_provider.add_span_processor(langfuse_processor)
 
-        tracer_provider = otel_trace_api.get_tracer_provider()
+        tracer_provider = cast(TracerProvider, otel_trace_api.get_tracer_provider())
         self._otel_tracer = tracer_provider.get_tracer(
             LANGFUSE_TRACER_NAME,
             langfuse_version,
@@ -195,7 +195,7 @@ def _initialize_instance(
             LANGFUSE_MEDIA_UPLOAD_ENABLED, "True"
         ).lower() not in ("false", "0")
 
-        self._media_upload_queue = Queue(100_000)
+        self._media_upload_queue: Queue[Any] = Queue(100_000)
         self._media_manager = MediaManager(
             api_client=self.api,
             media_upload_queue=self._media_upload_queue,
@@ -220,7 +220,7 @@ def _initialize_instance(
         self.prompt_cache = PromptCache()
 
         # Score ingestion
-        self._score_ingestion_queue = Queue(100_000)
+        self._score_ingestion_queue: Queue[Any] = Queue(100_000)
         self._ingestion_consumers = []
 
         ingestion_consumer = ScoreIngestionConsumer(

langfuse/_utils/serializer.py

@@ -21,14 +21,16 @@
 try:
     from langchain.load.serializable import Serializable
 except ImportError:
-    # If Serializable is not available, set it to NoneType
-    Serializable = type(None)
+    # If Serializable is not available, set it to a placeholder type
+    class Serializable:  # type: ignore
+        pass
+
 
 # Attempt to import numpy
 try:
     import numpy as np
 except ImportError:
-    np = None
+    np = None  # type: ignore
 
 logger = getLogger(__name__)
 

langfuse/langchain/CallbackHandler.py

@@ -60,8 +60,8 @@ def __init__(self, *, public_key: Optional[str] = None) -> None:
         self.client = get_client(public_key=public_key)
 
         self.runs: Dict[UUID, Union[LangfuseSpan, LangfuseGeneration]] = {}
-        self.prompt_to_parent_run_map = {}
-        self.updated_completion_start_time_memo = set()
+        self.prompt_to_parent_run_map: Dict[UUID, Any] = {}
+        self.updated_completion_start_time_memo: Set[UUID] = set()
 
     def on_llm_new_token(
         self,
@@ -166,19 +166,26 @@ def on_chain_start(
                 run_id=run_id, parent_run_id=parent_run_id, metadata=metadata
             )
 
-            content = {
-                "name": self.get_langchain_run_name(serialized, **kwargs),
-                "metadata": self.__join_tags_and_metadata(tags, metadata),
-                "input": inputs,
-                "level": "DEBUG" if tags and LANGSMITH_TAG_HIDDEN in tags else None,
-            }
+            span_name = self.get_langchain_run_name(serialized, **kwargs)
+            span_metadata = self.__join_tags_and_metadata(tags, metadata)
+            span_level = "DEBUG" if tags and LANGSMITH_TAG_HIDDEN in tags else None
 
             if parent_run_id is None:
-                self.runs[run_id] = self.client.start_span(**content)
+                self.runs[run_id] = self.client.start_span(
+                    name=span_name,
+                    metadata=span_metadata,
+                    input=inputs,
+                    level=span_level,
+                )
             else:
                 self.runs[run_id] = cast(
                     LangfuseSpan, self.runs[parent_run_id]
-                ).start_span(**content)
+                ).start_span(
+                    name=span_name,
+                    metadata=span_metadata,
+                    input=inputs,
+                    level=span_level,
+                )
 
         except Exception as e:
             langfuse_logger.exception(e)
@@ -431,23 +438,25 @@ def on_retriever_start(
             self._log_debug_event(
                 "on_retriever_start", run_id, parent_run_id, query=query
             )
+            span_name = self.get_langchain_run_name(serialized, **kwargs)
+            span_metadata = self.__join_tags_and_metadata(tags, metadata)
+            span_level = "DEBUG" if tags and LANGSMITH_TAG_HIDDEN in tags else None
+
             if parent_run_id is None:
-                content = {
-                    "name": self.get_langchain_run_name(serialized, **kwargs),
-                    "metadata": self.__join_tags_and_metadata(tags, metadata),
-                    "input": query,
-                    "level": "DEBUG" if tags and LANGSMITH_TAG_HIDDEN in tags else None,
-                }
-
-                self.runs[run_id] = self.client.start_span(**content)
+                self.runs[run_id] = self.client.start_span(
+                    name=span_name,
+                    metadata=span_metadata,
+                    input=query,
+                    level=span_level,
+                )
             else:
                 self.runs[run_id] = cast(
                     LangfuseSpan, self.runs[parent_run_id]
                 ).start_span(
-                    name=self.get_langchain_run_name(serialized, **kwargs),
+                    name=span_name,
                     input=query,
-                    metadata=self.__join_tags_and_metadata(tags, metadata),
-                    level="DEBUG" if tags and LANGSMITH_TAG_HIDDEN in tags else None,
+                    metadata=span_metadata,
+                    level=span_level,
                 )
 
         except Exception as e:

langfuse/langchain/utils.py

@@ -1,7 +1,7 @@
 """@private"""
 
 import re
-from typing import Any, Dict, List, Literal, Optional
+from typing import Any, Dict, List, Literal, Optional, cast
 
 # NOTE ON DEPENDENCIES:
 # - since Jan 2024, there is https://pypi.org/project/langchain-openai/ which is a separate package and imports openai models.
@@ -12,7 +12,7 @@
 def _extract_model_name(
     serialized: Optional[Dict[str, Any]],
     **kwargs: Any,
-):
+) -> Optional[str]:
     """Extracts the model name from the serialized or kwargs object. This is used to get the model names for Langfuse."""
     # In this function we return on the first match, so the order of operations is important
 
@@ -39,39 +39,54 @@ def _extract_model_name(
 
     for model_name, keys, select_from in models_by_id:
         model = _extract_model_by_path_for_id(
-            model_name, serialized, kwargs, keys, select_from
+            model_name,
+            serialized,
+            kwargs,
+            keys,
+            cast(Literal["serialized", "kwargs"], select_from),
         )
         if model:
             return model
 
     # Second, we match AzureOpenAI as we need to extract the model name, fdeployment version and deployment name
-    if serialized.get("id")[-1] == "AzureOpenAI":
-        if kwargs.get("invocation_params").get("model"):
-            return kwargs.get("invocation_params").get("model")
-
-        if kwargs.get("invocation_params").get("model_name"):
-            return kwargs.get("invocation_params").get("model_name")
-
-        deployment_name = None
-        deployment_version = None
-
-        if serialized.get("kwargs").get("openai_api_version"):
-            deployment_version = serialized.get("kwargs").get("deployment_version")
-
-        if serialized.get("kwargs").get("deployment_name"):
-            deployment_name = serialized.get("kwargs").get("deployment_name")
-
-        if not isinstance(deployment_name, str):
-            return None
-
-        if not isinstance(deployment_version, str):
-            return deployment_name
-
-        return (
-            deployment_name + "-" + deployment_version
-            if deployment_version not in deployment_name
-            else deployment_name
-        )
+    if serialized:
+        serialized_id = serialized.get("id")
+        if (
+            serialized_id
+            and isinstance(serialized_id, list)
+            and len(serialized_id) > 0
+            and serialized_id[-1] == "AzureOpenAI"
+        ):
+            invocation_params = kwargs.get("invocation_params")
+            if invocation_params and isinstance(invocation_params, dict):
+                if invocation_params.get("model"):
+                    return str(invocation_params.get("model"))
+
+                if invocation_params.get("model_name"):
+                    return str(invocation_params.get("model_name"))
+
+            deployment_name = None
+            deployment_version = None
+
+            serialized_kwargs = serialized.get("kwargs")
+            if serialized_kwargs and isinstance(serialized_kwargs, dict):
+                if serialized_kwargs.get("openai_api_version"):
+                    deployment_version = serialized_kwargs.get("deployment_version")
+
+                if serialized_kwargs.get("deployment_name"):
+                    deployment_name = serialized_kwargs.get("deployment_name")
+
+            if not isinstance(deployment_name, str):
+                return None
+
+            if not isinstance(deployment_version, str):
+                return deployment_name
+
+            return (
+                deployment_name + "-" + deployment_version
+                if deployment_version not in deployment_name
+                else deployment_name
+            )
 
     # Third, for some models, we are unable to extract the model by a path in an object. Langfuse provides us with a string representation of the model pbjects
     # We use regex to extract the model from the repr string
@@ -111,7 +126,9 @@ def _extract_model_name(
     ]
     for select in ["kwargs", "serialized"]:
         for path in random_paths:
-            model = _extract_model_by_path(serialized, kwargs, path, select)
+            model = _extract_model_by_path(
+                serialized, kwargs, path, cast(Literal["serialized", "kwargs"], select)
+            )
             if model:
                 return model
 
@@ -123,13 +140,20 @@ def _extract_model_from_repr_by_pattern(
     serialized: Optional[Dict[str, Any]],
     pattern: str,
     default: Optional[str] = None,
-):
+) -> Optional[str]:
     if serialized is None:
         return None
 
-    if serialized.get("id")[-1] == id:
-        if serialized.get("repr"):
-            extracted = _extract_model_with_regex(pattern, serialized.get("repr"))
+    serialized_id = serialized.get("id")
+    if (
+        serialized_id
+        and isinstance(serialized_id, list)
+        and len(serialized_id) > 0
+        and serialized_id[-1] == id
+    ):
+        repr_str = serialized.get("repr")
+        if repr_str and isinstance(repr_str, str):
+            extracted = _extract_model_with_regex(pattern, repr_str)
             return extracted if extracted else default if default else None
 
     return None
@@ -145,15 +169,24 @@ def _extract_model_with_regex(pattern: str, text: str):
 def _extract_model_by_path_for_id(
     id: str,
     serialized: Optional[Dict[str, Any]],
-    kwargs: dict,
+    kwargs: Dict[str, Any],
     keys: List[str],
     select_from: Literal["serialized", "kwargs"],
-):
+) -> Optional[str]:
     if serialized is None and select_from == "serialized":
         return None
 
-    if serialized.get("id")[-1] == id:
-        return _extract_model_by_path(serialized, kwargs, keys, select_from)
+    if serialized:
+        serialized_id = serialized.get("id")
+        if (
+            serialized_id
+            and isinstance(serialized_id, list)
+            and len(serialized_id) > 0
+            and serialized_id[-1] == id
+        ):
+            return _extract_model_by_path(serialized, kwargs, keys, select_from)
+
+    return None
 
 
 def _extract_model_by_path(
@@ -168,7 +201,10 @@ def _extract_model_by_path(
     current_obj = kwargs if select_from == "kwargs" else serialized
 
     for key in keys:
-        current_obj = current_obj.get(key)
+        if current_obj and isinstance(current_obj, dict):
+            current_obj = current_obj.get(key)
+        else:
+            return None
         if not current_obj:
             return None