azure_ai_search_vector_store_example.py

"""
Example: Using AzureAISearchVectorStore with custom document schemas.

This example demonstrates how to:
1. Define domain-specific documents by extending Document
2. Pass document_type to AzureAISearchVectorStore
3. Automatically get all fields indexed (universal + custom)
4. Filter on any field efficiently

All fields from your document class are automatically indexed in Azure Search.

Required .env variables:
- AZURE_AI_SEARCH_ENDPOINT: Your Azure Search service endpoint
- AZURE_AI_SEARCH_API_KEY: Your Azure Search API key
- SEARCH_SERVICE_NAME: Your Azure Search service name

Optional .env variables:
- AZURE_AI_SEARCH_EMBEDDING_DIMENSION: Embedding dimension (default: 1536)
"""

import os
import time
from dataclasses import dataclass
from datetime import datetime
from typing import Optional
from pathlib import Path

from dotenv import load_dotenv
from gmf_forge_ai_data.vector_stores import Document, AzureAISearchVectorStore
from gmf_forge_ai_data.indexing import AzureAISearchIndexBuilder
from gmf_forge_ai_data.embeddings import AzureOpenAIEmbeddings
from gmf_forge_ai_data.chunkers import RecursiveChunker
from gmf_forge_ai_shared_core.observability import BasicLogger
from gmf_forge_ai_shared_core.observability.tracing import get_tracer

# Load environment variables
env_path = Path(__file__).parent / '.env'
load_dotenv(env_path)
logger = BasicLogger(__name__)

# Corporate SSL certificate path (for corporate networks with SSL inspection)
WORKSPACE_ROOT = Path(__file__).parent.parent.parent.parent
CORPORATE_CERT = WORKSPACE_ROOT / "certs" / "gmf_and_public_cas.pem"


# ============================================================================
# Configuration from .env
# ============================================================================

def get_azure_search_config():
    """Load Azure Search configuration from environment variables.

    Authentication is selected by the ``AZURE_USE_MANAGED_IDENTITY`` env var:

    * **API key** (default): set ``AZURE_AI_SEARCH_API_KEY``.
    * **Managed identity**: set ``AZURE_USE_MANAGED_IDENTITY=true``.
      The token provider must request the **Azure AI Search** scope::

          https://search.azure.com/.default

      Note: Azure OpenAI / Cognitive Services uses a *different* scope
      (``https://cognitiveservices.azure.com/.default``) — each service
      requires its own token provider.
    """
    endpoint = os.getenv('AZURE_AI_SEARCH_ENDPOINT')
    api_key = os.getenv('AZURE_AI_SEARCH_API_KEY')
    service_name = os.getenv('SEARCH_SERVICE_NAME')
    embedding_dim = int(os.getenv('AZURE_AI_SEARCH_EMBEDDING_DIMENSION', '1536'))
    use_managed_identity = os.getenv('AZURE_USE_MANAGED_IDENTITY', '').lower() in ('1', 'true', 'yes')

    if not endpoint or (not api_key and not use_managed_identity):
        logger.warning(
            "Azure Search credentials not found in .env file",
            missing=["AZURE_AI_SEARCH_ENDPOINT",
                     "AZURE_AI_SEARCH_API_KEY (or set AZURE_USE_MANAGED_IDENTITY=true)"],
            note="Using placeholder values for demonstration",
        )
        return {
            'endpoint': 'https://your-search.search.windows.net',
            'api_key': 'your-api-key',
            'token_provider': None,
            'service_name': 'your-search-service',
            'embedding_dimension': embedding_dim,
            'configured': False
        }

    token_provider = None
    if use_managed_identity:
        # Managed identity — scope: https://search.azure.com/.default
        from azure.identity import DefaultAzureCredential, get_bearer_token_provider
        token_provider = get_bearer_token_provider(
            DefaultAzureCredential(),
            'https://search.azure.com/.default',
        )
        api_key = None

    return {
        'endpoint': endpoint,
        'api_key': api_key,
        'token_provider': token_provider,
        'service_name': service_name or 'default-search',
        'embedding_dimension': embedding_dim,
        'configured': True
    }


def get_ssl_cert_path() -> str:
    """Get SSL certificate path if it exists."""
    if CORPORATE_CERT.exists():
        return str(CORPORATE_CERT)
    return None


def _ensure_index(config: dict, index_name: str, document_type) -> None:
    """Provision the Azure AI Search index using AzureAISearchIndexBuilder."""
    AzureAISearchIndexBuilder(
        endpoint=config['endpoint'],
        api_key=config.get('api_key'),
        token_provider=config.get('token_provider'),
        index_name=index_name,
        embedding_dimension=config['embedding_dimension'],
        document_type=document_type,
        ssl_cert_path=get_ssl_cert_path(),
    ).create_index()


