rearrange spans

Author: nimarb

langfuse/_client/attributes.py

@@ -93,10 +93,10 @@ def create_span_attributes(
     level: Optional[SpanLevel] = None,
     status_message: Optional[str] = None,
     version: Optional[str] = None,
-    observation_type: Optional[str] = None,
+    observation_type: Optional[Literal["span", "guardrail"]] = "span",
 ) -> dict:
     attributes = {
-        LangfuseOtelSpanAttributes.OBSERVATION_TYPE: observation_type or "span",
+        LangfuseOtelSpanAttributes.OBSERVATION_TYPE: observation_type,
         LangfuseOtelSpanAttributes.OBSERVATION_LEVEL: level,
         LangfuseOtelSpanAttributes.OBSERVATION_STATUS_MESSAGE: status_message,
         LangfuseOtelSpanAttributes.VERSION: version,
@@ -123,9 +123,20 @@ 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",
 ) -> dict:
     attributes = {
-        LangfuseOtelSpanAttributes.OBSERVATION_TYPE: "generation",
+        LangfuseOtelSpanAttributes.OBSERVATION_TYPE: observation_type,
         LangfuseOtelSpanAttributes.OBSERVATION_LEVEL: level,
         LangfuseOtelSpanAttributes.OBSERVATION_STATUS_MESSAGE: status_message,
         LangfuseOtelSpanAttributes.VERSION: version,

langfuse/_client/client.py

@@ -848,22 +848,22 @@ def start_as_current_observation(
 
         if as_type == "generation":
             return self.start_as_current_generation(
-                    trace_context=trace_context,
-                    name=name,
-                    input=input,
-                    output=output,
-                    metadata=metadata,
-                    version=version,
-                    level=level,
-                    status_message=status_message,
-                    completion_start_time=completion_start_time,
-                    model=model,
-                    model_parameters=model_parameters,
-                    usage_details=usage_details,
-                    cost_details=cost_details,
-                    prompt=prompt,
-                    end_on_exit=end_on_exit,
-                )
+                trace_context=trace_context,
+                name=name,
+                input=input,
+                output=output,
+                metadata=metadata,
+                version=version,
+                level=level,
+                status_message=status_message,
+                completion_start_time=completion_start_time,
+                model=model,
+                model_parameters=model_parameters,
+                usage_details=usage_details,
+                cost_details=cost_details,
+                prompt=prompt,
+                end_on_exit=end_on_exit,
+            )
 
         if trace_context:
             trace_id = trace_context.get("trace_id", None)