cleanup

Author: nimarb

langfuse/_client/attributes.py

@@ -14,6 +14,8 @@
 from datetime import datetime
 from typing import Any, Dict, List, Literal, Optional, Union
 
+from langfuse._client.constants import ObservationTypeGenerationLike
+
 from langfuse._utils.serializer import EventSerializer
 from langfuse.model import PromptClient
 from langfuse.types import MapValue, SpanLevel
@@ -123,17 +125,7 @@ def create_generation_attributes(
     usage_details: Optional[Dict[str, int]] = None,
     cost_details: Optional[Dict[str, float]] = None,
     prompt: Optional[PromptClient] = None,
-    observation_type: Optional[
-        Literal[
-            "generation",
-            "agent",
-            "tool",
-            "chain",
-            "retriever",
-            "evaluator",
-            "embedding",
-        ]
-    ] = "generation",
+    observation_type: Optional[ObservationTypeGenerationLike] = "generation",
 ) -> dict:
     attributes = {
         LangfuseOtelSpanAttributes.OBSERVATION_TYPE: observation_type,

langfuse/_client/constants.py

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

langfuse/_client/span.py

@@ -45,7 +45,6 @@
     ObservationTypeLiteral,
     ObservationTypeGenerationLike,
     ObservationTypeLiteralNoEvent,
-    get_observation_types_list,
 )
 from langfuse.logger import langfuse_logger
 from langfuse.types import MapValue, ScoreDataType, SpanLevel
@@ -144,15 +143,7 @@ def __init__(
 
             attributes = {}
 
-            if as_type in [
-                "generation",
-                "agent",
-                "tool",
-                "chain",
-                "retriever",
-                "evaluator",
-                "embedding",
-            ]:
+            if as_type in ObservationTypeGenerationLike.__args__:
                 attributes = create_generation_attributes(
                     input=media_processed_input,
                     output=media_processed_output,
@@ -611,15 +602,7 @@ def update(
             self._otel_span.update_name(name)
 
         # Use same logic as __init__ to determine which attributes to create
-        if self._observation_type in [
-            "generation",
-            "agent",
-            "tool",
-            "chain",
-            "retriever",
-            "evaluator",
-            "embedding",
-        ]:
+        if self._observation_type in ObservationTypeGenerationLike.__args__:
             attributes = create_generation_attributes(
                 input=processed_input,
                 output=processed_output,