def get_embedder():
    """Initialize Azure OpenAI embeddings from environment variables.

    Authentication is selected by the ``AZURE_USE_MANAGED_IDENTITY`` env var:

    * **API key** (default): set ``AZURE_OPENAI_API_KEY``.
    * **Managed identity**: set ``AZURE_USE_MANAGED_IDENTITY=true``.
      The token provider must request the **Cognitive Services** scope::

          https://cognitiveservices.azure.com/.default

      Note: Azure AI Search uses a *different* scope
      (``https://search.azure.com/.default``) — each service requires
      its own token provider.
    """
    endpoint = os.getenv('AZURE_OPENAI_ENDPOINT')
    api_key = os.getenv('AZURE_OPENAI_API_KEY')
    deployment = os.getenv('AZURE_OPENAI_EMBEDDING_MODEL')
    api_version = os.getenv('AZURE_OPENAI_EMBEDDING_MODEL_VERSION', '2024-02-01')
    use_managed_identity = os.getenv('AZURE_USE_MANAGED_IDENTITY', '').lower() in ('1', 'true', 'yes')

    if not endpoint or not deployment or (not api_key and not use_managed_identity):
        logger.warning("Azure OpenAI credentials not found in .env file",
                       missing=["AZURE_OPENAI_ENDPOINT", "AZURE_OPENAI_EMBEDDING_MODEL",
                                "AZURE_OPENAI_API_KEY (or set AZURE_USE_MANAGED_IDENTITY=true)"])
        return None

    ssl_cert = get_ssl_cert_path()
    if ssl_cert:
        logger.info("Using SSL certificate", cert=CORPORATE_CERT.name)

    try:
        if use_managed_identity:
            # Managed identity — scope: https://cognitiveservices.azure.com/.default
            from azure.identity import DefaultAzureCredential, get_bearer_token_provider
            token_provider = get_bearer_token_provider(
                DefaultAzureCredential(),
                'https://cognitiveservices.azure.com/.default',
            )
            embedder = AzureOpenAIEmbeddings(
                endpoint=endpoint,
                deployment_name=deployment,
                token_provider=token_provider,
                api_version=api_version,
                ssl_cert_path=ssl_cert,
            )
        else:
            embedder = AzureOpenAIEmbeddings(
                endpoint=endpoint,
                deployment_name=deployment,
                api_key=api_key,
                api_version=api_version,
                ssl_cert_path=ssl_cert,
            )
        logger.info("Initialized embeddings", deployment=deployment)
        return embedder
    except Exception as e:
        logger.warning("Failed to initialize embedder", error=str(e))
        return None


# ============================================================================
# Example 1: Legal Document with Custom Fields
# ============================================================================

@dataclass
class LegalDocument(Document):
    """Legal document with case-specific fields."""
    case_number: str = ""
    court: str = ""
    jurisdiction: str = ""
    decision_date: Optional[datetime] = None
    case_type: str = ""
    source: str = ""
    page_number: Optional[int] = None


