diff --git a/python/openinference-instrumentation/examples/requirements.txt b/python/openinference-instrumentation/examples/requirements.txt new file mode 100644 index 000000000..5154c61ad --- /dev/null +++ b/python/openinference-instrumentation/examples/requirements.txt @@ -0,0 +1,4 @@ +jupyter +opentelemetry-sdk +opentelemetry-exporter-otlp +openinference-semantic-conventions diff --git a/python/openinference-instrumentation/examples/tracer.ipynb b/python/openinference-instrumentation/examples/tracer.ipynb new file mode 100644 index 000000000..082411427 --- /dev/null +++ b/python/openinference-instrumentation/examples/tracer.ipynb @@ -0,0 +1,593 @@ +{ + "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.semconv.trace 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, + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "from dataclasses import dataclass\n", + "from typing import Any, Dict\n", + "\n", + "import pydantic\n", + "from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter\n", + "from opentelemetry.sdk.resources import Resource\n", + "from opentelemetry.sdk.trace.export import ConsoleSpanExporter, SimpleSpanProcessor\n", + "from opentelemetry.trace import Status, StatusCode, set_tracer_provider\n", + "\n", + "from openinference.instrumentation import using_attributes\n", + "from openinference.instrumentation.config import (\n", + " OpenInferenceTracerProvider,\n", + " chain,\n", + " get_current_openinference_span,\n", + " get_input_value_and_mime_type,\n", + " get_openinference_span_kind,\n", + " get_output_value_and_mime_type,\n", + " get_tool_attributes,\n", + " suppress_tracing,\n", + " tool,\n", + ")\n", + "from openinference.semconv.resource import ResourceAttributes" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [], + "source": [ + "endpoint = \"http://127.0.0.1:6006/v1/traces\"\n", + "resource = Resource(attributes={ResourceAttributes.PROJECT_NAME: \"openinference-tracer\"})\n", + "tracer_provider = OpenInferenceTracerProvider(resource=resource)\n", + "tracer_provider.add_span_processor(SimpleSpanProcessor(OTLPSpanExporter(endpoint)))\n", + "tracer_provider.add_span_processor(SimpleSpanProcessor(ConsoleSpanExporter()))\n", + "set_tracer_provider(tracer_provider)\n", + "tracer = tracer_provider.get_tracer(__name__)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Chains" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Context Manager" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "with tracer.start_as_current_span(\n", + " \"chain-span-with-plain-text-io\",\n", + " openinference_span_kind=\"chain\",\n", + ") as span:\n", + " span.set_input(\"input\")\n", + " span.set_output(\"output\")\n", + " span.set_status(Status(StatusCode.OK))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "with tracer.start_as_current_span(\n", + " \"chain-span-with-json-io\",\n", + " openinference_span_kind=\"chain\",\n", + ") as span:\n", + " span.set_input(\n", + " {\"input-key\": \"input-value\"},\n", + " )\n", + " span.set_output(\n", + " json.dumps({\"output-key\": \"output-value\"}),\n", + " mime_type=\"application/json\",\n", + " )\n", + " span.set_status(Status(StatusCode.OK))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "with tracer.start_as_current_span(\n", + " \"chain-span-with-attribute-getters\",\n", + " attributes={\n", + " **get_openinference_span_kind(\"chain\"),\n", + " **get_input_value_and_mime_type(\"input\"),\n", + " },\n", + ") as span:\n", + " span.set_attributes(get_output_value_and_mime_type(\"output\"))\n", + " span.set_status(Status(StatusCode.OK))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "class InputModel(pydantic.BaseModel):\n", + " input: str\n", + "\n", + "\n", + "@dataclass\n", + "class OutputModel:\n", + " output: str\n", + "\n", + "\n", + "with tracer.start_as_current_span(\n", + " \"chain-span-with-pydantic-input-and-dataclass-output\",\n", + " openinference_span_kind=\"chain\",\n", + ") as span:\n", + " span.set_input(InputModel(input=\"input\"))\n", + " span.set_output(OutputModel(output=\"output\"))\n", + " span.set_status(Status(StatusCode.OK))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Decorator" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@chain\n", + "def decorated_chain_with_plain_text_output(input: str) -> str:\n", + " return \"output\"\n", + "\n", + "\n", + "decorated_chain_with_plain_text_output(\"input\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@chain\n", + "def decorated_chain_with_json_output(input: str) -> Dict[str, Any]:\n", + " return {\"output\": \"output\"}\n", + "\n", + "\n", + "decorated_chain_with_json_output(\"input\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@chain()\n", + "def decorated_chain_with_no_parameters(input: str) -> Dict[str, Any]:\n", + " return {\"output\": \"output\"}\n", + "\n", + "\n", + "decorated_chain_with_no_parameters(\"input\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@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", + "\n", + "this_name_should_be_overriden(\"input\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def chain_with_decorator_applied_as_function(input: str) -> Dict[str, Any]:\n", + " return {\"output\": \"output\"}\n", + "\n", + "\n", + "decorated = chain(chain_with_decorator_applied_as_function)\n", + "decorated(\"input\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def this_name_should_be_overriden_with_decorator_applied_as_function_with_parameters(\n", + " input: str,\n", + ") -> Dict[str, Any]:\n", + " 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(\"input\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@chain\n", + "async def decorated_async_chain(input: str) -> str:\n", + " return \"output\"\n", + "\n", + "\n", + "await decorated_async_chain(\"input\") # type: ignore[top-level-await]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@chain\n", + "def decorated_chain_with_error(input: str) -> str:\n", + " raise ValueError(\"error\")\n", + "\n", + "\n", + "try:\n", + " decorated_chain_with_error(\"input\")\n", + "except ValueError as e:\n", + " print(e)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@chain\n", + "def decorated_chain_with_child_span(input: str) -> str:\n", + " with tracer.start_as_current_span(\n", + " \"child-span\",\n", + " openinference_span_kind=\"chain\",\n", + " attributes=get_input_value_and_mime_type(\"child-span-input\"),\n", + " ) as child_span:\n", + " output = \"output\"\n", + " child_span.set_output(output)\n", + " child_span.set_status(Status(StatusCode.OK))\n", + " return output\n", + "\n", + "\n", + "decorated_chain_with_child_span(\"input\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@chain\n", + "def decorated_chain_with_child_span_error(input: str) -> str:\n", + " with tracer.start_as_current_span(\n", + " \"child-span\",\n", + " openinference_span_kind=\"chain\",\n", + " attributes=get_input_value_and_mime_type(\"child-span-input\"),\n", + " ):\n", + " raise ValueError(\"error\")\n", + "\n", + "\n", + "try:\n", + " decorated_chain_with_child_span_error(\"input\")\n", + "except ValueError as e:\n", + " print(e)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "class ChainRunner:\n", + " @chain\n", + " def decorated_chain_method(self, input1: str, input2: str) -> str:\n", + " return \"output\"\n", + "\n", + "\n", + "chain_runner = ChainRunner()\n", + "chain_runner.decorated_chain_method(\"input1\", \"input2\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@chain\n", + "def decorated_chain_with_input_and_output_set_inside_the_wrapped_function(input: str) -> str:\n", + " span = get_current_openinference_span()\n", + " span.set_input(\"overridden-input\")\n", + " span.set_output(\"overridden-output\")\n", + " return \"output\"\n", + "\n", + "\n", + "decorated_chain_with_input_and_output_set_inside_the_wrapped_function(\"input\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Suppress Tracing" + ] + }, + { + "cell_type": "code", + "execution_count": 19, + "metadata": {}, + "outputs": [], + "source": [ + "with suppress_tracing():\n", + " with tracer.start_as_current_span(\n", + " \"THIS-SPAN-SHOULD-NOT-BE-TRACED\",\n", + " openinference_span_kind=\"chain\",\n", + " ) as span:\n", + " span.set_input(\"input\")\n", + " span.set_output(\"output\")\n", + " span.set_status(Status(StatusCode.OK))" + ] + }, + { + "cell_type": "code", + "execution_count": 20, + "metadata": {}, + "outputs": [], + "source": [ + "@chain\n", + "def decorated_chain_with_suppress_tracing(input: str) -> str:\n", + " return \"output\"\n", + "\n", + "\n", + "with suppress_tracing():\n", + " decorated_chain_with_suppress_tracing(\"input\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Context Attributes" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "with using_attributes(session_id=\"123\"):\n", + " with tracer.start_as_current_span(\n", + " \"chain-span-with-context-attributes\",\n", + " openinference_span_kind=\"chain\",\n", + " ) as span:\n", + " span.set_input(\"input\")\n", + " span.set_output(\"output\")\n", + " span.set_status(Status(StatusCode.OK))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@chain\n", + "def decorated_chain_with_context_attributes(input: str) -> str:\n", + " return \"output\"\n", + "\n", + "\n", + "with using_attributes(session_id=\"123\"):\n", + " decorated_chain_with_context_attributes(\"input\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Tools" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Context Managers" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "with tracer.start_as_current_span(\n", + " \"tool-span\",\n", + " openinference_span_kind=\"tool\",\n", + ") as span:\n", + " span.set_input(\"input\")\n", + " span.set_output(\"output\")\n", + " span.set_tool(\n", + " name=\"tool-name\",\n", + " description=\"tool-description\",\n", + " parameters={\"input\": \"input\"},\n", + " )\n", + " span.set_status(Status(StatusCode.OK))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "with tracer.start_as_current_span(\n", + " \"tool-span-with-getter\",\n", + " openinference_span_kind=\"tool\",\n", + ") as span:\n", + " span.set_attributes(\n", + " get_tool_attributes(\n", + " name=\"tool-name\",\n", + " description=\"tool-description\",\n", + " parameters={\"input\": \"input\"},\n", + " )\n", + " )\n", + " span.set_status(Status(StatusCode.OK))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@tool\n", + "def decorated_tool(input1: str, input2: int) -> None:\n", + " \"\"\"\n", + " tool-description\n", + " \"\"\"\n", + "\n", + "\n", + "decorated_tool(\"input1\", 1)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@tool\n", + "async def decorated_tool_async(input1: str, input2: int) -> None:\n", + " \"\"\"\n", + " tool-description\n", + " \"\"\"\n", + "\n", + "\n", + "await decorated_tool_async(\"input1\", 1) # type: ignore[top-level-await]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@tool(\n", + " name=\"decorated-tool-with-overriden-name\",\n", + " description=\"overriden-tool-description\",\n", + ")\n", + "def this_tool_name_should_be_overriden(input1: str, input2: int) -> None:\n", + " \"\"\"\n", + " this tool description should be overriden\n", + " \"\"\"\n", + "\n", + "\n", + "this_tool_name_should_be_overriden(\"input1\", 1)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@tool\n", + "def tool_with_changes_inside_the_wrapped_function(input1: str, input2: int) -> str:\n", + " span = get_current_openinference_span()\n", + " print(type(span))\n", + " span.set_input(\"inside-input\")\n", + " span.set_output(\"inside-output\")\n", + " span.set_tool(\n", + " name=\"inside-tool-name\",\n", + " description=\"inside-tool-description\",\n", + " parameters={\"inside-input\": \"inside-input\"},\n", + " )\n", + " return \"output\"\n", + "\n", + "\n", + "tool_with_changes_inside_the_wrapped_function(\"input1\", 1)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.8.20" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/python/openinference-instrumentation/src/openinference/instrumentation/config.py b/python/openinference-instrumentation/src/openinference/instrumentation/config.py index dd52b8a30..e1b760641 100644 --- a/python/openinference-instrumentation/src/openinference/instrumentation/config.py +++ b/python/openinference-instrumentation/src/openinference/instrumentation/config.py @@ -1,49 +1,101 @@ +import asyncio +import inspect +import json import os +from collections.abc import Mapping, Sequence from contextlib import contextmanager -from dataclasses import dataclass, field, fields +from dataclasses import asdict, dataclass, field, fields +from datetime import datetime +from json import JSONEncoder from secrets import randbits +from types import ModuleType, TracebackType from typing import ( + TYPE_CHECKING, Any, + Awaitable, Callable, Dict, Iterator, + Literal, Optional, - Sequence, + Tuple, + Type, + TypeVar, Union, cast, get_args, ) -import wrapt +import wrapt # type: ignore[import-untyped] from opentelemetry.context import ( _SUPPRESS_INSTRUMENTATION_KEY, Context, attach, detach, + get_value, set_value, ) -from opentelemetry.sdk.trace import IdGenerator +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.id_generator import IdGenerator from opentelemetry.trace import ( + INVALID_SPAN, INVALID_SPAN_ID, INVALID_TRACE_ID, Link, Span, SpanKind, + Status, + StatusCode, Tracer, + get_current_span, + get_tracer, use_span, ) from opentelemetry.util.types import Attributes, AttributeValue +from typing_extensions import ParamSpec, TypeAlias, TypeGuard, overload from openinference.semconv.trace import ( EmbeddingAttributes, ImageAttributes, MessageAttributes, MessageContentAttributes, + OpenInferenceMimeTypeValues, + OpenInferenceSpanKindValues, SpanAttributes, ) +from .context_attributes import get_attributes_from_context from .logging import logger +if TYPE_CHECKING: + from _typeshed import DataclassInstance +pydantic: Optional[ModuleType] +try: + import pydantic # try to import without adding a dependency +except ImportError: + pydantic = None + + +OpenInferenceMimeType: TypeAlias = Union[ + Literal["application/json", "text/plain"], + OpenInferenceMimeTypeValues, +] +OpenInferenceSpanKind: TypeAlias = Union[ + Literal[ + "agent", + "chain", + "embedding", + "evaluator", + "guardrail", + "llm", + "reranker", + "retriever", + "tool", + "unknown", + ], + OpenInferenceSpanKindValues, +] + class suppress_tracing: """ @@ -63,10 +115,20 @@ def __aenter__(self) -> "suppress_tracing": self._token = attach(set_value(_SUPPRESS_INSTRUMENTATION_KEY, True)) return self - def __exit__(self, exc_type, exc_value, traceback) -> None: + def __exit__( + self, + exc_type: Optional[Type[BaseException]], + exc_value: Optional[BaseException], + traceback: Optional[TracebackType], + ) -> None: detach(self._token) - def __aexit__(self, exc_type, exc_value, traceback) -> None: + def __aexit__( + self, + exc_type: Optional[Type[BaseException]], + exc_value: Optional[BaseException], + traceback: Optional[TracebackType], + ) -> None: detach(self._token) @@ -218,23 +280,23 @@ def mask( value: Union[AttributeValue, Callable[[], AttributeValue]], ) -> Optional[AttributeValue]: if self.hide_llm_invocation_parameters and key == SpanAttributes.LLM_INVOCATION_PARAMETERS: - return + return None elif self.hide_inputs and key == SpanAttributes.INPUT_VALUE: value = REDACTED_VALUE elif self.hide_inputs and key == SpanAttributes.INPUT_MIME_TYPE: - return + return None elif self.hide_outputs and key == SpanAttributes.OUTPUT_VALUE: value = REDACTED_VALUE elif self.hide_outputs and key == SpanAttributes.OUTPUT_MIME_TYPE: - return + return None elif ( self.hide_inputs or self.hide_input_messages ) and SpanAttributes.LLM_INPUT_MESSAGES in key: - return + return None elif ( self.hide_outputs or self.hide_output_messages ) and SpanAttributes.LLM_OUTPUT_MESSAGES in key: - return + return None elif ( self.hide_input_text and SpanAttributes.LLM_INPUT_MESSAGES in key @@ -266,7 +328,7 @@ def mask( and SpanAttributes.LLM_INPUT_MESSAGES in key and MessageContentAttributes.MESSAGE_CONTENT_IMAGE in key ): - return + return None elif ( is_base64_url(value) # type:ignore and len(value) > self.base64_image_max_length # type:ignore @@ -280,7 +342,7 @@ def mask( and SpanAttributes.EMBEDDING_EMBEDDINGS in key and EmbeddingAttributes.EMBEDDING_VECTOR in key ): - return + return None return value() if callable(value) else value def _parse_value( @@ -318,7 +380,7 @@ def _cast_value( self, value: Any, cast_to: Any, - ) -> None: + ) -> Any: if cast_to is bool: if isinstance(value, str) and value.lower() == "true": return True @@ -334,7 +396,360 @@ def _cast_value( ] -class _WrappedSpan(wrapt.ObjectProxy): # type: ignore[misc] +def get_openinference_span_kind(kind: OpenInferenceSpanKind) -> Dict[str, AttributeValue]: + normalized_kind = _normalize_openinference_span_kind(kind) + return { + OPENINFERENCE_SPAN_KIND: normalized_kind.value, + } + + +def get_input_value_and_mime_type( + value: Any, + mime_type: Optional[OpenInferenceMimeType] = None, +) -> Dict[str, AttributeValue]: + normalized_mime_type: Optional[OpenInferenceMimeTypeValues] = None + if mime_type is not None: + normalized_mime_type = _normalize_mime_type(mime_type) + if normalized_mime_type is OpenInferenceMimeTypeValues.TEXT: + value = str(value) + elif normalized_mime_type is OpenInferenceMimeTypeValues.JSON: + if not isinstance(value, str): + value = safe_json_dumps_io_value(value) + else: + value, normalized_mime_type = _infer_serialized_io_value_and_mime_type(value) + attributes = { + INPUT_VALUE: value, + } + if normalized_mime_type is not None: + attributes[INPUT_MIME_TYPE] = normalized_mime_type.value + return attributes + + +def get_output_value_and_mime_type( + value: Any, + mime_type: Optional[OpenInferenceMimeType] = None, +) -> Dict[str, AttributeValue]: + normalized_mime_type: Optional[OpenInferenceMimeTypeValues] = None + if mime_type is not None: + normalized_mime_type = _normalize_mime_type(mime_type) + if normalized_mime_type is OpenInferenceMimeTypeValues.TEXT: + value = str(value) + elif normalized_mime_type is OpenInferenceMimeTypeValues.JSON: + if not isinstance(value, str): + value = safe_json_dumps_io_value(value) + else: + value, normalized_mime_type = _infer_serialized_io_value_and_mime_type(value) + attributes = { + OUTPUT_VALUE: value, + } + if normalized_mime_type is not None: + attributes[OUTPUT_MIME_TYPE] = normalized_mime_type.value + return attributes + + +def _infer_serialized_io_value_and_mime_type( + value: Any, +) -> Tuple[Any, Optional[OpenInferenceMimeTypeValues]]: + if isinstance(value, str): + return value, OpenInferenceMimeTypeValues.TEXT + if isinstance(value, (bool, int, float)): + return value, None + if isinstance(value, Sequence): + for element_type in (str, bool, int, float): + if all(isinstance(element, element_type) for element in value): + return value, None + return safe_json_dumps_io_value(value), OpenInferenceMimeTypeValues.JSON + if isinstance(value, Mapping): + return safe_json_dumps_io_value(value), OpenInferenceMimeTypeValues.JSON + if _is_dataclass_instance(value): + return safe_json_dumps_io_value(value), OpenInferenceMimeTypeValues.JSON + if pydantic is not None and isinstance(value, pydantic.BaseModel): + return safe_json_dumps_io_value(value), OpenInferenceMimeTypeValues.JSON + return str(value), OpenInferenceMimeTypeValues.TEXT + + +class IOValueJSONEncoder(JSONEncoder): + def default(self, obj: Any) -> Any: + try: + if _is_dataclass_instance(obj): + return asdict(obj) + if pydantic is not None and isinstance(obj, pydantic.BaseModel): + return obj.model_dump() + if isinstance(obj, datetime): + return obj.isoformat() + return super().default(obj) + except Exception: + return str(obj) + + +def safe_json_dumps_io_value(obj: Any, **kwargs: Any) -> str: + return json.dumps( + obj, + cls=IOValueJSONEncoder, + ensure_ascii=False, + ) + + +def get_tool_attributes( + *, + name: str, + description: Optional[str] = None, + parameters: Union[str, Dict[str, Any]], +) -> Dict[str, AttributeValue]: + if isinstance(parameters, str): + parameters_json = parameters + elif isinstance(parameters, Mapping): + parameters_json = safe_json_dumps_io_value(parameters) + else: + raise ValueError(f"Invalid parameters type: {type(parameters)}") + attributes = { + TOOL_NAME: name, + TOOL_PARAMETERS: parameters_json, + } + if description is not None: + attributes[TOOL_DESCRIPTION] = description + 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) self._self_config = config @@ -349,13 +764,13 @@ def set_attribute( key: str, value: Union[AttributeValue, Callable[[], AttributeValue]], ) -> None: - value = self._self_config.mask(key, value) - if value is not None: + masked_value = self._self_config.mask(key, value) + if masked_value is not None: if key in _IMPORTANT_ATTRIBUTES: - self._self_important_attributes[key] = value + self._self_important_attributes[key] = masked_value else: span = cast(Span, self.__wrapped__) - span.set_attribute(key, value) + span.set_attribute(key, masked_value) def end(self, end_time: Optional[int] = None) -> None: span = cast(Span, self.__wrapped__) @@ -363,6 +778,87 @@ def end(self, end_time: Optional[int] = None) -> None: span.set_attribute(k, v) span.end(end_time) + def set_input( + self, + value: Any, + mime_type: Optional[OpenInferenceMimeType] = None, + ) -> None: + self.set_attributes(get_input_value_and_mime_type(value, mime_type)) + + def set_output( + self, + value: Any, + mime_type: Optional[OpenInferenceMimeType] = None, + ) -> None: + self.set_attributes(get_output_value_and_mime_type(value, mime_type)) + + +class ChainSpan(OpenInferenceSpan): + def __init__(self, wrapped: Span, config: TraceConfig) -> None: + super().__init__(wrapped, config) + self.__wrapped__.set_attributes( + get_openinference_span_kind(OpenInferenceSpanKindValues.CHAIN) + ) + + +class ToolSpan(OpenInferenceSpan): + def __init__(self, wrapped: Span, config: TraceConfig) -> None: + super().__init__(wrapped, config) + self.__wrapped__.set_attributes( + get_openinference_span_kind(OpenInferenceSpanKindValues.TOOL) + ) + + def set_tool( + self, + *, + name: str, + description: Optional[str] = None, + parameters: Union[str, Dict[str, Any]], + ) -> None: + self.set_attributes( + get_tool_attributes( + name=name, + description=description, + parameters=parameters, + ) + ) + + +@overload +def get_current_openinference_span( + context: Optional[Context] = None, + *, + kind: Literal["chain"] = "chain", +) -> ChainSpan: ... + + +@overload +def get_current_openinference_span( + context: Optional[Context] = None, + *, + kind: Literal["tool"] = "tool", +) -> ToolSpan: ... + + +@overload +def get_current_openinference_span( + context: Optional[Context] = None, + *, + kind: None = None, +) -> OpenInferenceSpan: ... + + +def get_current_openinference_span( + context: Optional[Context] = None, + *, + kind: Optional[OpenInferenceSpanKind] = None, +) -> OpenInferenceSpan: + span_wrapper_cls = OpenInferenceSpan + if kind is not None: + normalized_span_kind = _normalize_openinference_span_kind(kind) + span_wrapper_cls = _get_span_wrapper_cls(normalized_span_kind) + return span_wrapper_cls(get_current_span(context), TraceConfig()) + class _IdGenerator(IdGenerator): """ @@ -391,6 +887,57 @@ def __init__(self, wrapped: Tracer, config: TraceConfig) -> None: def id_generator(self) -> IdGenerator: return self._self_id_generator + # @contextmanager + # @overload + # def start_as_current_span( + # self, + # name: str, + # context: Optional[Context] = None, + # kind: SpanKind = SpanKind.INTERNAL, + # attributes: Attributes = None, + # links: Optional["Sequence[Link]"] = (), + # start_time: Optional[int] = None, + # record_exception: bool = True, + # set_status_on_exception: bool = True, + # end_on_exit: bool = True, + # *, + # openinference_span_kind: Literal["chain"], + # ) -> Iterator[ChainSpan]: ... + + # @contextmanager + # @overload + # def start_as_current_span( + # self, + # name: str, + # context: Optional[Context] = None, + # kind: SpanKind = SpanKind.INTERNAL, + # attributes: Attributes = None, + # links: Optional["Sequence[Link]"] = (), + # start_time: Optional[int] = None, + # record_exception: bool = True, + # set_status_on_exception: bool = True, + # end_on_exit: bool = True, + # *, + # openinference_span_kind: Literal["tool"], + # ) -> Iterator[ToolSpan]: ... + + # @contextmanager + # @overload + # def start_as_current_span( + # self, + # name: str, + # context: Optional[Context] = None, + # kind: SpanKind = SpanKind.INTERNAL, + # attributes: Attributes = None, + # links: Optional["Sequence[Link]"] = (), + # start_time: Optional[int] = None, + # record_exception: bool = True, + # set_status_on_exception: bool = True, + # end_on_exit: bool = True, + # *, + # openinference_span_kind: Optional[OpenInferenceSpanKind] = None, + # ) -> Iterator[OpenInferenceSpan]: ... + @contextmanager def start_as_current_span( self, @@ -398,14 +945,17 @@ def start_as_current_span( context: Optional[Context] = None, kind: SpanKind = SpanKind.INTERNAL, attributes: Attributes = None, - links: Optional[Sequence[Link]] = (), + links: Optional["Sequence[Link]"] = (), start_time: Optional[int] = None, record_exception: bool = True, set_status_on_exception: bool = True, end_on_exit: bool = True, - ) -> Iterator[Span]: + *, + openinference_span_kind: Optional[OpenInferenceSpanKind] = None, + ) -> Iterator[OpenInferenceSpan]: span = self.start_span( name=name, + openinference_span_kind=openinference_span_kind, context=context, kind=kind, attributes=attributes, @@ -428,11 +978,19 @@ def start_span( context: Optional[Context] = None, kind: SpanKind = SpanKind.INTERNAL, attributes: Attributes = None, - links: Optional[Sequence[Link]] = (), + links: Optional["Sequence[Link]"] = (), start_time: Optional[int] = None, record_exception: bool = True, set_status_on_exception: bool = True, - ) -> Span: + *, + openinference_span_kind: Optional[OpenInferenceSpanKind] = None, + ) -> OpenInferenceSpan: + span_wrapper_cls = OpenInferenceSpan + if openinference_span_kind is not None: + normalized_span_kind = _normalize_openinference_span_kind(openinference_span_kind) + span_wrapper_cls = _get_span_wrapper_cls(normalized_span_kind) + if get_value(_SUPPRESS_INSTRUMENTATION_KEY): + return span_wrapper_cls(INVALID_SPAN, self._self_config) tracer = cast(Tracer, self.__wrapped__) span = tracer.__class__.start_span( self, @@ -445,13 +1003,75 @@ def start_span( record_exception=record_exception, set_status_on_exception=set_status_on_exception, ) - span = _WrappedSpan(span, config=self._self_config) + span = span_wrapper_cls(span, config=self._self_config) if attributes: - span.set_attributes(attributes) + span.set_attributes(dict(attributes)) + span.set_attributes(dict(get_attributes_from_context())) return span +class OpenInferenceTracerProvider(TracerProvider): + def __init__( + self, + *args: Any, + config: Optional[TraceConfig] = None, + **kwargs: Any, + ) -> None: + super().__init__(*args, **kwargs) + self._oi_trace_config = config or TraceConfig() + + def get_tracer( + self, + *args: Any, + **kwargs: Any, + ) -> OITracer: + tracer = super().get_tracer(*args, **kwargs) + return OITracer(tracer, config=self._oi_trace_config) + + def is_base64_url(url: str) -> bool: if not isinstance(url, str): return False return url.startswith("data:image/") and "base64" in url + + +def _normalize_mime_type(mime_type: OpenInferenceMimeType) -> OpenInferenceMimeTypeValues: + if isinstance(mime_type, OpenInferenceMimeTypeValues): + return mime_type + try: + return OpenInferenceMimeTypeValues(mime_type) + except ValueError: + raise ValueError(f"Invalid mime type: {mime_type}") + + +def _normalize_openinference_span_kind(kind: OpenInferenceSpanKind) -> OpenInferenceSpanKindValues: + if isinstance(kind, OpenInferenceSpanKindValues): + return kind + try: + return OpenInferenceSpanKindValues(kind.upper()) + except ValueError: + raise ValueError(f"Invalid OpenInference span kind: {kind}") + + +def _get_span_wrapper_cls(kind: OpenInferenceSpanKindValues) -> Type[OpenInferenceSpan]: + if kind is OpenInferenceSpanKindValues.CHAIN: + return ChainSpan + if kind is OpenInferenceSpanKindValues.TOOL: + return ToolSpan + raise NotImplementedError(f"Span kind {kind.value} is not yet supported") + + +def _is_dataclass_instance(obj: Any) -> TypeGuard["DataclassInstance"]: + cls = type(obj) + return hasattr(cls, "__dataclass_fields__") + + +# span attributes +INPUT_MIME_TYPE = SpanAttributes.INPUT_MIME_TYPE +INPUT_VALUE = SpanAttributes.INPUT_VALUE +OPENINFERENCE_SPAN_KIND = SpanAttributes.OPENINFERENCE_SPAN_KIND +OUTPUT_MIME_TYPE = SpanAttributes.OUTPUT_MIME_TYPE +OUTPUT_VALUE = SpanAttributes.OUTPUT_VALUE +TOOL_DESCRIPTION = SpanAttributes.TOOL_DESCRIPTION +TOOL_NAME = SpanAttributes.TOOL_NAME +TOOL_PARAMETERS = SpanAttributes.TOOL_PARAMETERS