diff --git a/langfuse/_client/client.py b/langfuse/_client/client.py index 2f1c8d783..8a241b463 100644 --- a/langfuse/_client/client.py +++ b/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 start_as_current_observation( def _get_span_class( self, - as_type: ObservationTypeLiteral, + as_type: str, ) -> Union[ Type[LangfuseAgent], Type[LangfuseTool], @@ -1108,6 +1107,21 @@ def _get_span_class( 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,7 +1392,10 @@ def update_current_span( 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, @@ -1431,11 +1448,9 @@ def set_current_trace_io( 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 @@ def set_current_trace_as_public(self) -> None: 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, diff --git a/tests/unit/test_observe.py b/tests/unit/test_observe.py index 24f79c3fc..5527be9b9 100644 --- a/tests/unit/test_observe.py +++ b/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: