feat: Pinecone Instrumentation

.github/workflows/ci-traceloop-sdk.yml

@@ -41,5 +41,7 @@ jobs:
       - name: Test with pytest
         env:
           OPENAI_API_KEY: ${{ secrets.OPEN_API_KEY }}
+          PINECONE_API_KEY: ${{ secrets.PINECONE_API_KEY }}
+          PINECONE_ENVIRONMENT: ${{ secrets.PINECONE_ENVIRONMENT }}
         working-directory: packages/traceloop-sdk
         run: poetry run pytest

README.md

@@ -85,7 +85,7 @@ OpenLLMetry can instrument everything that [OpenTelemetry already instruments](h
 
 ### Vector DBs
 
-- [ ] Pinecone
+- [x] Pinecone
 - [ ] Chroma
 - [ ] Weaviate
 

packages/opentelemetry-instrumentation-haystack/pyproject.toml

@@ -11,7 +11,7 @@ addopts = "--cov --cov-report html:'../../coverage/packages/opentelemetry-instru
 
 [tool.poetry]
 name = "opentelemetry-instrumentation-haystack"
-version = "0.0.1"
+version = "0.0.3"
 description = "OpenTelemetry Haystack instrumentation"
 authors = [
   "Gal Kleinman <gal@traceloop.com>",
@@ -28,7 +28,7 @@ include = "opentelemetry/instrumentation/haystack"
 python = ">=3.8.1,<3.12"
 opentelemetry-api = "^1.19.0"
 opentelemetry-instrumentation = "^0.40b0"
-opentelemetry-semantic-conventions-ai = "^0.0.1"
+opentelemetry-semantic-conventions-ai = "^0.0.3"
 
 [tool.poetry.group.dev.dependencies]
 autopep8 = "2.0.4"

packages/opentelemetry-instrumentation-openai/pyproject.toml

@@ -11,7 +11,7 @@ addopts = "--cov --cov-report html:'../../coverage/packages/opentelemetry-instru
 
 [tool.poetry]
 name = "opentelemetry-instrumentation-openai"
-version = "0.0.5"
+version = "0.0.8"
 description = "OpenTelemetry OpenAI instrumentation"
 authors = [
   "Gal Kleinman <gal@traceloop.com>",
@@ -27,7 +27,7 @@ include = "opentelemetry/instrumentation/openai"
 python = ">=3.8.1,<3.12"
 opentelemetry-api = "^1.19.0"
 opentelemetry-instrumentation = "^0.40b0"
-opentelemetry-semantic-conventions-ai = "^0.0.1"
+opentelemetry-semantic-conventions-ai = "^0.0.3"
 
 [tool.poetry.group.dev.dependencies]
 autopep8 = "2.0.4"

packages/opentelemetry-instrumentation-pinecone/opentelemetry/instrumentation/pinecone/__init__.py

@@ -1 +1,137 @@
 """OpenTelemetry Pinecone instrumentation"""
+
+import logging
+from typing import Collection
+from wrapt import wrap_function_wrapper
+import pinecone
+
+from opentelemetry import context as context_api
+from opentelemetry.trace import get_tracer, SpanKind
+from opentelemetry.trace.status import Status, StatusCode
+
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.instrumentation.utils import (
+    _SUPPRESS_INSTRUMENTATION_KEY,
+    unwrap,
+)
+
+from opentelemetry.semconv.ai import SpanAttributes
+
+logger = logging.getLogger(__name__)
+
+_instruments = ("pinecone-client ~= 2.2.2",)
+__version__ = "0.1.0"
+
+WRAPPED_METHODS = [
+    {
+        "object": "GRPCIndex",
+        "method": "query",
+        "span_name": "pinecone.query",
+    },
+    {
+        "object": "GRPCIndex",
+        "method": "upsert",
+        "span_name": "pinecone.upsert",
+    },
+    {
+        "object": "GRPCIndex",
+        "method": "delete",
+        "span_name": "pinecone.delete",
+    },
+]
+
+
+def _set_span_attribute(span, name, value):
+    if value is not None:
+        if value != "":
+            span.set_attribute(name, value)
+    return
+
+
+def _set_input_attributes(span, kwargs):
+    pass
+
+
+def _set_response_attributes(span, response):
+    pass
+
+
+def _with_tracer_wrapper(func):
+    """Helper for providing tracer for wrapper functions."""
+
+    def _with_tracer(tracer, to_wrap):
+        def wrapper(wrapped, instance, args, kwargs):
+            # prevent double wrapping
+            if hasattr(wrapped, "__wrapped__"):
+                return wrapped(*args, **kwargs)
+
+            return func(tracer, to_wrap, wrapped, instance, args, kwargs)
+
+        return wrapper
+
+    return _with_tracer
+
+
+@_with_tracer_wrapper
+def _wrap(tracer, to_wrap, wrapped, instance, args, kwargs):
+    """Instruments and calls every function defined in TO_WRAP."""
+    if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
+        return wrapped(*args, **kwargs)
+
+    name = to_wrap.get("span_name")
+    with tracer.start_as_current_span(
+        name,
+        kind=SpanKind.CLIENT,
+        attributes={
+            SpanAttributes.VECTOR_DB_VENDOR: "Pinecone",
+        },
+    ) as span:
+        try:
+            if span.is_recording():
+                _set_input_attributes(span, kwargs)
+
+        except Exception as ex:  # pylint: disable=broad-except
+            logger.warning(
+                "Failed to set input attributes for openai span, error: %s", str(ex)
+            )
+
+        response = wrapped(*args, **kwargs)
+
+        if response:
+            try:
+                if span.is_recording():
+                    _set_response_attributes(span, response)
+
+            except Exception as ex:  # pylint: disable=broad-except
+                logger.warning(
+                    "Failed to set response attributes for openai span, error: %s",
+                    str(ex),
+                )
+            if span.is_recording():
+                span.set_status(Status(StatusCode.OK))
+
+        return response
+
+
+class PineconeInstrumentor(BaseInstrumentor):
+    """An instrumentor for Pinecone's client library."""
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return _instruments
+
+    def _instrument(self, **kwargs):
+        tracer_provider = kwargs.get("tracer_provider")
+        tracer = get_tracer(__name__, __version__, tracer_provider)
+        for wrapped_method in WRAPPED_METHODS:
+            wrap_object = wrapped_method.get("object")
+            wrap_method = wrapped_method.get("method")
+            wrap_function_wrapper(
+                "pinecone",
+                f"{wrap_object}.{wrap_method}",
+                _wrap(tracer, wrapped_method),
+            )
+
+    def _uninstrument(self, **kwargs):
+        for wrapped_method in WRAPPED_METHODS:
+            wrap_object = wrapped_method.get("object")
+            unwrap(f"pinecone.{wrap_object}", wrapped_method.get("method"))