feat: LanceDB Integration

.cz.toml

@@ -22,6 +22,8 @@ version_files = [
     "packages/opentelemetry-instrumentation-haystack/opentelemetry/instrumentation/haystack/version.py",
     "packages/opentelemetry-instrumentation-langchain/pyproject.toml:^version",
     "packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/version.py",
+    "packages/opentelemetry-instrumentation-lancedb/pyproject.toml:^version",
+    "packages/opentelemetry-instrumentation-lancedb/opentelemetry/instrumentation/lancedb/version.py",
     "packages/opentelemetry-instrumentation-milvus/pyproject.toml:^version",
     "packages/opentelemetry-instrumentation-milvus/opentelemetry/instrumentation/milvus/version.py",
     "packages/opentelemetry-instrumentation-mistralai/pyproject.toml:^version",

README.md

@@ -128,6 +128,8 @@ OpenLLMetry can instrument everything that [OpenTelemetry already instruments](h
 - ✅ Weaviate
 - ✅ Milvus
 - ✅ Marqo
+- ✅ LanceDB
+
 
 ### Frameworks
 

packages/opentelemetry-instrumentation-lancedb/.flake8

@@ -0,0 +1,11 @@
+[flake8]
+exclude =
+	.git,
+	__pycache__,
+	build,
+	dist,
+	.tox,
+	venv,
+	.venv,
+	.pytest_cache
+max-line-length = 120

packages/opentelemetry-instrumentation-lancedb/.python-version

@@ -0,0 +1 @@
+3.9.5

packages/opentelemetry-instrumentation-lancedb/README.md

@@ -0,0 +1,21 @@
+# OpenTelemetry LanceDB Instrumentation
+
+<a href="https://pypi.org/project/opentelemetry-instrumentation-lancedb/">
+    <img src="https://badge.fury.io/py/opentelemetry-instrumentation-lancedb.svg">
+</a>
+
+This library allows tracing client-side calls to LanceDB sent with the official [LanceDB library](https://github.com/lancedb/lancedb).
+
+## Installation
+
+```bash
+pip install opentelemetry-instrumentation-lancedb
+```
+
+## Example usage
+
+```python
+from opentelemetry.instrumentation.lancedb import LanceInstrumentor
+
+LanceInstrumentor().instrument()
+```

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

@@ -0,0 +1,75 @@
+"""OpenTelemetry LanceDB instrumentation"""
+
+import logging
+import lancedb.table
+
+from typing import Collection
+
+from opentelemetry.instrumentation.lancedb.config import Config
+from opentelemetry.trace import get_tracer
+from wrapt import wrap_function_wrapper
+
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.instrumentation.utils import unwrap
+
+from opentelemetry.instrumentation.lancedb.wrapper import _wrap
+from opentelemetry.instrumentation.lancedb.version import __version__
+
+logger = logging.getLogger(__name__)
+
+_instruments = ("lancedb >= 0.9.0",)
+
+WRAPPED_METHODS = [
+    {
+        "package": lancedb.table,
+        "object": "LanceTable",
+        "method": "add",
+        "span_name": "lancedb.add"
+    },
+    {
+        "package": lancedb.table,
+        "object": "LanceTable",
+        "method": "search",
+        "span_name": "lancedb.search"
+    },
+    {
+        "package": lancedb.table,
+        "object": "LanceTable",
+        "method": "delete",
+        "span_name": "lancedb.delete"
+    },
+]
+
+
+class LanceInstrumentor(BaseInstrumentor):
+    """An instrumentor for Lance DB's client library."""
+
+    def __init__(self, exception_logger=None):
+        super().__init__()
+        Config.exception_logger = exception_logger
+
+    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_package = wrapped_method.get("package")
+            wrap_object = wrapped_method.get("object")
+            wrap_method = wrapped_method.get("method")
+            if getattr(wrap_package, wrap_object, None):
+                wrap_function_wrapper(
+                    wrap_package,
+                    f"{wrap_object}.{wrap_method}",
+                    _wrap(tracer, wrapped_method),
+                )
+
+    def _uninstrument(self, **kwargs):
+        for wrapped_method in WRAPPED_METHODS:
+            wrap_package = wrapped_method.get("package")
+            wrap_object = wrapped_method.get("object")
+
+            wrapped = getattr(wrap_package, wrap_object, None)
+            if wrapped:
+                unwrap(wrapped, wrapped_method.get("method"))

packages/opentelemetry-instrumentation-lancedb/opentelemetry/instrumentation/lancedb/config.py

@@ -0,0 +1,2 @@
+class Config:
+    exception_logger = None

packages/opentelemetry-instrumentation-lancedb/opentelemetry/instrumentation/lancedb/utils.py

@@ -0,0 +1,28 @@
+import logging
+import traceback
+from opentelemetry.instrumentation.lancedb.config import Config
+
+
+def dont_throw(func):
+    """
+    A decorator that wraps the passed in function and logs exceptions instead of throwing them.
+
+    @param func: The function to wrap
+    @return: The wrapper function
+    """
+    # Obtain a logger specific to the function's module
+    logger = logging.getLogger(func.__module__)
+
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except Exception as e:
+            logger.debug(
+                "OpenLLMetry failed to trace in %s, error: %s",
+                func.__name__,
+                traceback.format_exc(),
+            )
+            if Config.exception_logger:
+                Config.exception_logger(e)
+
+    return wrapper

packages/opentelemetry-instrumentation-lancedb/opentelemetry/instrumentation/lancedb/version.py

@@ -0,0 +1 @@
+__version__ = "0.26.3"

packages/opentelemetry-instrumentation-lancedb/opentelemetry/instrumentation/lancedb/wrapper.py

@@ -0,0 +1,86 @@
+from opentelemetry.instrumentation.lancedb.utils import dont_throw
+from opentelemetry.semconv.trace import SpanAttributes
+
+from opentelemetry import context as context_api
+from opentelemetry.instrumentation.utils import (
+    _SUPPRESS_INSTRUMENTATION_KEY,
+)
+from opentelemetry.semconv_ai import SpanAttributes as AISpanAttributes
+
+
+def _with_tracer_wrapper(func):
+    """Helper for providing tracer for wrapper functions."""
+
+    def _with_tracer(tracer, to_wrap):
+        def wrapper(wrapped, instance, args, kwargs):
+            return func(tracer, to_wrap, wrapped, instance, args, kwargs)
+
+        return wrapper
+
+    return _with_tracer
+
+
+def _set_span_attribute(span, name, value):
+    if value is not None:
+        if value != "":
+            span.set_attribute(name, value)
+    return
+
+
+@_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) as span:
+        span.set_attribute(SpanAttributes.DB_SYSTEM, "lancedb")
+        span.set_attribute(SpanAttributes.DB_OPERATION, to_wrap.get("method"))
+
+        if to_wrap.get("method") == "add":
+            _set_add_attributes(span, kwargs)
+        elif to_wrap.get("method") == "search":
+            _set_search_attributes(span, kwargs)
+        elif to_wrap.get("method") == "delete":
+            _set_delete_attributes(span, kwargs)
+
+        return_value = wrapped(*args, **kwargs)
+
+    return return_value
+
+
+def _encode_query(_query):
+    _query_str = None
+    if _query:
+        _query_str = str(_query)
+
+    return _query_str
+
+
+def _count_or_none(obj):
+    if obj:
+        return len(obj)
+
+    return None
+
+
+@dont_throw
+def _set_add_attributes(span, kwargs):
+    _set_span_attribute(
+        span, AISpanAttributes.MILVUS_INSERT_DATA_COUNT, _count_or_none(kwargs.get("data"))
+    )
+
+
+@dont_throw
+def _set_search_attributes(span, kwargs):
+    _set_span_attribute(
+        span, AISpanAttributes.MILVUS_SEARCH_FILTER, _encode_query(kwargs.get("query"))
+    )
+
+
+@dont_throw
+def _set_delete_attributes(span, kwargs):
+    _set_span_attribute(
+        span, AISpanAttributes.CHROMADB_DELETE_WHERE, kwargs.get("where")
+    )