""" rag_with_ai_search.py — Full RAG Pipeline with Azure AI Search =============================================================== End-to-end Retrieval-Augmented Generation: 1. Create (or update) an Azure AI Search index 2. Generate embeddings for sample documents 3. Upload documents + embeddings to the index 4. Perform a hybrid search (keyword + vector) 5. Feed results to Azure OpenAI for grounded generation Prerequisites: pip install azure-search-documents openai azure-identity python-dotenv rich Usage: python rag_with_ai_search.py """ import json import os import sys from dotenv import load_dotenv from azure.identity import DefaultAzureCredential, get_bearer_token_provider from azure.search.documents import SearchClient from azure.search.documents.indexes import SearchIndexClient from azure.search.documents.indexes.models import ( SearchIndex, SimpleField, SearchableField, SearchField, SearchFieldDataType, VectorSearch, HnswAlgorithmConfiguration, VectorSearchProfile, SemanticConfiguration, SemanticSearch, SemanticPrioritizedFields, SemanticField, ) from azure.search.documents.models import VectorizedQuery from openai import AzureOpenAI from rich.console import Console load_dotenv() console = Console() # --------------------------------------------------------------------------- # Sample documents to index # --------------------------------------------------------------------------- DOCUMENTS = [ { "id": "1", "title": "Azure AI Foundry Overview", "content": ( "Azure AI Foundry is a unified platform for building, evaluating, and " "deploying AI solutions. It provides a hub-and-project model that lets " "teams share resources while maintaining isolation." ), "category": "platform", }, { "id": "2", "title": "RAG Pattern with Azure AI Search", "content": ( "Retrieval-Augmented Generation (RAG) combines information retrieval with " "large language models. Azure AI Search provides vector, keyword, and hybrid " "search capabilities that make it ideal for RAG workloads." ), "category": "patterns", }, { "id": "3", "title": "Prompt Flow in AI Foundry", "content": ( "Prompt Flow is a development tool inside Azure AI Foundry that lets you " "build LLM-powered workflows visually. You can chain prompts, tools, and " "code together, then evaluate and deploy the resulting flow." ), "category": "tools", }, { "id": "4", "title": "AI Evaluation Metrics", "content": ( "Azure AI Foundry includes built-in evaluation metrics such as groundedness, " "relevance, coherence, fluency, and similarity. These metrics help you " "measure the quality of generative AI outputs systematically." ), "category": "evaluation", }, ] def get_env(key: str) -> str: value = os.getenv(key) if not value: console.print(f"[red]ERROR:[/red] '{key}' not set. See .env.example.") sys.exit(1) return value def get_openai_client() -> AzureOpenAI: endpoint = get_env("AZURE_OPENAI_ENDPOINT") api_version = get_env("AZURE_OPENAI_API_VERSION") api_key = os.getenv("AZURE_OPENAI_API_KEY") if api_key: return AzureOpenAI(azure_endpoint=endpoint, api_key=api_key, api_version=api_version) token_provider = get_bearer_token_provider( DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default" ) return AzureOpenAI(azure_endpoint=endpoint, azure_ad_token_provider=token_provider, api_version=api_version) # --------------------------------------------------------------------------- # Step 1 — Create the search index # --------------------------------------------------------------------------- def create_index(index_client: SearchIndexClient, index_name: str) -> None: """Create or update the search index with vector and semantic configuration.""" console.print("\n[bold]Step 1:[/bold] Creating search index...", style="cyan") fields = [ SimpleField(name="id", type=SearchFieldDataType.String, key=True, filterable=True), SearchableField(name="title", type=SearchFieldDataType.String), 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, vector_search_dimensions=1536, 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")], ), ) semantic_search = SemanticSearch(configurations=[semantic_config]) index = SearchIndex( name=index_name, fields=fields, vector_search=vector_search, semantic_search=semantic_search, ) index_client.create_or_update_index(index) console.print(f" ✅ Index '{index_name}' ready.") # --------------------------------------------------------------------------- # Step 2 — Generate embeddings # --------------------------------------------------------------------------- def generate_embeddings(client: AzureOpenAI, texts: list[str], deployment: str) -> list[list[float]]: """Generate embeddings for a batch of texts.""" console.print("\n[bold]Step 2:[/bold] Generating embeddings...", style="cyan") response = client.embeddings.create(input=texts, model=deployment) embeddings = [item.embedding for item in response.data] console.print(f" ✅ Generated {len(embeddings)} embeddings (dim={len(embeddings[0])}).") return embeddings # --------------------------------------------------------------------------- # Step 3 — Upload documents # --------------------------------------------------------------------------- def upload_documents(search_client: SearchClient, docs: list[dict], embeddings: list[list[float]]) -> None: """Merge documents with their embedding vectors and upload to the index.""" console.print("\n[bold]Step 3:[/bold] Uploading documents...", style="cyan") for doc, emb in zip(docs, embeddings): doc["content_vector"] = emb result = search_client.upload_documents(documents=docs) succeeded = sum(1 for r in result if r.succeeded) console.print(f" ✅ Uploaded {succeeded}/{len(docs)} documents.") # --------------------------------------------------------------------------- # Step 4 — Hybrid search # --------------------------------------------------------------------------- def hybrid_search( search_client: SearchClient, openai_client: AzureOpenAI, query: str, embedding_deployment: str, top_k: int = 3, ) -> list[dict]: """Run a hybrid (keyword + vector) search and return the top results.""" console.print(f"\n[bold]Step 4:[/bold] Searching for: '{query}'", style="cyan") # Embed the query query_embedding = openai_client.embeddings.create( input=[query], model=embedding_deployment ).data[0].embedding vector_query = VectorizedQuery( vector=query_embedding, k_nearest_neighbors=top_k, fields="content_vector", ) results = search_client.search( search_text=query, vector_queries=[vector_query], select=["id", "title", "content", "category"], top=top_k, ) hits = [] for r in results: hits.append({"title": r["title"], "content": r["content"], "score": r["@search.score"]}) console.print(f" 📄 {r['title']} (score: {r['@search.score']:.4f})") return hits # --------------------------------------------------------------------------- # Step 5 — Generate a grounded answer (RAG) # --------------------------------------------------------------------------- def generate_answer(client: AzureOpenAI, chat_deployment: str, query: str, context_docs: list[dict]) -> str: """Send the search results as context to the LLM for a grounded answer.""" console.print("\n[bold]Step 5:[/bold] Generating grounded answer...", style="cyan") context = "\n\n".join( f"### {doc['title']}\n{doc['content']}" for doc in context_docs ) response = client.chat.completions.create( model=chat_deployment, messages=[ { "role": "system", "content": ( "You are an AI assistant. Answer the user's question using ONLY " "the provided context. If the context doesn't contain the answer, " "say so. Cite the document titles in your answer." ), }, { "role": "user", "content": f"Context:\n{context}\n\nQuestion: {query}", }, ], temperature=0.3, max_tokens=512, ) return response.choices[0].message.content # --------------------------------------------------------------------------- # Main # --------------------------------------------------------------------------- def main() -> None: # Config search_endpoint = get_env("AZURE_SEARCH_ENDPOINT") index_name = os.getenv("AZURE_SEARCH_INDEX_NAME", "rag-demo-index") embedding_deployment = get_env("AZURE_OPENAI_EMBEDDING_DEPLOYMENT") chat_deployment = get_env("AZURE_OPENAI_CHAT_DEPLOYMENT") credential = DefaultAzureCredential() openai_client = get_openai_client() # Clients index_client = SearchIndexClient(endpoint=search_endpoint, credential=credential) search_client = SearchClient(endpoint=search_endpoint, index_name=index_name, credential=credential) # Pipeline create_index(index_client, index_name) texts = [doc["content"] for doc in DOCUMENTS] embeddings = generate_embeddings(openai_client, texts, embedding_deployment) upload_documents(search_client, DOCUMENTS, embeddings) query = "What is RAG and how does Azure AI Search support it?" hits = hybrid_search(search_client, openai_client, query, embedding_deployment) answer = generate_answer(openai_client, chat_deployment, query, hits) console.print("\n" + "=" * 60) console.print("[bold green]🤖 Grounded Answer:[/bold green]") console.print(answer) console.print("=" * 60) if __name__ == "__main__": main()