Grounded Search Agents: Microsoft Foundry + Azure AI Search + Cosmos DB + Function Calling
DEV Community

Grounded Search Agents: Microsoft Foundry + Azure AI Search + Cosmos DB + Function Calling

Grounded Search Agents: Microsoft Foundry + Azure AI Search + Cosmos DB + Function Calling

Why "search" isn't enough anymore

A plain RAG pipeline - embed a query, hit a vector index, stuff the top-k chunks into a prompt - works fine for a single-turn FAQ bot. It falls apart the moment you need:

  • Multi-turn memory ("what about last quarter?" needs to know what "last quarter" refers to)
  • A decision about whether to search at all (some questions are structured lookups, not document retrieval)
  • Citations the caller can verify

This post builds a small but complete grounded search agent: Azure OpenAI (via Microsoft Foundry) decides when to call a search tool versus a structured data tool, Azure AI Search handles retrieval and semantic ranking over unstructured documents, and Azure Cosmos DB stores conversation state and structured order/customer metadata. The agent uses OpenAI function calling to route between the two data sources and always returns citations. We'll build this from the ground up in Python.

Architecture overview

The system has two phases: ingestion (getting documents into a searchable index) and runtime (the agent loop that answers questions). At runtime, the agent sits between the user and two tools, using function calling to decide which one (or both) to invoke.

Step 1 - Provision and index documents in Azure AI Search

Create the index with both a vector field (for semantic similarity) and a semantic configuration (for ranking), since combining the two consistently outperforms either alone.

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

endpoint = "https://<your-search-service>.search.windows.net"
credential = AzureKeyCredential("<admin-key>")
index_client = SearchIndexClient(endpoint, credential)

fields = [
    SimpleField(name="id", type=SearchFieldDataType.String, key=True),
    SearchableField(name="content", type=SearchFieldDataType.String),
    SimpleField(name="source_url", type=SearchFieldDataType.String),
    SearchField(
        name="content_vector",
        type=SearchFieldDataType.Collection(SearchFieldDataType.Single),
        searchable=True,
        vector_search_dimensions=1536,
        vector_search_profile_name="default-profile",
    ),
]

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

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

index = SearchIndex(
    name="support-docs-index",
    fields=fields,
    vector_search=vector_search,
    semantic_search=SemanticSearch(configurations=[semantic_config]),
)

index_client.create_or_update_index(index)
print("Index created.")

Now embed and upload chunks. Chunking matters more than most tutorials admit - 512-token chunks with ~15% overlap is a reasonable default for policy/support documents.

# ingest.py
from openai import AzureOpenAI
from azure.search.documents import SearchClient
from azure.core.credentials import AzureKeyCredential
import uuid

aoai = AzureOpenAI(
    azure_endpoint="https://<your-foundry-resource>.openai.azure.com",
    api_key="<key>",
    api_version="2024-10-21",
)

search_client = SearchClient(
    endpoint=endpoint,
    index_name="support-docs-index",
    credential=AzureKeyCredential("<admin-key>"),
)

def chunk_text(text: str, max_tokens: int = 512, overlap: int = 80):
    words = text.split()
    step = max_tokens - overlap
    return [" ".join(words[i:i + max_tokens]) for i in range(0, len(words), step)]

def embed(text: str) -> list[float]:
    resp = aoai.embeddings.create(model="text-embedding-3-small", input=text)
    return resp.data[0].embedding

def ingest_document(text: str, source_url: str):
    docs = []
    for chunk in chunk_text(text):
        docs.append({
            "id": str(uuid.uuid4()),
            "content": chunk,
            "source_url": source_url,
            "content_vector": embed(chunk),
        })
    search_client.upload_documents(documents=docs)

with open("return_policy.txt") as f:
    ingest_document(f.read(), source_url="policies/return_policy.txt")

Step 2 - Structured lookups in Cosmos DB

Not every question is a document-retrieval problem. "Where's my order?" is a point lookup. Store operational data in Cosmos DB, partitioned by customer_id for predictable query performance.

# cosmos_setup.py
from azure.cosmos import CosmosClient, PartitionKey

client = CosmosClient("<cosmos-endpoint>", credential="<cosmos-key>")
db = client.create_database_if_not_exists("retail")
orders = db.create_container_if_not_exists(
    id="orders",
    partition_key=PartitionKey(path="/customer_id"),
    offer_throughput=400,
)

orders.upsert_item({
    "id": "4471",
    "customer_id": "cust_9001",
    "status": "delayed",
    "delay_reason": "carrier weather disruption",
    "eta": "2026-08-02",
})

Step 3 - The agent loop with function calling

