From 05d0f815e0b4a909a5982bd66d354c5bc645ee13 Mon Sep 17 00:00:00 2001 From: Alexander Song Date: Tue, 21 Jan 2025 11:20:35 -0800 Subject: [PATCH] attributes --- .../openinference/instrumentation/config.py | 4 +++ .../tests/test_manual_instrumentation.py | 36 +++++++++++++++++++ 2 files changed, 40 insertions(+) diff --git a/python/openinference-instrumentation/src/openinference/instrumentation/config.py b/python/openinference-instrumentation/src/openinference/instrumentation/config.py index 76e38f6ae..4365cb7bf 100644 --- a/python/openinference-instrumentation/src/openinference/instrumentation/config.py +++ b/python/openinference-instrumentation/src/openinference/instrumentation/config.py @@ -550,6 +550,8 @@ def set_input( *, mime_type: Optional[OpenInferenceMimeType] = None, ) -> None: + if OPENINFERENCE_SPAN_KIND not in self._self_important_attributes: + raise ValueError("Cannot set input attributes on a non-OpenInference span") self.set_attributes(get_input_attributes(value, mime_type=mime_type)) def set_output( @@ -558,6 +560,8 @@ def set_output( *, mime_type: Optional[OpenInferenceMimeType] = None, ) -> None: + if OPENINFERENCE_SPAN_KIND not in self._self_important_attributes: + raise ValueError("Cannot set output attributes on a non-OpenInference span") self.set_attributes(get_output_attributes(value, mime_type=mime_type)) def set_tool( diff --git a/python/openinference-instrumentation/tests/test_manual_instrumentation.py b/python/openinference-instrumentation/tests/test_manual_instrumentation.py index f94cdbc7b..a83f0a1bf 100644 --- a/python/openinference-instrumentation/tests/test_manual_instrumentation.py +++ b/python/openinference-instrumentation/tests/test_manual_instrumentation.py @@ -165,6 +165,42 @@ def test_set_tool_attributes_on_non_tool_span_raises_exception( ) assert str(exc_info.value) == "Cannot set tool attributes on a non-tool span" + def test_non_openinference_span( + self, + in_memory_span_exporter: InMemorySpanExporter, + tracer: OITracer, + ) -> None: + with tracer.start_as_current_span("non-openinference-span") as non_openinference_span: + non_openinference_span.set_attribute("custom.attribute", "value") + + spans = in_memory_span_exporter.get_finished_spans() + assert len(spans) == 1 + span = spans[0] + assert span.name == "non-openinference-span" + attributes = dict(span.attributes or {}) + assert attributes.pop("custom.attribute") == "value" + assert not attributes + + def test_cannot_set_input_on_non_openinference_span( + self, + in_memory_span_exporter: InMemorySpanExporter, + tracer: OITracer, + ) -> None: + with tracer.start_as_current_span("non-openinference-span") as span: + with pytest.raises(ValueError) as exc_info: + span.set_input("input") + assert str(exc_info.value) == "Cannot set input attributes on a non-OpenInference span" + + def test_cannot_set_output_on_non_openinference_span( + self, + in_memory_span_exporter: InMemorySpanExporter, + tracer: OITracer, + ) -> None: + with tracer.start_as_current_span("non-openinference-span") as span: + with pytest.raises(ValueError) as exc_info: + span.set_output("output") + assert str(exc_info.value) == "Cannot set output attributes on a non-OpenInference span" + def test_tool( self, in_memory_span_exporter: InMemorySpanExporter,