def example_legal_documents():
    """Legal document vector store with automatic field indexing."""
    logger.info("Starting example", example="Legal Documents")
    
    # Load config from .env
    config = get_azure_search_config()
    
    if not config['configured']:
        logger.info("Skipping Azure Search operations - no credentials", note="Showing schema and usage examples only")
        return
    
    # Initialize embedder
    embedder = get_embedder()
    if not embedder:
        logger.info("Skipping document indexing - no embedder available")
        return
    
    # Create store with custom document type
    # ALL fields from LegalDocument are automatically indexed
    store = AzureAISearchVectorStore(
        endpoint=config['endpoint'],
        index_name="legal_documents",
        api_key=config.get('api_key'),
        token_provider=config.get('token_provider'),
        embedding_dimension=config['embedding_dimension'],
        document_type=LegalDocument,  # <-- Pass your document type
    )
    _ensure_index(config, "legal_documents", LegalDocument)
    
    logger.info("Azure Search index created",
        universal_fields=["id", "content", "embedding", "timestamp"],
        custom_fields=["case_number", "court", "jurisdiction", "decision_date", "case_type", "source", "page_number"],
        note="All fields are indexed and filterable",
    )
    
    # Create sample legal documents
    docs = [
        LegalDocument(
            id="case_001",
            content="In the matter of Smith v. Jones, the court finds...",
            case_number="2024-CV-12345",
            court="Supreme Court",
            jurisdiction="Federal",
            decision_date=datetime(2024, 1, 15),
            case_type="civil",
            source="court_records.pdf",
            page_number=42,
            timestamp=datetime.now()
        ),
        LegalDocument(
            id="case_002",
            content="United States v. Johnson, regarding antitrust violations...",
            case_number="2024-CR-67890",
            court="District Court",
            jurisdiction="Federal",
            decision_date=datetime(2024, 2, 20),
            case_type="criminal",
            source="court_records.pdf",
            page_number=105,
            timestamp=datetime.now()
        ),
        LegalDocument(
            id="case_003",
            content="State of California v. Tech Corp, patent infringement...",
            case_number="2024-CV-11111",
            court="State Supreme Court",
            jurisdiction="State",
            decision_date=datetime(2024, 3, 10),
            case_type="civil",
            source="state_records.pdf",
            page_number=1,
            timestamp=datetime.now()
        )
    ]
    
    # Generate embeddings for all documents
    logger.info("Generating embeddings")
    for doc in docs:
        doc.embedding = embedder.embed_text(doc.content)
    logger.info("Generated embeddings", count=len(docs))
    
    # Add documents — embeddings already generated above
    store.add_documents(docs)
    logger.info("Added documents to index", count=len(docs))

    # ── Vector search with filter ─────────────────────────────────────────
    query = "patent infringement technology"
    logger.info("Vector search with filter", query=query, filter="jurisdiction=Federal")
    query_embedding = embedder.embed_text(query)
    results = store.search(
        query_embedding=query_embedding,
        top_k=3,
        filters={"jurisdiction": "Federal"},
        search_type="vector",
    )
    for r in results:
        doc: LegalDocument = r.document
        logger.info("Result", rank=r.rank, id=doc.id, score=round(r.score, 4),
                    case_number=doc.case_number, court=doc.court, jurisdiction=doc.jurisdiction)

    # ── get_document ──────────────────────────────────────────────────────
    fetched = store.get_document("case_001")
    if fetched:
        logger.info("get_document", id=fetched.id, case_number=fetched.case_number,
                    court=fetched.court, case_type=fetched.case_type)

    logger.info("Supported filter syntax",
                equality="{'jurisdiction': 'Federal'}",
                date_range="{'decision_date': ('>=', datetime(2024, 1, 1))}",
                multiple="{'jurisdiction': 'Federal', 'case_type': 'civil'}")


# ============================================================================
# Example 2: Product Catalog
# ============================================================================

@dataclass
class ProductDocument(Document):
    """Product document for e-commerce."""
    sku: str = ""
    category: str = ""
    price: float = 0.0
    in_stock: bool = True
    brand: str = ""
    rating: float = 0.0
    review_count: int = 0


def example_product_catalog():
    """Product catalog with automatic field indexing."""
    logger.info("Starting example", example="Product Catalog")
    
    # Load config from .env
    config = get_azure_search_config()
    
    if not config['configured']:
        return
    
    # Initialize embedder
    embedder = get_embedder()
    if not embedder:
        return
    
    # Create store - all ProductDocument fields automatically indexed
    store = AzureAISearchVectorStore(
            endpoint=config['endpoint'],
            index_name="products",
            api_key=config.get('api_key'),
            token_provider=config.get('token_provider'),
            embedding_dimension=config['embedding_dimension'],
            document_type=ProductDocument,  # <-- Your custom type
        )
    _ensure_index(config, "products", ProductDocument)
    
    logger.info("Azure Search index created",
        universal_fields=["id", "content", "embedding", "timestamp"],
        custom_fields=["sku", "category", "price", "in_stock", "brand", "rating", "review_count"],
    )
    
    # Sample products
    products = [
        ProductDocument(
            id="prod_001",
            content="Premium wireless headphones with noise cancellation...",
            sku="WH-1000XM5",
            category="Electronics",
            price=399.99,
            in_stock=True,
            brand="Sony",
            rating=4.8,
            review_count=1250,
            timestamp=datetime.now()
        ),
        ProductDocument(
            id="prod_002",
            content="Professional DSLR camera with 24MP sensor...",
            sku="EOS-R5",
            category="Electronics",
            price=3899.00,
            in_stock=False,
            brand="Canon",
            rating=4.9,
            review_count=890,
            timestamp=datetime.now()
        )
    ]
    
    # Generate embeddings for products
    logger.info("Generating embeddings")
    for product in products:
        product.embedding = embedder.embed_text(product.content)
    logger.info("Generated embeddings", count=len(products))
    
    # Add documents — embeddings already generated above
    store.add_documents(products)
    logger.info("Added products to index", count=len(products))

    # ── Keyword search ────────────────────────────────────────────────────
    query = "wireless headphones noise cancellation"
    logger.info("Keyword search", query=query, search_type="keyword")
    results = store.search(query=query, top_k=3, search_type="keyword")
    for r in results:
        prod: ProductDocument = r.document
        logger.info("Result", rank=r.rank, id=prod.id, score=round(r.score, 4),
                    sku=prod.sku, brand=prod.brand, price=prod.price)

    # ── count ─────────────────────────────────────────────────────────────
    logger.info("Document count", total=store.count())

    logger.info("Supported filter syntax",
                numeric_range="{'price': ('<=', 500)}",
                boolean="{'in_stock': True}",
                combined="{'category': 'Electronics', 'in_stock': True, 'rating': ('>=', 4.5)}")