This is the core piece: Azure OpenAI decides, per turn, whether to call search_docs, get_order, both, or neither. We define the tools as JSON schemas and let the model pick.

# agent.py
import json
from openai import AzureOpenAI
from azure.search.documents import SearchClient
from azure.cosmos import CosmosClient

aoai = AzureOpenAI(
    azure_endpoint="https://<your-foundry-resource>.openai.azure.com",
    api_key="<key>",
    api_version="2024-10-21",
)

search_client = SearchClient(endpoint, "support-docs-index", AzureKeyCredential("<admin-key>"))
cosmos = CosmosClient("<cosmos-endpoint>", credential="<cosmos-key>")
orders_container = cosmos.get_database_client("retail").get_container_client("orders")

tools = [
    {
        "type": "function",
        "function": {
            "name": "search_docs",
            "description": "Semantic search over policy and support documents. Returns passages with citations.",
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {"type": "string"}
                },
                "required": ["query"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_order",
            "description": "Look up a customer order by ID for status, ETA, and delay reason.",
            "parameters": {
                "type": "object",
                "properties": {
                    "order_id": {"type": "string"},
                    "customer_id": {"type": "string"},
                },
                "required": ["order_id", "customer_id"],
            },
        },
    },
]

def search_docs(query: str):
    results = search_client.search(
        search_text=query,
        query_type="semantic",
        semantic_configuration_name="default-semantic",
        top=3,
    )
    return [{"content": r["content"], "source": r["source_url"]} for r in results]

def get_order(order_id: str, customer_id: str):
    item = orders_container.read_item(item=order_id, partition_key=customer_id)
    return {"status": item["status"], "reason": item.get("delay_reason"), "eta": item.get("eta")}

def run_agent(user_message: str, customer_id: str, history: list):
    messages = history + [{"role": "user", "content": user_message}]
    response = aoai.chat.completions.create(
        model="gpt-4o",
        messages=messages,
        tools=tools,
        tool_choice="auto",
    )
    msg = response.choices[0].message
    if msg.tool_calls:
        messages.append(msg)
        for call in msg.tool_calls:
            args = json.loads(call.function.arguments)
            if call.function.name == "search_docs":
                result = search_docs(**args)
            elif call.function.name == "get_order":
                args["customer_id"] = customer_id  # enforce, don't trust model-provided id
                result = get_order(**args)
            messages.append({
                "role": "tool",
                "tool_call_id": call.id,
                "content": json.dumps(result),
            })
        final = aoai.chat.completions.create(model="gpt-4o", messages=messages)
        return final.choices[0].message.content
    return msg.content

A few things worth calling out:

  • Never trust model-supplied identity fields. customer_id is overwritten server-side before hitting Cosmos DB, even though the model could technically supply one - this prevents cross-tenant data leaks via prompt injection.
  • tool_choice="auto" lets the model skip tools entirely for small talk, avoiding unnecessary search calls.
  • The second chat.completions.create call (after tool results are appended) is what produces the grounded, cited final answer.

Step 4 - Session state in Cosmos DB

Conversation history should not live in application memory - store it in Cosmos DB so any instance of the service can resume a session.

# session_store.py
def load_history(session_id: str) -> list:
    container = cosmos.get_database_client("retail").get_container_client("sessions")
    try:
        item = container.read_item(item=session_id, partition_key=session_id)
        return item["messages"]
    except Exception:
        return []

def save_history(session_id: str, messages: list):
    container = cosmos.get_database_client("retail").get_container_client("sessions")
    container.upsert_item({"id": session_id, "messages": messages})

Comparing retrieval strategies

Choosing between search modes is a real architectural decision, not just a config toggle:

Retrieval Mode Best for Weakness
Keyword (BM25) Exact terms, SKUs, error codes Misses paraphrased queries
Pure vector search Conceptual / paraphrased questions Can surface plausible-but-wrong matches
Hybrid (keyword + vector) General-purpose document Q&A Slightly higher latency than either alone
Semantic ranking on top of hybrid Conversational, nuanced queries Adds a reranking pass; not needed for simple lookups
Structured query (Cosmos DB point read) Known-ID lookups (orders, accounts) Not applicable to unstructured content

For most support/document agents, hybrid retrieval with semantic ranking is the right default, reserving structured lookups for anything with a stable ID.

Production considerations

  • Grounding failures are silent failures. If retrieval returns weak matches, the model can still generate a fluent but wrong answer - always surface citations so users (and evaluators) can catch this.
  • Access control matters. If your index mixes content from different permission levels, filter search results by the caller's identity before they reach the model, not after.
  • Rising input tokens. Retrieved passages add to the prompt on every call - budget for this in both cost and latency.

References

Comments

No comments yet. Start the discussion.