diff --git a/python/openinference-instrumentation/examples/tracer.ipynb b/python/openinference-instrumentation/examples/tracer.ipynb index 8e9a5df75..9ca61998f 100644 --- a/python/openinference-instrumentation/examples/tracer.ipynb +++ b/python/openinference-instrumentation/examples/tracer.ipynb @@ -1,30 +1,5 @@ { "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Notes:\n", - "\n", - "- imagine `phoenix.otel.register` returns `OpenInferenceTracerProvider`\n", - "- imagine there is an `@agent` decorator that behaves the same as `@chain`\n", - "- imagine decorators can be also be used from a particular tracer with `@tracer.chain`\n", - "- import paths will change\n", - "\n", - "\n", - "Proposed imports:\n", - "\n", - "```python\n", - "from openinference.instrumentation import (\n", - " get_input_value_and_mime_type,\n", - " get_openinference_span_kind,\n", - " get_output_value_and_mime_type,\n", - " get_tool_attributes,\n", - ")\n", - "from openinference.instrumentation import tool, chain, OpenInferenceTracerProvider\n", - "```" - ] - }, { "cell_type": "code", "execution_count": 1, @@ -43,14 +18,12 @@ "\n", "from openinference.instrumentation import (\n", " TracerProvider,\n", - " chain,\n", " get_current_span,\n", " get_input_value_and_mime_type,\n", " get_output_value_and_mime_type,\n", " get_span_kind,\n", " get_tool_attributes,\n", " suppress_tracing,\n", - " tool,\n", " using_attributes,\n", ")\n", "from openinference.semconv.resource import ResourceAttributes" @@ -174,7 +147,7 @@ "metadata": {}, "outputs": [], "source": [ - "@chain\n", + "@tracer.chain\n", "def decorated_chain_with_plain_text_output(input: str) -> str:\n", " return \"output\"\n", "\n", @@ -188,7 +161,7 @@ "metadata": {}, "outputs": [], "source": [ - "@chain\n", + "@tracer.chain\n", "def decorated_chain_with_json_output(input: str) -> Dict[str, Any]:\n", " return {\"output\": \"output\"}\n", "\n", @@ -202,7 +175,7 @@ "metadata": {}, "outputs": [], "source": [ - "@chain()\n", + "@tracer.chain()\n", "def decorated_chain_with_no_parameters(input: str) -> Dict[str, Any]:\n", " return {\"output\": \"output\"}\n", "\n", @@ -216,7 +189,7 @@ "metadata": {}, "outputs": [], "source": [ - "@chain(name=\"decorated-chain-with-overriden-name\")\n", + "@tracer.chain(name=\"decorated-chain-with-overriden-name\")\n", "def this_name_should_be_overriden(input: str) -> Dict[str, Any]:\n", " return {\"output\": \"output\"}\n", "\n", @@ -234,7 +207,7 @@ " return {\"output\": \"output\"}\n", "\n", "\n", - "decorated = chain(chain_with_decorator_applied_as_function)\n", + "decorated = tracer.chain(chain_with_decorator_applied_as_function)\n", "decorated(\"input\")" ] }, @@ -250,9 +223,9 @@ " return {\"output\": \"output\"}\n", "\n", "\n", - "decorated = chain(name=\"decorated-chain-with-decorator-applied-as-function-with-overriden-name\")(\n", - " this_name_should_be_overriden_with_decorator_applied_as_function_with_parameters\n", - ")\n", + "decorated = tracer.chain(\n", + " name=\"decorated-chain-with-decorator-applied-as-function-with-overriden-name\"\n", + ")(this_name_should_be_overriden_with_decorator_applied_as_function_with_parameters)\n", "decorated(\"input\")" ] }, @@ -262,7 +235,7 @@ "metadata": {}, "outputs": [], "source": [ - "@chain\n", + "@tracer.chain\n", "async def decorated_async_chain(input: str) -> str:\n", " return \"output\"\n", "\n", @@ -276,7 +249,7 @@ "metadata": {}, "outputs": [], "source": [ - "@chain\n", + "@tracer.chain\n", "def decorated_chain_with_error(input: str) -> str:\n", " raise ValueError(\"error\")\n", "\n", @@ -293,7 +266,7 @@ "metadata": {}, "outputs": [], "source": [ - "@chain\n", + "@tracer.chain\n", "def decorated_chain_with_child_span(input: str) -> str:\n", " with tracer.start_as_current_span(\n", " \"child-span\",\n", @@ -315,7 +288,7 @@ "metadata": {}, "outputs": [], "source": [ - "@chain\n", + "@tracer.chain\n", "def decorated_chain_with_child_span_error(input: str) -> str:\n", " with tracer.start_as_current_span(\n", " \"child-span\",\n", @@ -338,7 +311,7 @@ "outputs": [], "source": [ "class ChainRunner:\n", - " @chain\n", + " @tracer.chain\n", " def decorated_chain_method(self, input1: str, input2: str) -> str:\n", " return \"output\"\n", "\n", @@ -353,7 +326,7 @@ "metadata": {}, "outputs": [], "source": [ - "@chain\n", + "@tracer.chain\n", "def decorated_chain_with_input_and_output_set_inside_the_wrapped_function(input: str) -> str:\n", " span = get_current_span()\n", " span.set_input(\"overridden-input\")\n", @@ -393,7 +366,7 @@ "metadata": {}, "outputs": [], "source": [ - "@chain\n", + "@tracer.chain\n", "def decorated_chain_with_suppress_tracing(input: str) -> str:\n", " return \"output\"\n", "\n", @@ -431,7 +404,7 @@ "metadata": {}, "outputs": [], "source": [ - "@chain\n", + "@tracer.chain\n", "def decorated_chain_with_context_attributes(input: str) -> str:\n", " return \"output\"\n", "\n", @@ -500,7 +473,7 @@ "metadata": {}, "outputs": [], "source": [ - "@tool\n", + "@tracer.tool\n", "def decorated_tool(input1: str, input2: int) -> None:\n", " \"\"\"\n", " tool-description\n", @@ -516,7 +489,7 @@ "metadata": {}, "outputs": [], "source": [ - "@tool\n", + "@tracer.tool\n", "async def decorated_tool_async(input1: str, input2: int) -> None:\n", " \"\"\"\n", " tool-description\n", @@ -532,7 +505,7 @@ "metadata": {}, "outputs": [], "source": [ - "@tool(\n", + "@tracer.tool(\n", " name=\"decorated-tool-with-overriden-name\",\n", " description=\"overriden-tool-description\",\n", ")\n", @@ -551,7 +524,7 @@ "metadata": {}, "outputs": [], "source": [ - "@tool\n", + "@tracer.tool\n", "def tool_with_changes_inside_the_wrapped_function(input1: str, input2: int) -> str:\n", " span = get_current_span()\n", " print(type(span))\n", diff --git a/python/openinference-instrumentation/src/openinference/instrumentation/__init__.py b/python/openinference-instrumentation/src/openinference/instrumentation/__init__.py index e1b0fcfcb..809f4c195 100644 --- a/python/openinference-instrumentation/src/openinference/instrumentation/__init__.py +++ b/python/openinference-instrumentation/src/openinference/instrumentation/__init__.py @@ -3,14 +3,12 @@ OITracer, TraceConfig, TracerProvider, - chain, get_current_span, get_input_value_and_mime_type, get_output_value_and_mime_type, get_span_kind, get_tool_attributes, suppress_tracing, - tool, ) from .context_attributes import ( get_attributes_from_context, @@ -42,11 +40,9 @@ "OITracer", "REDACTED_VALUE", "TracerProvider", - "chain", "get_current_span", "get_input_value_and_mime_type", "get_output_value_and_mime_type", "get_span_kind", "get_tool_attributes", - "tool", ] diff --git a/python/openinference-instrumentation/src/openinference/instrumentation/config.py b/python/openinference-instrumentation/src/openinference/instrumentation/config.py index 802d73ea6..00d63d83a 100644 --- a/python/openinference-instrumentation/src/openinference/instrumentation/config.py +++ b/python/openinference-instrumentation/src/openinference/instrumentation/config.py @@ -95,6 +95,8 @@ ], OpenInferenceSpanKindValues, ] +ParametersType = ParamSpec("ParametersType") +ReturnType = TypeVar("ReturnType") class suppress_tracing: @@ -511,244 +513,6 @@ def get_tool_attributes( return attributes -ParametersType = ParamSpec("ParametersType") -ReturnType = TypeVar("ReturnType") - - -# overload for @chain usage (no parameters) -@overload -def chain( - wrapped_function: Callable[ParametersType, ReturnType], - /, - *, - name: None = None, -) -> Callable[ParametersType, ReturnType]: ... - - -# overload for @chain(name="name") usage (with parameters) -@overload -def chain( - wrapped_function: None = None, - /, - *, - name: Optional[str] = None, -) -> Callable[[Callable[ParametersType, ReturnType]], Callable[ParametersType, ReturnType]]: ... - - -def chain( - wrapped_function: Optional[Callable[ParametersType, ReturnType]] = None, - /, - *, - name: Optional[str] = None, -) -> Union[ - Callable[ParametersType, ReturnType], - Callable[[Callable[ParametersType, ReturnType]], Callable[ParametersType, ReturnType]], -]: - @wrapt.decorator # type: ignore[misc] - def sync_wrapper( - wrapped: Callable[ParametersType, ReturnType], - instance: Any, - args: Tuple[Any, ...], - kwargs: Dict[str, Any], - ) -> ReturnType: - tracer = OITracer(get_tracer(__name__), config=TraceConfig()) - span_name = name or wrapped.__name__ - bound_args = inspect.signature(wrapped).bind(*args, **kwargs) - bound_args.apply_defaults() - arguments = bound_args.arguments - - if len(arguments) == 1: - argument = next(iter(arguments.values())) - input_attributes = get_input_value_and_mime_type(value=argument) - else: - input_attributes = get_input_value_and_mime_type(value=arguments) - with tracer.start_as_current_span( - span_name, - openinference_span_kind=OpenInferenceSpanKindValues.CHAIN, - attributes=input_attributes, - ) as span: - output = wrapped(*args, **kwargs) - span.set_status(Status(StatusCode.OK)) - attributes = getattr( - span, "attributes", {} - ) # INVALID_SPAN does not have the attributes property - has_output = OUTPUT_VALUE in attributes - if has_output: - return output # don't overwrite if the output is set inside the wrapped function - span.set_output(value=output) - return output - - @wrapt.decorator # type: ignore[misc] - async def async_wrapper( - wrapped: Callable[ParametersType, Awaitable[ReturnType]], - instance: Any, - args: Tuple[Any, ...], - kwargs: Dict[str, Any], - ) -> ReturnType: - tracer = OITracer(get_tracer(__name__), config=TraceConfig()) - span_name = name or wrapped.__name__ - bound_args = inspect.signature(wrapped).bind(*args, **kwargs) - bound_args.apply_defaults() - arguments = bound_args.arguments - - if len(arguments) == 1: - argument = next(iter(arguments.values())) - input_attributes = get_input_value_and_mime_type(value=argument) - else: - input_attributes = get_input_value_and_mime_type(value=arguments) - with tracer.start_as_current_span( - span_name, - openinference_span_kind=OpenInferenceSpanKindValues.CHAIN, - attributes=input_attributes, - ) as span: - output = await wrapped(*args, **kwargs) - span.set_status(Status(StatusCode.OK)) - attributes = getattr( - span, "attributes", {} - ) # INVALID_SPAN does not have the attributes property - has_output = OUTPUT_VALUE in attributes - if has_output: - return output # don't overwrite if the output is set inside the wrapped function - span.set_output(value=output) - return output - - if wrapped_function is not None: - 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) - - -# overload for @tool usage (no parameters) -@overload -def tool( - wrapped_function: Callable[ParametersType, ReturnType], - /, - *, - name: None = None, - description: Optional[str] = None, -) -> Callable[ParametersType, ReturnType]: ... - - -# overload for @tool(name="name") usage (with parameters) -@overload -def tool( - wrapped_function: None = None, - /, - *, - name: Optional[str] = None, - description: Optional[str] = None, -) -> Callable[[Callable[ParametersType, ReturnType]], Callable[ParametersType, ReturnType]]: ... - - -def tool( - wrapped_function: Optional[Callable[ParametersType, ReturnType]] = None, - /, - *, - name: Optional[str] = None, - description: Optional[str] = None, -) -> Union[ - Callable[ParametersType, ReturnType], - Callable[[Callable[ParametersType, ReturnType]], Callable[ParametersType, ReturnType]], -]: - @wrapt.decorator # type: ignore[misc] - def sync_wrapper( - wrapped: Callable[ParametersType, ReturnType], - instance: Any, - args: Tuple[Any, ...], - kwargs: Dict[str, Any], - ) -> ReturnType: - tracer = OITracer(get_tracer(__name__), config=TraceConfig()) - span_name = name or wrapped.__name__ - bound_args = inspect.signature(wrapped).bind(*args, **kwargs) - bound_args.apply_defaults() - arguments = bound_args.arguments - - if len(arguments) == 1: - argument = next(iter(arguments.values())) - input_attributes = get_input_value_and_mime_type(value=argument) - else: - input_attributes = get_input_value_and_mime_type(value=arguments) - tool_parameters = safe_json_dumps_io_value(arguments) - tool_attributes = get_tool_attributes( - name=name or wrapped.__name__, - description=description or wrapped.__doc__, - parameters=tool_parameters, - ) - with tracer.start_as_current_span( - span_name, - openinference_span_kind=OpenInferenceSpanKindValues.TOOL, - attributes={ - **input_attributes, - **tool_attributes, - }, - ) as span: - output = wrapped(*args, **kwargs) - span.set_status(Status(StatusCode.OK)) - attributes = getattr( - span, "attributes", {} - ) # INVALID_SPAN does not have the attributes property - has_output = OUTPUT_VALUE in attributes - if not has_output: - span.set_output(value=output) - return output - - @wrapt.decorator # type: ignore[misc] - async def async_wrapper( - wrapped: Callable[ParametersType, Awaitable[ReturnType]], - instance: Any, - args: Tuple[Any, ...], - kwargs: Dict[str, Any], - ) -> ReturnType: - tracer = OITracer(get_tracer(__name__), config=TraceConfig()) - span_name = name or wrapped.__name__ - bound_args = inspect.signature(wrapped).bind(*args, **kwargs) - bound_args.apply_defaults() - arguments = bound_args.arguments - - if len(arguments) == 1: - argument = next(iter(arguments.values())) - input_attributes = get_input_value_and_mime_type(value=argument) - else: - input_attributes = get_input_value_and_mime_type(value=arguments) - tool_parameters = safe_json_dumps_io_value(arguments) - tool_description: Optional[str] = None - if (docstring := wrapped.__doc__) is not None: - tool_description = docstring.strip() - tool_attributes = get_tool_attributes( - name=name or wrapped.__name__, - description=tool_description, - parameters=tool_parameters, - ) - with tracer.start_as_current_span( - span_name, - openinference_span_kind=OpenInferenceSpanKindValues.TOOL, - attributes={ - **input_attributes, - **tool_attributes, - }, - ) as span: - output = await wrapped(*args, **kwargs) - span.set_status(Status(StatusCode.OK)) - attributes = getattr( - span, "attributes", {} - ) # INVALID_SPAN does not have the attributes property - has_output = OUTPUT_VALUE in attributes - if not has_output: # don't overwrite if the output is set inside the wrapped function - span.set_output(value=output) - return output - - if wrapped_function is not None: - 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) - - class OpenInferenceSpan(wrapt.ObjectProxy): # type: ignore[misc] def __init__(self, wrapped: Span, config: TraceConfig) -> None: super().__init__(wrapped) @@ -1005,6 +769,246 @@ def start_span( span.set_attributes(dict(get_attributes_from_context())) return span + # overload for @tracer.chain usage (no parameters) + @overload + def chain( + self, + wrapped_function: Callable[ParametersType, ReturnType], + /, + *, + name: None = None, + ) -> Callable[ParametersType, ReturnType]: ... + + # overload for @tracer.chain(name="name") usage (with parameters) + @overload + def chain( + self, + wrapped_function: None = None, + /, + *, + name: Optional[str] = None, + ) -> Callable[[Callable[ParametersType, ReturnType]], Callable[ParametersType, ReturnType]]: ... + + def chain( + self, + wrapped_function: Optional[Callable[ParametersType, ReturnType]] = None, + /, + *, + name: Optional[str] = None, + ) -> Union[ + Callable[ParametersType, ReturnType], + Callable[[Callable[ParametersType, ReturnType]], Callable[ParametersType, ReturnType]], + ]: + @wrapt.decorator # type: ignore[misc] + def sync_wrapper( + wrapped: Callable[ParametersType, ReturnType], + instance: Any, + args: Tuple[Any, ...], + kwargs: Dict[str, Any], + ) -> ReturnType: + tracer = self + span_name = name or wrapped.__name__ + bound_args = inspect.signature(wrapped).bind(*args, **kwargs) + bound_args.apply_defaults() + arguments = bound_args.arguments + + if len(arguments) == 1: + argument = next(iter(arguments.values())) + input_attributes = get_input_value_and_mime_type(value=argument) + else: + input_attributes = get_input_value_and_mime_type(value=arguments) + with tracer.start_as_current_span( + span_name, + openinference_span_kind=OpenInferenceSpanKindValues.CHAIN, + attributes=input_attributes, + ) as span: + output = wrapped(*args, **kwargs) + span.set_status(Status(StatusCode.OK)) + attributes = getattr( + span, "attributes", {} + ) # INVALID_SPAN does not have the attributes property + has_output = OUTPUT_VALUE in attributes + if has_output: + return ( + output # don't overwrite if the output is set inside the wrapped function + ) + span.set_output(value=output) + return output + + @wrapt.decorator # type: ignore[misc] + async def async_wrapper( + wrapped: Callable[ParametersType, Awaitable[ReturnType]], + instance: Any, + args: Tuple[Any, ...], + kwargs: Dict[str, Any], + ) -> ReturnType: + tracer = self + span_name = name or wrapped.__name__ + bound_args = inspect.signature(wrapped).bind(*args, **kwargs) + bound_args.apply_defaults() + arguments = bound_args.arguments + + if len(arguments) == 1: + argument = next(iter(arguments.values())) + input_attributes = get_input_value_and_mime_type(value=argument) + else: + input_attributes = get_input_value_and_mime_type(value=arguments) + with tracer.start_as_current_span( + span_name, + openinference_span_kind=OpenInferenceSpanKindValues.CHAIN, + attributes=input_attributes, + ) as span: + output = await wrapped(*args, **kwargs) + span.set_status(Status(StatusCode.OK)) + attributes = getattr( + span, "attributes", {} + ) # INVALID_SPAN does not have the attributes property + has_output = OUTPUT_VALUE in attributes + if has_output: + return ( + output # don't overwrite if the output is set inside the wrapped function + ) + span.set_output(value=output) + return output + + if wrapped_function is not None: + 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) + + # overload for @tool usage (no parameters) + @overload + def tool( + self, + wrapped_function: Callable[ParametersType, ReturnType], + /, + *, + name: None = None, + description: Optional[str] = None, + ) -> Callable[ParametersType, ReturnType]: ... + + # overload for @tool(name="name") usage (with parameters) + @overload + def tool( + self, + wrapped_function: None = None, + /, + *, + name: Optional[str] = None, + description: Optional[str] = None, + ) -> Callable[[Callable[ParametersType, ReturnType]], Callable[ParametersType, ReturnType]]: ... + + def tool( + self, + wrapped_function: Optional[Callable[ParametersType, ReturnType]] = None, + /, + *, + name: Optional[str] = None, + description: Optional[str] = None, + ) -> Union[ + Callable[ParametersType, ReturnType], + Callable[[Callable[ParametersType, ReturnType]], Callable[ParametersType, ReturnType]], + ]: + @wrapt.decorator # type: ignore[misc] + def sync_wrapper( + wrapped: Callable[ParametersType, ReturnType], + instance: Any, + args: Tuple[Any, ...], + kwargs: Dict[str, Any], + ) -> ReturnType: + tracer = self + span_name = name or wrapped.__name__ + bound_args = inspect.signature(wrapped).bind(*args, **kwargs) + bound_args.apply_defaults() + arguments = bound_args.arguments + + if len(arguments) == 1: + argument = next(iter(arguments.values())) + input_attributes = get_input_value_and_mime_type(value=argument) + else: + input_attributes = get_input_value_and_mime_type(value=arguments) + tool_parameters = safe_json_dumps_io_value(arguments) + tool_attributes = get_tool_attributes( + name=name or wrapped.__name__, + description=description or wrapped.__doc__, + parameters=tool_parameters, + ) + with tracer.start_as_current_span( + span_name, + openinference_span_kind=OpenInferenceSpanKindValues.TOOL, + attributes={ + **input_attributes, + **tool_attributes, + }, + ) as span: + output = wrapped(*args, **kwargs) + span.set_status(Status(StatusCode.OK)) + attributes = getattr( + span, "attributes", {} + ) # INVALID_SPAN does not have the attributes property + has_output = OUTPUT_VALUE in attributes + if not has_output: + span.set_output(value=output) + return output + + @wrapt.decorator # type: ignore[misc] + async def async_wrapper( + wrapped: Callable[ParametersType, Awaitable[ReturnType]], + instance: Any, + args: Tuple[Any, ...], + kwargs: Dict[str, Any], + ) -> ReturnType: + tracer = self + span_name = name or wrapped.__name__ + bound_args = inspect.signature(wrapped).bind(*args, **kwargs) + bound_args.apply_defaults() + arguments = bound_args.arguments + + if len(arguments) == 1: + argument = next(iter(arguments.values())) + input_attributes = get_input_value_and_mime_type(value=argument) + else: + input_attributes = get_input_value_and_mime_type(value=arguments) + tool_parameters = safe_json_dumps_io_value(arguments) + tool_description: Optional[str] = None + if (docstring := wrapped.__doc__) is not None: + tool_description = docstring.strip() + tool_attributes = get_tool_attributes( + name=name or wrapped.__name__, + description=tool_description, + parameters=tool_parameters, + ) + with tracer.start_as_current_span( + span_name, + openinference_span_kind=OpenInferenceSpanKindValues.TOOL, + attributes={ + **input_attributes, + **tool_attributes, + }, + ) as span: + output = await wrapped(*args, **kwargs) + span.set_status(Status(StatusCode.OK)) + attributes = getattr( + span, "attributes", {} + ) # INVALID_SPAN does not have the attributes property + has_output = OUTPUT_VALUE in attributes + if ( + not has_output + ): # don't overwrite if the output is set inside the wrapped function + span.set_output(value=output) + return output + + if wrapped_function is not None: + 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) + class TracerProvider(OTelTracerProvider): def __init__(