# Agentic Retrieval with Knowledge Bases

Agentic retrieval integrates an LLM to process queries, retrieve content, and generate grounded answers.

## Architecture

```
Knowledge Base (wraps LLM + sources)
    ├── Knowledge Source 1 → Search Index A
    ├── Knowledge Source 2 → Search Index B
    └── Azure OpenAI Model (query planning + answer synthesis)
```

## Setup Workflow

### 1. Create Index with Semantic Configuration

```python
from azure.search.documents.indexes import SearchIndexClient
from azure.search.documents.indexes.models import (
    SearchIndex, SearchField, VectorSearch, VectorSearchProfile,
    HnswAlgorithmConfiguration, AzureOpenAIVectorizer,
    AzureOpenAIVectorizerParameters, SemanticSearch,
    SemanticConfiguration, SemanticPrioritizedFields, SemanticField
)

index = SearchIndex(
    name="my-index",
    fields=[
        SearchField(name="id", type="Edm.String", key=True, filterable=True),
        SearchField(name="content", type="Edm.String", filterable=False),
        SearchField(name="embedding", type="Collection(Edm.Single)",
                   stored=False, vector_search_dimensions=3072,
                   vector_search_profile_name="hnsw-profile"),
        SearchField(name="category", type="Edm.String", filterable=True, facetable=True)
    ],
    vector_search=VectorSearch(
        profiles=[VectorSearchProfile(
            name="hnsw-profile",
            algorithm_configuration_name="hnsw-algo",
            vectorizer_name="aoai-vectorizer"
        )],
        algorithms=[HnswAlgorithmConfiguration(name="hnsw-algo")],
        vectorizers=[AzureOpenAIVectorizer(
            vectorizer_name="aoai-vectorizer",
            parameters=AzureOpenAIVectorizerParameters(
                resource_url=aoai_endpoint,
                deployment_name="text-embedding-3-large",
                model_name="text-embedding-3-large"
            )
        )]
    ),
    # REQUIRED for agentic retrieval
    semantic_search=SemanticSearch(
        default_configuration_name="semantic-config",
        configurations=[SemanticConfiguration(
            name="semantic-config",
            prioritized_fields=SemanticPrioritizedFields(
                content_fields=[SemanticField(field_name="content")]
            )
        )]
    )
)

index_client = SearchIndexClient(endpoint, credential)
index_client.create_or_update_index(index)
```

### 2. Create Knowledge Source

```python
from azure.search.documents.indexes.models import (
    SearchIndexKnowledgeSource,
    SearchIndexKnowledgeSourceParameters,
    SearchIndexFieldReference
)

knowledge_source = SearchIndexKnowledgeSource(
    name="my-knowledge-source",
    description="Knowledge source for document retrieval",
    search_index_parameters=SearchIndexKnowledgeSourceParameters(
        search_index_name="my-index",
        source_data_fields=[
            SearchIndexFieldReference(name="id"),
            SearchIndexFieldReference(name="category")
        ]
    )
)

index_client.create_or_update_knowledge_source(knowledge_source)
```

### 3. Create Knowledge Base

```python
from azure.search.documents.indexes.models import (
    KnowledgeBase, KnowledgeBaseAzureOpenAIModel,
    KnowledgeSourceReference, AzureOpenAIVectorizerParameters,
    KnowledgeRetrievalOutputMode
)

knowledge_base = KnowledgeBase(
    name="my-knowledge-base",
    models=[KnowledgeBaseAzureOpenAIModel(
        azure_open_ai_parameters=AzureOpenAIVectorizerParameters(
            resource_url=aoai_endpoint,
            deployment_name="gpt-4o-mini",
            model_name="gpt-4o-mini"
        )
    )],
    knowledge_sources=[KnowledgeSourceReference(name="my-knowledge-source")],
    output_mode=KnowledgeRetrievalOutputMode.ANSWER_SYNTHESIS,
    answer_instructions="Provide concise, well-cited answers based on retrieved documents."
)

index_client.create_or_update_knowledge_base(knowledge_base)
```

## Querying the Knowledge Base

