DEV Community

Cover image for Grounded Search Agents: Microsoft Foundry + Azure AI Search + Cosmos DB + Function Calling
Jubin Soni
Jubin Soni Subscriber

Posted on

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).

system description

At runtime, the agent sits between the user and two tools, using function calling to decide which one (or both) to invoke:

agent description

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.")
Enter fullscreen mode Exit fullscreen mode

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")
Enter fullscreen mode Exit fullscreen mode

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",
})
Enter fullscreen mode Exit fullscreen mode

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
Enter fullscreen mode Exit fullscreen mode

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})
Enter fullscreen mode Exit fullscreen mode

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

  1. Retrieval augmented generation (RAG) and indexes in Microsoft Foundry — Microsoft Learn
  2. RAG and Generative AI overview for Azure AI Search — Microsoft Learn
  3. RAG with Agent Framework and Semantic Kernel — Microsoft Learn
  4. RAG using Semantic Kernel with Azure OpenAI and Azure Cosmos DB for MongoDB vCore — Microsoft Learn Shows
  5. Step-by-Step Guide on Building Agentic RAG with Semantic Kernel and Azure AI Search — Data Science Collective

Top comments (0)