Skip to content

Commit

Permalink
tmp
Browse files Browse the repository at this point in the history
  • Loading branch information
axiomofjoy committed Jan 14, 2025
1 parent 4bef033 commit bef939f
Show file tree
Hide file tree
Showing 3 changed files with 208 additions and 37 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
jupyter
opentelemetry-sdk
opentelemetry-exporter-otlp
openinference-semantic-conventions
108 changes: 108 additions & 0 deletions python/openinference-instrumentation/examples/tracer.ipynb
Original file line number Diff line number Diff line change
@@ -0,0 +1,108 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"import json\n",
"\n",
"from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter\n",
"from opentelemetry.sdk.resources import Resource\n",
"from opentelemetry.sdk.trace import TracerProvider\n",
"from opentelemetry.sdk.trace.export import SimpleSpanProcessor\n",
"from opentelemetry.trace import Status, StatusCode\n",
"\n",
"from openinference.instrumentation import OITracerProvider, TraceConfig\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 = OITracerProvider(wrapped=TracerProvider(resource=resource), config=TraceConfig())\n",
"tracer_provider.add_span_processor(SimpleSpanProcessor(OTLPSpanExporter(endpoint)))\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": 3,
"metadata": {},
"outputs": [],
"source": [
"with tracer.start_as_current_chain_span(\n",
" \"plain-text-input-plain-text-output\",\n",
" 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": 4,
"metadata": {},
"outputs": [],
"source": [
"with tracer.start_as_current_chain_span(\n",
" \"json-input-json-output\",\n",
") as span:\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": []
}
],
"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
}
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
Union,
cast,
get_args,
overload,
)

import wrapt
Expand All @@ -37,6 +38,7 @@
from opentelemetry.util.types import Attributes, AttributeValue
from typing_extensions import TypeAlias

from openinference.instrumentation.helpers import safe_json_dumps
from openinference.semconv.trace import (
EmbeddingAttributes,
ImageAttributes,
Expand All @@ -53,6 +55,10 @@
Literal[tuple(mime_type.value for mime_type in OpenInferenceMimeTypeValues)],
OpenInferenceMimeTypeValues,
]
OpenInferenceSpanKind: TypeAlias = Union[
Literal[tuple(kind.value.lower() for kind in OpenInferenceSpanKindValues)],
OpenInferenceSpanKindValues,
]


class suppress_tracing:
Expand Down Expand Up @@ -373,28 +379,52 @@ def end(self, end_time: Optional[int] = None) -> None:
span.set_attribute(k, v)
span.end(end_time)

def set_input_value(self, input_value: str, /) -> None:
if not isinstance(input_value, str):
raise TypeError(f"input value must be a string, got {type(input_value).__name__}")
self.set_attribute(INPUT_VALUE, input_value)

def set_input_mime_type(self, input_mime_type: OpenInferenceMimeType, /) -> None:
self.set_attribute(INPUT_MIME_TYPE, _normalize_mime_type(input_mime_type))

def set_output_value(self, output_value: str, /) -> None:
if not isinstance(output_value, str):
raise TypeError(f"output value must be a string, got {type(output_value).__name__}")
self.set_attribute(OUTPUT_VALUE, output_value)

def set_output_mime_type(self, output_mime_type: OpenInferenceMimeType, /) -> None:
self.set_attribute(OUTPUT_MIME_TYPE, _normalize_mime_type(output_mime_type))
def set_input(
self, value: str, mime_type: OpenInferenceMimeType = OpenInferenceMimeTypeValues.TEXT
) -> None:
if not isinstance(value, str):
raise TypeError(f"input value must be a string, got {type(value).__name__}")
mime_type = _normalize_mime_type(mime_type)
self.set_attribute(INPUT_VALUE, value)
self.set_attribute(INPUT_MIME_TYPE, mime_type)

def set_output(
self, value: str, mime_type: OpenInferenceMimeType = OpenInferenceMimeTypeValues.TEXT
) -> None:
if not isinstance(value, str):
raise TypeError(f"output value must be a string, got {type(value).__name__}")
mime_type = _normalize_mime_type(mime_type)
self.set_attribute(OUTPUT_VALUE, value)
self.set_attribute(OUTPUT_MIME_TYPE, mime_type)


class OpenInferenceChainSpan(OpenInferenceSpan):
class ChainSpan(OpenInferenceSpan):
def __init__(self, wrapped: Span, config: TraceConfig) -> None:
super().__init__(wrapped, config)
self.__wrapped__.set_attribute(OPENINFERENCE_SPAN_KIND, CHAIN)

def set_input(self, value: Any, mime_type: Optional[OpenInferenceMimeType] = None) -> None:
normalized_mime_type: OpenInferenceMimeTypeValues
if isinstance(value, str):
normalized_mime_type = (
_normalize_mime_type(mime_type) if mime_type is not None else TEXT
)
else:
if mime_type == TEXT:
raise ValueError('input value must be a string for "text/plain" mime type')
value = safe_json_dumps(value)
normalized_mime_type = JSON
self.set_attribute(INPUT_VALUE, value)
self.set_attribute(INPUT_MIME_TYPE, normalized_mime_type)