```python
from azure.search.documents.knowledgebases import KnowledgeBaseRetrievalClient
from azure.search.documents.knowledgebases.models import (
    KnowledgeBaseRetrievalRequest,
    KnowledgeRetrievalSemanticIntent,
    KnowledgeRetrievalMinimalReasoningEffort,
)

client = KnowledgeBaseRetrievalClient(
    endpoint=endpoint,
    credential=credential,
)

# Build retrieval request with semantic intents
request = KnowledgeBaseRetrievalRequest(
    intents=[KnowledgeRetrievalSemanticIntent(search="What is vector search?")]
)

result = client.retrieve(
    knowledge_base_name="my-knowledge-base",
    retrieval_request=request,
)
```

## Processing Results

```python
import json

# Extract response content
response_parts = []
for resp in result.response or []:
    for content in resp.content or []:
        if hasattr(content, "text"):
            response_parts.append(content.text)

if response_parts:
    response_content = "\n\n".join(response_parts)
    print(response_content)

# Extract references (source documents)
if result.references:
    for ref in result.references:
        print(f"Reference ID: {ref.id}")
        if hasattr(ref, 'reranker_score'):
            print(f"  Score: {ref.reranker_score}")
        if ref.source_data:
            print(f"  Content: {ref.source_data.get('content', '')[:200]}")
```

## Multi-turn Conversations

```python
from azure.search.documents.knowledgebases.models import (
    KnowledgeBaseRetrievalRequest,
    KnowledgeRetrievalSemanticIntent,
)

def ask(question: str) -> str:
    """Ask a question against the knowledge base."""
    request = KnowledgeBaseRetrievalRequest(
        intents=[KnowledgeRetrievalSemanticIntent(search=question)]
    )
    
    result = client.retrieve(
        knowledge_base_name="my-knowledge-base",
        retrieval_request=request,
    )
    
    # Extract response
    response_text = "\n\n".join(
        content.text
        for resp in (result.response or [])
        for content in (resp.content or [])
        if hasattr(content, "text")
    )
    
    return response_text
```

## Output Modes

| Mode | Description |
|------|-------------|
| `EXTRACTIVE_DATA` | Return raw chunks from knowledge sources |
| `ANSWER_SYNTHESIS` | LLM generates answers citing retrieved content |

## Reasoning Effort Levels

| Level | Class |
|-------|-------|
| `KnowledgeRetrievalMinimalReasoningEffort` | No query planning or iterative search |
| `KnowledgeRetrievalLowReasoningEffort` | Basic query decomposition |
| `KnowledgeRetrievalMediumReasoningEffort` | More sophisticated reasoning |

**Usage:**
```python
from azure.search.documents.knowledgebases.models import KnowledgeRetrievalMinimalReasoningEffort

knowledge_base = KnowledgeBase(
    name="my-knowledge-base",
    knowledge_sources=[KnowledgeSourceReference(name="my-knowledge-source")],
    retrieval_reasoning_effort=KnowledgeRetrievalMinimalReasoningEffort(),
)
```

## Async Pattern

```python
from azure.search.documents.knowledgebases.aio import KnowledgeBaseRetrievalClient

async with KnowledgeBaseRetrievalClient(endpoint, credential=credential) as client:
    result = await client.retrieve(
        knowledge_base_name="my-knowledge-base",
        retrieval_request=request,
    )
```

## Clean Up

```python
# Delete in reverse order of creation
index_client.delete_knowledge_base("my-knowledge-base")
index_client.delete_knowledge_source("my-knowledge-source")
index_client.delete_index("my-index")
```

# Semantic Ranking Patterns

Semantic ranking uses machine learning to re-rank search results for better relevance.

## Semantic Configuration

### Basic Semantic Configuration

```python
from azure.search.documents.indexes.models import (
    SemanticSearch,
    SemanticConfiguration,
    SemanticPrioritizedFields,
    SemanticField,
)

semantic_search = SemanticSearch(
    default_configuration_name="my-semantic-config",
    configurations=[
        SemanticConfiguration(
            name="my-semantic-config",
            prioritized_fields=SemanticPrioritizedFields(
                title_field=SemanticField(field_name="title"),
                content_fields=[
                    SemanticField(field_name="content"),
                    SemanticField(field_name="summary")
                ],
                keywords_fields=[
                    SemanticField(field_name="tags")
                ]
            )
        )
    ]
)
```

