fix(observe): preserve current observation type on span update

langfuse/_client/client.py

@@ -47,7 +47,6 @@
 from langfuse._client.constants import (
     LANGFUSE_SDK_EXPERIMENT_ENVIRONMENT,
     ObservationTypeGenerationLike,
-    ObservationTypeLiteral,
     ObservationTypeLiteralNoEvent,
     ObservationTypeSpanLike,
     get_observation_types_list,
@@ -1069,7 +1068,7 @@
 
     def _get_span_class(
         self,
-        as_type: ObservationTypeLiteral,
+        as_type: str,
     ) -> Union[
         Type[LangfuseAgent],
         Type[LangfuseTool],
@@ -1108,6 +1107,21 @@
         else:
             return LangfuseSpan
 
+    @staticmethod
+    def _get_observation_type_from_otel_span(otel_span: otel_trace_api.Span) -> str:
+        if not otel_span.is_recording():
+            return "span"
+
+        attributes = getattr(otel_span, "attributes", None)
+        if attributes is None or not hasattr(attributes, "get"):
+            return "span"
+
+        observation_type = attributes.get(
+            LangfuseOtelSpanAttributes.OBSERVATION_TYPE, "span"
+        )
+
+        return observation_type if isinstance(observation_type, str) else "span"
+
     @_agnosticcontextmanager
     def _create_span_with_parent_context(
         self,
@@ -1378,10 +1392,13 @@
         current_otel_span = self._get_current_otel_span()
 
         if current_otel_span is not None:
-            span = LangfuseSpan(
+            span_class = self._get_span_class(
+                self._get_observation_type_from_otel_span(current_otel_span)
+            )
+            span = span_class(
                 otel_span=current_otel_span,
                 langfuse_client=self,
                 environment=self._environment,
                 release=self._release,
             )
 
@@ -1431,11 +1448,9 @@
         current_otel_span = self._get_current_otel_span()
 
         if current_otel_span is not None and current_otel_span.is_recording():
-            existing_observation_type = current_otel_span.attributes.get(  # type: ignore[attr-defined]
-                LangfuseOtelSpanAttributes.OBSERVATION_TYPE, "span"
+            span_class = self._get_span_class(
+                self._get_observation_type_from_otel_span(current_otel_span)
             )
-            # We need to preserve the class to keep the correct observation type
-            span_class = self._get_span_class(existing_observation_type)
             span = span_class(
                 otel_span=current_otel_span,
                 langfuse_client=self,
@@ -1468,11 +1483,9 @@
         current_otel_span = self._get_current_otel_span()
 
         if current_otel_span is not None and current_otel_span.is_recording():
-            existing_observation_type = current_otel_span.attributes.get(  # type: ignore[attr-defined]
-                LangfuseOtelSpanAttributes.OBSERVATION_TYPE, "span"
+            span_class = self._get_span_class(
+                self._get_observation_type_from_otel_span(current_otel_span)
             )
-            # We need to preserve the class to keep the correct observation type
-            span_class = self._get_span_class(existing_observation_type)
             span = span_class(
                 otel_span=current_otel_span,
                 langfuse_client=self,

tests/unit/test_observe.py

@@ -1,6 +1,7 @@
 import asyncio
 import contextvars
 import gc
+import json
 import sys
 from typing import Any, AsyncGenerator, Generator, cast
 
@@ -32,6 +33,28 @@ def _finished_spans_by_name(memory_exporter: Any, name: str) -> list[Any]:
     return [span for span in memory_exporter.get_finished_spans() if span.name == name]
 
 
+@pytest.mark.asyncio
+async def test_capture_output_false_preserves_type_when_current_span_is_updated(
+    langfuse_memory_client: Any, memory_exporter: Any
+) -> None:
+    @observe(name="guardrail_check", as_type="guardrail", capture_output=False)
+    async def guardrail_check() -> bool:
+        langfuse_memory_client.update_current_span(output={"verdict": "manually set"})
+        return True
+
+    assert await guardrail_check() is True
+
+    langfuse_memory_client.flush()
+
+    guardrail_span = _finished_spans_by_name(memory_exporter, "guardrail_check")[0]
+    attributes = guardrail_span.attributes
+
+    assert attributes[LangfuseOtelSpanAttributes.OBSERVATION_TYPE] == "guardrail"
+    assert json.loads(attributes[LangfuseOtelSpanAttributes.OBSERVATION_OUTPUT]) == {
+        "verdict": "manually set"
+    }
+
+
 def test_sync_generator_preserves_context_without_output_capture(
     langfuse_memory_client: Any, memory_exporter: Any
 ) -> None: