🍏 Basic Retrieval-Augmented Generation (RAG) with AIProjectClient 🍎¶

import os
import time
import json
from dotenv import load_dotenv

# azure-ai-projects
from azure.ai.projects import AIProjectClient
from azure.identity import DefaultAzureCredential

# We'll embed with azure-ai-inference
from azure.ai.inference import EmbeddingsClient, ChatCompletionsClient
from azure.ai.inference.models import UserMessage, SystemMessage

# For vector search or hybrid search
from azure.search.documents import SearchClient
from azure.search.documents.indexes import SearchIndexClient
from azure.core.credentials import AzureKeyCredential

load_dotenv()

conn_string = os.environ.get("PROJECT_CONNECTION_STRING")
chat_model = os.environ.get("MODEL_DEPLOYMENT_NAME", "gpt-4o-mini")
search_index_name = os.environ.get("SEARCH_INDEX_NAME", "healthtips-index")

try:
    project_client = AIProjectClient.from_connection_string(
        credential=DefaultAzureCredential(),
        conn_str=conn_string,
    )
    print("✅ AIProjectClient created successfully!")
except Exception as e:
    print("❌ Error creating AIProjectClient:", e)

health_tips = [
    {
        "id": "doc1",
        "content": "Daily 30-minute walks help maintain a healthy weight and reduce stress.",
        "source": "General Fitness"
    },
    {
        "id": "doc2",
        "content": "Stay hydrated by drinking 8-10 cups of water per day.",
        "source": "General Fitness"
    },
    {
        "id": "doc3",
        "content": "Consistent sleep patterns (7-9 hours) improve muscle recovery.",
        "source": "General Fitness"
    },
    {
        "id": "doc4",
        "content": "For cardio endurance, try interval training like HIIT.",
        "source": "Workout Advice"
    },
    {
        "id": "doc5",
        "content": "Warm up with dynamic stretches before running to reduce injury risk.",
        "source": "Workout Advice"
    },
    {
        "id": "doc6",
        "content": "Balanced diets typically include protein, whole grains, fruits, vegetables, and healthy fats.",
        "source": "Nutrition"
    },
]
print("Created a small list of health tips.")

from azure.ai.projects.models import ConnectionType

# Get the default search connection, with credentials
search_conn = project_client.connections.get_default(
    connection_type=ConnectionType.AZURE_AI_SEARCH, include_credentials=True
)
if not search_conn:
    raise RuntimeError("❌ No default Azure AI Search connection found!")

search_client = SearchClient(
    endpoint=search_conn.endpoint_url,
    index_name=search_index_name,
    credential=AzureKeyCredential(search_conn.key)
)
print("✅ Created Azure SearchClient.")

# Now create embeddings client
embeddings_client = project_client.inference.get_embeddings_client()
print("✅ Created embeddings client.")

search_docs = []
for doc in health_tips:
    # embed doc content
    emb_response = embeddings_client.embed(
        input=[doc["content"]]
    )
    emb_vec = emb_response.data[0].embedding

    # We'll build a doc with 'id', 'content', 'source', 'embedding'
    search_docs.append(
        {
            "id": doc["id"],
            "content": doc["content"],
            "source": doc["source"],
            "embedding": emb_vec,
        }
    )

result = search_client.upload_documents(documents=search_docs)
print(f"Uploaded {len(search_docs)} docs to Search index '{search_index_name}'.")

def rag_chat(query: str, top_k: int = 3) -> str:
    # 1) Embed user query
    user_vec = embeddings_client.embed(input=[query]).data[0].embedding

    # 2) Vector search
    # We'll assume your index has the vector field named 'embedding',
    # using vector search or hybrid. We'll do a minimal example.

    results = search_client.search(
        search_text="",  # for vector search we don't rely on textual search
        vector=user_vec,
        vector_fields="embedding",
        top=top_k
    )

    # gather the top docs
    top_docs_content = []
    for r in results:
        # Each doc is a SearchResult object with doc: dict.
        c = r["content"]
        s = r["source"]
        top_docs_content.append(f"Source: {s} => {c}")

    # 3) Chat with retrieved docs.
    # We'll pass a system message instructing the model to use them.
    # We'll just do a minimal approach.

    system_text = (
        "You are a health & fitness assistant.\n"
        "Answer user questions using ONLY the text from these docs.\n"
        "Docs:\n"
        + "\n".join(top_docs_content)
        + "\nIf unsure, say 'I'm not sure'.\n"
    )

    with project_client.inference.get_chat_completions_client() as chat_client:
        response = chat_client.complete(
            model=chat_model,
            messages=[
                SystemMessage(content=system_text),
                UserMessage(content=query)
            ]
        )
    return response.choices[0].message.content

user_query = "What's a good short cardio routine for me if I'm busy?"
answer = rag_chat(user_query)
print("🗣️ User Query:", user_query)
print("🤖 RAG Answer:", answer)