# ============================================================================
# Example 3: Financial Documents
# ============================================================================

@dataclass
class FinancialDocument(Document):
    """Financial report document."""
    document_type: str = ""  # 10-K, 10-Q, 8-K, etc.
    fiscal_year: int = 0
    quarter: str = ""  # Q1, Q2, Q3, Q4, FY
    company_ticker: str = ""
    sector: str = ""
    report_date: Optional[datetime] = None
    source: str = ""
    page_number: Optional[int] = None


def example_financial_documents():
    """Financial documents with automatic field indexing."""
    logger.info("Starting example", example="Financial Documents")
    
    # Load config from .env
    config = get_azure_search_config()
    
    if not config['configured']:
        return
    
    # Initialize embedder
    embedder = get_embedder()
    if not embedder:
        return
    
    store = AzureAISearchVectorStore(
            endpoint=config['endpoint'],
            index_name="financial_reports",
            api_key=config.get('api_key'),
            token_provider=config.get('token_provider'),
            embedding_dimension=config['embedding_dimension'],
            document_type=FinancialDocument,
        )
    _ensure_index(config, "financial_reports", FinancialDocument)
    
    logger.info("Azure Search index created",
        universal_fields=["id", "content", "embedding", "timestamp"],
        custom_fields=["document_type", "fiscal_year", "quarter", "company_ticker", "sector", "report_date", "source", "page_number"],
        note="All fields are indexed and filterable",
    )
    
    # Create sample financial documents
    fin_docs = [
        FinancialDocument(
            id="fin_001",
            content="Microsoft Q1 2024 earnings report shows strong cloud growth...",
            document_type="10-Q",
            fiscal_year=2024,
            quarter="Q1",
            company_ticker="MSFT",
            sector="Technology",
            report_date=datetime(2024, 1, 30),
            source="sec_filings.pdf",
            page_number=1,
            timestamp=datetime.now()
        ),
        FinancialDocument(
            id="fin_002",
            content="Apple annual 10-K filing for fiscal year 2024...",
            document_type="10-K",
            fiscal_year=2024,
            quarter="FY",
            company_ticker="AAPL",
            sector="Technology",
            report_date=datetime(2024, 10, 28),
            source="sec_filings.pdf",
            page_number=15,
            timestamp=datetime.now()
        )
    ]
    
    # Generate embeddings for financial documents
    logger.info("Generating embeddings")
    for doc in fin_docs:
        doc.embedding = embedder.embed_text(doc.content)
    logger.info("Generated embeddings", count=len(fin_docs))
    
    # Add documents — embeddings already generated above
    store.add_documents(fin_docs)
    logger.info("Added financial documents to index", count=len(fin_docs))
    time.sleep(2)  # allow Azure Search to index before searching

    # ── Hybrid search with filter ─────────────────────────────────────────
    query = "quarterly earnings cloud growth"
    logger.info("Hybrid search with filter", query=query, filter="sector=Technology",
                search_type="hybrid")
    query_embedding = embedder.embed_text(query)
    results = store.search(
        query=query,
        query_embedding=query_embedding,
        top_k=3,
        filters={"sector": "Technology"},
        search_type="hybrid",
    )
    for r in results:
        doc: FinancialDocument = r.document
        logger.info("Result", rank=r.rank, id=doc.id, score=round(r.score, 4),
                    document_type=doc.document_type, company_ticker=doc.company_ticker)

    # ── CRUD operations ───────────────────────────────────────────────────
    logger.info("CRUD demonstration")
    fetched = store.get_document("fin_001")
    if fetched:
        logger.info("get_document", id=fetched.id, company_ticker=fetched.company_ticker,
                    quarter=fetched.quarter, sector=fetched.sector)

    updated = store.update_document("fin_001", metadata={"reviewed": True, "reviewer": "compliance_team"})
    logger.info("update_document", id="fin_001", success=updated)

    logger.info("count", total=store.count())

    deleted = store.delete_documents(["fin_002"])
    time.sleep(2)  # Azure Search index updates are eventually consistent
    logger.info("delete_documents", deleted=deleted, remaining=store.count())


# ============================================================================
# Document type for AI/ML Knowledge Base
# ============================================================================

@dataclass
class AIMLDocument(Document):
    """AI/ML knowledge base document."""
    topic: str = ""  # Specific topic like "deep_learning", "CNNs", "transformers"
    category: str = ""  # Broader category like "ML", "NLP", "CV"


