Tutorial 08: RAG with Azure AI Search¶

import os
from azure.search.documents.indexes import SearchIndexClient
from azure.search.documents.indexes.models import (
    SearchIndex,
    SearchField,
    SearchFieldDataType,
    SimpleField,
    SearchableField,
    VectorSearch,
    HnswAlgorithmConfiguration,
    VectorSearchProfile,
    SemanticConfiguration,
    SemanticSearch,
    SemanticPrioritizedFields,
    SemanticField,
    SearchIndex,
)
from azure.core.credentials import AzureKeyCredential

client = SearchIndexClient(
    endpoint=os.environ["SEARCH_ENDPOINT"],
    credential=AzureKeyCredential(os.environ["SEARCH_KEY"]),
)

fields = [
    SimpleField(name="id", type=SearchFieldDataType.String, key=True, filterable=True),
    SearchableField(name="title", type=SearchFieldDataType.String),
    SearchableField(name="content", type=SearchFieldDataType.String),
    SearchableField(name="source", type=SearchFieldDataType.String, filterable=True),
    SearchableField(name="category", type=SearchFieldDataType.String, filterable=True, facetable=True),
    SimpleField(name="last_updated", type=SearchFieldDataType.DateTimeOffset, filterable=True, sortable=True),
    SearchField(
        name="content_vector",
        type=SearchFieldDataType.Collection(SearchFieldDataType.Single),
        searchable=True,
        vector_search_dimensions=3072,
        vector_search_profile_name="vector-profile",
    ),
]

vector_search = VectorSearch(
    algorithms=[HnswAlgorithmConfiguration(name="hnsw-config")],
    profiles=[VectorSearchProfile(name="vector-profile", algorithm_configuration_name="hnsw-config")],
)

semantic_config = SemanticConfiguration(
    name="semantic-config",
    prioritized_fields=SemanticPrioritizedFields(
        title_field=SemanticField(field_name="title"),
        content_fields=[SemanticField(field_name="content")],
    ),
)

index = SearchIndex(
    name=os.environ["INDEX_NAME"],
    fields=fields,
    vector_search=vector_search,
    semantic_search=SemanticSearch(configurations=[semantic_config]),
)

result = client.create_or_update_index(index)
print(f"Index '{result.name}' created/updated successfully")

import os
import hashlib
from openai import AzureOpenAI
from azure.search.documents import SearchClient
from azure.core.credentials import AzureKeyCredential

oai = AzureOpenAI(
    azure_endpoint=os.environ["AOAI_ENDPOINT"],
    api_key=os.environ["AOAI_KEY"],
    api_version="2025-04-01-preview",
)

search = SearchClient(
    endpoint=os.environ["SEARCH_ENDPOINT"],
    index_name=os.environ["INDEX_NAME"],
    credential=AzureKeyCredential(os.environ["SEARCH_KEY"]),
)


def chunk_text(text: str, chunk_size: int = 1000, overlap: int = 200) -> list[str]:
    """Split text into overlapping chunks."""
    chunks = []
    start = 0
    while start < len(text):
        end = start + chunk_size
        chunks.append(text[start:end])
        start += chunk_size - overlap
    return chunks


def embed_text(text: str) -> list[float]:
    """Generate embedding vector for text."""
    resp = oai.embeddings.create(
        model="text-embedding-3-large",
        input=text,
    )
    return resp.data[0].embedding


def index_document(title: str, content: str, source: str, category: str):
    """Chunk, embed, and upload a document to the search index."""
    chunks = chunk_text(content)
    documents = []

    for i, chunk in enumerate(chunks):
        doc_id = hashlib.md5(f"{title}_{i}".encode()).hexdigest()
        vector = embed_text(chunk)
        documents.append({
            "id": doc_id,
            "title": f"{title} (part {i+1})",
            "content": chunk,
            "source": source,
            "category": category,
            "content_vector": vector,
        })

    result = search.upload_documents(documents)
    succeeded = sum(1 for r in result if r.succeeded)
    print(f"  Indexed {succeeded}/{len(documents)} chunks for '{title}'")
    return succeeded

import os
import glob
from embed_documents import index_document

# Index CSA platform documentation
docs_path = "docs/"
md_files = glob.glob(os.path.join(docs_path, "**/*.md"), recursive=True)

print(f"Found {len(md_files)} markdown files to index")

total = 0
for filepath in md_files:
    with open(filepath, "r", encoding="utf-8") as f:
        content = f.read()
    if len(content) < 50:  # Skip tiny files
        continue
    title = os.path.basename(filepath).replace(".md", "").replace("_", " ").title()
    total += index_document(
        title=title,
        content=content,
        source=filepath,
        category="documentation",
    )

