Skip to content

Commit

Permalink
Browse files Browse the repository at this point in the history
Added support for Azure search
Fixes traceloop#2303
  • Loading branch information
Mr-Imperium committed Nov 19, 2024
1 parent 61a4bf9 commit fd4845d
Show file tree
Hide file tree
Showing 11 changed files with 944 additions and 0 deletions.
15 changes: 15 additions & 0 deletions packages/opentelemetry-instrumentation-azureai/.flake8
Original file line number Diff line number Diff line change
@@ -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
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
3.8.10
3.9.16
3.10.11
3.11.4
50 changes: 50 additions & 0 deletions packages/opentelemetry-instrumentation-azureai/README.md
Original file line number Diff line number Diff line change
@@ -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
Original file line number Diff line number Diff line change
@@ -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
Original file line number Diff line number Diff line change
@@ -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"
)
Loading

0 comments on commit fd4845d

Please sign in to comment.