diff --git a/langfuse/_client/propagation.py b/langfuse/_client/propagation.py index cf06dcd33..d97f33317 100644 --- a/langfuse/_client/propagation.py +++ b/langfuse/_client/propagation.py @@ -6,7 +6,19 @@ """ import re -from typing import Any, Dict, Generator, List, Literal, Optional, TypedDict, Union, cast +from typing import ( + Any, + Dict, + Generator, + List, + Literal, + Mapping, + Optional, + Tuple, + TypedDict, + Union, + cast, +) from opentelemetry import ( baggage, @@ -29,6 +41,7 @@ from langfuse._client.attributes import LangfuseOtelSpanAttributes from langfuse._client.constants import LANGFUSE_SDK_EXPERIMENT_ENVIRONMENT from langfuse.logger import langfuse_logger +from langfuse.model import PromptClient PropagatedKeys = Literal[ "user_id", @@ -38,6 +51,8 @@ "tags", "trace_name", "environment", + "prompt_name", + "prompt_version", ] InternalPropagatedKeys = Literal[ @@ -58,6 +73,8 @@ "tags", "trace_name", "environment", + "prompt_name", + "prompt_version", "experiment_id", "experiment_name", "experiment_metadata", @@ -103,6 +120,7 @@ def propagate_attributes( tags: Optional[List[str]] = None, trace_name: Optional[str] = None, environment: Optional[str] = None, + prompt: Optional[Union[PromptClient, Mapping[str, Any]]] = None, as_baggage: bool = False, ) -> _AgnosticContextManager[Any]: """Propagate trace-level attributes to all spans created within this context. @@ -139,6 +157,17 @@ def propagate_attributes( tags: List of tags to categorize the group of observations trace_name: Name to assign to the trace. Must be US-ASCII string, ≤200 characters. Use this to set a consistent trace name for all spans created within this context. + prompt: Langfuse prompt to link to generations created within this context. + Accepts a `PromptClient` returned by `langfuse.get_prompt(...)` or any + object/dict exposing `name` (string) and `version` (integer) — e.g. + `{"name": "my-prompt", "version": 3}`. This is the recommended way to + link prompts to generations emitted by auto-instrumentation libraries + (e.g. LiteLLM's `langfuse_otel`, OpenAI Agents SDK, OpenInference) + where you don't create the generation via the Langfuse SDK yourself. + The prompt link is only applied to generation-type observations by the + Langfuse backend. Fallback prompts are never linked. An explicit + `prompt` passed to `start_observation` / `update_current_generation` + takes precedence over the propagated one. environment: Langfuse environment to assign to spans created in this context. Must be a lowercase alphanumeric string with optional hyphens or underscores, must be ≤40 characters, and must not start with "langfuse". This maps to @@ -184,6 +213,24 @@ def propagate_attributes( ... ``` + Prompt linking with auto-instrumented libraries: + + ```python + from langfuse import Langfuse + + langfuse = Langfuse() + prompt = langfuse.get_prompt("my-prompt") + + with langfuse.propagate_attributes(prompt=prompt): + # Generations emitted by auto-instrumentation (LiteLLM langfuse_otel, + # OpenAI Agents SDK, OpenInference, ...) within this context are + # linked to the prompt version. + completion = litellm.completion( + model="gpt-4o", + messages=prompt.compile(topic="chickens"), + ) + ``` + Late propagation (anti-pattern): ```python @@ -244,6 +291,7 @@ def propagate_attributes( tags=tags, trace_name=trace_name, environment=environment, + prompt=prompt, as_baggage=as_baggage, ) @@ -258,6 +306,7 @@ def _propagate_attributes( tags: Optional[List[str]] = None, trace_name: Optional[str] = None, environment: Optional[str] = None, + prompt: Optional[Union[PromptClient, Mapping[str, Any]]] = None, as_baggage: bool = False, experiment: Optional[PropagatedExperimentAttributes] = None, ) -> Generator[Any, Any, Any]: @@ -273,6 +322,25 @@ def _propagate_attributes( "environment": environment, } + prompt_info = _extract_propagated_prompt(prompt) if prompt is not None else None + if prompt_info is not None: + prompt_name, prompt_version = prompt_info + + context = _set_propagated_attribute( + key="prompt_name", + value=prompt_name, + context=context, + span=current_span, + as_baggage=as_baggage, + ) + context = _set_propagated_attribute( + key="prompt_version", + value=prompt_version, + context=context, + span=current_span, + as_baggage=as_baggage, + ) + propagated_metadata_attributes: Dict[str, Optional[Dict[str, Any]]] = { "metadata": metadata, } @@ -336,10 +404,51 @@ def _propagate_attributes( _detach_context_token_safely(token) +def _extract_propagated_prompt( + prompt: Union[PromptClient, Mapping[str, Any]], +) -> Optional[Tuple[str, int]]: + """Extract and validate (name, version) from a prompt-like value. + + Accepts a PromptClient or any mapping/object exposing `name` and `version`. + Returns None (with a warning) if the value is invalid or a fallback prompt. + """ + if isinstance(prompt, Mapping): + name = prompt.get("name") + version = prompt.get("version") + is_fallback = bool(prompt.get("is_fallback", False)) + else: + name = getattr(prompt, "name", None) + version = getattr(prompt, "version", None) + is_fallback = bool(getattr(prompt, "is_fallback", False)) + + if is_fallback: + langfuse_logger.debug( + "Propagated prompt is a fallback prompt. Skipping prompt linking." + ) + return None + + if not isinstance(name, str) or not name: + langfuse_logger.warning( + "Propagated 'prompt' has no valid 'name' (non-empty string required). Dropping prompt link." + ) + return None + + if isinstance(version, str) and version.isdigit(): + version = int(version) + + if not isinstance(version, int) or isinstance(version, bool): + langfuse_logger.warning( + "Propagated 'prompt' has no valid 'version' (integer required). Dropping prompt link." + ) + return None + + return name, version + + def _get_propagated_attributes_from_context( context: otel_context_api.Context, -) -> Dict[str, Union[str, List[str]]]: - propagated_attributes: Dict[str, Union[str, List[str]]] = {} +) -> Dict[str, Union[str, int, List[str]]]: + propagated_attributes: Dict[str, Union[str, int, List[str]]] = {} # Handle baggage baggage_entries = baggage.get_all(context=context) @@ -362,6 +471,14 @@ def _get_propagated_attributes_from_context( propagated_attributes[span_key] = validated_environment continue + if ( + span_key == LangfuseOtelSpanAttributes.OBSERVATION_PROMPT_VERSION + and isinstance(baggage_value, str) + and baggage_value.isdigit() + ): + propagated_attributes[span_key] = int(baggage_value) + continue + propagated_attributes[span_key] = ( baggage_value if isinstance(baggage_value, (str, list)) @@ -398,7 +515,7 @@ def _get_propagated_attributes_from_context( span_key = _get_propagated_span_key(key) propagated_attributes[span_key] = ( - value if isinstance(value, (str, list)) else str(value) + value if isinstance(value, (str, int, list)) else str(value) ) if ( @@ -415,7 +532,7 @@ def _get_propagated_attributes_from_context( def _set_propagated_attribute( *, key: str, - value: Union[str, List[str], Dict[str, str]], + value: Union[str, int, List[str], Dict[str, str]], context: otel_context_api.Context, span: otel_trace_api.Span, as_baggage: bool, @@ -472,7 +589,9 @@ def _set_propagated_attribute( ) else: context = otel_baggage_api.set_baggage( - name=baggage_key, value=value, context=context + name=baggage_key, + value=str(value) if isinstance(value, int) else value, + context=context, ) return context @@ -622,6 +741,8 @@ def _get_propagated_span_key(key: str) -> str: "trace_name": LangfuseOtelSpanAttributes.TRACE_NAME, "environment": LangfuseOtelSpanAttributes.ENVIRONMENT, "metadata": LangfuseOtelSpanAttributes.TRACE_METADATA, + "prompt_name": LangfuseOtelSpanAttributes.OBSERVATION_PROMPT_NAME, + "prompt_version": LangfuseOtelSpanAttributes.OBSERVATION_PROMPT_VERSION, "experiment_id": LangfuseOtelSpanAttributes.EXPERIMENT_ID, "experiment_name": LangfuseOtelSpanAttributes.EXPERIMENT_NAME, "experiment_metadata": LangfuseOtelSpanAttributes.EXPERIMENT_METADATA, diff --git a/langfuse/_client/span_processor.py b/langfuse/_client/span_processor.py index bfebf7f87..aee15d17b 100644 --- a/langfuse/_client/span_processor.py +++ b/langfuse/_client/span_processor.py @@ -144,6 +144,19 @@ def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None context = parent_context or context_api.get_current() propagated_attributes = _get_propagated_attributes_from_context(context) + # An explicit prompt set at span creation takes precedence over a propagated one + existing_attributes = span.attributes or {} + if LangfuseOtelSpanAttributes.OBSERVATION_PROMPT_NAME in existing_attributes: + propagated_attributes = { + key: value + for key, value in propagated_attributes.items() + if key + not in ( + LangfuseOtelSpanAttributes.OBSERVATION_PROMPT_NAME, + LangfuseOtelSpanAttributes.OBSERVATION_PROMPT_VERSION, + ) + } + if propagated_attributes: span.set_attributes(propagated_attributes) diff --git a/tests/unit/test_propagate_attributes.py b/tests/unit/test_propagate_attributes.py index 0ef92ac10..985c65af8 100644 --- a/tests/unit/test_propagate_attributes.py +++ b/tests/unit/test_propagate_attributes.py @@ -3368,3 +3368,283 @@ def test_trace_name_with_baggage(self, langfuse_client, memory_exporter): self.verify_span_attribute( child_span, LangfuseOtelSpanAttributes.TRACE_USER_ID, "user_123" ) + + +class TestPropagateAttributesPrompt(TestPropagateAttributesBase): + """Tests for prompt linking via propagate_attributes.""" + + def _make_prompt_client(self, name="test-prompt", version=3, is_fallback=False): + from langfuse.api import Prompt_Text + from langfuse.model import TextPromptClient + + prompt = Prompt_Text( + name=name, + version=version, + prompt="Make me laugh", + type="text", + labels=[], + config={}, + tags=[], + ) + + return TextPromptClient(prompt, is_fallback=is_fallback) + + def test_prompt_client_propagates_to_child_spans( + self, langfuse_client, memory_exporter + ): + """Verify a PromptClient propagates name and version to all child spans.""" + prompt = self._make_prompt_client(name="test-prompt", version=3) + + with langfuse_client.start_as_current_observation(name="parent-span"): + with propagate_attributes(prompt=prompt): + child1 = langfuse_client.start_observation(name="child-span-1") + child1.end() + + child2 = langfuse_client.start_observation( + name="child-span-2", as_type="generation" + ) + child2.end() + + for name in ("child-span-1", "child-span-2"): + child_span = self.get_span_by_name(memory_exporter, name) + self.verify_span_attribute( + child_span, + LangfuseOtelSpanAttributes.OBSERVATION_PROMPT_NAME, + "test-prompt", + ) + self.verify_span_attribute( + child_span, + LangfuseOtelSpanAttributes.OBSERVATION_PROMPT_VERSION, + 3, + ) + + def test_prompt_dict_propagates_to_child_spans( + self, langfuse_client, memory_exporter + ): + """Verify a plain dict with name and version keys is supported.""" + with langfuse_client.start_as_current_observation(name="parent-span"): + with propagate_attributes(prompt={"name": "dict-prompt", "version": 7}): + child = langfuse_client.start_observation(name="child-span") + child.end() + + child_span = self.get_span_by_name(memory_exporter, "child-span") + self.verify_span_attribute( + child_span, + LangfuseOtelSpanAttributes.OBSERVATION_PROMPT_NAME, + "dict-prompt", + ) + self.verify_span_attribute( + child_span, + LangfuseOtelSpanAttributes.OBSERVATION_PROMPT_VERSION, + 7, + ) + + def test_prompt_duck_typed_object_supported(self, langfuse_client, memory_exporter): + """Verify any object exposing name and version attributes is supported.""" + + class MyPrompt: + name = "duck-prompt" + version = 2 + + with langfuse_client.start_as_current_observation(name="parent-span"): + with propagate_attributes(prompt=MyPrompt()): + child = langfuse_client.start_observation(name="child-span") + child.end() + + child_span = self.get_span_by_name(memory_exporter, "child-span") + self.verify_span_attribute( + child_span, + LangfuseOtelSpanAttributes.OBSERVATION_PROMPT_NAME, + "duck-prompt", + ) + self.verify_span_attribute( + child_span, + LangfuseOtelSpanAttributes.OBSERVATION_PROMPT_VERSION, + 2, + ) + + def test_prompt_version_digit_string_coerced_to_int( + self, langfuse_client, memory_exporter + ): + """Verify a numeric string version is coerced to an integer.""" + with propagate_attributes(prompt={"name": "prompt", "version": "5"}): + child = langfuse_client.start_observation(name="child-span") + child.end() + + child_span = self.get_span_by_name(memory_exporter, "child-span") + self.verify_span_attribute( + child_span, + LangfuseOtelSpanAttributes.OBSERVATION_PROMPT_VERSION, + 5, + ) + + def test_fallback_prompt_not_linked(self, langfuse_client, memory_exporter): + """Verify fallback prompts are never linked.""" + prompt = self._make_prompt_client(is_fallback=True) + + with propagate_attributes(prompt=prompt): + child = langfuse_client.start_observation(name="child-span") + child.end() + + child_span = self.get_span_by_name(memory_exporter, "child-span") + self.verify_missing_attribute( + child_span, LangfuseOtelSpanAttributes.OBSERVATION_PROMPT_NAME + ) + self.verify_missing_attribute( + child_span, LangfuseOtelSpanAttributes.OBSERVATION_PROMPT_VERSION + ) + + def test_prompt_without_name_dropped(self, langfuse_client, memory_exporter): + """Verify a prompt without a valid name is dropped entirely.""" + with propagate_attributes(prompt={"version": 3}): + child = langfuse_client.start_observation(name="child-span") + child.end() + + child_span = self.get_span_by_name(memory_exporter, "child-span") + self.verify_missing_attribute( + child_span, LangfuseOtelSpanAttributes.OBSERVATION_PROMPT_NAME + ) + self.verify_missing_attribute( + child_span, LangfuseOtelSpanAttributes.OBSERVATION_PROMPT_VERSION + ) + + def test_prompt_without_valid_version_dropped( + self, langfuse_client, memory_exporter + ): + """Verify a prompt without an integer version is dropped entirely.""" + with propagate_attributes(prompt={"name": "prompt", "version": "latest"}): + child = langfuse_client.start_observation(name="child-span") + child.end() + + child_span = self.get_span_by_name(memory_exporter, "child-span") + self.verify_missing_attribute( + child_span, LangfuseOtelSpanAttributes.OBSERVATION_PROMPT_NAME + ) + self.verify_missing_attribute( + child_span, LangfuseOtelSpanAttributes.OBSERVATION_PROMPT_VERSION + ) + + def test_explicit_prompt_takes_precedence_over_propagated( + self, langfuse_client, memory_exporter + ): + """Verify an explicit prompt on an observation wins over the propagated one.""" + propagated = self._make_prompt_client(name="propagated-prompt", version=1) + explicit = self._make_prompt_client(name="explicit-prompt", version=9) + + with propagate_attributes(prompt=propagated): + child = langfuse_client.start_observation( + name="explicit-child", as_type="generation", prompt=explicit + ) + child.end() + + other = langfuse_client.start_observation(name="propagated-child") + other.end() + + explicit_span = self.get_span_by_name(memory_exporter, "explicit-child") + self.verify_span_attribute( + explicit_span, + LangfuseOtelSpanAttributes.OBSERVATION_PROMPT_NAME, + "explicit-prompt", + ) + self.verify_span_attribute( + explicit_span, + LangfuseOtelSpanAttributes.OBSERVATION_PROMPT_VERSION, + 9, + ) + + propagated_span = self.get_span_by_name(memory_exporter, "propagated-child") + self.verify_span_attribute( + propagated_span, + LangfuseOtelSpanAttributes.OBSERVATION_PROMPT_NAME, + "propagated-prompt", + ) + + def test_nested_prompt_inner_overwrites(self, langfuse_client, memory_exporter): + """Verify a nested propagated prompt shadows the outer one within its scope.""" + with propagate_attributes(prompt={"name": "outer-prompt", "version": 1}): + with propagate_attributes(prompt={"name": "inner-prompt", "version": 2}): + inner = langfuse_client.start_observation(name="inner-span") + inner.end() + + outer = langfuse_client.start_observation(name="outer-span") + outer.end() + + inner_span = self.get_span_by_name(memory_exporter, "inner-span") + self.verify_span_attribute( + inner_span, + LangfuseOtelSpanAttributes.OBSERVATION_PROMPT_NAME, + "inner-prompt", + ) + self.verify_span_attribute( + inner_span, + LangfuseOtelSpanAttributes.OBSERVATION_PROMPT_VERSION, + 2, + ) + + outer_span = self.get_span_by_name(memory_exporter, "outer-span") + self.verify_span_attribute( + outer_span, + LangfuseOtelSpanAttributes.OBSERVATION_PROMPT_NAME, + "outer-prompt", + ) + + def test_prompt_propagates_via_baggage(self, langfuse_client, memory_exporter): + """Verify prompt propagates through baggage with version restored as int.""" + with propagate_attributes( + prompt={"name": "baggage-prompt", "version": 4}, + as_baggage=True, + ): + child = langfuse_client.start_observation(name="child-span") + child.end() + + child_span = self.get_span_by_name(memory_exporter, "child-span") + self.verify_span_attribute( + child_span, + LangfuseOtelSpanAttributes.OBSERVATION_PROMPT_NAME, + "baggage-prompt", + ) + self.verify_span_attribute( + child_span, + LangfuseOtelSpanAttributes.OBSERVATION_PROMPT_VERSION, + 4, + ) + + def test_prompt_composes_with_outer_propagated_attributes( + self, langfuse_client, memory_exporter + ): + """Verify a nested prompt-only context preserves outer attributes.""" + with propagate_attributes(session_id="session_abc", user_id="user_123"): + with propagate_attributes(prompt={"name": "my-prompt", "version": 3}): + inner = langfuse_client.start_observation(name="inner-span") + inner.end() + + after = langfuse_client.start_observation(name="after-span") + after.end() + + # Inner span has outer session/user attributes AND the prompt + inner_span = self.get_span_by_name(memory_exporter, "inner-span") + self.verify_span_attribute( + inner_span, LangfuseOtelSpanAttributes.TRACE_SESSION_ID, "session_abc" + ) + self.verify_span_attribute( + inner_span, LangfuseOtelSpanAttributes.TRACE_USER_ID, "user_123" + ) + self.verify_span_attribute( + inner_span, + LangfuseOtelSpanAttributes.OBSERVATION_PROMPT_NAME, + "my-prompt", + ) + self.verify_span_attribute( + inner_span, + LangfuseOtelSpanAttributes.OBSERVATION_PROMPT_VERSION, + 3, + ) + + # After the inner context exits, the prompt is gone but outer attributes remain + after_span = self.get_span_by_name(memory_exporter, "after-span") + self.verify_span_attribute( + after_span, LangfuseOtelSpanAttributes.TRACE_SESSION_ID, "session_abc" + ) + self.verify_missing_attribute( + after_span, LangfuseOtelSpanAttributes.OBSERVATION_PROMPT_NAME + )