feat(azureai-search)

packages/opentelemetry-instrumentation-azureai/.flake8

@@ -0,0 +1,15 @@
+[flake8]
+max-line-length = 120
+exclude = 
+    .git,
+    __pycache__,
+    build,
+    dist,
+    *.pyc,
+    *.egg-info,
+    .cache,
+    .eggs
+ignore = 
+    E203,  # whitespace before ':' (black compatibility)
+    W503,  # line break before binary operator
+    F401   # unused imports

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

@@ -0,0 +1,4 @@
+3.8.10
+3.9.16
+3.10.11
+3.11.4

packages/opentelemetry-instrumentation-azureai/README.md

@@ -0,0 +1,50 @@
+# OpenLLMetry Azure AI Search Integration
+
+## Overview
+
+This package provides an integration for Azure AI Search within the OpenLLMetry observability framework. It enables detailed tracing and metrics collection for search operations.
+
+## Installation
+
+```bash
+pip install openllmetry-azure-search
+```
+
+Or with Poetry:
+
+```bash
+poetry add openllmetry-azure-search
+```
+
+## Usage
+
+```python
+from openllmetry.integrations.azure_search import AzureSearchIntegration
+
+# Initialize the integration
+search = AzureSearchIntegration(
+    endpoint="https://your-search-service.search.windows.net",
+    key="your-api-key",
+    index_name="your-index-name"
+)
+
+# Perform a search
+results = search.search(
+    query="your search query",
+    filter="category eq 'books'",
+    top=10
+)
+```
+
+## Features
+
+- Automatic OpenTelemetry span creation for search operations
+- Comprehensive metrics collection
+- Error tracking and status reporting
+- Flexible configuration options
+
+## Requirements
+
+- Python 3.8+
+- Azure AI Search
+- OpenTelemetry

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

@@ -0,0 +1,69 @@
+from typing import Dict, List, Optional, Any
+from opentelemetry.base import BaseIntegration
+from azure.search.documents import SearchClient
+from azure.core.credentials import AzureKeyCredential
+
+class AzureSearchIntegration(BaseIntegration):
+    """Integration for Azure AI Search in OpenLLMetry"""
+
+    def __init__(
+        self,
+        endpoint: str,
+        key: str,
+        index_name: str,
+        service_name: str = "azure-search"
+    ):
+        super().__init__(service_name)
+        self.endpoint = endpoint
+        self.credential = AzureKeyCredential(key)
+        self.index_name = index_name
+        self.client = SearchClient(
+            endpoint=endpoint,
+            index_name=index_name,
+            credential=self.credential
+        )
+
+    def search(
+        self,
+        query: str,
+        filter: Optional[str] = None,
+        top: Optional[int] = None,
+        **kwargs
+    ) -> Dict[str, Any]:
+        """
+        Perform a search operation and collect metrics
+        """
+        with self.start_span("azure.search.query") as span:
+            # Add relevant attributes to span
+            span.set_attribute("azure.search.query", query)
+            span.set_attribute("azure.search.index", self.index_name)
+            if filter:
+                span.set_attribute("azure.search.filter", filter)
+            if top:
+                span.set_attribute("azure.search.top", top)
+
+            try:
+                results = list(self.client.search(
+                    search_text=query,
+                    filter=filter,
+                    top=top,
+                    **kwargs
+                ))
+
+                # Record metrics
+                self.record_metric(
+                    "azure.search.results.count",
+                    len(results),
+                    {"index": self.index_name}
+                )
+
+                return {
+                    "results": results,
+                    "count": len(results)
+                }
+
+            except Exception as e:
+                # Record error in span
+                span.record_exception(e)
+                span.set_status(Status(StatusCode.ERROR))
+                raise

packages/opentelemetry-instrumentation-azureai/opentelemetry/instrumentation/azureai/metrics.py

@@ -0,0 +1,24 @@
+from opentelemetry.metrics import Counter, Histogram
+from typing import Dict
+
+class AzureSearchMetrics:
+    """Metrics definitions for Azure Search integration"""
+
+    def __init__(self):
+        self.search_count = Counter(
+            name="azure.search.requests",
+            description="Number of search requests made",
+            unit="1"
+        )
+
+        self.search_latency = Histogram(
+            name="azure.search.latency",
+            description="Search request latency",
+            unit="ms"
+        )
+
+        self.results_count = Counter(
+            name="azure.search.results.total",
+            description="Total number of search results returned",
+            unit="1"
+        )