push

Author: hassiebp

langfuse/_client/attributes.py

@@ -12,7 +12,7 @@
 
 import json
 from datetime import datetime
-from typing import Any, Dict, List, Literal, Optional
+from typing import Any, Dict, List, Literal, Optional, Union
 
 from langfuse._utils.serializer import EventSerializer
 from langfuse.model import PromptClient
@@ -164,14 +164,14 @@ def _flatten_and_serialize_metadata(
         else LangfuseOtelSpanAttributes.TRACE_METADATA
     )
 
-    metadata_attributes = {}
+    metadata_attributes: Dict[str, Union[str, int, None]] = {}
 
     if not isinstance(metadata, dict):
         metadata_attributes[prefix] = _serialize(metadata)
     else:
         for key, value in metadata.items():
             metadata_attributes[f"{prefix}.{key}"] = (
-                str(value)
+                value
                 if isinstance(value, str) or isinstance(value, int)
                 else _serialize(value)
             )

langfuse/_client/client.py

@@ -147,6 +147,7 @@ class Langfuse:
 
     _resources: Optional[LangfuseResourceManager] = None
     _mask: Optional[MaskFunction] = None
+    _otel_tracer: otel_trace_api.Tracer
 
     def __init__(
         self,
@@ -167,10 +168,14 @@ def __init__(
         mask: Optional[MaskFunction] = None,
         blocked_instrumentation_scopes: Optional[List[str]] = None,
         additional_headers: Optional[Dict[str, str]] = None,
-        tracer_provider: Optional[otel_trace_api.TracerProvider] = None,
+        tracer_provider: Optional[TracerProvider] = None,
     ):
-        self._host = host or os.environ.get(LANGFUSE_HOST, "https://cloud.langfuse.com")
-        self._environment = environment or os.environ.get(LANGFUSE_TRACING_ENVIRONMENT)
+        self._host = host or cast(
+            str, os.environ.get(LANGFUSE_HOST, "https://cloud.langfuse.com")
+        )
+        self._environment = environment or cast(
+            str, os.environ.get(LANGFUSE_TRACING_ENVIRONMENT)
+        )
         self._project_id: Optional[str] = None
         sample_rate = sample_rate or float(os.environ.get(LANGFUSE_SAMPLE_RATE, 1.0))
         if not 0.0 <= sample_rate <= 1.0:
@@ -218,7 +223,7 @@ def __init__(
         self._resources = LangfuseResourceManager(
             public_key=public_key,
             secret_key=secret_key,
-            host=self._host or "https://cloud.langfuse.com",
+            host=self._host,
             timeout=timeout,
             environment=environment,
             release=release,
@@ -231,15 +236,14 @@ def __init__(
             tracing_enabled=self._tracing_enabled,
             blocked_instrumentation_scopes=blocked_instrumentation_scopes,
             additional_headers=additional_headers,
-            tracer_provider=cast(Optional[TracerProvider], tracer_provider),
+            tracer_provider=tracer_provider,
         )
         self._mask = self._resources.mask
 
-        self._otel_tracer = cast(
-            Any,
+        self._otel_tracer = (
             self._resources.tracer
             if self._tracing_enabled and self._resources.tracer is not None
-            else otel_trace_api.NoOpTracer(),
+            else otel_trace_api.NoOpTracer()
         )
         self.api = self._resources.api
         self.async_api = self._resources.async_api
@@ -660,8 +664,8 @@ def _create_span_with_parent_context(
         self,
         *,
         name: str,
-        parent: Any,
-        remote_parent_span: Any,
+        parent: Optional[otel_trace_api.Span] = None,
+        remote_parent_span: Optional[otel_trace_api.Span] = None,
         as_type: Literal["generation", "span"],
         end_on_exit: Optional[bool] = None,
         input: Optional[Any] = None,

langfuse/_client/datasets.py

@@ -1,5 +1,6 @@
 import datetime as dt
 import logging
+from .span import LangfuseSpan
 from typing import TYPE_CHECKING, Any, Generator, List, Optional
 
 from opentelemetry.util._decorator import _agnosticcontextmanager
@@ -91,7 +92,7 @@ def run(
         run_name: str,
         run_metadata: Optional[Any] = None,
         run_description: Optional[str] = None,
-    ) -> Generator[Any, None, None]:
+    ) -> Generator[LangfuseSpan, None, None]:
         """Create a context manager for the dataset item run that links the execution to a Langfuse trace.
 
         This method is a context manager that creates a trace for the dataset run and yields a span

langfuse/_client/observe.py

@@ -207,9 +207,11 @@ def _async_observe(
         transform_to_string: Optional[Callable[[Iterable], str]] = None,
     ) -> F:
         @wraps(func)
-        async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
-            trace_id = kwargs.pop("langfuse_trace_id", None)
-            parent_observation_id = kwargs.pop("langfuse_parent_observation_id", None)
+        async def async_wrapper(*args: Tuple[Any], **kwargs: Dict[str, Any]) -> Any:
+            trace_id = cast(str, kwargs.pop("langfuse_trace_id", None))
+            parent_observation_id = cast(
+                str, kwargs.pop("langfuse_parent_observation_id", None)
+            )
             trace_context: Optional[TraceContext] = (
                 {
                     "trace_id": trace_id,
@@ -228,7 +230,7 @@ async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
                 if capture_input
                 else None
             )
-            public_key = kwargs.pop("langfuse_public_key", None)
+            public_key = cast(str, kwargs.pop("langfuse_public_key", None))
             langfuse_client = get_client(public_key=public_key)
             context_manager: Optional[
                 Union[

langfuse/_client/span.py

@@ -1071,7 +1071,7 @@ def update(
         usage_details: Optional[Dict[str, int]] = None,
         cost_details: Optional[Dict[str, float]] = None,
         prompt: Optional[PromptClient] = None,
-        **kwargs: Any,
+        **kwargs: Dict[str, Any],
     ) -> "LangfuseGeneration":
         """Update this generation span with new information.
 

langfuse/_task_manager/media_manager.py

@@ -170,18 +170,18 @@ def _process_media(
         trace_id: str,
         observation_id: Optional[str],
         field: str,
-    ) -> str:
+    ) -> None:
         if (
             media._content_length is None
             or media._content_type is None
             or media._content_sha256_hash is None
             or media._content_bytes is None
         ):
-            return ""
+            return
 
         if media._media_id is None:
             self._log.error("Media ID is None. Skipping upload.")
-            return ""
+            return
 
         try:
             upload_media_job = UploadMediaJob(
@@ -212,9 +212,6 @@ def _process_media(
             self._log.error(
                 f"Media processing error: Failed to process media_id={media._media_id} for trace_id={trace_id}. Error: {str(e)}"
             )
-            return ""
-
-        return media._reference_string or ""
 
     def _process_upload_media_job(
         self,

langfuse/_utils/parse_error.py

@@ -63,7 +63,7 @@ def generate_error_message_fern(error: Error) -> str:
             else defaultErrorResponse
         )
 
-    # This line is unreachable but kept for completeness
+    return defaultErrorResponse  # type: ignore
 
 
 def handle_fern_exception(exception: Error) -> None:

langfuse/model.py

@@ -2,7 +2,7 @@
 
 import re
 from abc import ABC, abstractmethod
-from typing import Any, Dict, List, Literal, Optional, Tuple, TypedDict, Union
+from typing import Any, Dict, List, Literal, Optional, Sequence, Tuple, TypedDict, Union
 from langfuse.logger import langfuse_logger
 
 from langfuse.api.resources.commons.types.dataset import (
@@ -161,7 +161,12 @@ def __init__(self, prompt: Prompt, is_fallback: bool = False):
         self.is_fallback = is_fallback
 
     @abstractmethod
-    def compile(self, **kwargs: Any) -> Union[str, List[Dict[str, Any]]]:
+    def compile(
+        self, **kwargs: Union[str, Any]
+    ) -> Union[
+        str,
+        Sequence[Union[ChatMessageDict, ChatMessageWithPlaceholdersDict_Placeholder]],
+    ]:
         pass
 
     @property
@@ -255,7 +260,7 @@ def __init__(self, prompt: Prompt_Text, is_fallback: bool = False):
         super().__init__(prompt, is_fallback)
         self.prompt = prompt.prompt
 
-    def compile(self, **kwargs: Any) -> str:
+    def compile(self, **kwargs: Union[str, Any]) -> str:
         return TemplateParser.compile_template(self.prompt, kwargs)
 
     @property
@@ -274,7 +279,7 @@ def __eq__(self, other: object) -> bool:
 
         return False
 
-    def get_langchain_prompt(self, **kwargs: Any) -> str:
+    def get_langchain_prompt(self, **kwargs: Union[str, Any]) -> str:
         """Convert Langfuse prompt into string compatible with Langchain PromptTemplate.
 
         This method adapts the mustache-style double curly braces {{variable}} used in Langfuse
@@ -314,12 +319,15 @@ def __init__(self, prompt: Prompt_Chat, is_fallback: bool = False):
                 self.prompt.append(
                     ChatMessageWithPlaceholdersDict_Message(
                         type="message",
-                        role=p.role,
-                        content=p.content,
+                        role=p.role,  # type: ignore
+                        content=p.content,  # type: ignore
                     ),
                 )
 
-    def compile(self, **kwargs: Any) -> List[Dict[str, Any]]:
+    def compile(
+        self,
+        **kwargs: Union[str, Any],
+    ) -> Sequence[Union[ChatMessageDict, ChatMessageWithPlaceholdersDict_Placeholder]]:
         """Compile the prompt with placeholders and variables.
 
         Args:
@@ -358,9 +366,9 @@ def compile(self, **kwargs: Any) -> List[Dict[str, Any]]:
                             ):
                                 compiled_messages.append(
                                     {
-                                        "role": msg["role"],
+                                        "role": msg["role"],  # type: ignore
                                         "content": TemplateParser.compile_template(
-                                            msg["content"],
+                                            msg["content"],  # type: ignore
                                             kwargs,
                                         ),
                                     },
@@ -392,7 +400,7 @@ def compile(self, **kwargs: Any) -> List[Dict[str, Any]]:
             unresolved_placeholders_message = f"Placeholders {unresolved_placeholders} have not been resolved. Pass them as keyword arguments to compile()."
             langfuse_logger.warning(unresolved_placeholders_message)
 
-        return compiled_messages
+        return compiled_messages  # type: ignore
 
     @property
     def variables(self) -> List[str]:
@@ -434,7 +442,9 @@ def __eq__(self, other: object) -> bool:
 
         return False
 
-    def get_langchain_prompt(self, **kwargs: Any) -> List[Union[Tuple[str, str], Any]]:
+    def get_langchain_prompt(
+        self, **kwargs: Union[str, Any]
+    ) -> List[Union[Tuple[str, str], Any]]:
         """Convert Langfuse prompt into string compatible with Langchain ChatPromptTemplate.
 
         It specifically adapts the mustache-style double curly braces {{variable}} used in Langfuse
@@ -455,9 +465,9 @@ def get_langchain_prompt(self, **kwargs: Any) -> List[Union[Tuple[str, str], Any
         langchain_messages: List[Union[Tuple[str, str], Any]] = []
 
         for msg in compiled_messages:
-            if isinstance(msg, dict) and "type" in msg and msg["type"] == "placeholder":
+            if isinstance(msg, dict) and "type" in msg and msg["type"] == "placeholder":  # type: ignore
                 # unresolved placeholder -> add LC MessagesPlaceholder
-                placeholder_name = msg["name"]
+                placeholder_name = msg["name"]  # type: ignore
                 try:
                     from langchain_core.prompts.chat import MessagesPlaceholder  # noqa: PLC0415, I001
 
@@ -471,8 +481,8 @@ def get_langchain_prompt(self, **kwargs: Any) -> List[Union[Tuple[str, str], Any
                 if isinstance(msg, dict) and "role" in msg and "content" in msg:
                     langchain_messages.append(
                         (
-                            msg["role"],
-                            self._get_langchain_prompt_string(msg["content"]),
+                            msg["role"],  # type: ignore
+                            self._get_langchain_prompt_string(msg["content"]),  # type: ignore
                         ),
                     )