change generation like and span like

Author: nimarb

langfuse/_client/attributes.py

@@ -14,7 +14,10 @@
 from datetime import datetime
 from typing import Any, Dict, List, Literal, Optional, Union
 
-from langfuse._client.constants import ObservationTypeGenerationLike
+from langfuse._client.constants import (
+    ObservationTypeGenerationLike,
+    ObservationTypeSpanLike,
+)
 
 from langfuse._utils.serializer import EventSerializer
 from langfuse.model import PromptClient
@@ -95,7 +98,9 @@ def create_span_attributes(
     level: Optional[SpanLevel] = None,
     status_message: Optional[str] = None,
     version: Optional[str] = None,
-    observation_type: Optional[Literal["span", "guardrail", "event"]] = "span",
+    observation_type: Optional[
+        Union[ObservationTypeSpanLike, Literal["event"]]
+    ] = "span",
 ) -> dict:
     attributes = {
         LangfuseOtelSpanAttributes.OBSERVATION_TYPE: observation_type,

langfuse/_client/constants.py

@@ -12,20 +12,22 @@
 """Note: this type is used with .__args__ / get_args in some cases and therefore must remain flat"""
 ObservationTypeGenerationLike: TypeAlias = Literal[
     "generation",
+    "embedding",
+]
+
+ObservationTypeSpanLike: TypeAlias = Literal[
+    "span",
     "agent",
     "tool",
     "chain",
     "retriever",
     "evaluator",
-    "embedding",
+    "guardrail",
 ]
 
 ObservationTypeLiteralNoEvent: TypeAlias = Union[
     ObservationTypeGenerationLike,
-    Literal[
-        "span",
-        "guardrail",
-    ],
+    ObservationTypeSpanLike,
 ]
 
 ObservationTypeLiteral: TypeAlias = Union[

langfuse/_client/span.py

@@ -47,7 +47,9 @@
 from langfuse._client.constants import (
     ObservationTypeLiteral,
     ObservationTypeGenerationLike,
+    ObservationTypeSpanLike,
     ObservationTypeLiteralNoEvent,
+    get_observation_types_list,
 )
 from langfuse.logger import langfuse_logger
 from langfuse.types import MapValue, ScoreDataType, SpanLevel
@@ -142,7 +144,7 @@ def __init__(
 
             attributes = {}
 
-            if as_type in get_args(ObservationTypeGenerationLike):
+            if as_type in get_observation_types_list(ObservationTypeGenerationLike):
                 attributes = create_generation_attributes(
                     input=media_processed_input,
                     output=media_processed_output,
@@ -157,21 +159,13 @@ def __init__(
                     cost_details=cost_details,
                     prompt=prompt,
                     observation_type=cast(
-                        Literal[
-                            "generation",
-                            "agent",
-                            "tool",
-                            "chain",
-                            "retriever",
-                            "evaluator",
-                            "embedding",
-                        ],
+                        ObservationTypeGenerationLike,
                         as_type,
                     ),
                 )
 
             else:
-                # For span-like types: "span", "guardrail", "event"
+                # For span-like types and events
                 attributes = create_span_attributes(
                     input=media_processed_input,
                     output=media_processed_output,
@@ -180,8 +174,12 @@ def __init__(
                     level=level,
                     status_message=status_message,
                     observation_type=cast(
-                        Optional[Literal["span", "guardrail", "event"]],
-                        as_type if as_type in ["span", "guardrail", "event"] else None,
+                        Optional[Union[ObservationTypeSpanLike, Literal["event"]]],
+                        as_type
+                        if as_type
+                        in get_observation_types_list(ObservationTypeSpanLike)
+                        or as_type == "event"
+                        else None,
                     ),
                 )
 
@@ -597,7 +595,9 @@ def update(
         if name:
             self._otel_span.update_name(name)
 
-        if self._observation_type in get_args(ObservationTypeGenerationLike):
+        if self._observation_type in get_observation_types_list(
+            ObservationTypeGenerationLike
+        ):
             attributes = create_generation_attributes(
                 input=processed_input,
                 output=processed_output,
@@ -617,7 +617,7 @@ def update(
                 prompt=prompt,
             )
         else:
-            # For span-like types: "span", "guardrail", "event"
+            # For span-like types and events
             attributes = create_span_attributes(
                 input=processed_input,
                 output=processed_output,
@@ -626,9 +626,11 @@ def update(
                 level=level,
                 status_message=status_message,
                 observation_type=cast(
-                    Optional[Literal["span", "guardrail", "event"]],
+                    Optional[Union[ObservationTypeSpanLike, Literal["event"]]],
                     self._observation_type
-                    if self._observation_type in ["span", "guardrail", "event"]
+                    if self._observation_type
+                    in get_observation_types_list(ObservationTypeSpanLike)
+                    or self._observation_type == "event"
                     else None,
                 ),
             )
@@ -683,12 +685,6 @@ def start_observation(
         version: Optional[str] = None,
         level: Optional[SpanLevel] = None,
         status_message: Optional[str] = None,
-        completion_start_time: Optional[datetime] = None,
-        model: Optional[str] = None,
-        model_parameters: Optional[Dict[str, MapValue]] = None,
-        usage_details: Optional[Dict[str, int]] = None,
-        cost_details: Optional[Dict[str, float]] = None,
-        prompt: Optional[PromptClient] = None,
     ) -> "LangfuseAgent": ...
 
     @overload
@@ -703,12 +699,6 @@ def start_observation(
         version: Optional[str] = None,
         level: Optional[SpanLevel] = None,
         status_message: Optional[str] = None,
-        completion_start_time: Optional[datetime] = None,
-        model: Optional[str] = None,
-        model_parameters: Optional[Dict[str, MapValue]] = None,
-        usage_details: Optional[Dict[str, int]] = None,
-        cost_details: Optional[Dict[str, float]] = None,
-        prompt: Optional[PromptClient] = None,
     ) -> "LangfuseTool": ...
 
     @overload
@@ -723,12 +713,6 @@ def start_observation(
         version: Optional[str] = None,
         level: Optional[SpanLevel] = None,
         status_message: Optional[str] = None,
-        completion_start_time: Optional[datetime] = None,
-        model: Optional[str] = None,
-        model_parameters: Optional[Dict[str, MapValue]] = None,
-        usage_details: Optional[Dict[str, int]] = None,
-        cost_details: Optional[Dict[str, float]] = None,
-        prompt: Optional[PromptClient] = None,
     ) -> "LangfuseChain": ...
 
     @overload
@@ -743,12 +727,6 @@ def start_observation(
         version: Optional[str] = None,
         level: Optional[SpanLevel] = None,
         status_message: Optional[str] = None,
-        completion_start_time: Optional[datetime] = None,
-        model: Optional[str] = None,
-        model_parameters: Optional[Dict[str, MapValue]] = None,
-        usage_details: Optional[Dict[str, int]] = None,
-        cost_details: Optional[Dict[str, float]] = None,
-        prompt: Optional[PromptClient] = None,
     ) -> "LangfuseRetriever": ...
 
     @overload
@@ -763,12 +741,6 @@ def start_observation(
         version: Optional[str] = None,
         level: Optional[SpanLevel] = None,
         status_message: Optional[str] = None,
-        completion_start_time: Optional[datetime] = None,
-        model: Optional[str] = None,
-        model_parameters: Optional[Dict[str, MapValue]] = None,
-        usage_details: Optional[Dict[str, int]] = None,
-        cost_details: Optional[Dict[str, float]] = None,
-        prompt: Optional[PromptClient] = None,
     ) -> "LangfuseEvaluator": ...
 
     @overload
@@ -912,7 +884,7 @@ def start_observation(
             "status_message": status_message,
         }
 
-        if as_type in get_args(ObservationTypeGenerationLike):
+        if as_type in get_observation_types_list(ObservationTypeGenerationLike):
             common_args.update(
                 {
                     "completion_start_time": completion_start_time,
@@ -960,15 +932,81 @@ def start_as_current_observation(
         prompt: Optional[PromptClient] = None,
     ) -> _AgnosticContextManager[
         Union[
-            "LangfuseAgent",
-            "LangfuseTool",
-            "LangfuseChain",
-            "LangfuseRetriever",
-            "LangfuseEvaluator",
+            "LangfuseGeneration",
             "LangfuseEmbedding",
         ]
     ]: ...
 
+    @overload
+    def start_as_current_observation(
+        self,
+        *,
+        name: str,
+        as_type: Literal["agent"],
+        input: Optional[Any] = None,
+        output: Optional[Any] = None,
+        metadata: Optional[Any] = None,
+        version: Optional[str] = None,
+        level: Optional[SpanLevel] = None,
+        status_message: Optional[str] = None,
+    ) -> _AgnosticContextManager["LangfuseAgent"]: ...
+
+    @overload
+    def start_as_current_observation(
+        self,
+        *,
+        name: str,
+        as_type: Literal["tool"],
+        input: Optional[Any] = None,
+        output: Optional[Any] = None,
+        metadata: Optional[Any] = None,
+        version: Optional[str] = None,
+        level: Optional[SpanLevel] = None,
+        status_message: Optional[str] = None,
+    ) -> _AgnosticContextManager["LangfuseTool"]: ...
+
+    @overload
+    def start_as_current_observation(
+        self,
+        *,
+        name: str,
+        as_type: Literal["chain"],
+        input: Optional[Any] = None,
+        output: Optional[Any] = None,
+        metadata: Optional[Any] = None,
+        version: Optional[str] = None,
+        level: Optional[SpanLevel] = None,
+        status_message: Optional[str] = None,
+    ) -> _AgnosticContextManager["LangfuseChain"]: ...
+
+    @overload
+    def start_as_current_observation(
+        self,
+        *,
+        name: str,
+        as_type: Literal["retriever"],
+        input: Optional[Any] = None,
+        output: Optional[Any] = None,
+        metadata: Optional[Any] = None,
+        version: Optional[str] = None,
+        level: Optional[SpanLevel] = None,
+        status_message: Optional[str] = None,
+    ) -> _AgnosticContextManager["LangfuseRetriever"]: ...
+
+    @overload
+    def start_as_current_observation(
+        self,
+        *,
+        name: str,
+        as_type: Literal["evaluator"],
+        input: Optional[Any] = None,
+        output: Optional[Any] = None,
+        metadata: Optional[Any] = None,
+        version: Optional[str] = None,
+        level: Optional[SpanLevel] = None,
+        status_message: Optional[str] = None,
+    ) -> _AgnosticContextManager["LangfuseEvaluator"]: ...
+
     @overload
     def start_as_current_observation(
         self,