Building a Knowledge Graph on Cosmos DB for AI

Knowledge Graphs Are Having a Moment

Retrieval-Augmented Generation (RAG) has become the standard pattern for grounding large language models in real data. The typical approach is to chunk documents, embed them into vectors, and retrieve relevant passages at query time. It works — but it has a ceiling. Vector similarity alone cannot capture the structured relationships between entities that make domain knowledge truly useful.

That is where knowledge graphs come in. By representing domain knowledge as entities and typed relationships, a knowledge graph gives an LLM something that flat vector stores cannot: context with structure. And when you build that graph on Azure Cosmos DB, you get global distribution, elastic scale, and single-digit millisecond reads — everything a production AI application demands.

What Is a Knowledge Graph?

A knowledge graph is a network of entities (vertices) connected by typed relationships (edges), with properties on both. A few examples make the concept concrete:

• Alice -[WORKS_AT]-> Contoso
• Q3 Report -[MENTIONS]-> Revenue Growth
• Product A -[RELATED_TO]-> Product B
• Dr. Chen -[AUTHORED]-> Scalability Whitepaper

Each vertex has a label (Person, Document, Concept) and properties (name, created date, confidence score). Each edge carries a label that describes the relationship and can hold its own properties — a weight, a timestamp, a source reference.

This is fundamentally more expressive than a bag of embedded text chunks. When an LLM asks “who are the experts on revenue growth?”, a knowledge graph can traverse from the concept to the documents that mention it, then to the authors of those documents — returning a precise, structured answer rather than a ranked list of vaguely similar paragraphs.

Why Azure Cosmos DB for Knowledge Graphs?

Azure Cosmos DB’s Gremlin API is purpose-built for this kind of work:

• Global distribution — Replicate your knowledge graph across regions so multi-region AI applications get low-latency reads everywhere.
• Single-digit millisecond point reads — RAG retrieval needs to be fast. Cosmos DB delivers consistent sub-10ms reads at any scale.
• Gremlin traversals — The Apache TinkerPop Gremlin language lets you express multi-hop relationship queries that would be painful in SQL or impossible in a vector store.
• Elastic throughput — Scale from a development workload to enterprise traffic by adjusting Request Units. No re-architecture required.
• Azure ecosystem integration — Connect directly to Azure OpenAI, Azure AI Search, and Azure Functions to build end-to-end RAG pipelines.

Designing Your Graph Schema

Before writing any Gremlin, plan your schema. A well-designed knowledge graph makes retrieval straightforward.

Vertex labels represent entity types. For a corporate knowledge graph, you might use:

Label	Purpose	Key Properties
Person	Employees, authors, experts	name, role, department
Document	Reports, articles, wikis	title, publishedDate, source
Concept	Topics, themes, keywords	name, description, category
Product	Products, services, offerings	name, version, status
Team	Organizational units	name, division

Edge labels represent relationship types: AUTHORED, MENTIONS, WORKS_ON, RELATED_TO, BELONGS_TO, EXPERT_IN.

Partition key strategy is critical in Cosmos DB. Use the entity type or a domain identifier as your partition key. A property like pk set to the vertex label keeps related entities co-located for efficient queries.

Building the Graph with Gremlin

With your schema in hand, start populating the graph. Here are the fundamental patterns.

Adding Entities

g.addV('document')
  .property('pk', 'document')
  .property('title', 'Q3 Revenue Report')
  .property('publishedDate', '2026-01-15')
  .property('source', 'finance-team')

g.addV('concept')
  .property('pk', 'concept')
  .property('name', 'Revenue Growth')
  .property('category', 'finance')

g.addV('person')
  .property('pk', 'person')
  .property('name', 'Dr. Sarah Chen')
  .property('role', 'Senior Analyst')

Adding Relationships

g.V().has('document', 'title', 'Q3 Revenue Report')
  .addE('mentions')
  .property('confidence', 0.95)
  .to(g.V().has('concept', 'name', 'Revenue Growth'))

g.V().has('person', 'name', 'Dr. Sarah Chen')
  .addE('authored')
  .to(g.V().has('document', 'title', 'Q3 Revenue Report'))

Bulk Loading Tips

When loading data at scale, keep each batch under 1000 RUs to avoid throttling. Group your addV and addE queries into batches of 50-100 operations. Load vertices first, then edges — edges reference existing vertices by their properties.

Querying for RAG Context

This is where the knowledge graph earns its keep. Instead of retrieving loosely related text chunks, you retrieve precise, structured context to inject into LLM prompts.

Find Everything Related to a Concept

g.V().has('concept', 'name', 'Revenue Growth')
  .bothE().otherV()
  .path()

This traversal returns the concept, every edge connected to it, and every vertex on the other end — documents that mention it, people who are experts, related products.

Find Documents Related to a Topic (Multi-Hop)

g.V().has('concept', 'name', 'Revenue Growth')
  .in('related_to')
  .in('mentions')
  .hasLabel('document')
  .dedup()
  .valueMap('title', 'publishedDate')

This walks from a concept to related concepts, then finds all documents that mention any of them — casting a wider net than a single-hop query.

Find the Experts on a Topic

g.V().has('concept', 'name', 'Revenue Growth')
  .in('mentions')
  .hasLabel('document')
  .in('authored')
  .hasLabel('person')
  .dedup()
  .valueMap('name', 'role')

Traverse from concept to documents to authors. The result is a list of people who have written about the topic — exactly the kind of structured answer that makes RAG responses more grounded and trustworthy.

Injecting Graph Context into LLM Prompts

The pattern is straightforward. Run your Gremlin traversal, format the results as structured text, and prepend it to your LLM prompt:

System: You are an analyst assistant. Use the following knowledge graph context to answer the user’s question.

Context: The concept “Revenue Growth” is mentioned in: Q3 Revenue Report (authored by Dr. Sarah Chen), Annual Strategy Deck (authored by James Wu). Related concepts: Market Expansion, Customer Retention.

User: Who should I talk to about revenue growth trends?

The LLM now has structured, relationship-aware context instead of a wall of text fragments.

Using GremlinStudio to Build and Explore Your Knowledge Graph

Building a knowledge graph is iterative. You need to see what you have, find gaps, and refine your schema. GremlinStudio is designed for exactly this workflow:

• Visual graph exploration — Run a traversal and see the result as an interactive force-directed graph. Spot disconnected nodes, find unexpected clusters, and understand the shape of your data at a glance.
• WYSIWYG CRUD — Right-click the graph canvas to add vertices and edges directly. Edit properties inline. No need to write addV queries for one-off corrections.
• AI-powered natural language queries — Describe what you want to find in plain English. GremlinStudio translates it to Gremlin and executes it. Useful when you are exploring an unfamiliar graph.
• Schema discovery — Auto-detect all vertex labels, edge labels, and properties in your graph. Understand what is in your knowledge graph before writing a single query.
• Step-by-step debugger — Complex multi-hop RAG traversals can be hard to reason about. The debugger lets you execute one Gremlin step at a time, watching traversers move through the graph and inspecting intermediate results.

Start Building Your Knowledge Graph

The combination of Azure Cosmos DB’s scale and Gremlin’s expressive traversals makes it an excellent foundation for knowledge graphs that power AI applications. Whether you are building a RAG pipeline, a recommendation engine, or a domain-specific search system, a well-designed knowledge graph gives your LLM the structured context it needs to deliver accurate, grounded answers.

Download GremlinStudio to start building and exploring your knowledge graph today. If you are new to Cosmos DB, the Connection Setup guide will get you connected in minutes.