### Field Priority Rules

| Field Type | Max Fields | Purpose |
|------------|------------|---------|
| `title_field` | 1 | Most important for ranking |
| `content_fields` | 10 | Main content, ordered by priority |
| `keywords_fields` | 10 | Keywords/tags for matching |

## Semantic Queries

### Basic Semantic Query

```python
from azure.search.documents.models import QueryType

def semantic_search(client, query: str):
    """Execute a semantic search query."""
    results = client.search(
        search_text=query,
        query_type=QueryType.SEMANTIC,
        semantic_configuration_name="my-semantic-config",
        select=["id", "title", "content"]
    )
    
    return list(results)
```

### Semantic Search with Captions

Captions provide extractive summaries highlighting relevant passages.

```python
from azure.search.documents.models import QueryType, QueryCaptionType

def semantic_search_with_captions(client, query: str):
    """Semantic search with extractive captions."""
    results = client.search(
        search_text=query,
        query_type=QueryType.SEMANTIC,
        semantic_configuration_name="my-semantic-config",
        query_caption=QueryCaptionType.EXTRACTIVE,
        select=["id", "title", "content"]
    )
    
    for result in results:
        print(f"Title: {result['title']}")
        print(f"Score: {result['@search.score']}")
        print(f"Reranker Score: {result.get('@search.reranker_score', 'N/A')}")
        
        # Extract captions
        captions = result.get("@search.captions", [])
        for caption in captions:
            print(f"Caption: {caption.text}")
            print(f"Highlights: {caption.highlights}")
```

### Semantic Search with Answers

Answers extract specific text that directly answers the query.

```python
from azure.search.documents.models import QueryType, QueryCaptionType, QueryAnswerType

def semantic_search_with_answers(client, query: str):
    """Semantic search with extractive answers."""
    results = client.search(
        search_text=query,
        query_type=QueryType.SEMANTIC,
        semantic_configuration_name="my-semantic-config",
        query_caption=QueryCaptionType.EXTRACTIVE,
        query_answer=QueryAnswerType.EXTRACTIVE,
        query_answer_count=3,  # Request up to 3 answers
        select=["id", "title", "content"]
    )
    
    # Get semantic answers (top-level, not per-document)
    answers = results.get_answers()
    if answers:
        for answer in answers:
            print(f"Answer: {answer.text}")
            print(f"Highlights: {answer.highlights}")
            print(f"Score: {answer.score}")
            print(f"Key: {answer.key}")
    
    # Process documents
    for result in results:
        print(f"Document: {result['title']}")
```

## Hybrid Search with Semantic Ranking

Combine keyword, vector, and semantic ranking for best results.

```python
from azure.search.documents.models import (
    VectorizedQuery,
    QueryType,
    QueryCaptionType,
)

def hybrid_semantic_search(
    client,
    query: str,
    query_vector: list[float],
    top: int = 10
):
    """Hybrid search with semantic re-ranking."""
    vector_query = VectorizedQuery(
        vector=query_vector,
        k_nearest_neighbors=50,  # Over-fetch for re-ranking
        fields="content_vector"
    )
    
    results = client.search(
        search_text=query,
        vector_queries=[vector_query],
        query_type=QueryType.SEMANTIC,
        semantic_configuration_name="my-semantic-config",
        query_caption=QueryCaptionType.EXTRACTIVE,
        top=top,
        select=["id", "title", "content"]
    )
    
    return list(results)
```

## Semantic Configuration for Different Content Types

### Document Search

```python
doc_semantic_config = SemanticConfiguration(
    name="document-semantic-config",
    prioritized_fields=SemanticPrioritizedFields(
        title_field=SemanticField(field_name="document_title"),
        content_fields=[
            SemanticField(field_name="body_text"),
            SemanticField(field_name="abstract")
        ],
        keywords_fields=[
            SemanticField(field_name="categories"),
            SemanticField(field_name="authors")
        ]
    )
)
```

### Product Search

