From a36fb0f2521d81dad9de649508a5670bb2224cd8 Mon Sep 17 00:00:00 2001 From: sanyamk23 Date: Sat, 1 Aug 2026 00:32:33 +0530 Subject: [PATCH] feat(azure-search): add Azure AI Search instrumentation and semantic conventions --- .../.python-version | 1 + .../README.md | 80 +++++ .../opentelemetry/__init__.py | 0 .../opentelemetry/instrumentation/__init__.py | 0 .../instrumentation/azure_search/__init__.py | 327 ++++++++++++++++++ .../instrumentation/azure_search/config.py | 2 + .../instrumentation/azure_search/utils.py | 42 +++ .../instrumentation/azure_search/version.py | 1 + .../poetry.toml | 2 + .../project.json | 75 ++++ .../pyproject.toml | 74 ++++ .../tests/__init__.py | 0 .../tests/conftest.py | 27 ++ .../test_azure_search_instrumentation.py | 326 +++++++++++++++++ .../opentelemetry/semconv_ai/__init__.py | 18 + 15 files changed, 975 insertions(+) create mode 100644 packages/opentelemetry-instrumentation-azure-search/.python-version create mode 100644 packages/opentelemetry-instrumentation-azure-search/README.md create mode 100644 packages/opentelemetry-instrumentation-azure-search/opentelemetry/__init__.py create mode 100644 packages/opentelemetry-instrumentation-azure-search/opentelemetry/instrumentation/__init__.py create mode 100644 packages/opentelemetry-instrumentation-azure-search/opentelemetry/instrumentation/azure_search/__init__.py create mode 100644 packages/opentelemetry-instrumentation-azure-search/opentelemetry/instrumentation/azure_search/config.py create mode 100644 packages/opentelemetry-instrumentation-azure-search/opentelemetry/instrumentation/azure_search/utils.py create mode 100644 packages/opentelemetry-instrumentation-azure-search/opentelemetry/instrumentation/azure_search/version.py create mode 100644 packages/opentelemetry-instrumentation-azure-search/poetry.toml create mode 100644 packages/opentelemetry-instrumentation-azure-search/project.json create mode 100644 packages/opentelemetry-instrumentation-azure-search/pyproject.toml create mode 100644 packages/opentelemetry-instrumentation-azure-search/tests/__init__.py create mode 100644 packages/opentelemetry-instrumentation-azure-search/tests/conftest.py create mode 100644 packages/opentelemetry-instrumentation-azure-search/tests/test_azure_search_instrumentation.py diff --git a/packages/opentelemetry-instrumentation-azure-search/.python-version b/packages/opentelemetry-instrumentation-azure-search/.python-version new file mode 100644 index 0000000000..c8cfe39591 --- /dev/null +++ b/packages/opentelemetry-instrumentation-azure-search/.python-version @@ -0,0 +1 @@ +3.10 diff --git a/packages/opentelemetry-instrumentation-azure-search/README.md b/packages/opentelemetry-instrumentation-azure-search/README.md new file mode 100644 index 0000000000..e83586f7bf --- /dev/null +++ b/packages/opentelemetry-instrumentation-azure-search/README.md @@ -0,0 +1,80 @@ +# OpenTelemetry Azure AI Search Instrumentation + +[![License](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](https://opensource.org/licenses/Apache-2.0) + +This library allows tracing client-side calls to Azure AI Search using OpenTelemetry. + +## Installation + +```bash +pip install opentelemetry-instrumentation-azure-search +``` + +## Supported Operations + +### SearchClient + +| Method | Span Name | Attributes Captured | +|--------|-----------|-------------------| +| `search()` | `azure_search.search` | search_text, top, filter, index_name | +| `get_document()` | `azure_search.get_document` | index_name | +| `autocomplete()` | `azure_search.autocomplete` | autocomplete_text, index_name | +| `suggest()` | `azure_search.suggest` | suggest_text, index_name | +| `index_documents()` | `azure_search.index_documents` | documents_count, succeeded_count | +| `upload_documents()` | `azure_search.upload_documents` | documents_count, succeeded_count | +| `merge_documents()` | `azure_search.merge_documents` | documents_count, succeeded_count | +| `merge_or_upload_documents()` | `azure_search.merge_or_upload_documents` | documents_count, succeeded_count | +| `delete_documents()` | `azure_search.delete_documents` | documents_count, succeeded_count | +| `get_document_count()` | `azure_search.get_document_count` | documents_count | + +### SearchIndexClient + +| Method | Span Name | Attributes Captured | +|--------|-----------|-------------------| +| `create_index()` | `azure_search.create_index` | index_name | +| `create_or_update_index()` | `azure_search.create_or_update_index` | index_name | +| `delete_index()` | `azure_search.delete_index` | index_name | +| `get_index()` | `azure_search.get_index` | index_name | +| `list_indexes()` | `azure_search.list_indexes` | — | +| `get_index_statistics()` | `azure_search.get_index_statistics` | index_doc_count, index_size_bytes | +| `analyze_text()` | `azure_search.analyze_text` | index_name | +| `get_service_statistics()` | `azure_search.get_service_statistics` | service_usage, service_limit | + +### SearchIndexerClient + +| Method | Span Name | Attributes Captured | +|--------|-----------|-------------------| +| `create_indexer()` | `azure_search.create_indexer` | indexer_name | +| `create_or_update_indexer()` | `azure_search.create_or_update_indexer` | indexer_name | +| `delete_indexer()` | `azure_search.delete_indexer` | indexer_name | +| `get_indexer()` | `azure_search.get_indexer` | indexer_name | +| `get_indexers()` | `azure_search.get_indexers` | — | +| `get_indexer_status()` | `azure_search.get_indexer_status` | indexer_status | +| `run_indexer()` | `azure_search.run_indexer` | indexer_name | +| `reset_indexer()` | `azure_search.reset_indexer` | indexer_name | +| `create_data_source_connection()` | `azure_search.create_data_source_connection` | — | +| `create_skillset()` | `azure_search.create_skillset` | — | +| `get_skillset()` | `azure_search.get_skillset` | — | +| `delete_skillset()` | `azure_search.delete_skillset` | — | + +## Usage + +### Auto-instrumentation via Traceloop SDK + +```python +from traceloop.sdk import Traceloop + +Traceloop.init(app_name="my-rag-app") +``` + +### Manual instrumentation + +```python +from opentelemetry.instrumentation.azure_search import AzureSearchInstrumentor + +AzureSearchInstrumentor().instrument() +``` + +## License + +This project is licensed under the Apache License 2.0 — see the [LICENSE](LICENSE) file for details. diff --git a/packages/opentelemetry-instrumentation-azure-search/opentelemetry/__init__.py b/packages/opentelemetry-instrumentation-azure-search/opentelemetry/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/packages/opentelemetry-instrumentation-azure-search/opentelemetry/instrumentation/__init__.py b/packages/opentelemetry-instrumentation-azure-search/opentelemetry/instrumentation/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/packages/opentelemetry-instrumentation-azure-search/opentelemetry/instrumentation/azure_search/__init__.py b/packages/opentelemetry-instrumentation-azure-search/opentelemetry/instrumentation/azure_search/__init__.py new file mode 100644 index 0000000000..aa1d69015f --- /dev/null +++ b/packages/opentelemetry-instrumentation-azure-search/opentelemetry/instrumentation/azure_search/__init__.py @@ -0,0 +1,327 @@ +"""OpenTelemetry Azure AI Search instrumentation""" + +import importlib +import logging +from typing import Collection + +from opentelemetry import context as context_api +from opentelemetry.instrumentation.instrumentor import BaseInstrumentor +from opentelemetry.instrumentation.utils import ( + _SUPPRESS_INSTRUMENTATION_KEY, + unwrap, +) +from opentelemetry.semconv_ai import SpanAttributes +from opentelemetry.trace import SpanKind, get_tracer +from opentelemetry.trace.status import Status, StatusCode +from wrapt import wrap_function_wrapper + +from opentelemetry.instrumentation.azure_search.config import Config +from opentelemetry.instrumentation.azure_search.utils import ( + dont_throw, + set_span_attribute, +) +from opentelemetry.instrumentation.azure_search.version import __version__ + +logger = logging.getLogger(__name__) + +_instruments = ("azure-search-documents >= 11.4.0",) + +WRAPPED_SEARCH_CLIENT_METHODS = [ + {"method": "search", "span_name": "azure_search.search"}, + {"method": "get_document", "span_name": "azure_search.get_document"}, + {"method": "autocomplete", "span_name": "azure_search.autocomplete"}, + {"method": "suggest", "span_name": "azure_search.suggest"}, + {"method": "index_documents", "span_name": "azure_search.index_documents"}, + {"method": "upload_documents", "span_name": "azure_search.upload_documents"}, + {"method": "merge_documents", "span_name": "azure_search.merge_documents"}, + {"method": "merge_or_upload_documents", "span_name": "azure_search.merge_or_upload_documents"}, + {"method": "delete_documents", "span_name": "azure_search.delete_documents"}, + {"method": "get_document_count", "span_name": "azure_search.get_document_count"}, +] + +WRAPPED_INDEX_CLIENT_METHODS = [ + {"method": "create_index", "span_name": "azure_search.create_index"}, + {"method": "create_or_update_index", "span_name": "azure_search.create_or_update_index"}, + {"method": "delete_index", "span_name": "azure_search.delete_index"}, + {"method": "get_index", "span_name": "azure_search.get_index"}, + {"method": "list_indexes", "span_name": "azure_search.list_indexes"}, + {"method": "get_index_statistics", "span_name": "azure_search.get_index_statistics"}, + {"method": "analyze_text", "span_name": "azure_search.analyze_text"}, + {"method": "get_service_statistics", "span_name": "azure_search.get_service_statistics"}, +] + +WRAPPED_INDEXER_CLIENT_METHODS = [ + {"method": "create_indexer", "span_name": "azure_search.create_indexer"}, + {"method": "create_or_update_indexer", "span_name": "azure_search.create_or_update_indexer"}, + {"method": "delete_indexer", "span_name": "azure_search.delete_indexer"}, + {"method": "get_indexer", "span_name": "azure_search.get_indexer"}, + {"method": "get_indexers", "span_name": "azure_search.get_indexers"}, + {"method": "get_indexer_status", "span_name": "azure_search.get_indexer_status"}, + {"method": "run_indexer", "span_name": "azure_search.run_indexer"}, + {"method": "reset_indexer", "span_name": "azure_search.reset_indexer"}, + {"method": "create_data_source_connection", "span_name": "azure_search.create_data_source_connection"}, + {"method": "create_skillset", "span_name": "azure_search.create_skillset"}, + {"method": "get_skillset", "span_name": "azure_search.get_skillset"}, + {"method": "delete_skillset", "span_name": "azure_search.delete_skillset"}, +] + +WRAPPED_METHODS = ( + WRAPPED_SEARCH_CLIENT_METHODS + + WRAPPED_INDEX_CLIENT_METHODS + + WRAPPED_INDEXER_CLIENT_METHODS +) + + +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 + + +@_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) + + span_name = to_wrap.get("span_name") + with tracer.start_as_current_span( + span_name, + kind=SpanKind.CLIENT, + attributes={ + SpanAttributes.VECTOR_DB_VENDOR: "azure_search", + }, + record_exception=False, + set_status_on_exception=False, + ) as span: + _set_input_attributes(span, instance, to_wrap, args, kwargs) + + try: + response = wrapped(*args, **kwargs) + except Exception as e: + span.record_exception(e) + span.set_status(Status(StatusCode.ERROR, str(e))) + raise + + if span.is_recording(): + _set_response_attributes(span, to_wrap, response) + span.set_status(Status(StatusCode.OK)) + + return response + + +@dont_throw +def _set_input_attributes(span, instance, to_wrap, args, kwargs): + method = to_wrap.get("method") + + endpoint = getattr(instance, "_endpoint", None) + if endpoint: + set_span_attribute(span, "server.address", endpoint) + + if method == "search": + set_span_attribute( + span, SpanAttributes.AZURE_SEARCH_SEARCH_TEXT, kwargs.get("search_text"), + ) + set_span_attribute(span, SpanAttributes.AZURE_SEARCH_TOP, kwargs.get("top")) + set_span_attribute( + span, SpanAttributes.AZURE_SEARCH_FILTER, kwargs.get("filter"), + ) + elif method == "autocomplete": + set_span_attribute( + span, + SpanAttributes.AZURE_SEARCH_AUTOCOMPLETE_TEXT, + kwargs.get("search_text"), + ) + elif method == "suggest": + set_span_attribute( + span, + SpanAttributes.AZURE_SEARCH_SUGGEST_TEXT, + kwargs.get("search_text"), + ) + elif method in ( + "index_documents", + "upload_documents", + "merge_documents", + "merge_or_upload_documents", + "delete_documents", + ): + documents = kwargs.get("documents") or (args[0] if args else []) + set_span_attribute( + span, SpanAttributes.AZURE_SEARCH_DOCUMENTS_COUNT, len(documents), + ) + + index_name = getattr(instance, "_index_name", None) + if index_name: + set_span_attribute(span, SpanAttributes.AZURE_SEARCH_INDEX_NAME, index_name) + + if method in ("get_index", "delete_index", "create_index", "create_or_update_index"): + index = kwargs.get("index") or kwargs.get("index_name") or (args[0] if args else None) + if hasattr(index, "name"): + set_span_attribute(span, SpanAttributes.AZURE_SEARCH_INDEX_NAME, index.name) + elif isinstance(index, str): + set_span_attribute(span, SpanAttributes.AZURE_SEARCH_INDEX_NAME, index) + + if method in ("get_indexer", "delete_indexer", "run_indexer", "reset_indexer", "create_indexer"): + indexer = kwargs.get("indexer") or (args[0] if args else None) + if hasattr(indexer, "name"): + set_span_attribute(span, SpanAttributes.AZURE_SEARCH_INDEXER_NAME, indexer.name) + elif isinstance(indexer, str): + set_span_attribute(span, SpanAttributes.AZURE_SEARCH_INDEXER_NAME, indexer) + elif method in ("get_skillset", "delete_skillset", "create_skillset"): + skillset = kwargs.get("skillset") or (args[0] if args else None) + if hasattr(skillset, "name"): + set_span_attribute(span, SpanAttributes.AZURE_SEARCH_INDEXER_NAME, skillset.name) + + +@dont_throw +def _set_response_attributes(span, to_wrap, response): + method = to_wrap.get("method") + + if method == "search" and response is not None: + set_span_attribute(span, SpanAttributes.VECTOR_DB_OPERATION, "search") + elif method == "get_document_count" and response is not None: + set_span_attribute(span, SpanAttributes.AZURE_SEARCH_DOCUMENTS_COUNT, response) + elif method in ( + "index_documents", + "upload_documents", + "merge_documents", + "merge_or_upload_documents", + "delete_documents", + ): + if response is not None and hasattr(response, "results"): + results = list(response.results) + succeeded = sum(1 for r in results if not r.error) + set_span_attribute( + span, SpanAttributes.AZURE_SEARCH_SUCCEEDED_COUNT, succeeded, + ) + set_span_attribute( + span, SpanAttributes.AZURE_SEARCH_DOCUMENTS_COUNT, len(results), + ) + elif method == "get_indexer_status" and response is not None: + status = getattr(response, "status", None) + set_span_attribute(span, SpanAttributes.AZURE_SEARCH_INDEXER_STATUS, status) + elif method == "get_index_statistics" and response is not None: + set_span_attribute( + span, + SpanAttributes.AZURE_SEARCH_INDEX_DOC_COUNT, + getattr(response, "document_count", None), + ) + set_span_attribute( + span, + SpanAttributes.AZURE_SEARCH_INDEX_SIZE_BYTES, + getattr(response, "storage_size", None), + ) + elif method == "get_service_statistics" and response is not None: + counters = getattr(response, "counters", None) + if counters: + set_span_attribute( + span, + SpanAttributes.AZURE_SEARCH_SERVICE_USAGE, + getattr(counters, "search_service_usage", None), + ) + set_span_attribute( + span, + SpanAttributes.AZURE_SEARCH_SERVICE_LIMIT, + getattr(counters, "search_service_limit", None), + ) + elif method in ("create_index", "get_index", "create_or_update_index") and response is not None: + if hasattr(response, "name"): + set_span_attribute(span, SpanAttributes.AZURE_SEARCH_INDEX_NAME, response.name) + + +class AzureSearchInstrumentor(BaseInstrumentor): + """An instrumentor for the Azure AI Search client library. + + Instruments SearchClient, SearchIndexClient, and SearchIndexerClient + to emit OpenTelemetry spans for search, indexing, and indexer operations. + """ + + 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_SEARCH_CLIENT_METHODS: + _instrument_method( + "azure.search.documents", + "SearchClient", + wrapped_method, + tracer, + ) + + for wrapped_method in WRAPPED_INDEX_CLIENT_METHODS: + _instrument_method( + "azure.search.documents.indexes", + "SearchIndexClient", + wrapped_method, + tracer, + ) + + for wrapped_method in WRAPPED_INDEXER_CLIENT_METHODS: + _instrument_method( + "azure.search.documents.indexes", + "SearchIndexerClient", + wrapped_method, + tracer, + ) + + def _uninstrument(self, **kwargs): + for wrapped_method in WRAPPED_SEARCH_CLIENT_METHODS: + _uninstrument_method( + "azure.search.documents", "SearchClient", wrapped_method, + ) + + for wrapped_method in WRAPPED_INDEX_CLIENT_METHODS: + _uninstrument_method( + "azure.search.documents.indexes", + "SearchIndexClient", + wrapped_method, + ) + + for wrapped_method in WRAPPED_INDEXER_CLIENT_METHODS: + _uninstrument_method( + "azure.search.documents.indexes", + "SearchIndexerClient", + wrapped_method, + ) + + +def _instrument_method(module, class_name, wrapped_method, tracer): + method = wrapped_method.get("method") + try: + mod = importlib.import_module(module) + cls = getattr(mod, class_name, None) + if cls and hasattr(cls, method) and callable(getattr(cls, method)): + wrap_function_wrapper( + module, + f"{class_name}.{method}", + _wrap(tracer, wrapped_method), + ) + except (ImportError, ModuleNotFoundError): + pass + + +def _uninstrument_method(module, class_name, wrapped_method): + method = wrapped_method.get("method") + try: + unwrap(f"{module}.{class_name}", method) + except Exception as e: + logger.debug( + "Failed to unwrap %s.%s.%s: %s", + module, + class_name, + method, + e, + ) diff --git a/packages/opentelemetry-instrumentation-azure-search/opentelemetry/instrumentation/azure_search/config.py b/packages/opentelemetry-instrumentation-azure-search/opentelemetry/instrumentation/azure_search/config.py new file mode 100644 index 0000000000..4689e9292f --- /dev/null +++ b/packages/opentelemetry-instrumentation-azure-search/opentelemetry/instrumentation/azure_search/config.py @@ -0,0 +1,2 @@ +class Config: + exception_logger = None diff --git a/packages/opentelemetry-instrumentation-azure-search/opentelemetry/instrumentation/azure_search/utils.py b/packages/opentelemetry-instrumentation-azure-search/opentelemetry/instrumentation/azure_search/utils.py new file mode 100644 index 0000000000..9c1af72d61 --- /dev/null +++ b/packages/opentelemetry-instrumentation-azure-search/opentelemetry/instrumentation/azure_search/utils.py @@ -0,0 +1,42 @@ +import logging +import os +import traceback +from functools import wraps + +from opentelemetry.instrumentation.azure_search.config import Config + + +def dont_throw(func): + """Decorator that wraps the passed in function and logs exceptions instead of throwing them. + + Args: + func: The function to wrap. + + Returns: + The wrapper function. + """ + logger = logging.getLogger(func.__module__) + + @wraps(func) + 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 + + +def set_span_attribute(span, name, value): + if value is not None and value != "": + span.set_attribute(name, value) + + +def is_metrics_enabled() -> bool: + return (os.getenv("TRACELOOP_METRICS_ENABLED") or "true").lower() == "true" diff --git a/packages/opentelemetry-instrumentation-azure-search/opentelemetry/instrumentation/azure_search/version.py b/packages/opentelemetry-instrumentation-azure-search/opentelemetry/instrumentation/azure_search/version.py new file mode 100644 index 0000000000..b7b61f3032 --- /dev/null +++ b/packages/opentelemetry-instrumentation-azure-search/opentelemetry/instrumentation/azure_search/version.py @@ -0,0 +1 @@ +__version__ = "0.62.1" diff --git a/packages/opentelemetry-instrumentation-azure-search/poetry.toml b/packages/opentelemetry-instrumentation-azure-search/poetry.toml new file mode 100644 index 0000000000..ab1033bd37 --- /dev/null +++ b/packages/opentelemetry-instrumentation-azure-search/poetry.toml @@ -0,0 +1,2 @@ +[virtualenvs] +in-project = true diff --git a/packages/opentelemetry-instrumentation-azure-search/project.json b/packages/opentelemetry-instrumentation-azure-search/project.json new file mode 100644 index 0000000000..cc422769f1 --- /dev/null +++ b/packages/opentelemetry-instrumentation-azure-search/project.json @@ -0,0 +1,75 @@ +{ + "name": "opentelemetry-instrumentation-azure-search", + "$schema": "../../node_modules/nx/schemas/project-schema.json", + "projectType": "library", + "sourceRoot": "packages/opentelemetry-instrumentation-azure-search/opentelemetry/instrumentation/azure_search", + "targets": { + "lock": { + "executor": "nx:run-commands", + "options": { + "command": "uv lock", + "cwd": "packages/opentelemetry-instrumentation-azure-search" + } + }, + "add": { + "executor": "@nxlv/python:add", + "options": {} + }, + "update": { + "executor": "@nxlv/python:update", + "options": {} + }, + "remove": { + "executor": "@nxlv/python:remove", + "options": {} + }, + "build": { + "executor": "@nxlv/python:build", + "outputs": [ + "{projectRoot}/dist" + ], + "options": { + "outputPath": "packages/opentelemetry-instrumentation-azure-search/dist", + "publish": false, + "lockedVersions": true, + "bundleLocalDependencies": true + } + }, + "install": { + "executor": "nx:run-commands", + "options": { + "command": "uv sync --all-groups", + "cwd": "packages/opentelemetry-instrumentation-azure-search" + } + }, + "lint": { + "executor": "nx:run-commands", + "options": { + "command": "uv run ruff check .", + "cwd": "packages/opentelemetry-instrumentation-azure-search" + } + }, + "test": { + "executor": "nx:run-commands", + "outputs": [ + "{workspaceRoot}/reports/packages/opentelemetry-instrumentation-azure-search/unittests", + "{workspaceRoot}/coverage/packages/opentelemetry-instrumentation-azure-search" + ], + "options": { + "command": "uv run pytest tests/", + "cwd": "packages/opentelemetry-instrumentation-azure-search" + } + }, + "build-release": { + "executor": "nx:run-commands", + "commands": [ + "chmod +x ../../scripts/build-release.sh", + "../../scripts/build-release.sh" + ], + "cwd": "packages/opentelemetry-instrumentation-azure-search" + } + }, + "tags": [ + "instrumentation" + ] +} diff --git a/packages/opentelemetry-instrumentation-azure-search/pyproject.toml b/packages/opentelemetry-instrumentation-azure-search/pyproject.toml new file mode 100644 index 0000000000..84a7e856e1 --- /dev/null +++ b/packages/opentelemetry-instrumentation-azure-search/pyproject.toml @@ -0,0 +1,74 @@ +[project] +name = "opentelemetry-instrumentation-azure-search" +version = "0.62.1" +description = "OpenTelemetry Azure AI Search instrumentation" +authors = [ + { name = "Contributors", email = "openllmetry@traceloop.com" }, +] +license = "Apache-2.0" +readme = "README.md" +requires-python = ">=3.10,<4" +dependencies = [ + "opentelemetry-api>=1.38.0,<2", + "opentelemetry-instrumentation>=0.59b0", + "opentelemetry-semantic-conventions-ai>=0.5.2,<0.6.0", + "opentelemetry-semantic-conventions>=0.63b1", + "wrapt>=1.14,<2", +] + +[project.urls] +Repository = "https://github.com/traceloop/openllmetry/tree/main/packages/opentelemetry-instrumentation-azure-search" + +[project.optional-dependencies] +instruments = ["azure-search-documents>=11.4.0"] + +[project.entry-points."opentelemetry_instrumentor"] +azure_search_client = "opentelemetry.instrumentation.azure_search:AzureSearchInstrumentor" + +[dependency-groups] +dev = [ + "autopep8>=2.2.0,<3", + "ruff>=0.4.0", +] +test = [ + "azure-search-documents>=11.4.0", + "opentelemetry-sdk>=1.38.0,<2", + "pytest-sugar==1.0.0", + "pytest>=8.2.2,<9", + "vcrpy>=8.0.0,<9", +] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["opentelemetry"] + +[tool.coverage.run] +branch = true +source = ["opentelemetry/instrumentation/azure_search"] + +[tool.coverage.report] +exclude_lines = ["if TYPE_CHECKING:"] +show_missing = true + +[tool.ruff] +line-length = 120 +exclude = [ + ".git", + "__pycache__", + "build", + "dist", + ".venv", + ".pytest_cache", +] + +[tool.ruff.lint] +select = ["E", "F", "W"] + +[tool.uv] +constraint-dependencies = ["urllib3>=2.6.3", "pip>=25.3"] + +[tool.uv.sources] +opentelemetry-semantic-conventions-ai = { path = "../opentelemetry-semantic-conventions-ai", editable = true } diff --git a/packages/opentelemetry-instrumentation-azure-search/tests/__init__.py b/packages/opentelemetry-instrumentation-azure-search/tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/packages/opentelemetry-instrumentation-azure-search/tests/conftest.py b/packages/opentelemetry-instrumentation-azure-search/tests/conftest.py new file mode 100644 index 0000000000..98f8790660 --- /dev/null +++ b/packages/opentelemetry-instrumentation-azure-search/tests/conftest.py @@ -0,0 +1,27 @@ +"""Unit tests configuration module.""" + +import pytest +from opentelemetry import trace +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.instrumentation.azure_search import AzureSearchInstrumentor + + +@pytest.fixture(scope="session") +def exporter(): + exporter = InMemorySpanExporter() + processor = SimpleSpanProcessor(exporter) + + provider = TracerProvider() + provider.add_span_processor(processor) + trace.set_tracer_provider(provider) + + AzureSearchInstrumentor().instrument() + + return exporter + + +@pytest.fixture(autouse=True) +def clear_exporter(exporter): + exporter.clear() diff --git a/packages/opentelemetry-instrumentation-azure-search/tests/test_azure_search_instrumentation.py b/packages/opentelemetry-instrumentation-azure-search/tests/test_azure_search_instrumentation.py new file mode 100644 index 0000000000..356dc09a9e --- /dev/null +++ b/packages/opentelemetry-instrumentation-azure-search/tests/test_azure_search_instrumentation.py @@ -0,0 +1,326 @@ +"""Unit tests for Azure AI Search instrumentation. + +Uses mock clients since Azure Search requires a cloud endpoint. +Spans are verified via the InMemorySpanExporter from conftest.py. +""" + +from unittest.mock import MagicMock + +import pytest +from opentelemetry.semconv_ai import SpanAttributes + + +def _make_search_client(index_name="test-index"): + """Create a mock SearchClient with patched endpoint and index.""" + from azure.search.documents import SearchClient + + client = MagicMock(spec=SearchClient) + client._endpoint = "https://test.search.windows.net" + client._index_name = index_name + return client + + +def _make_index_client(): + """Create a mock SearchIndexClient.""" + from azure.search.documents.indexes import SearchIndexClient + + client = MagicMock(spec=SearchIndexClient) + client._endpoint = "https://test.search.windows.net" + return client + + +def _make_indexer_client(): + """Create a mock SearchIndexerClient.""" + from azure.search.documents.indexes import SearchIndexerClient + + client = MagicMock(spec=SearchIndexerClient) + client._endpoint = "https://test.search.windows.net" + return client + + +# --------------------------------------------------------------------------- +# SearchClient tests +# --------------------------------------------------------------------------- + + +def test_search_creates_span(exporter): + client = _make_search_client() + client.search.return_value = MagicMock() + + client.search(search_text="hello world", top=5, filter="category eq 'docs'") + + spans = exporter.get_finished_spans() + search_spans = [s for s in spans if s.name == "azure_search.search"] + assert len(search_spans) == 1 + + span = search_spans[0] + assert span.attributes.get(SpanAttributes.VECTOR_DB_VENDOR) == "azure_search" + assert span.attributes.get(SpanAttributes.AZURE_SEARCH_SEARCH_TEXT) == "hello world" + assert span.attributes.get(SpanAttributes.AZURE_SEARCH_TOP) == 5 + assert span.attributes.get(SpanAttributes.AZURE_SEARCH_FILTER) == "category eq 'docs'" + assert span.attributes.get(SpanAttributes.AZURE_SEARCH_INDEX_NAME) == "test-index" + assert span.attributes.get("server.address") == "https://test.search.windows.net" + + +def test_get_document_creates_span(exporter): + client = _make_search_client() + client.get_document.return_value = {"id": "1", "title": "Test"} + + client.get_document(document_id="1") + + spans = exporter.get_finished_spans() + get_doc_spans = [s for s in spans if s.name == "azure_search.get_document"] + assert len(get_doc_spans) == 1 + + span = get_doc_spans[0] + assert span.attributes.get(SpanAttributes.AZURE_SEARCH_INDEX_NAME) == "test-index" + + +def test_autocomplete_creates_span(exporter): + client = _make_search_client() + client.autocomplete.return_value = MagicMock() + + client.autocomplete(search_text="hel", autocomplete_mode="twoTerms") + + spans = exporter.get_finished_spans() + ac_spans = [s for s in spans if s.name == "azure_search.autocomplete"] + assert len(ac_spans) == 1 + + span = ac_spans[0] + assert span.attributes.get(SpanAttributes.AZURE_SEARCH_AUTOCOMPLETE_TEXT) == "hel" + + +def test_suggest_creates_span(exporter): + client = _make_search_client() + client.suggest.return_value = MagicMock() + + client.suggest(search_text="hel", suggester_name="sg") + + spans = exporter.get_finished_spans() + suggest_spans = [s for s in spans if s.name == "azure_search.suggest"] + assert len(suggest_spans) == 1 + + span = suggest_spans[0] + assert span.attributes.get(SpanAttributes.AZURE_SEARCH_SUGGEST_TEXT) == "hel" + + +def test_index_documents_creates_span(exporter): + client = _make_search_client() + response = MagicMock() + response.results = [MagicMock(error=None), MagicMock(error=None)] + client.index_documents.return_value = response + + docs = [{"id": "1", "title": "A"}, {"id": "2", "title": "B"}] + client.index_documents(documents=docs) + + spans = exporter.get_finished_spans() + idx_spans = [s for s in spans if s.name == "azure_search.index_documents"] + assert len(idx_spans) == 1 + + span = idx_spans[0] + assert span.attributes.get(SpanAttributes.AZURE_SEARCH_DOCUMENTS_COUNT) == 2 + assert span.attributes.get(SpanAttributes.AZURE_SEARCH_SUCCEEDED_COUNT) == 2 + + +def test_upload_documents_creates_span(exporter): + client = _make_search_client() + response = MagicMock() + response.results = [MagicMock(error=None)] + client.upload_documents.return_value = response + + client.upload_documents(documents=[{"id": "1"}]) + + spans = exporter.get_finished_spans() + spans = [s for s in spans if s.name == "azure_search.upload_documents"] + assert len(spans) == 1 + assert spans[0].attributes.get(SpanAttributes.AZURE_SEARCH_DOCUMENTS_COUNT) == 1 + + +def test_delete_documents_creates_span(exporter): + client = _make_search_client() + response = MagicMock() + response.results = [ + MagicMock(error=None), + MagicMock(error=None), + MagicMock(error=None), + ] + client.delete_documents.return_value = response + + client.delete_documents(documents=[{"id": "1"}, {"id": "2"}, {"id": "3"}]) + + spans = exporter.get_finished_spans() + del_spans = [s for s in spans if s.name == "azure_search.delete_documents"] + assert len(del_spans) == 1 + assert del_spans[0].attributes.get(SpanAttributes.AZURE_SEARCH_DOCUMENTS_COUNT) == 3 + + +def test_get_document_count_creates_span(exporter): + client = _make_search_client() + client.get_document_count.return_value = 42 + + client.get_document_count() + + spans = exporter.get_finished_spans() + count_spans = [s for s in spans if s.name == "azure_search.get_document_count"] + assert len(count_spans) == 1 + assert count_spans[0].attributes.get(SpanAttributes.AZURE_SEARCH_DOCUMENTS_COUNT) == 42 + + +# --------------------------------------------------------------------------- +# SearchIndexClient tests +# --------------------------------------------------------------------------- + + +def test_create_index_creates_span(exporter): + client = _make_index_client() + index = MagicMock() + index.name = "my-index" + client.create_index.return_value = index + + client.create_index(index=index) + + spans = exporter.get_finished_spans() + create_spans = [s for s in spans if s.name == "azure_search.create_index"] + assert len(create_spans) == 1 + assert create_spans[0].attributes.get(SpanAttributes.AZURE_SEARCH_INDEX_NAME) == "my-index" + + +def test_delete_index_creates_span(exporter): + client = _make_index_client() + client.delete_index.return_value = None + + client.delete_index(index_name="my-index") + + spans = exporter.get_finished_spans() + del_spans = [s for s in spans if s.name == "azure_search.delete_index"] + assert len(del_spans) == 1 + assert del_spans[0].attributes.get(SpanAttributes.AZURE_SEARCH_INDEX_NAME) == "my-index" + + +def test_get_index_creates_span(exporter): + client = _make_index_client() + index = MagicMock() + index.name = "my-index" + client.get_index.return_value = index + + client.get_index(index_name="my-index") + + spans = exporter.get_finished_spans() + get_spans = [s for s in spans if s.name == "azure_search.get_index"] + assert len(get_spans) == 1 + assert get_spans[0].attributes.get(SpanAttributes.AZURE_SEARCH_INDEX_NAME) == "my-index" + + +def test_get_index_statistics_creates_span(exporter): + client = _make_index_client() + stats = MagicMock() + stats.document_count = 1000 + stats.storage_size = 1048576 + client.get_index_statistics.return_value = stats + + client.get_index_statistics(index_name="my-index") + + spans = exporter.get_finished_spans() + stat_spans = [s for s in spans if s.name == "azure_search.get_index_statistics"] + assert len(stat_spans) == 1 + assert stat_spans[0].attributes.get(SpanAttributes.AZURE_SEARCH_INDEX_DOC_COUNT) == 1000 + assert stat_spans[0].attributes.get(SpanAttributes.AZURE_SEARCH_INDEX_SIZE_BYTES) == 1048576 + + +def test_get_service_statistics_creates_span(exporter): + client = _make_index_client() + counters = MagicMock() + counters.search_service_usage = "85%" + counters.search_service_limit = "100%" + stats = MagicMock() + stats.counters = counters + client.get_service_statistics.return_value = stats + + client.get_service_statistics() + + spans = exporter.get_finished_spans() + svc_spans = [s for s in spans if s.name == "azure_search.get_service_statistics"] + assert len(svc_spans) == 1 + assert svc_spans[0].attributes.get(SpanAttributes.AZURE_SEARCH_SERVICE_USAGE) == "85%" + assert svc_spans[0].attributes.get(SpanAttributes.AZURE_SEARCH_SERVICE_LIMIT) == "100%" + + +# --------------------------------------------------------------------------- +# SearchIndexerClient tests +# --------------------------------------------------------------------------- + + +def test_create_indexer_creates_span(exporter): + client = _make_indexer_client() + indexer = MagicMock() + indexer.name = "my-indexer" + client.create_indexer.return_value = indexer + + client.create_indexer(indexer=indexer) + + spans = exporter.get_finished_spans() + idx_spans = [s for s in spans if s.name == "azure_search.create_indexer"] + assert len(idx_spans) == 1 + assert idx_spans[0].attributes.get(SpanAttributes.AZURE_SEARCH_INDEXER_NAME) == "my-indexer" + + +def test_get_indexer_status_creates_span(exporter): + client = _make_indexer_client() + status = MagicMock() + status.status = "running" + client.get_indexer_status.return_value = status + + client.get_indexer_status(indexer_name="my-indexer") + + spans = exporter.get_finished_spans() + status_spans = [s for s in spans if s.name == "azure_search.get_indexer_status"] + assert len(status_spans) == 1 + assert status_spans[0].attributes.get(SpanAttributes.AZURE_SEARCH_INDEXER_STATUS) == "running" + + +def test_run_indexer_creates_span(exporter): + client = _make_indexer_client() + client.run_indexer.return_value = None + + client.run_indexer(indexer_name="my-indexer") + + spans = exporter.get_finished_spans() + run_spans = [s for s in spans if s.name == "azure_search.run_indexer"] + assert len(run_spans) == 1 + assert run_spans[0].attributes.get(SpanAttributes.AZURE_SEARCH_INDEXER_NAME) == "my-indexer" + + +# --------------------------------------------------------------------------- +# Error handling +# --------------------------------------------------------------------------- + + +def test_search_exception_records_exception(exporter): + client = _make_search_client() + client.search.side_effect = RuntimeError("Connection refused") + + with pytest.raises(RuntimeError, match="Connection refused"): + client.search(search_text="fail") + + spans = exporter.get_finished_spans() + search_spans = [s for s in spans if s.name == "azure_search.search"] + assert len(search_spans) == 1 + + span = search_spans[0] + assert span.status.status_code.name == "ERROR" + assert len(span.events) == 1 + assert span.events[0].name == "exception" + + +# --------------------------------------------------------------------------- +# Instrumentor API +# --------------------------------------------------------------------------- + + +def test_instrumentor_has_correct_interface(): + from opentelemetry.instrumentation.azure_search import AzureSearchInstrumentor + + instrumentor = AzureSearchInstrumentor() + assert hasattr(instrumentor, "_instrument") + assert hasattr(instrumentor, "_uninstrument") + assert "azure-search-documents" in instrumentor.instrumentation_dependencies()[0] diff --git a/packages/opentelemetry-semantic-conventions-ai/opentelemetry/semconv_ai/__init__.py b/packages/opentelemetry-semantic-conventions-ai/opentelemetry/semconv_ai/__init__.py index 24ea52decb..ced1813979 100644 --- a/packages/opentelemetry-semantic-conventions-ai/opentelemetry/semconv_ai/__init__.py +++ b/packages/opentelemetry-semantic-conventions-ai/opentelemetry/semconv_ai/__init__.py @@ -302,6 +302,24 @@ class SpanAttributes: QDRANT_UPSERT_COLLECTION_NAME = "qdrant.upsert.collection_name" QDRANT_UPSERT_POINTS_COUNT = "qdrant.upsert.points_count" + # Azure AI Search + AZURE_SEARCH_INDEX_NAME = "db.azure_search.index_name" + AZURE_SEARCH_SEARCH_TEXT = "db.azure_search.search_text" + AZURE_SEARCH_TOP = "db.azure_search.top" + AZURE_SEARCH_FILTER = "db.azure_search.filter" + AZURE_SEARCH_RESULT_COUNT = "db.azure_search.result_count" + AZURE_SEARCH_DOCUMENTS_COUNT = "db.azure_search.documents_count" + AZURE_SEARCH_SUCCEEDED_COUNT = "db.azure_search.succeeded_count" + AZURE_SEARCH_INDEXER_NAME = "db.azure_search.indexer_name" + AZURE_SEARCH_INDEXER_STATUS = "db.azure_search.indexer_status" + AZURE_SEARCH_INDEX_DIMENSIONS = "db.azure_search.index_dimensions" + AZURE_SEARCH_INDEX_DOC_COUNT = "db.azure_search.index_document_count" + AZURE_SEARCH_INDEX_SIZE_BYTES = "db.azure_search.index_size_bytes" + AZURE_SEARCH_SERVICE_LIMIT = "db.azure_search.service_limit" + AZURE_SEARCH_SERVICE_USAGE = "db.azure_search.service_usage" + AZURE_SEARCH_AUTOCOMPLETE_TEXT = "db.azure_search.autocomplete_text" + AZURE_SEARCH_SUGGEST_TEXT = "db.azure_search.suggest_text" + # Marqo MARQO_SEARCH_QUERY = "db.marqo.search.query" MARQO_SEARCH_PROCESSING_TIME = "db.marqo.search.processing_time"