# ============================================================================
# Example 4: AI/ML Knowledge Base (for retrieval examples)
# ============================================================================

def example_ai_ml_knowledge_base():
    """Create AI/ML knowledge base for retrieval strategy demonstrations."""
    logger.info("Starting example", example="AI/ML Knowledge Base")
    
    # Load config from .env
    config = get_azure_search_config()
    
    if not config['configured']:
        return
    
    # Initialize embedder
    embedder = get_embedder()
    if not embedder:
        return
    
    store = AzureAISearchVectorStore(
        endpoint=config['endpoint'],
        index_name="ai_ml_knowledge",
        api_key=config.get('api_key'),
        token_provider=config.get('token_provider'),
        embedding_dimension=config['embedding_dimension'],
        document_type=AIMLDocument,
    )
    _ensure_index(config, "ai_ml_knowledge", AIMLDocument)
    
    logger.info("Azure Search index created",
        universal_fields=["id", "content", "embedding", "timestamp"],
        custom_fields=["topic", "category"],
    )
    
    # Create AI/ML knowledge base documents.
    # Each document is a multi-sentence paragraph so that context processing
    # examples (reranking, compression, windowing) have realistic content to
    # work with.  Sentences within a chunk intentionally vary in relevance to
    # the shared query "How do neural networks learn and what role does
    # backpropagation play?" — this lets ContextCompressor demonstrate
    # meaningful sentence-level extraction and token reduction.
    documents_data = [
        {
            "id": "doc_001",
            "content": (
                "Machine learning is a subset of artificial intelligence that enables systems "
                "to learn and improve automatically from experience without being explicitly programmed. "
                "Rather than relying on hand-crafted rules, ML models generalise patterns from training "
                "data to make predictions on new, unseen inputs. "
                "Supervised learning trains models on labelled examples by minimising a loss function "
                "that measures the gap between predictions and ground truth. "
                "This optimisation process relies on computing gradients — the direction and magnitude "
                "in which to adjust model parameters — making calculus a foundational tool in machine learning. "
                "Once training is complete, the model's parameters remain fixed and predictions are made "
                "through a single forward pass."
            ),
            "category": "ML",
            "topic": "introduction"
        },
        {
            "id": "doc_002",
            "content": (
                "Deep learning uses neural networks with multiple stacked layers to learn hierarchical "
                "representations of data directly from raw inputs. "
                "A neural network learns by adjusting its weights through backpropagation, an algorithm "
                "that computes the gradient of the loss with respect to every weight using the chain rule "
                "of calculus. "
                "During forward propagation, each layer transforms its input through a weighted sum "
                "followed by a non-linear activation function, producing the final prediction. "
                "Backward propagation reverses this process: the loss gradient flows from the output layer "
                "back through each layer, enabling the network to attribute responsibility for the "
                "prediction error to each individual weight. "
                "Gradient descent then uses these gradients to nudge every weight in the direction that "
                "reduces the overall loss, and this cycle repeats over many mini-batches until convergence."
            ),
            "category": "ML",
            "topic": "deep_learning"
        },
        {
            "id": "doc_003",
            "content": (
                "Natural language processing (NLP) allows computers to understand, interpret, and generate "
                "human language in a meaningful way. "
                "It combines rule-based approaches with statistical and neural techniques to tackle tasks "
                "such as text classification, machine translation, and question answering. "
                "Words are typically represented as dense vectors — word embeddings — that encode semantic "
                "relationships so that similar words appear close to each other in the vector space. "
                "Neural language models are pre-trained on massive text corpora and learn to predict masked "
                "or next tokens, developing rich internal representations of grammar, facts, and reasoning. "
                "Fine-tuning adapts these pre-trained weights to specific downstream tasks using small "
                "task-specific labelled datasets."
            ),
            "category": "NLP",
            "topic": "introduction"
        },
        {
            "id": "doc_004",
            "content": (
                "The transformer architecture replaced recurrent networks as the dominant model for NLP "
                "by processing entire sequences in parallel using self-attention. "
                "Self-attention computes a weighted combination of all token representations simultaneously, "
                "allowing each position to directly attend to any other without the vanishing-gradient "
                "bottleneck of sequential processing. "
                "Unlike recurrent networks that propagate gradients through time steps, transformers propagate "
                "gradients through attention and feed-forward sub-layers, making training of very deep models "
                "more stable and allowing better gradient flow during backpropagation. "
                "Pre-trained models such as BERT, GPT, and T5 distil knowledge from billions of text tokens "
                "into their weights through iterative gradient-based updates. "
                "Positional encodings inject sequence-order information into the attention computation, "
                "since the self-attention mechanism itself is inherently order-invariant."
            ),
            "category": "NLP",
            "topic": "transformers"
        },
        {
            "id": "doc_005",
            "content": (
                "Optimising a machine learning model means finding the weight values that minimise a chosen "
                "loss function on the training data. "
                "Gradient descent is the workhorse algorithm: it iteratively moves the parameters in the "
                "direction of the negative gradient, taking small steps controlled by a learning rate. "
                "In neural networks, these gradients are efficiently calculated end-to-end by backpropagation, "
                "which applies the chain rule to propagate error signals layer by layer from the output "
                "back to the earliest weights. "
                "Adaptive optimisers such as Adam and RMSProp maintain per-parameter learning rates that "
                "adjust based on the history of gradient magnitudes, often converging faster than vanilla SGD. "
                "Learning rate schedules including warm-up phases and cosine annealing further stabilise "
                "training of large models."
            ),
            "category": "ML",
            "topic": "optimisation"
        },
        {
            "id": "doc_006",
            "content": (
                "Computer vision is the field of AI that enables machines to extract meaning from images, "
                "videos, and other visual inputs. "
                "Convolutional neural networks (CNNs) became the standard for vision tasks after "
                "demonstrating breakthrough performance on the ImageNet benchmark in 2012. "
                "A CNN stacks convolutional layers — which detect local spatial patterns using learnable "
                "filters — with pooling layers that progressively reduce spatial resolution. "
                "CNNs are trained end-to-end with backpropagation: gradients flow from the classification "
                "loss back through pooling and convolutional layers, updating the filter weights so the "
                "network learns to recognise task-relevant visual features. "
                "Modern architectures like ResNet use residual connections to allow gradients to bypass "
                "individual layers, enabling stable training of networks hundreds of layers deep."
            ),
            "category": "CV",
            "topic": "CNNs"
        },
        {
            "id": "doc_007",
            "content": (
                "Activation functions introduce the non-linearities that allow neural networks to model "
                "complex, high-dimensional relationships beyond what a linear model can capture. "
                "Early networks used sigmoid and tanh activations, but both suffer from vanishing gradients: "
                "when inputs fall in the saturated tails the derivative approaches zero and backpropagation "
                "cannot effectively update the earlier layers of the network. "
                "The Rectified Linear Unit (ReLU) addresses this by having a constant gradient of 1 for "
                "positive inputs, greatly accelerating training in deep architectures. "
                "Variants such as Leaky ReLU, GELU, and Swish make subtle adjustments to improve gradient "
                "flow in specific model families. "
                "Choosing an appropriate activation function is therefore a practical consideration for "
                "ensuring that backpropagation can reliably train all layers of a deep network."
            ),
            "category": "ML",
            "topic": "activation_functions"
        },
        {
            "id": "doc_008",
            "content": (
                "Reinforcement learning (RL) trains an agent to take actions in an environment by "
                "maximising a cumulative reward signal over time, rather than learning from labelled "
                "input–output pairs. "
                "The agent interacts episodically with the environment, observing states, selecting actions "
                "according to its current policy, and receiving scalar reward or penalty feedback. "
                "Deep reinforcement learning uses neural networks as function approximators for the policy "
                "or value function; these networks are updated through gradient-based methods derived from "
                "the Bellman equation, with backpropagation computing the necessary weight gradients. "
                "Stabilisation techniques such as experience replay and target networks are critical because "
                "RL training signals are non-stationary and temporally correlated, making naive gradient "
                "updates unstable. "
                "Prominent algorithms including DQN, PPO, and SAC have achieved superhuman performance on "
                "games and continuous control benchmarks using these deep RL principles."
            ),
            "category": "ML",
            "topic": "reinforcement"
        },
    ]
    
    logger.info("Creating AI/ML knowledge base documents")
    documents = []
    for doc_data in documents_data:
        doc = AIMLDocument(
            id=doc_data["id"],
            content=doc_data["content"],
            category=doc_data["category"],
            topic=doc_data["topic"],
            timestamp=datetime.now()
        )
        documents.append(doc)
    
    # Generate embeddings
    logger.info("Generating embeddings")
    for doc in documents:
        doc.embedding = embedder.embed_text(doc.content)
    logger.info("Generated embeddings", count=len(documents))
    
    # Add documents — embeddings already generated above
    store.add_documents(documents)
    logger.info("Added AI/ML documents to index", count=len(documents))
    
    logger.info("Knowledge base created",
        ml_docs=5, ml_topics=["introduction", "deep_learning", "optimisation", "activation_functions", "reinforcement"],
        nlp_docs=2, nlp_topics=["introduction", "transformers"],
        cv_docs=1, cv_topics=["CNNs"],
        note="Each document is a multi-sentence paragraph with mixed relevance to neural network learning and backpropagation",
    )
    logger.info("Index used by downstream examples",
        retrieval_example="Vector, Keyword, Hybrid, MMR, Ensemble, Parent Document, Hierarchical, Graph retrievers",
        query_example="QueryExpansion, HyDE, SubQuestion, StepBack",
        context_example="RelevanceFilter, ContextDeduplicator, ContextReranker, ContextCompressor, ContextWindowManager",
    )