```python
product_semantic_config = SemanticConfiguration(
    name="product-semantic-config",
    prioritized_fields=SemanticPrioritizedFields(
        title_field=SemanticField(field_name="product_name"),
        content_fields=[
            SemanticField(field_name="description"),
            SemanticField(field_name="features")
        ],
        keywords_fields=[
            SemanticField(field_name="brand"),
            SemanticField(field_name="category")
        ]
    )
)
```

### FAQ Search

```python
faq_semantic_config = SemanticConfiguration(
    name="faq-semantic-config",
    prioritized_fields=SemanticPrioritizedFields(
        title_field=SemanticField(field_name="question"),
        content_fields=[
            SemanticField(field_name="answer")
        ],
        keywords_fields=[
            SemanticField(field_name="topic")
        ]
    )
)
```

## Multiple Semantic Configurations

Define different configurations for different query patterns.

```python
semantic_search = SemanticSearch(
    default_configuration_name="general-config",
    configurations=[
        SemanticConfiguration(
            name="general-config",
            prioritized_fields=SemanticPrioritizedFields(
                title_field=SemanticField(field_name="title"),
                content_fields=[SemanticField(field_name="content")]
            )
        ),
        SemanticConfiguration(
            name="qa-config",
            prioritized_fields=SemanticPrioritizedFields(
                title_field=SemanticField(field_name="question"),
                content_fields=[SemanticField(field_name="answer")]
            )
        )
    ]
)

# Use specific config at query time
results = client.search(
    search_text=query,
    query_type=QueryType.SEMANTIC,
    semantic_configuration_name="qa-config"  # Override default
)
```

## Best Practices

1. **Title field**: Always specify a title field for best ranking
2. **Content field order**: List most important content fields first
3. **Over-fetch**: Request more results than needed, let semantic re-ranking select the best
4. **Captions for UI**: Use captions to show relevant snippets in search results
5. **Answers for Q&A**: Use answers for question-answering scenarios
6. **Combine with hybrid**: Semantic ranking works best with hybrid (keyword + vector) search

# Vector Search Patterns

Detailed patterns for vector search with Azure AI Search.

## Vector Search Configuration

### HNSW Algorithm Configuration

HNSW (Hierarchical Navigable Small World) is the recommended algorithm for large-scale vector search.

```python
from azure.search.documents.indexes.models import (
    VectorSearch,
    HnswAlgorithmConfiguration,
    VectorSearchProfile,
    HnswParameters,
)

vector_search = VectorSearch(
    algorithms=[
        HnswAlgorithmConfiguration(
            name="hnsw-config",
            parameters=HnswParameters(
                m=4,              # Bi-directional links per node (default: 4)
                ef_construction=400,  # Size of dynamic candidate list during indexing
                ef_search=500,    # Size of dynamic candidate list during search
                metric="cosine"   # Distance metric: cosine, euclidean, dotProduct
            )
        )
    ],
    profiles=[
        VectorSearchProfile(
            name="vector-profile",
            algorithm_configuration_name="hnsw-config"
        )
    ]
)
```

### Parameter Tuning Guide

| Parameter | Range | Trade-off |
|-----------|-------|-----------|
| `m` | 4-16 | Higher = better recall, more memory |
| `ef_construction` | 100-1000 | Higher = better index quality, slower indexing |
| `ef_search` | 100-1000 | Higher = better recall, slower queries |

## Integrated Vectorization

Let Azure AI Search generate embeddings automatically using Azure OpenAI.

### Vectorizer Configuration

```python
from azure.search.documents.indexes.models import (
    AzureOpenAIVectorizer,
    AzureOpenAIVectorizerParameters,
)

vectorizer = AzureOpenAIVectorizer(
    vectorizer_name="openai-vectorizer",
    parameters=AzureOpenAIVectorizerParameters(
        resource_url="https://<resource>.openai.azure.com",
        deployment_name="text-embedding-3-large",
        model_name="text-embedding-3-large"
    )
)
```

### Vector Field with Integrated Vectorization

```python
from azure.search.documents.indexes.models import (
    SearchField,
    SearchFieldDataType,
)

vector_field = SearchField(
    name="content_vector",
    type=SearchFieldDataType.Collection(SearchFieldDataType.Single),
    searchable=True,
    stored=False,  # Don't store vectors to save space
    vector_search_dimensions=3072,  # text-embedding-3-large
    vector_search_profile_name="vector-profile"
)
```