print(f"\nTotal chunks indexed: {total}")

import os
from openai import AzureOpenAI
from azure.search.documents import SearchClient
from azure.search.documents.models import (
    VectorizableTextQuery,
    QueryType,
    QueryCaptionType,
    QueryAnswerType,
)
from azure.core.credentials import AzureKeyCredential

oai = AzureOpenAI(
    azure_endpoint=os.environ["AOAI_ENDPOINT"],
    api_key=os.environ["AOAI_KEY"],
    api_version="2025-04-01-preview",
)

search = SearchClient(
    endpoint=os.environ["SEARCH_ENDPOINT"],
    index_name=os.environ["INDEX_NAME"],
    credential=AzureKeyCredential(os.environ["SEARCH_KEY"]),
)


def embed_query(query: str) -> list[float]:
    resp = oai.embeddings.create(model="text-embedding-3-large", input=query)
    return resp.data[0].embedding


def hybrid_search(query: str, top_k: int = 5) -> list[dict]:
    """Perform hybrid search: vector + keyword + semantic reranking."""
    vector = embed_query(query)

    results = search.search(
        search_text=query,
        vector_queries=[
            VectorizableTextQuery(
                text=query,
                k_nearest_neighbors=top_k,
                fields="content_vector",
            )
        ],
        query_type=QueryType.SEMANTIC,
        semantic_configuration_name="semantic-config",
        query_caption=QueryCaptionType.EXTRACTIVE,
        query_answer=QueryAnswerType.EXTRACTIVE,
        top=top_k,
    )

    hits = []
    for r in results:
        hits.append({
            "title": r["title"],
            "content": r["content"],
            "source": r["source"],
            "score": r["@search.score"],
            "reranker_score": r.get("@search.reranker_score", 0),
        })
    return hits


if __name__ == "__main__":
    query = "How does the medallion architecture work?"
    print(f"Query: {query}\n")
    results = hybrid_search(query)
    for i, r in enumerate(results, 1):
        print(f"{i}. [{r['reranker_score']:.2f}] {r['title']}")
        print(f"   {r['content'][:120]}...")
        print()

import os
from openai import AzureOpenAI
from search import hybrid_search

oai = AzureOpenAI(
    azure_endpoint=os.environ["AOAI_ENDPOINT"],
    api_key=os.environ["AOAI_KEY"],
    api_version="2025-04-01-preview",
)

SYSTEM_PROMPT = (
    "You are a knowledgeable assistant for the CSA-in-a-Box data platform. "
    "Answer questions using ONLY the provided context. "
    "If the context does not contain the answer, say so. "
    "Cite your sources by referencing the document title."
)


def rag_chat(question: str, history: list[dict] | None = None) -> str:
    # 1. Retrieve relevant context
    results = hybrid_search(question, top_k=5)
    context = "\n\n---\n\n".join(
        f"**{r['title']}** (source: {r['source']})\n{r['content']}"
        for r in results
    )

    # 2. Build prompt with context
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    messages.extend(history or [])
    messages.append({
        "role": "user",
        "content": f"Context:\n{context}\n\nQuestion: {question}",
    })

    # 3. Generate answer
    resp = oai.chat.completions.create(
        model="gpt-54", messages=messages, temperature=0.2, max_tokens=1024
    )
    return resp.choices[0].message.content


if __name__ == "__main__":
    print("RAG Chatbot (type 'quit' to exit)")
    print("-" * 50)
    history = []
    while True:
        q = input("\nYou: ").strip()
        if q.lower() in ("quit", "exit"):
            break
        reply = rag_chat(q, history)
        print(f"\nAssistant: {reply}")
        history += [{"role": "user", "content": q}, {"role": "assistant", "content": reply}]

from azure.cosmos import CosmosClient
import uuid
from datetime import datetime

cosmos = CosmosClient(os.environ["COSMOS_ENDPOINT"], os.environ["COSMOS_KEY"])
db = cosmos.get_database_client("chatbot")
container = db.get_container_client("conversations")

def save_conversation(user_id: str, messages: list[dict]):
    container.upsert_item({
        "id": str(uuid.uuid4()),
        "user_id": user_id,
        "messages": messages,
        "timestamp": datetime.utcnow().isoformat(),
    })

def load_history(user_id: str, limit: int = 10) -> list[dict]:
    query = f"SELECT TOP {limit} * FROM c WHERE c.user_id = @uid ORDER BY c.timestamp DESC"
    items = list(container.query_items(query, parameters=[{"name": "@uid", "value": user_id}]))
    if items:
        return items[0].get("messages", [])
    return []