feat(python): openai instrumentator (#35)

.github/workflows/python-CI.yaml

@@ -0,0 +1,26 @@
+name: Python CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    paths:
+      - "python/**"
+
+defaults:
+  run:
+    working-directory: ./python
+
+jobs:
+  ci:
+    name: CI Python
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: |
+            3.8
+            3.11
+      - run: pip install tox==4.11.4
+      - run: tox run-parallel

README.md

@@ -20,9 +20,10 @@ OpenInference provides a set of instrumentations for popular machine learning SD
 
 ## Python
 
-| Package                                                                                       | Description                                   |
-| --------------------------------------------------------------------------------------------- | --------------------------------------------- |
-| [`openinference-semantic-conventions`](./python/openinference-semantic-conventions/README.md) | Semantic conventions for tracing of LLM Apps. |
+| Package                                                                                                            | Description                                   |
+|--------------------------------------------------------------------------------------------------------------------|-----------------------------------------------|
+| [`openinference-semantic-conventions`](./python/openinference-semantic-conventions/README.md)                      | Semantic conventions for tracing of LLM Apps. |
+| [`openinference-instrumentation-openai`](./python/instrumentation/openinference-instrumentation-openai/README.rst) | OpenInference Instrumentation for OpenAI SDK. |
 
 ## JavaScript
 

python/dev-requirements.txt

@@ -0,0 +1,3 @@
+pytest == 7.4.4
+ruff == 0.1.11
+mypy == 1.8.0

python/instrumentation/openinference-instrumentation-openai/examples/chat_completions.py

@@ -0,0 +1,25 @@
+import openai
+from openinference.instrumentation.openai import OpenAIInstrumentor
+from opentelemetry import trace as trace_api
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk import trace as trace_sdk
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+
+resource = Resource(attributes={})
+tracer_provider = trace_sdk.TracerProvider(resource=resource)
+span_exporter = OTLPSpanExporter(endpoint="http://127.0.0.1:6006/v1/traces")
+span_processor = SimpleSpanProcessor(span_exporter=span_exporter)
+tracer_provider.add_span_processor(span_processor=span_processor)
+trace_api.set_tracer_provider(tracer_provider=tracer_provider)
+
+OpenAIInstrumentor().instrument()
+
+
+if __name__ == "__main__":
+    response = openai.OpenAI().chat.completions.create(
+        model="gpt-3.5-turbo",
+        messages=[{"role": "user", "content": "Write a haiku."}],
+        max_tokens=20,
+    )
+    print(response.choices[0].message.content)

python/instrumentation/openinference-instrumentation-openai/examples/chat_completions_async_stream.py

@@ -0,0 +1,37 @@
+import asyncio
+
+import openai
+from openinference.instrumentation.openai import OpenAIInstrumentor
+from opentelemetry import trace as trace_api
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk import trace as trace_sdk
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+
+resource = Resource(attributes={})
+tracer_provider = trace_sdk.TracerProvider(resource=resource)
+span_exporter = OTLPSpanExporter(endpoint="http://127.0.0.1:6006/v1/traces")
+span_processor = SimpleSpanProcessor(span_exporter=span_exporter)
+tracer_provider.add_span_processor(span_processor=span_processor)
+trace_api.set_tracer_provider(tracer_provider=tracer_provider)
+
+OpenAIInstrumentor().instrument()
+
+
+async def chat_completions(**kwargs):
+    client = openai.AsyncOpenAI()
+    async for chunk in await client.chat.completions.create(**kwargs):
+        if content := chunk.choices[0].delta.content:
+            print(content, end="")
+    print()
+
+
+if __name__ == "__main__":
+    asyncio.run(
+        chat_completions(
+            model="gpt-3.5-turbo",
+            messages=[{"role": "user", "content": "Write a haiku."}],
+            max_tokens=20,
+            stream=True,
+        ),
+    )

python/instrumentation/openinference-instrumentation-openai/examples/embeddings.py

@@ -0,0 +1,24 @@
+import openai
+from openinference.instrumentation.openai import OpenAIInstrumentor
+from opentelemetry import trace as trace_api
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk import trace as trace_sdk
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+
+resource = Resource(attributes={})
+tracer_provider = trace_sdk.TracerProvider(resource=resource)
+span_exporter = OTLPSpanExporter(endpoint="http://127.0.0.1:6006/v1/traces")
+span_processor = SimpleSpanProcessor(span_exporter=span_exporter)
+tracer_provider.add_span_processor(span_processor=span_processor)
+trace_api.set_tracer_provider(tracer_provider=tracer_provider)
+
+OpenAIInstrumentor().instrument()
+
+
+if __name__ == "__main__":
+    response = openai.OpenAI().embeddings.create(
+        model="text-embedding-ada-002",
+        input="hello world",
+    )
+    print(response.data[0].embedding)

python/instrumentation/openinference-instrumentation-openai/examples/with_httpx_instrumentor.py

@@ -0,0 +1,29 @@
+from importlib import import_module
+
+from openinference.instrumentation.openai import OpenAIInstrumentor
+from opentelemetry import trace as trace_api
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor
+from opentelemetry.sdk import trace as trace_sdk
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+
+resource = Resource(attributes={})
+tracer_provider = trace_sdk.TracerProvider(resource=resource)
+span_exporter = OTLPSpanExporter(endpoint="http://127.0.0.1:6006/v1/traces")
+span_processor = SimpleSpanProcessor(span_exporter=span_exporter)
+tracer_provider.add_span_processor(span_processor=span_processor)
+trace_api.set_tracer_provider(tracer_provider=tracer_provider)
+
+HTTPXClientInstrumentor().instrument()
+OpenAIInstrumentor().instrument()
+
+
+if __name__ == "__main__":
+    openai = import_module("openai")
+    response = openai.OpenAI().chat.completions.create(
+        model="gpt-3.5-turbo",
+        messages=[{"role": "user", "content": "Write a haiku."}],
+        max_tokens=20,
+    )
+    print(response.choices[0].message.content)

python/instrumentation/openinference-instrumentation-openai/pyproject.toml

@@ -10,30 +10,37 @@ readme = "README.rst"
 license = "Apache-2.0"
 requires-python = ">=3.8, <3.12"
 authors = [
-    { name = "OpenInference Authors", email = "[email protected]" },
+  { name = "OpenInference Authors", email = "[email protected]" },
 ]
 classifiers = [
-    "Development Status :: 4 - Beta",
-    "Intended Audience :: Developers",
-    "License :: OSI Approved :: Apache Software License",
-    "Programming Language :: Python",
-    "Programming Language :: Python :: 3",
-    "Programming Language :: Python :: 3.8",
-    "Programming Language :: Python :: 3.9",
-    "Programming Language :: Python :: 3.10",
-    "Programming Language :: Python :: 3.11",
+  "Development Status :: 4 - Beta",
+  "Intended Audience :: Developers",
+  "License :: OSI Approved :: Apache Software License",
+  "Programming Language :: Python",
+  "Programming Language :: Python :: 3",
+  "Programming Language :: Python :: 3.8",
+  "Programming Language :: Python :: 3.9",
+  "Programming Language :: Python :: 3.10",
+  "Programming Language :: Python :: 3.11",
 ]
 dependencies = [
-    "opentelemetry-api",
-    "opentelemetry-instrumentation",
-    "opentelemetry-semantic-conventions",
-    "openinference-semantic-conventions",
-    "wrapt",
+  "opentelemetry-api",
+  "opentelemetry-instrumentation",
+  "opentelemetry-semantic-conventions",
+  "openinference-semantic-conventions",
+  "wrapt",
 ]
 
 [project.optional-dependencies]
+instruments = [
+  "openai >= 1.0.0",
+]
 test = [
-    "openai == 1.0.0",
+  "openai == 1.0.0",
+  "opentelemetry-sdk",
+  "opentelemetry-instrumentation-httpx",
+  "respx",
+  "numpy",
 ]
 
 [project.urls]
@@ -44,8 +51,8 @@ path = "src/openinference/instrumentation/openai/version.py"
 
 [tool.hatch.build.targets.sdist]
 include = [
-    "/src",
-    "/tests",
+  "/src",
+  "/tests",
 ]
 
 [tool.hatch.build.targets.wheel]

python/instrumentation/openinference-instrumentation-openai/src/openinference/instrumentation/openai/__init__.py

@@ -0,0 +1,55 @@
+import logging
+from importlib import import_module
+from typing import Any, Collection
+
+from openinference.instrumentation.openai._request import (
+    _AsyncRequest,
+    _Request,
+)
+from openinference.instrumentation.openai.package import _instruments
+from openinference.instrumentation.openai.version import __version__
+from opentelemetry import trace as trace_api
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor  # type: ignore
+from wrapt import wrap_function_wrapper
+
+logger = logging.getLogger(__name__)
+logger.addHandler(logging.NullHandler())
+
+_MODULE = "openai"
+
+
+class OpenAIInstrumentor(BaseInstrumentor):  # type: ignore
+    """
+    An instrumentor for openai
+    """
+
+    __slots__ = (
+        "_original_request",
+        "_original_async_request",
+    )
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return _instruments
+
+    def _instrument(self, **kwargs: Any) -> None:
+        if not (tracer_provider := kwargs.get("tracer_provider")):
+            tracer_provider = trace_api.get_tracer_provider()
+        tracer = trace_api.get_tracer(__name__, __version__, tracer_provider)
+        openai = import_module(_MODULE)
+        self._original_request = openai.OpenAI.request
+        self._original_async_request = openai.AsyncOpenAI.request
+        wrap_function_wrapper(
+            module=_MODULE,
+            name="OpenAI.request",
+            wrapper=_Request(tracer=tracer, openai=openai),
+        )
+        wrap_function_wrapper(
+            module=_MODULE,
+            name="AsyncOpenAI.request",
+            wrapper=_AsyncRequest(tracer=tracer, openai=openai),
+        )
+
+    def _uninstrument(self, **kwargs: Any) -> None:
+        openai = import_module(_MODULE)
+        openai.OpenAI.request = self._original_request
+        openai.AsyncOpenAI.request = self._original_async_request