## Vector Query Patterns

### Basic Vector Query

```python
from azure.search.documents.models import VectorizedQuery

def vector_search(client, query_vector: list[float], k: int = 10):
    """Execute a pure vector search."""
    vector_query = VectorizedQuery(
        vector=query_vector,
        k_nearest_neighbors=k,
        fields="content_vector"
    )
    
    results = client.search(
        search_text=None,
        vector_queries=[vector_query],
        select=["id", "title", "content"]
    )
    
    return list(results)
```

### Multi-Vector Query

Search across multiple vector fields simultaneously.

```python
from azure.search.documents.models import VectorizedQuery

def multi_vector_search(client, title_vector: list[float], content_vector: list[float]):
    """Search across title and content vectors."""
    results = client.search(
        search_text=None,
        vector_queries=[
            VectorizedQuery(
                vector=title_vector,
                k_nearest_neighbors=10,
                fields="title_vector"
            ),
            VectorizedQuery(
                vector=content_vector,
                k_nearest_neighbors=10,
                fields="content_vector"
            )
        ],
        select=["id", "title", "content"]
    )
    
    return list(results)
```

### Vector Search with Filters

Pre-filter documents before vector search for better performance.

```python
def filtered_vector_search(
    client,
    query_vector: list[float],
    category: str,
    min_rating: float = 4.0
):
    """Vector search with pre-filtering."""
    vector_query = VectorizedQuery(
        vector=query_vector,
        k_nearest_neighbors=10,
        fields="content_vector"
    )
    
    results = client.search(
        search_text=None,
        vector_queries=[vector_query],
        filter=f"category eq '{category}' and rating ge {min_rating}",
        select=["id", "title", "content", "category", "rating"]
    )
    
    return list(results)
```

## Exhaustive KNN Search

For exact nearest neighbor search (smaller indexes or when precision is critical).

```python
from azure.search.documents.indexes.models import (
    VectorSearch,
    ExhaustiveKnnAlgorithmConfiguration,
    ExhaustiveKnnParameters,
    VectorSearchProfile,
)

vector_search = VectorSearch(
    algorithms=[
        ExhaustiveKnnAlgorithmConfiguration(
            name="exhaustive-knn",
            parameters=ExhaustiveKnnParameters(
                metric="cosine"
            )
        )
    ],
    profiles=[
        VectorSearchProfile(
            name="exhaustive-profile",
            algorithm_configuration_name="exhaustive-knn"
        )
    ]
)
```

## Scalar Quantization (Preview)

Reduce vector storage size with minimal quality loss.

```python
from azure.search.documents.indexes.models import (
    VectorSearch,
    VectorSearchProfile,
    ScalarQuantizationCompression,
    ScalarQuantizationParameters,
)

vector_search = VectorSearch(
    profiles=[
        VectorSearchProfile(
            name="quantized-profile",
            algorithm_configuration_name="hnsw-config",
            compression_name="scalar-quantization"
        )
    ],
    compressions=[
        ScalarQuantizationCompression(
            compression_name="scalar-quantization",
            parameters=ScalarQuantizationParameters(
                quantized_data_type="int8"
            )
        )
    ]
)
```

## Async Vector Search

```python
from azure.search.documents.aio import SearchClient
from azure.search.documents.models import VectorizedQuery

async def async_vector_search(client: SearchClient, query_vector: list[float]):
    """Async vector search."""
    vector_query = VectorizedQuery(
        vector=query_vector,
        k_nearest_neighbors=10,
        fields="content_vector"
    )
    
    results = client.search(
        search_text=None,
        vector_queries=[vector_query],
        select=["id", "title", "content"]
    )
    
    docs = []
    async for result in results:
        docs.append(result)
    
    return docs
```

## Best Practices

1. **Dimensions**: Match your embedding model (text-embedding-3-large = 3072, ada-002 = 1536)
2. **Stored vectors**: Set `stored=False` to save storage if you don't need to retrieve vectors
3. **Pre-filtering**: Use filters to narrow search space before vector comparison
4. **Hybrid search**: Combine with keyword search for best relevance (see hybrid-search.md)
5. **Compression**: Use scalar quantization for large indexes to reduce costs

