From 143c56bcf83ce1b2af574bef5650625181877ede Mon Sep 17 00:00:00 2001 From: Alexander Song Date: Sat, 18 Jan 2025 19:56:55 -0800 Subject: [PATCH] fix tool bug --- .../openinference/instrumentation/config.py | 4 +-- .../tests/test_manual_instrumentation.py | 29 +++++++++++++++++++ 2 files changed, 30 insertions(+), 3 deletions(-) diff --git a/python/openinference-instrumentation/src/openinference/instrumentation/config.py b/python/openinference-instrumentation/src/openinference/instrumentation/config.py index b7b8582f2..1eba54167 100644 --- a/python/openinference-instrumentation/src/openinference/instrumentation/config.py +++ b/python/openinference-instrumentation/src/openinference/instrumentation/config.py @@ -1050,9 +1050,7 @@ async def async_wrapper( if asyncio.iscoroutinefunction(wrapped_function): return async_wrapper(wrapped_function) # type: ignore[no-any-return] return sync_wrapper(wrapped_function) # type: ignore[no-any-return] - if asyncio.iscoroutinefunction(wrapped_function): - return lambda x: async_wrapper(x) - return lambda x: sync_wrapper(x) + return lambda f: async_wrapper(f) if asyncio.iscoroutinefunction(f) else sync_wrapper(f) class TracerProvider(OTelTracerProvider): diff --git a/python/openinference-instrumentation/tests/test_manual_instrumentation.py b/python/openinference-instrumentation/tests/test_manual_instrumentation.py index aba1bb0a9..decfed8ac 100644 --- a/python/openinference-instrumentation/tests/test_manual_instrumentation.py +++ b/python/openinference-instrumentation/tests/test_manual_instrumentation.py @@ -888,6 +888,35 @@ async def decorated_async_tool(input: str) -> None: assert json.loads(tool_parameters) == {} assert not attributes + async def test_async_tool_with_overridden_name( + self, + in_memory_span_exporter: InMemorySpanExporter, + tracer: OITracer, + ) -> None: + @tracer.tool(name="overridden-name") + async def decorated_async_tool(input: str) -> None: + pass + + await decorated_async_tool("input") + + spans = in_memory_span_exporter.get_finished_spans() + assert len(spans) == 1 + span = spans[0] + assert span.name == "overridden-name" + assert span.status.is_ok + assert not span.events + attributes = dict(span.attributes or {}) + assert attributes.pop(OPENINFERENCE_SPAN_KIND) == TOOL + assert attributes.pop(INPUT_MIME_TYPE) == JSON + assert isinstance(input_value := attributes.pop(INPUT_VALUE), str) + assert json.loads(input_value) == {"input": "input"} + assert attributes.pop(OUTPUT_MIME_TYPE) == TEXT + assert attributes.pop(OUTPUT_VALUE) == "None" + assert attributes.pop(TOOL_NAME) == "overridden-name" + assert isinstance(tool_parameters := attributes.pop(TOOL_PARAMETERS), str) + assert json.loads(tool_parameters) == {} + assert not attributes + def test_tool_with_zero_arguments_and_overridden_name_and_description( self, in_memory_span_exporter: InMemorySpanExporter,