Building a Grounded Search Agent With Foundry, Azure AI Search, and Cosmos DB, Containerized With Docker and Observed With Grafana
Build a grounded AI search agent with Azure OpenAI, AI Search, and Cosmos DB that combines RAG, structured lookups, memory, and citations.
CORE
·
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. We’ll then package the agent in a Docker container and instrument it with Prometheus metrics and Grafana dashboards, so you can see what it’s doing in production.
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:
The runtime is a stateless service packaged as a Docker image, which is why session state lives in Cosmos DB, and it exposes metrics that Grafana visualizes.
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.contentA few things worth calling out:
- Never trust model-supplied identity fields.
customer_idis 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.createcall (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, including multiple Docker containers behind a load balancer, 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.
- Observability in Grafana. Track tool-call rates, retrieval latency, and token usage in Grafana, and alert on spikes.
References
- Retrieval-augmented generation (RAG) and indexes in Microsoft Foundry — Microsoft Learn
- RAG and Generative AI overview for Azure AI Search — Microsoft Learn
- RAG with Agent Framework and Semantic Kernel — Microsoft Learn
- RAG using Semantic Kernel with Azure OpenAI and Azure Cosmos DB for MongoDB vCore — Microsoft Learn Shows
- Step-by-Step Guide on Building Agentic RAG with Semantic Kernel and Azure AI Search — Data Science Collective
Opinions expressed by DZone contributors are their own.