#!/usr/bin/env python3
"""
Create Azure AI Search agentic retrieval infrastructure.

Usage:
    python setup_agentic_retrieval.py --index-name <name> --kb-name <name> [options]

Environment variables required:
    SEARCH_ENDPOINT: Azure AI Search endpoint
    AOAI_ENDPOINT: Azure OpenAI endpoint
    AOAI_EMBEDDING_MODEL: Embedding model name (default: text-embedding-3-large)
    AOAI_EMBEDDING_DEPLOYMENT: Embedding deployment name
    AOAI_GPT_MODEL: GPT model name (default: gpt-4o-mini)
    AOAI_GPT_DEPLOYMENT: GPT deployment name
"""

import argparse
import os
from azure.identity import DefaultAzureCredential
from azure.search.documents.indexes import SearchIndexClient
from azure.search.documents.indexes.models import (
    SearchIndex, SearchField, VectorSearch, VectorSearchProfile,
    HnswAlgorithmConfiguration, AzureOpenAIVectorizer,
    AzureOpenAIVectorizerParameters, SemanticSearch,
    SemanticConfiguration, SemanticPrioritizedFields, SemanticField,
    SearchIndexKnowledgeSource, SearchIndexKnowledgeSourceParameters,
    SearchIndexFieldReference, KnowledgeBase, KnowledgeBaseAzureOpenAIModel,
    KnowledgeSourceReference, KnowledgeRetrievalOutputMode
)


def create_index(client: SearchIndexClient, name: str, aoai_endpoint: str,
                 embedding_deployment: str, embedding_model: str,
                 dimensions: int = 3072) -> SearchIndex:
    """Create a search index with vector and semantic search."""
    index = SearchIndex(
        name=name,
        fields=[
            SearchField(name="id", type="Edm.String", key=True,
                       filterable=True, sortable=True),
            SearchField(name="content", type="Edm.String",
                       filterable=False, sortable=False),
            SearchField(name="embedding", type="Collection(Edm.Single)",
                       stored=False, vector_search_dimensions=dimensions,
                       vector_search_profile_name="hnsw-profile"),
            SearchField(name="metadata", type="Edm.String",
                       filterable=True, sortable=False)
        ],
        vector_search=VectorSearch(
            profiles=[VectorSearchProfile(
                name="hnsw-profile",
                algorithm_configuration_name="hnsw-algo",
                vectorizer_name="aoai-vectorizer"
            )],
            algorithms=[HnswAlgorithmConfiguration(name="hnsw-algo")],
            vectorizers=[AzureOpenAIVectorizer(
                vectorizer_name="aoai-vectorizer",
                parameters=AzureOpenAIVectorizerParameters(
                    resource_url=aoai_endpoint,
                    deployment_name=embedding_deployment,
                    model_name=embedding_model
                )
            )]
        ),
        semantic_search=SemanticSearch(
            default_configuration_name="semantic-config",
            configurations=[SemanticConfiguration(
                name="semantic-config",
                prioritized_fields=SemanticPrioritizedFields(
                    content_fields=[SemanticField(field_name="content")]
                )
            )]
        )
    )
    return client.create_or_update_index(index)


def create_knowledge_source(client: SearchIndexClient, name: str,
                           index_name: str, description: str = "") -> None:
    """Create a knowledge source pointing to an index."""
    ks = SearchIndexKnowledgeSource(
        name=name,
        description=description or f"Knowledge source for {index_name}",
        search_index_parameters=SearchIndexKnowledgeSourceParameters(
            search_index_name=index_name,
            source_data_fields=[
                SearchIndexFieldReference(name="id"),
                SearchIndexFieldReference(name="metadata")
            ]
        )
    )
    client.create_or_update_knowledge_source(ks)


def create_knowledge_base(client: SearchIndexClient, name: str,
                         knowledge_source_name: str, aoai_endpoint: str,
                         gpt_deployment: str, gpt_model: str,
                         answer_instructions: str = "") -> None:
    """Create a knowledge base with LLM integration."""
    kb = KnowledgeBase(
        name=name,
        models=[KnowledgeBaseAzureOpenAIModel(
            azure_open_ai_parameters=AzureOpenAIVectorizerParameters(
                resource_url=aoai_endpoint,
                deployment_name=gpt_deployment,
                model_name=gpt_model
            )
        )],
        knowledge_sources=[KnowledgeSourceReference(name=knowledge_source_name)],
        output_mode=KnowledgeRetrievalOutputMode.ANSWER_SYNTHESIS,
        answer_instructions=answer_instructions or
            "Provide concise, accurate answers citing the retrieved documents."
    )
    client.create_or_update_knowledge_base(kb)


