# Azure AI Search Plugin — Semantic Search & Cosmos DB Retriever
 
The **Azure AI Search Plugin** is a unified, lightweight, and reusable **Python utility suite** built to connect, query, and retrieve data from **Azure Cognitive Search** which is "main.py" file and **Azure Cosmos DB** which is in "cosmos_data_retriever.py" using the official Azure SDKs.  
It is part of the `azure_ai_search_plugin` package and designed for seamless integration with **AI agents**, **Semantic Kernel**, or **intelligent retrieval systems**.
 
Together, these retrievers enable **semantic search**, **structured querying**, and **data enrichment** from multiple Azure sources — all with **async support**, **robust error handling**, and **clean structured responses**.
 
---
 
## Features
 
### 🔹 Azure Cognitive Search (Semantic Search Retriever) (main.py file)
- Fully **asynchronous operations** using the `aio` SDK  
- **Semantic & keyword search** support  
- **Top matching document retrieval** by confidence score  
- **Dynamic parameter handling** (no hardcoded values)  
- **Safe score conversion** and **robust exception handling**  
- **Semantic Kernel compatible** (`@kernel_function` ready)  
 
### 🔹 Azure Cosmos DB (Data Retriever) (cosmos_data_retriever.py)
- **Simple connection setup** via `.env` or direct parameters  
- **Fetch all items** or filter by specific field  
- **Retrieve by ID and Partition Key**  
- **Custom SQL-like query execution**  
- **Integrated logging** and **custom error handling**
 
---
 
## Requirements
 
| Dependency | Minimum Version |
|-------------|----------------|
| Python | 3.8+ |
| azure-search-documents | ≥ 11.4.0 |
| azure-cosmos | ≥ 4.0.0 |
| semantic-kernel | ≥ 0.5.0 |
| python-dotenv | ≥ 1.0 |
 
---
 
## Installation
 
```bash
pip install azure-search-documents azure-cosmos semantic-kernel python-dotenv
```

---
# Note
While installing the dependencies for this package, the build/ folder and its related files will be automatically created as part of the packaging process.
---
 
## Environment Setup
 
Create a `.env` file at your project root with both Cognitive Search and Cosmos DB configuration:
 
```bash
# Azure Cognitive Search
AZURE_SEARCH_ENDPOINT="https://your-search-service.search.windows.net"
AZURE_SEARCH_API_KEY="your-api-key"
AZURE_SEARCH_INDEX="your-index-name"
AZURE_SEMANTIC_CONFIG="your-semantic-config"  # Optional
 
# Azure Cosmos DB
COSMOS_ACCOUNT_URI="your-cosmosdb-endpoint"
COSMOS_ACCOUNT_KEY="your-cosmosdb-key"
COSMOS_DATABASE_NAME="your-database-name"
COSMOS_CONTAINER_NAME="your-container-name"
```
 
---
 
## Usage Examples
 
### 1 Semantic Search (Azure Cognitive Search)
 
```python
import asyncio
import os
from dotenv import load_dotenv
from azure_ai_search_plugin.main import AzureSearchPlugin
 
load_dotenv()
 
async def main():
    plugin = AzureSearchPlugin(
        endpoint=os.getenv("AZURE_SEARCH_ENDPOINT"),
        api_key=os.getenv("AZURE_SEARCH_API_KEY"),
        index_name=os.getenv("AZURE_SEARCH_INDEX"),
        semantic_config=os.getenv("AZURE_SEMANTIC_CONFIG")
    )
 
    query = "Explain Azure Cognitive Search"
    result = await plugin.search_top(query=query, top_k=3)
 
    if result:
        print("Top matching document:")
        print(result)
    else:
        print("No results found.")
 
if __name__ == "__main__":
    asyncio.run(main())
```
 
---
 
## Sample Outputs
 
### Azure Cognitive Search
```
Top matching document:
{
  "id": "article-001",
  "title": "Introduction to Azure Cognitive Search",
  "content": "Azure Cognitive Search is a fully managed service that provides indexing and querying capabilities...",
  "_confidence": 0.982
}
```
 
### 2 Cosmos DB Retriever
 
```python
import os
from dotenv import load_dotenv
from azure_ai_search_plugin.cosmos_data_retriever import CosmosDBRetriever
 
load_dotenv()
 
retriever = CosmosDBRetriever(
    database_name=os.getenv("COSMOS_DATABASE_NAME"),
    container_name=os.getenv("COSMOS_CONTAINER_NAME"),
    endpoint=os.getenv("COSMOS_ACCOUNT_URI"),
    key=os.getenv("COSMOS_ACCOUNT_KEY")
)
 
# Fetch all items
items = retriever.fetch_all_items()
print(f"Total items fetched: {len(items)}")
 
# Fetch by field
if items:
    sample_message_id = items[0].get("message_id")
    filtered = retriever.fetch_by_field("message_id", sample_message_id)
    print(f"Fetched {len(filtered)} items where message_id={sample_message_id}")
 
# Get item by ID
if items:
    sample_id = items[0].get("id")
    partition_key = items[0].get("thread_id")
    item = retriever.get_item_by_id(sample_id, partition_key)
    print(item)
 
# Run custom query
query = "SELECT TOP 5 c.id, c.message_id, c.subject, c.body FROM c WHERE IS_DEFINED(c.subject)"
result = retriever.query_knowledge_source(query)
print(f"Query returned {result['count']} items.")
```
 
---
 
## Sample Outputs
 
### Cosmos DB Retriever
```
=== Get item by ID ===
[INFO] azure_ai_search_plugin.cosmos_data_retriever: Retrieved item with ID: 54554614-c7fc-49a1-a9f0-7b1bd34bfeae
{'id': '54554614...', 'subject': 'test', 'body': 'I have a query regarding how to create a new vendor'}
```
 
---
 
## Class Overview
 
### `AzureSearchPlugin`
| Method | Description |
|--------|--------------|
| `search_top(query: str, top_k: int)` | Performs semantic or keyword search and returns the top matching document |
 
### `CosmosDBRetriever`
| Method | Description |
|--------|--------------|
| `fetch_all_items()` | Retrieves all documents from the container |
| `fetch_by_field(field_name, field_value)` | Fetches items matching a given field |
| `get_item_by_id(item_id, partition_key)` | Retrieves a single item by ID and partition key |
| `query_knowledge_source(sql_query, parameters=None)` | Executes a custom SQL-like query |
 
---
 
## Error Handling
 
Both retrievers include **robust exception management**:
 
| Error Type | Description |
|-------------|-------------|
| `ValueError` | Missing or invalid configuration |
| `TypeError / ValueError` | Invalid confidence score or query parameter |
| `CosmosConnectionError` | Cosmos DB connection or authentication issue |
| `CosmosQueryError` | Query or data retrieval error |
| `CosmosQuerySyntaxError` | Query contains invalid syntax or unsafe operations|
 
---
 
## Async Execution (Search Plugin Only)
 
The **Semantic Search Retriever** is fully asynchronous:
- Non-blocking I/O for faster parallel operations  
- Ideal for async frameworks like **FastAPI**, **Semantic Kernel**, or **LangChain**
 
---
 
## Integration with Semantic Kernel
 
```python
from semantic_kernel import Kernel
 
kernel = Kernel()
kernel.import_skill(AzureSearchPlugin(...), skill_name="azure_ai_search")
result = await kernel.skills.azure_ai_search.search_top("What is Azure AI Search?")
```
 
---
 
## Logging
 
- All Cosmos DB operations are logged for transparency.  
- The Search Plugin can also be extended to include custom logging.  
- Integrate Python’s built-in `logging` module for full observability.
 
---