def set_output(self, value: str, mime_type: Optional[OpenInferenceMimeType] = None) -> None:
if not isinstance(value, str):
raise TypeError(f"output value must be a string, got {type(value).__name__}")
mime_type = _normalize_mime_type(mime_type)
self.set_attribute(OUTPUT_VALUE, value)
self.set_attribute(OUTPUT_MIME_TYPE, mime_type)
self.set_attributes()


class _IdGenerator(IdGenerator):
"""
Expand Down Expand Up @@ -424,7 +454,7 @@ def id_generator(self) -> IdGenerator:
return self._self_id_generator

@contextmanager
def start_as_current_span(
def _start_as_current_span(
self,
name: str,
context: Optional[Context] = None,
Expand Down Expand Up @@ -487,32 +517,36 @@ def start_chain_span(
name: str,
*,
input_value: Optional[str] = None,
input_mime_type: Optional[OpenInferenceMimeType] = None,
input_mime_type: OpenInferenceMimeType = OpenInferenceMimeTypeValues.TEXT,
**kwargs: Any,
) -> OpenInferenceChainSpan:
) -> ChainSpan:
span = self.start_span(name, **kwargs)
span = OpenInferenceChainSpan(span)
if input_value is not None:
span.set_input_value(input_value)
if input_mime_type is not None:
span.set_input_mime_type(input_mime_type)
span = ChainSpan(span)
input_mime_type = _normalize_mime_type(input_mime_type)
span.set_input(input_value, mime_type=input_mime_type)
return span

@overload
def start_as_current_span(
self,
name: str,
*,
kind: Union[Literal["chain"], Literal[OpenInferenceSpanKindValues.CHAIN]],
**kwargs: Any,
) -> Iterator[ChainSpan]: ...

@contextmanager
def start_as_current_chain_span(
def start_as_current_span(
self,
name: str,
*,
input_value: Optional[str] = None,
input_mime_type: Optional[OpenInferenceMimeType] = None,
kind: Optional[OpenInferenceSpanKind] = None,
**kwargs: Any,
) -> Iterator[OpenInferenceChainSpan]:
with self.start_as_current_span(name, **kwargs) as span:
span = OpenInferenceChainSpan(span, config=self._self_config)
if input_value is not None:
span.set_input_value(input_value)
if input_mime_type is not None:
span.set_input_mime_type(input_mime_type)
) -> Iterator[Span]:
kind = _normalize_span_kind(kind)
with self._start_as_current_span(name, **kwargs) as span:
if kind is not None:
span = _as_openinference_span(span, self._self_config, kind)
yield span


Expand All @@ -532,18 +566,39 @@ def is_base64_url(url: str) -> bool:
return url.startswith("data:image/") and "base64" in url


def _normalize_mime_type(mime_type: OpenInferenceMimeType) -> OpenInferenceMimeTypeValues:
def _normalize_mime_type(mime_type: OpenInferenceMimeType) -> str:
if not isinstance(mime_type, OpenInferenceMimeTypeValues):
try:
return OpenInferenceMimeTypeValues(mime_type)
mime_type = OpenInferenceMimeTypeValues(mime_type)
except ValueError:
valid_mime_types = ", ".join(
f'"{mime_type.value}"' for mime_type in OpenInferenceMimeTypeValues
)
raise ValueError(
f"Invalid mime type: {mime_type}. Valid mime types are: {valid_mime_types}"
)
return mime_type
return mime_type.value


def _normalize_span_kind(kind: OpenInferenceSpanKind) -> OpenInferenceSpanKindValues:
try:
return OpenInferenceSpanKindValues(kind)
except ValueError:
valid_span_kinds = ", ".join(f'"{kind.value}"' for kind in OpenInferenceSpanKindValues)
raise ValueError(f"Invalid span kind: {kind}. Valid span kinds are: {valid_span_kinds}")


def _infer_mime_type(value: Any) -> OpenInferenceMimeType:
if isinstance(value, str):
return TEXT


def _as_openinference_span(
span: Span, config: TraceConfig, kind: OpenInferenceSpanKindValues
) -> OpenInferenceSpan:
if kind == CHAIN:
return ChainSpan(span, config=config)
raise NotImplementedError(f"Span kind {kind.value} is not yet supported")


# span attributes
Expand All @@ -553,5 +608,9 @@ def _normalize_mime_type(mime_type: OpenInferenceMimeType) -> OpenInferenceMimeT
OUTPUT_MIME_TYPE = SpanAttributes.OUTPUT_MIME_TYPE
OUTPUT_VALUE = SpanAttributes.OUTPUT_VALUE

# mime types
JSON = OpenInferenceMimeTypeValues.JSON.value
TEXT = OpenInferenceMimeTypeValues.TEXT.value

# span kinds
CHAIN = OpenInferenceSpanKindValues.CHAIN.value

0 comments on commit bef939f

Please sign in to comment.