From e8a0839b15b654234c8514d5e0c00f1a0f07947d Mon Sep 17 00:00:00 2001 From: Hassieb Pakzad <68423100+hassiebp@users.noreply.github.com> Date: Mon, 14 Jul 2025 15:36:25 +0200 Subject: [PATCH] feat(client): allow updating span/generation name --- langfuse/_client/client.py | 10 ++++++ langfuse/_client/span.py | 10 ++++++ tests/test_otel.py | 74 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 94 insertions(+) diff --git a/langfuse/_client/client.py b/langfuse/_client/client.py index 755184619..8311d0872 100644 --- a/langfuse/_client/client.py +++ b/langfuse/_client/client.py @@ -773,6 +773,7 @@ def _get_current_otel_span(self) -> Optional[otel_trace_api.Span]: def update_current_generation( self, *, + name: Optional[str] = None, input: Optional[Any] = None, output: Optional[Any] = None, metadata: Optional[Any] = None, @@ -793,6 +794,7 @@ def update_current_generation( details that become available during or after model generation. Args: + name: The generation name input: Updated input data for the model output: Output from the model (e.g., completions) metadata: Additional metadata to associate with the generation @@ -835,6 +837,9 @@ def update_current_generation( otel_span=current_otel_span, langfuse_client=self ) + if name: + current_otel_span.update_name(name) + generation.update( input=input, output=output, @@ -853,6 +858,7 @@ def update_current_generation( def update_current_span( self, *, + name: Optional[str] = None, input: Optional[Any] = None, output: Optional[Any] = None, metadata: Optional[Any] = None, @@ -867,6 +873,7 @@ def update_current_span( that become available during execution. Args: + name: The span name input: Updated input data for the operation output: Output data from the operation metadata: Additional metadata to associate with the span @@ -905,6 +912,9 @@ def update_current_span( environment=self._environment, ) + if name: + current_otel_span.update_name(name) + span.update( input=input, output=output, diff --git a/langfuse/_client/span.py b/langfuse/_client/span.py index f3c5ef920..f61de550c 100644 --- a/langfuse/_client/span.py +++ b/langfuse/_client/span.py @@ -561,6 +561,7 @@ def __init__( def update( self, *, + name: Optional[str] = None, input: Optional[Any] = None, output: Optional[Any] = None, metadata: Optional[Any] = None, @@ -575,6 +576,7 @@ def update( during execution, such as outputs, metadata, or status changes. Args: + name: Span name input: Updated input data for the operation output: Output data from the operation metadata: Additional metadata to associate with the span @@ -607,6 +609,9 @@ def update( data=metadata, field="metadata", span=self._otel_span ) + if name: + self._otel_span.update_name(name) + attributes = create_span_attributes( input=processed_input, output=processed_output, @@ -1048,6 +1053,7 @@ def __init__( def update( self, *, + name: Optional[str] = None, input: Optional[Any] = None, output: Optional[Any] = None, metadata: Optional[Any] = None, @@ -1069,6 +1075,7 @@ def update( token usage statistics, or cost details. Args: + name: The generation name input: Updated input data for the model output: Output from the model (e.g., completions) metadata: Additional metadata to associate with the generation @@ -1123,6 +1130,9 @@ def update( data=metadata, field="metadata", span=self._otel_span ) + if name: + self._otel_span.update_name(name) + attributes = create_generation_attributes( input=processed_input, output=processed_output, diff --git a/tests/test_otel.py b/tests/test_otel.py index 8434f3703..3e52b10fc 100644 --- a/tests/test_otel.py +++ b/tests/test_otel.py @@ -309,6 +309,23 @@ def test_span_hierarchy(self, langfuse_client, memory_exporter): # All spans should have the same trace ID assert len(set(s["trace_id"] for s in spans)) == 1 + def test_update_current_span_name(self, langfuse_client, memory_exporter): + """Test updating current span name via update_current_span method.""" + # Create a span using context manager + with langfuse_client.start_as_current_span(name="original-current-span"): + # Update the current span name + langfuse_client.update_current_span(name="updated-current-span") + + # Verify the span name was updated + spans = self.get_spans_by_name(memory_exporter, "updated-current-span") + assert len(spans) == 1, "Expected one span with updated name" + + # Also verify no spans exist with the original name + original_spans = self.get_spans_by_name( + memory_exporter, "original-current-span" + ) + assert len(original_spans) == 0, "Expected no spans with original name" + def test_span_attributes(self, langfuse_client, memory_exporter): """Test that span attributes are correctly set and updated.""" # Create a span with attributes @@ -360,6 +377,23 @@ def test_span_attributes(self, langfuse_client, memory_exporter): == "Test status" ) + def test_span_name_update(self, langfuse_client, memory_exporter): + """Test updating span name via update method.""" + # Create a span with initial name + span = langfuse_client.start_span(name="original-span-name") + + # Update the span name + span.update(name="updated-span-name") + span.end() + + # Verify the span name was updated + spans = self.get_spans_by_name(memory_exporter, "updated-span-name") + assert len(spans) == 1, "Expected one span with updated name" + + # Also verify no spans exist with the original name + original_spans = self.get_spans_by_name(memory_exporter, "original-span-name") + assert len(original_spans) == 0, "Expected no spans with original name" + def test_generation_span(self, langfuse_client, memory_exporter): """Test creating a generation span with model-specific attributes.""" # Create a generation @@ -394,6 +428,27 @@ def test_generation_span(self, langfuse_client, memory_exporter): ) assert usage == {"input": 10, "output": 5, "total": 15} + def test_generation_name_update(self, langfuse_client, memory_exporter): + """Test updating generation name via update method.""" + # Create a generation with initial name + generation = langfuse_client.start_generation( + name="original-generation-name", model="gpt-4" + ) + + # Update the generation name + generation.update(name="updated-generation-name") + generation.end() + + # Verify the generation name was updated + spans = self.get_spans_by_name(memory_exporter, "updated-generation-name") + assert len(spans) == 1, "Expected one generation with updated name" + + # Also verify no spans exist with the original name + original_spans = self.get_spans_by_name( + memory_exporter, "original-generation-name" + ) + assert len(original_spans) == 0, "Expected no generations with original name" + def test_trace_update(self, langfuse_client, memory_exporter): """Test updating trace level attributes.""" # Create a span and update trace attributes @@ -514,6 +569,25 @@ def test_complex_scenario(self, langfuse_client, memory_exporter): assert llm_input == {"prompt": "Summarize this text"} assert llm_output == {"text": "This is a summary"} + def test_update_current_generation_name(self, langfuse_client, memory_exporter): + """Test updating current generation name via update_current_generation method.""" + # Create a generation using context manager + with langfuse_client.start_as_current_generation( + name="original-current-generation", model="gpt-4" + ): + # Update the current generation name + langfuse_client.update_current_generation(name="updated-current-generation") + + # Verify the generation name was updated + spans = self.get_spans_by_name(memory_exporter, "updated-current-generation") + assert len(spans) == 1, "Expected one generation with updated name" + + # Also verify no spans exist with the original name + original_spans = self.get_spans_by_name( + memory_exporter, "original-current-generation" + ) + assert len(original_spans) == 0, "Expected no generations with original name" + def test_custom_trace_id(self, langfuse_client, memory_exporter): """Test setting a custom trace ID.""" # Create a custom trace ID