cleanup

Author: nimarb

langfuse/_client/client.py

@@ -10,7 +10,18 @@
 from datetime import datetime
 from hashlib import sha256
 from time import time_ns
-from typing import Any, Dict, List, Literal, Optional, Union, cast, overload
+from typing import (
+    Any,
+    Dict,
+    List,
+    Literal,
+    Optional,
+    Union,
+    Type,
+    cast,
+    overload,
+    get_args,
+)
 
 import backoff
 import httpx
@@ -39,6 +50,7 @@
 from langfuse._client.constants import (
     ObservationTypeLiteral,
     ObservationTypeLiteralNoEvent,
+    ObservationTypeGenerationLike,
 )
 from langfuse._client.resource_manager import LangfuseResourceManager
 from langfuse._client.span import (
@@ -656,10 +668,14 @@ def _create_observation_from_otel_span(
         LangfuseGuardrail,
     ]:
         """Create the appropriate observation type from an OTEL span."""
-        if as_type == "generation":
-            return LangfuseGeneration(
+        if as_type in get_args(ObservationTypeGenerationLike):
+            observation_class = self._get_span_class(as_type)
+            # Type ignore to prevent overloads of internal _get_span_class function,
+            # issue is that LangfuseEvent could be returned
+            return observation_class(  # type: ignore[return-value]
                 otel_span=otel_span,
                 langfuse_client=self,
+                environment=self._environment,
                 input=input,
                 output=output,
                 metadata=metadata,
@@ -673,21 +689,12 @@ def _create_observation_from_otel_span(
                 cost_details=cost_details,
                 prompt=prompt,
             )
-        elif as_type == "span":
-            return LangfuseSpan(
-                otel_span=otel_span,
-                langfuse_client=self,
-                environment=self._environment,
-                input=input,
-                output=output,
-                metadata=metadata,
-                version=version,
-                level=level,
-                status_message=status_message,
-            )
         else:
-            # For all other observation types (agent, tool, etc.), create a span and set the type
-            span = LangfuseSpan(
+            # For other types (e.g. span, guardrail), create appropriate class without generation properties
+            observation_class = self._get_span_class(as_type)
+            # Type ignore to prevent overloads of internal _get_span_class function,
+            # issue is that LangfuseEvent could be returned
+            return observation_class(  # type: ignore[return-value]
                 otel_span=otel_span,
                 langfuse_client=self,
                 environment=self._environment,
@@ -698,10 +705,9 @@ def _create_observation_from_otel_span(
                 level=level,
                 status_message=status_message,
             )
-            # Set the observation type on the span
-            span._observation_type = as_type
-            span._otel_span.set_attribute("langfuse.observation.type", as_type)
-            return span
+            # span._observation_type = as_type
+            # span._otel_span.set_attribute("langfuse.observation.type", as_type)
+            # return span
 
     def start_generation(
         self,
@@ -1075,8 +1081,8 @@ def start_as_current_observation(
             status_message: Optional status message for the observation
             end_on_exit (default: True): Whether to end the span automatically when leaving the context manager. If False, the span must be manually ended to avoid memory leaks.
 
-            # TODO: also add the generation like types here!
-            The following parameters are only available when as_type="generation":
+            The following parameters are available when as_type is: "generation", "agent",
+            "tool", "chain", "retriever", "evaluator" or "embedding".
             completion_start_time: When the model started generating the response
             model: Name/identifier of the AI model used (e.g., "gpt-4")
             model_parameters: Parameters used for the model (e.g., temperature, max_tokens)
@@ -1285,7 +1291,18 @@ def start_as_current_observation(
     def _get_span_class(
         self,
         as_type: ObservationTypeLiteral,
-    ) -> type:
+    ) -> Union[
+        Type[LangfuseAgent],
+        Type[LangfuseTool],
+        Type[LangfuseChain],
+        Type[LangfuseRetriever],
+        Type[LangfuseEvaluator],
+        Type[LangfuseEmbedding],
+        Type[LangfuseGuardrail],
+        Type[LangfuseGeneration],
+        Type[LangfuseEvent],
+        Type[LangfuseSpan],
+    ]:
         """Get the appropriate span class based on as_type."""
         normalized_type = as_type.lower()
 
@@ -1400,7 +1417,15 @@ def _start_as_current_otel_span_with_processed_media(
                 "status_message": status_message,
             }
 
-            if span_class == LangfuseGeneration:
+            if span_class in [
+                LangfuseGeneration,
+                LangfuseAgent,
+                LangfuseTool,
+                LangfuseChain,
+                LangfuseRetriever,
+                LangfuseEvaluator,
+                LangfuseEmbedding,
+            ]:
                 common_args.update(
                     {
                         "completion_start_time": completion_start_time,
@@ -1411,19 +1436,7 @@ def _start_as_current_otel_span_with_processed_media(
                         "prompt": prompt,
                     }
                 )
-            # TODO: for some, create generation-like classes with those props.
-            elif span_class in [
-                LangfuseSpan,
-                LangfuseAgent,
-                LangfuseTool,
-                LangfuseChain,
-                LangfuseRetriever,
-                LangfuseEvaluator,
-                LangfuseEmbedding,
-                LangfuseGuardrail,
-            ]:
-                # set their type internally in the class
-                pass
+            # For span and guardrail types, no generation properties needed
 
             yield span_class(**common_args)
 

langfuse/_client/constants.py

@@ -9,7 +9,7 @@
 LANGFUSE_TRACER_NAME = "langfuse-sdk"
 
 
-"""Note: this type is used with .__args__ in some cases and therefore must remain flat"""
+"""Note: this type is used with .__args__ / get_args in some cases and therefore must remain flat"""
 ObservationTypeGenerationLike: TypeAlias = Literal[
     "generation",
     "agent",

langfuse/_client/observe.py

@@ -109,8 +109,9 @@ def observe(
             name (Optional[str]): Custom name for the created trace or span. If not provided, the function name is used.
             as_type (Optional[Literal]): Set the observation type. Supported values:
                     "generation", "span", "agent", "tool", "chain", "retriever", "embedding", "evaluator", "guardrail".
-                    When set to "generation", creates a specialized LLM generation span with model metrics support.
-                    Other types create spans with the specified type for graph visualization and filtering in the Langfuse UI.
+                    Observation types are highlighted in the Langfuse UI for filtering and visualization.
+                    The types "generation", "agent", "tool", "chain", "retriever", "embedding", "evaluator" create a span which
+                    allows to set additional attributes such as model metrics..
 
         Returns:
             Callable: A wrapped version of the original function that automatically creates and manages Langfuse spans.
@@ -136,13 +137,6 @@ async def generate_answer(query):
                 return response.choices[0].message.content
             ```
 
-            For automatic graph instrumentation with agent workflows:
-            ```python
-            @observe(as_type="agent")
-            def planning_agent():
-                return create_plan()
-            ```
-
             For trace context propagation between functions:
             ```python
             @observe()
@@ -369,7 +363,6 @@ def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
             # Set public key in execution context for nested decorated functions
             with _set_current_public_key(public_key):
                 langfuse_client = get_client(public_key=public_key)
-
                 context_manager: Optional[
                     Union[
                         _AgnosticContextManager[LangfuseGeneration],

langfuse/_client/span.py

@@ -116,7 +116,6 @@ def __init__(
             LangfuseOtelSpanAttributes.OBSERVATION_TYPE, as_type
         )
         self._langfuse_client = langfuse_client
-
         self._observation_type = as_type
 
         self.trace_id = self._langfuse_client._get_otel_trace_id(otel_span)
@@ -606,15 +605,7 @@ def update(
                 level=level,
                 status_message=status_message,
                 observation_type=cast(
-                    Literal[
-                        "generation",
-                        "agent",
-                        "tool",
-                        "chain",
-                        "retriever",
-                        "evaluator",
-                        "embedding",
-                    ],
+                    ObservationTypeGenerationLike,
                     self._observation_type,
                 ),
                 completion_start_time=completion_start_time,
@@ -1072,7 +1063,8 @@ class LangfuseSpan(LangfuseObservationWrapper):
     This class represents a general-purpose span that can be used to trace
     any operation in your application. It extends the base LangfuseObservationWrapper
     with specific methods for creating child spans, generations, and updating
-    span-specific attributes.
+    span-specific attributes. If possible, use a more specific type for
+    better observability and insights.
     """
 
     def __init__(
@@ -1602,7 +1594,7 @@ def update(
 
 
 class LangfuseAgent(LangfuseObservationWrapper):
-    """Specialized span for agent observations in agentic workflows."""
+    """Agent observation for reasoning blocks that act on tools using LLM guidance."""
 
     def __init__(self, **kwargs: Any) -> None:
         """Initialize a new LangfuseAgent span."""
@@ -1611,7 +1603,7 @@ def __init__(self, **kwargs: Any) -> None:
 
 
 class LangfuseTool(LangfuseObservationWrapper):
-    """Specialized span for tool observations in agentic workflows."""
+    """Tool observation representing external tool calls, e.g., calling a weather API."""
 
     def __init__(self, **kwargs: Any) -> None:
         """Initialize a new LangfuseTool span."""
@@ -1620,7 +1612,7 @@ def __init__(self, **kwargs: Any) -> None:
 
 
 class LangfuseChain(LangfuseObservationWrapper):
-    """Specialized span for chain observations in agentic workflows."""
+    """Chain observation for connecting LLM application steps, e.g. passing context from retriever to LLM."""
 
     def __init__(self, **kwargs: Any) -> None:
         """Initialize a new LangfuseChain span."""
@@ -1629,7 +1621,7 @@ def __init__(self, **kwargs: Any) -> None:
 
 
 class LangfuseRetriever(LangfuseObservationWrapper):
-    """Specialized span for retriever observations in agentic workflows."""
+    """Retriever observation for data retrieval steps, e.g. vector store or database queries."""
 
     def __init__(self, **kwargs: Any) -> None:
         """Initialize a new LangfuseRetriever span."""
@@ -1638,7 +1630,7 @@ def __init__(self, **kwargs: Any) -> None:
 
 
 class LangfuseEmbedding(LangfuseObservationWrapper):
-    """Specialized span for embedding observations in agentic workflows."""
+    """Embedding observation for LLM embedding calls, typically used before retrieval."""
 
     def __init__(self, **kwargs: Any) -> None:
         """Initialize a new LangfuseEmbedding span."""
@@ -1647,7 +1639,7 @@ def __init__(self, **kwargs: Any) -> None:
 
 
 class LangfuseEvaluator(LangfuseObservationWrapper):
-    """Specialized span for evaluator observations in agentic workflows."""
+    """Evaluator observation for assessing relevance, correctness, or helpfulness of LLM outputs."""
 
     def __init__(self, **kwargs: Any) -> None:
         """Initialize a new LangfuseEvaluator span."""
@@ -1656,7 +1648,7 @@ def __init__(self, **kwargs: Any) -> None:
 
 
 class LangfuseGuardrail(LangfuseObservationWrapper):
-    """Specialized span for guardrail observations in agentic workflows."""
+    """Guardrail observation for protection e.g. against jailbreaks or offensive content."""
 
     def __init__(self, **kwargs: Any) -> None:
         """Initialize a new LangfuseGuardrail span."""