def main():
    parser = argparse.ArgumentParser(
        description="Set up Azure AI Search agentic retrieval infrastructure"
    )
    parser.add_argument("--index-name", required=True, help="Search index name")
    parser.add_argument("--kb-name", required=True, help="Knowledge base name")
    parser.add_argument("--ks-name", help="Knowledge source name (default: <index>-source)")
    parser.add_argument("--dimensions", type=int, default=3072,
                       help="Vector dimensions (default: 3072)")
    parser.add_argument("--answer-instructions", default="",
                       help="Custom answer synthesis instructions")
    args = parser.parse_args()

    # Load environment
    search_endpoint = os.environ["SEARCH_ENDPOINT"]
    aoai_endpoint = os.environ["AOAI_ENDPOINT"]
    embedding_model = os.environ.get("AOAI_EMBEDDING_MODEL", "text-embedding-3-large")
    embedding_deployment = os.environ.get("AOAI_EMBEDDING_DEPLOYMENT", embedding_model)
    gpt_model = os.environ.get("AOAI_GPT_MODEL", "gpt-4o-mini")
    gpt_deployment = os.environ.get("AOAI_GPT_DEPLOYMENT", gpt_model)

    ks_name = args.ks_name or f"{args.index_name}-source"

    credential = DefaultAzureCredential()
    client = SearchIndexClient(endpoint=search_endpoint, credential=credential)

    print(f"Creating index '{args.index_name}'...")
    create_index(client, args.index_name, aoai_endpoint,
                embedding_deployment, embedding_model, args.dimensions)
    print(f"  ✓ Index created")

    print(f"Creating knowledge source '{ks_name}'...")
    create_knowledge_source(client, ks_name, args.index_name)
    print(f"  ✓ Knowledge source created")

    print(f"Creating knowledge base '{args.kb_name}'...")
    create_knowledge_base(client, args.kb_name, ks_name, aoai_endpoint,
                         gpt_deployment, gpt_model, args.answer_instructions)
    print(f"  ✓ Knowledge base created")

    print(f"\nAgentic retrieval setup complete!")
    print(f"  Index: {args.index_name}")
    print(f"  Knowledge Source: {ks_name}")
    print(f"  Knowledge Base: {args.kb_name}")


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
Create an Azure AI Search index with vector search capabilities.

Usage:
    python setup_vector_index.py --index-name <name> [options]

Environment variables required:
    AZURE_SEARCH_ENDPOINT: Azure AI Search endpoint
    AZURE_OPENAI_ENDPOINT: Azure OpenAI endpoint (for integrated vectorization)
    AZURE_OPENAI_EMBEDDING_DEPLOYMENT: Embedding model deployment name
