forked from traceloop/openllmetry
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Added support for Azure search Fixes traceloop#2303
- Loading branch information
1 parent
61a4bf9
commit fd4845d
Showing
11 changed files
with
944 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 |
4 changes: 4 additions & 0 deletions
4
packages/opentelemetry-instrumentation-azureai/.python-version
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 |
69 changes: 69 additions & 0 deletions
69
...s/opentelemetry-instrumentation-azureai/opentelemetry/instrumentation/azureai/__init__.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 |
24 changes: 24 additions & 0 deletions
24
...es/opentelemetry-instrumentation-azureai/opentelemetry/instrumentation/azureai/metrics.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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" | ||
) |
Oops, something went wrong.