# ============================================================================
# Document type for the RAG Pipeline example
# ============================================================================

@dataclass
class KnowledgeBaseArticle(Document):
    """Knowledge base article with domain classification and versioning fields."""
    domain: str = ""      # e.g. "ai", "security", "infrastructure"
    language: str = "en"  # ISO 639-1 language code
    version: str = "1.0"  # document version for change tracking


# ============================================================================
# Example 5: End-to-End RAG Pipeline (Chunk → Embed → Index → Retrieve)
# ============================================================================

def example_rag_pipeline():
    """Demonstrate a complete RAG ingestion pipeline: chunk → embed → index → retrieve."""
    logger.info("Starting example", example="End-to-End RAG Pipeline")
    logger.info("Scenario",
        description="Ingest multiple knowledge base articles through the full RAG pipeline",
        pipeline="articles → RecursiveChunker → embed_batch → add_documents(KnowledgeBaseArticle) → vector_search",
        note="Each article carries its own domain, language, and version — enabling filtered retrieval",
    )

    config = get_azure_search_config()
    if not config['configured']:
        return

    embedder = get_embedder()
    if not embedder:
        return

    store = AzureAISearchVectorStore(
        endpoint=config['endpoint'],
        index_name="rag_pipeline_demo",
        api_key=config.get('api_key'),
        token_provider=config.get('token_provider'),
        embedding_dimension=config['embedding_dimension'],
        document_type=KnowledgeBaseArticle,
    )
    _ensure_index(config, "rag_pipeline_demo", KnowledgeBaseArticle)
    logger.info("Azure Search index schema",
        document_type="KnowledgeBaseArticle",
        custom_fields=["domain", "language", "version"],
        note="Each chunk carries these fields — enables filtering by domain, language, or version")

    # ── Source articles (each with its own domain, language, and version) ───────
    articles = [
        {
            "domain": "ai",
            "language": "en",
            "version": "2.0",
            "text": (
                "Retrieval-augmented generation (RAG) is a technique that enhances large language "
                "model responses by grounding them in factual, up-to-date information retrieved "
                "from an external knowledge base. Rather than relying solely on the information "
                "encoded in model weights during training, a RAG system retrieves relevant passages "
                "at inference time and supplies them as context to the model, dramatically reducing "
                "hallucinations and enabling answers to reflect recent events. The ingestion pipeline "
                "is the foundation of every RAG system: raw documents are loaded by a connector, "
                "split into overlapping chunks by a chunker, converted into dense embedding vectors, "
                "and stored in a vector database alongside the original text and metadata."
            ),
        },
        {
            "domain": "security",
            "language": "en",
            "version": "1.3",
            "text": (
                "Zero-trust architecture is a security model that assumes no implicit trust for any "
                "user, device, or network segment. Every access request is verified, validated, and "
                "encrypted before being granted, regardless of whether the request originates from "
                "inside or outside the corporate network. Micro-segmentation divides the network into "
                "small zones, each requiring separate authentication. Continuous monitoring of user "
                "behaviour and device posture helps detect anomalies and respond to potential breaches "
                "in real time. Implementing zero-trust requires identity-aware proxies, multi-factor "
                "authentication, and robust logging across all access points."
            ),
        },
        {
            "domain": "infrastructure",
            "language": "en",
            "version": "1.0",
            "text": (
                "Infrastructure as Code (IaC) lets teams manage cloud resources through declarative "
                "configuration files instead of manual provisioning. Tools such as Terraform, Bicep, "
                "and Pulumi enable version-controlled, repeatable deployments across development, staging, "
                "and production environments. Drift detection compares the desired state in code against "
                "the actual state in the cloud provider, flagging any manual changes that may have been "
                "introduced. Modular IaC patterns encourage reusable modules that encapsulate best "
                "practices for networking, compute, and storage, reducing duplication and enforcing "
                "organisational standards at scale."
            ),
        },
    ]

    logger.info("Source articles", count=len(articles),
        domains=[a["domain"] for a in articles])

    query = "How does RAG reduce hallucinations?"
    tracer = get_tracer()
    with tracer.trace("rag_pipeline", input=query) as trace:
        with trace.span("chunking") as span:
            # ── Step 1: Chunk ──────────────────────────────────────────────────────
            logger.info("Step 1 - Chunking",
                strategy="RecursiveChunker", chunk_size=400, chunk_overlap=50)
            chunker = RecursiveChunker(chunk_size=400, chunk_overlap=50)
            all_chunks_per_article: list = []
            for article in articles:
                chunks = chunker.chunk(
                    article["text"],
                    metadata={"source": f"kb_{article['domain']}", "domain": article["domain"]},
                )
                logger.info("Chunked article",
                    domain=article["domain"], num_chunks=len(chunks),
                    avg_chars=sum(len(c.text) for c in chunks) // max(len(chunks), 1))
                all_chunks_per_article.append((article, chunks))
            span.set_output({"documents": len(articles), "total_chunks": sum(len(c) for _, c in all_chunks_per_article)})

        # ── Step 2: Embed ──────────────────────────────────────────────────────────
        with trace.span("embedding") as span:
            all_docs: list = []
            for article, chunks in all_chunks_per_article:
                texts = [c.text for c in chunks]
                logger.info("Step 2 - Embedding", domain=article["domain"], chunks=len(texts))
                embeddings = embedder.embed_batch(texts)

                # ── Build typed documents ──────────────────────────────────────────
                for i, (chunk, emb) in enumerate(zip(chunks, embeddings)):
                    all_docs.append(
                        KnowledgeBaseArticle(
                            id=f"{article['domain']}_{chunk.chunk_id}",
                            content=chunk.text,
                            embedding=emb,
                            domain=article["domain"],
                            language=article["language"],
                            version=article["version"],
                            metadata={**chunk.metadata, "chunk_index": i,
                                      "char_start": chunk.start_pos, "char_end": chunk.end_pos},
                        )
                    )
            span.set_output({"chunks": len(all_docs)})

        # ── Step 3: Index ──────────────────────────────────────────────────────────
        with trace.span("indexing") as span:
            logger.info("Step 3 - Indexing", index="rag_pipeline_demo",
                document_type="KnowledgeBaseArticle", total_chunks=len(all_docs))
            store.add_documents(all_docs, generate_embeddings=False)
            logger.info("Chunks indexed", count=len(all_docs))
            span.set_output({"chunks_indexed": len(all_docs)})

        # ── Step 4: Retrieve ───────────────────────────────────────────────────────
        with trace.span("retrieval", input=query) as span:
            logger.info("Step 4 - Retrieval", query=query, top_k=3)
            query_embedding = embedder.embed_text(query)
            results = store.search(query_embedding=query_embedding, top_k=3)
            logger.info("Retrieved chunks", returned=len(results))
            for r in results:
                article: KnowledgeBaseArticle = r.document
                snippet = article.content[:100].replace("\n", " ")
                logger.info("Retrieved chunk",
                    rank=r.rank, id=article.id, score=round(r.score, 4),
                    domain=article.domain, language=article.language, version=article.version,
                    snippet=snippet)
            span.set_output({"results": len(results), "top_score": round(results[0].score, 4) if results else 0})

        trace.set_output({"chunks_indexed": len(all_docs), "query_results": len(results)})

    logger.info("RAG pipeline complete",
        chunks_indexed=len(all_docs),
        query_results=len(results),
        tip="Filter by domain='security' or version='2.0' to narrow results at query time")


# ============================================================================
# Key Benefits Summary
# ============================================================================

def print_benefits():
    """Log key benefits of this approach."""
    logger.info("Key Benefits: Simple API",
        detail="Pass document_type parameter - no need to manually define search fields")
    logger.info("Key Benefits: Automatic Indexing",
        detail="All document fields automatically indexed with appropriate data types")
    logger.info("Key Benefits: Complete Filtering",
        detail="Filter on ANY field - string equality, numeric ranges, date ranges, boolean, multiple conditions")
    logger.info("Key Benefits: Type Safety",
        detail="Documents returned as custom type with IntelliSense support")
    logger.info("Key Benefits: Production Ready",
        detail="Azure Search scalability, HNSW vector search, hybrid search (vector + BM25)")


# ============================================================================
# Main
# ============================================================================

if __name__ == "__main__":
    logger.info("Custom Azure Vector Store Examples",
        note="Using AzureAISearchVectorStore with custom document schemas - all fields automatically indexed")
    
    example_legal_documents()
    logger.info("")
    example_product_catalog()
    logger.info("")
    example_financial_documents()
    logger.info("")
    example_ai_ml_knowledge_base()
    logger.info("")
    example_rag_pipeline()
    logger.info("")
    print_benefits()
    
    logger.info("Usage in production",
        step_1="Define your domain-specific Document subclass",
        step_2="Pass document_type to AzureAISearchVectorStore",
        step_3="All fields automatically indexed",
        step_4="Filter on any field efficiently",
    )