"""

import argparse
import os
from azure.identity import DefaultAzureCredential
from azure.search.documents.indexes import SearchIndexClient
from azure.search.documents.indexes.models import (
    SearchIndex,
    SearchField,
    SearchFieldDataType,
    VectorSearch,
    VectorSearchProfile,
    HnswAlgorithmConfiguration,
    HnswParameters,
    AzureOpenAIVectorizer,
    AzureOpenAIVectorizerParameters,
    SemanticSearch,
    SemanticConfiguration,
    SemanticPrioritizedFields,
    SemanticField,
    SearchableField,
    SimpleField,
)


def create_vector_index(
    client: SearchIndexClient,
    index_name: str,
    aoai_endpoint: str | None = None,
    embedding_deployment: str | None = None,
    dimensions: int = 1536,
    enable_semantic: bool = True,
) -> SearchIndex:
    """Create a search index with vector and optional semantic search."""

    # Define fields
    fields = [
        SimpleField(
            name="id",
            type=SearchFieldDataType.String,
            key=True,
            filterable=True,
            sortable=True,
        ),
        SearchableField(
            name="title",
            type=SearchFieldDataType.String,
            filterable=True,
            sortable=True,
        ),
        SearchableField(
            name="content",
            type=SearchFieldDataType.String,
        ),
        SimpleField(
            name="category",
            type=SearchFieldDataType.String,
            filterable=True,
            facetable=True,
        ),
        SearchField(
            name="content_vector",
            type=SearchFieldDataType.Collection(SearchFieldDataType.Single),
            searchable=True,
            stored=False,
            vector_search_dimensions=dimensions,
            vector_search_profile_name="vector-profile",
        ),
    ]

    # Configure vector search
    vectorizers = []
    if aoai_endpoint and embedding_deployment:
        vectorizers.append(
            AzureOpenAIVectorizer(
                vectorizer_name="openai-vectorizer",
                parameters=AzureOpenAIVectorizerParameters(
                    resource_url=aoai_endpoint,
                    deployment_name=embedding_deployment,
                    model_name=embedding_deployment,
                ),
            )
        )

    vector_search = VectorSearch(
        algorithms=[
            HnswAlgorithmConfiguration(
                name="hnsw-algo",
                parameters=HnswParameters(
                    m=4,
                    ef_construction=400,
                    ef_search=500,
                    metric="cosine",
                ),
            )
        ],
        profiles=[
            VectorSearchProfile(
                name="vector-profile",
                algorithm_configuration_name="hnsw-algo",
                vectorizer_name="openai-vectorizer" if vectorizers else None,
            )
        ],
        vectorizers=vectorizers if vectorizers else None,
    )

    # Configure semantic search
    semantic_search = None
    if enable_semantic:
        semantic_search = SemanticSearch(
            default_configuration_name="semantic-config",
            configurations=[
                SemanticConfiguration(
                    name="semantic-config",
                    prioritized_fields=SemanticPrioritizedFields(
                        title_field=SemanticField(field_name="title"),
                        content_fields=[SemanticField(field_name="content")],
                        keywords_fields=[SemanticField(field_name="category")],
                    ),
                )
            ],
        )

    # Create index
    index = SearchIndex(
        name=index_name,
        fields=fields,
        vector_search=vector_search,
        semantic_search=semantic_search,
    )

    return client.create_or_update_index(index)


def main():
    parser = argparse.ArgumentParser(
        description="Create an Azure AI Search index with vector search"
    )
    parser.add_argument("--index-name", required=True, help="Search index name")
    parser.add_argument(
        "--dimensions",
        type=int,
        default=1536,
        help="Vector dimensions (default: 1536 for ada-002, use 3072 for text-embedding-3-large)",
    )
    parser.add_argument(
        "--no-semantic",
        action="store_true",
        help="Disable semantic search configuration",
    )
    parser.add_argument(
        "--no-vectorizer",
        action="store_true",
        help="Skip integrated vectorization (provide vectors manually)",
    )
    args = parser.parse_args()

    # Load environment
    search_endpoint = os.environ["AZURE_SEARCH_ENDPOINT"]
    aoai_endpoint = (
        os.environ.get("AZURE_OPENAI_ENDPOINT") if not args.no_vectorizer else None
    )
    embedding_deployment = (
        os.environ.get("AZURE_OPENAI_EMBEDDING_DEPLOYMENT")
        if not args.no_vectorizer
        else None
    )

    credential = DefaultAzureCredential()
    client = SearchIndexClient(endpoint=search_endpoint, credential=credential)

    print(f"Creating index '{args.index_name}'...")
    index = create_vector_index(
        client=client,
        index_name=args.index_name,
        aoai_endpoint=aoai_endpoint,
        embedding_deployment=embedding_deployment,
        dimensions=args.dimensions,
        enable_semantic=not args.no_semantic,
    )

    print(f"  Index created: {index.name}")
    print(f"  Fields: {[f.name for f in index.fields]}")
    print(f"  Vector dimensions: {args.dimensions}")
    print(f"  Semantic search: {'enabled' if not args.no_semantic else 'disabled'}")
    print(f"  Integrated vectorization: {'enabled' if aoai_endpoint else 'disabled'}")


if __name__ == "__main__":
    main()