Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions experimental/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ See the root [README](../README.md) for details on the stable install path.
### 🤖 AI & Agents
- **databricks-ai-functions** - Built-in AI Functions (ai_classify, ai_extract, ai_summarize, ai_query, ai_forecast, ai_parse_document, and more) with SQL and PySpark patterns, function selection guidance, document processing pipelines, and custom RAG (parse → chunk → index → query)
- **databricks-agent-bricks** - Knowledge Assistants, Genie Spaces, Supervisor Agents
- **databricks-graphrag** - GraphRAG and knowledge-graph augmented retrieval architecture: compare GraphRAG vs traditional RAG, combine AI Search with graph traversal, and evaluate relationship-aware retrieval
- **databricks-mlflow-evaluation** - End-to-end agent evaluation workflow
- **databricks-unstructured-pdf-generation** - Generate synthetic PDFs for RAG

Expand Down
89 changes: 89 additions & 0 deletions experimental/databricks-graphrag/SKILL.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,89 @@
---
name: databricks-graphrag
description: "Design GraphRAG and knowledge-graph augmented retrieval on Databricks. Use when comparing GraphRAG vs traditional RAG, combining AI Search with graph traversal, designing entity/relationship models, evaluating hybrid retrieval quality, or integrating Neo4j/Text2Cypher patterns with Unity Catalog, Databricks Apps, Model Serving, and MLflow."
compatibility: Requires databricks CLI (>= v1.0.0)
metadata:
version: "0.1.0"
parent: databricks-core
---

# GraphRAG on Databricks

GraphRAG augments retrieval with explicit entities and relationships. Use it when
answers require multi-hop reasoning, exact relationship paths, explainable
provenance, or structured graph constraints that plain chunk retrieval misses.

This skill is an architecture and evaluation guide. Databricks provides the
governed platform spine — Unity Catalog, AI Search, Databricks Apps, Model
Serving, and MLflow evaluation — while graph stores and Text2Cypher libraries are
integration choices, not Databricks-native APIs.

## Critical Rules

- **Start with AI Search hybrid retrieval first.** Establish a measured baseline
with keyword + vector retrieval, filters, and reranking before adding graph
complexity.
- **Use GraphRAG only when relationships matter.** If the query only needs
semantically similar passages, use `databricks-vector-search` / AI Search.
- **Keep graph queries read-only.** Text2Cypher and generated graph traversals
must run against allow-listed labels, relationship types, and result limits.
- **Evaluate before promoting.** Compare ANN, hybrid, reranker, graph traversal,
and Text2Cypher with MLflow metrics and groundedness judges.
- **Govern inputs and outputs in Unity Catalog.** Source Delta tables, derived
entity/relationship tables, AI Search indexes, and model/agent registrations
should have UC ownership, permissions, lineage, and auditability.

## When to Use

| User asks for | Use this skill? | Why |
|---|---:|---|
| "Should I use GraphRAG or normal RAG?" | Yes | Architecture decisioning and trade-offs |
| "Find related policies, owners, systems, and exceptions" | Yes | Relationship traversal is part of the answer |
| "Build document Q&A over PDFs" | Usually no | Start with Knowledge Assistant or AI Search RAG |
| "Improve retrieval precision for SKUs/error codes" | Maybe | Try AI Search hybrid + filters + reranker first |
| "Generate Cypher from questions" | Yes | Requires Text2Cypher safety guidance |

## Workflow

1. **Define answer shape.** Identify whether the answer needs passages,
entities, paths, summaries, or calculations.
2. **Build the baseline.** Use AI Search hybrid retrieval with metadata filters
and optional reranker; measure recall and groundedness.
3. **Model the graph.** If the baseline misses relationships, design entity and
relationship tables in Unity Catalog before loading an external graph store.
4. **Choose retrieval pattern.** Pick vector-to-graph expansion, full-text to
graph, Text2Cypher, or a router that chooses among them.
5. **Assemble the app/agent.** Prefer Databricks Apps for custom UX/API agents;
use Model Serving when deploying custom models, embeddings, or agent endpoints.
6. **Evaluate and operate.** Use MLflow traces, retrieval metrics, groundedness,
sufficiency, and latency budgets before recommending GraphRAG over simpler RAG.

## Reference Files

| File | Purpose |
|---|---|
| [references/architecture-decisions.md](references/architecture-decisions.md) | Decide when GraphRAG is worth the added complexity |
| [references/retrieval-patterns.md](references/retrieval-patterns.md) | Retrieval patterns for AI Search, Neo4j, Text2Cypher, and routers |
| [references/evaluation-and-operations.md](references/evaluation-and-operations.md) | MLflow evaluation, safety checks, observability, and operations |

## Related Skills

- **databricks-vector-search** — AI Search / Vector Search indexes, hybrid
search, filters, reranking, and baseline RAG.
- **databricks-apps** — Databricks Apps for custom GraphRAG UX and APIs.
- **databricks-model-serving** — serving LLMs, embeddings, custom rerankers, and
model/agent endpoints.
- **[databricks-mlflow-evaluation](../databricks-mlflow-evaluation/SKILL.md)** —
MLflow GenAI traces, datasets, scorers, and production monitoring.
- **[databricks-unity-catalog](../databricks-unity-catalog/SKILL.md)** — source
tables, governance, access control, lineage, and system tables.
- **[databricks-agent-bricks](../databricks-agent-bricks/SKILL.md)** — managed
Knowledge Assistant baseline and supervisor-agent routing comparison.

## Official References

- AI Search: https://docs.databricks.com/aws/en/ai-search/ai-search
- AI Search retrieval quality: https://docs.databricks.com/aws/en/ai-search/retrieval-quality
- Databricks Apps resources: https://docs.databricks.com/aws/en/dev-tools/databricks-apps/resources
- Model Serving: https://docs.databricks.com/aws/en/machine-learning/model-serving/
- MLflow GenAI evaluation: https://docs.databricks.com/aws/en/mlflow3/genai/eval-monitor/
7 changes: 7 additions & 0 deletions experimental/databricks-graphrag/agents/openai.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
interface:
display_name: "Databricks Graphrag"
short_description: "Design GraphRAG and knowledge-graph augmented retrieval on Databricks."
icon_small: "./assets/databricks.svg"
icon_large: "./assets/databricks.png"
brand_color: "#FF3621"
default_prompt: "Use $databricks-graphrag for design graphrag and knowledge-graph augmented retrieval on databricks."
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
3 changes: 3 additions & 0 deletions experimental/databricks-graphrag/assets/databricks.svg
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Original file line number Diff line number Diff line change
@@ -0,0 +1,58 @@
# GraphRAG Architecture Decisions

Use GraphRAG when explicit relationships are part of the answer, not because it
sounds more advanced than RAG. The safest Databricks path is to start with a
measured AI Search baseline, then add graph retrieval only when evaluation shows
relationship context improves quality enough to justify the extra operational
cost.

## Decision Tree

1. **Can AI Search answer it with hybrid retrieval?**
- Use AI Search hybrid when questions need both semantic recall and exact
identifiers such as SKUs, policy IDs, entity names, error codes, or product
names.
- Add metadata filters and a reranker before introducing graph traversal.
2. **Does the user need relationship paths?**
- Use GraphRAG when answers depend on graph facts: ownership chains,
dependency paths, customer-product-symptom links, policy exceptions, or
multi-hop relationships.
3. **Is the graph schema stable enough?**
- GraphRAG works best with curated labels, relationship types, primary keys,
and provenance columns. If the schema changes every week, keep the graph as
a derived index built from governed Unity Catalog tables.
4. **Can the question be routed safely?**
- Use a router that chooses among AI Search, graph traversal, Text2Cypher,
SQL/Genie, or a managed Knowledge Assistant.

## Architecture Options

| Pattern | Use when | Databricks spine | Graph component |
|---|---|---|---|
| Databricks-native RAG | Documents and metadata are enough | Unity Catalog -> AI Search -> Databricks Apps or Model Serving -> MLflow | None |
| Vector-to-graph expansion | Similar chunks identify seed entities, then relationships add context | AI Search for seeds; UC for source/provenance; MLflow for eval | Neo4j or graph DB traversal |
| Full-text + graph traversal | Entity names and exact identifiers dominate | UC source tables; App/API orchestration | Full-text index plus graph relationships |
| Text2Cypher | Users ask graph-schema questions directly | Databricks App or serving endpoint with prompt safety | Read-only Cypher generation |
| Router / supervisor | Different question types need different retrievers | Databricks Apps or Agent Bricks supervisor; MLflow traces | Optional retriever tools |

## Databricks Integration Spine

- **Unity Catalog** governs source Delta tables, entity and relationship tables,
AI Search indexes, registered models, permissions, lineage, and auditability.
- **AI Search** is the default retrieval baseline. Use hybrid search, filters,
and reranking before graph traversal.
- **Databricks Apps** are the preferred place to host custom GraphRAG UX, query
routers, and API endpoints because app resources avoid hardcoded credentials.
- **Model Serving** is appropriate when serving custom models, embeddings,
rerankers, or an agent endpoint outside an App.
- **MLflow** records traces, retrieval spans, scorer output, and versioned
comparisons across ANN, hybrid, reranker, graph traversal, and Text2Cypher.

## Anti-Patterns

- Building a graph before defining answer shapes and evaluation metrics.
- Using Text2Cypher on an unconstrained schema or allowing write/delete Cypher.
- Treating Neo4j or a LangChain GraphRAG demo as a Databricks-native API.
- Skipping the AI Search baseline; most retrieval quality issues are fixed with
parsing, chunking, metadata filters, hybrid search, or a reranker.
- Returning graph facts without source documents, entity IDs, or provenance.
Original file line number Diff line number Diff line change
@@ -0,0 +1,76 @@
# GraphRAG Evaluation and Operations

Do not recommend GraphRAG unless it beats a simpler AI Search baseline on
measured quality, latency, cost, or explainability. GraphRAG adds pipelines,
schemas, indexes, graph-store operations, and prompt-safety risk.

## Evaluation Protocol

1. **Create an evaluation set** with representative questions, expected answer
notes, expected entities/relationships, and source documents.
2. **Run retrieval variants**: ANN, AI Search hybrid, hybrid + reranker, graph
traversal, Text2Cypher, and router combinations.
3. **Score retrieval** with Recall@k, Precision@k, MRR, DCG@10, entity/path
recall, and latency p95.
4. **Score answer quality** with MLflow judges for groundedness, context
relevance, context sufficiency, safety, and task-specific correctness.
5. **Inspect traces** for missed entities, bad graph paths, hallucinated Cypher,
stale edges, and over-broad graph expansion.

## MLflow Trace Tags

Record enough metadata to debug retrieval choices:

```python
with mlflow.start_span(name="graphrag_retrieve") as span:
span.set_attribute("retriever", "hybrid_plus_graph")
span.set_attribute("query_type", query_type)
span.set_attribute("candidate_count", len(candidates))
span.set_attribute("graph_hops", max_hops)
span.set_attribute("reranker", reranker_name)
```

Track at least:

- retriever selected by the router
- AI Search index name and query type
- graph store and schema version
- Text2Cypher prompt version and generated Cypher
- candidate counts before/after reranker
- answer citations, entity IDs, and relationship IDs

## Operational Checklist

- **Govern source data in Unity Catalog.** Keep entity and relationship source
tables in UC, with owners and lineage back to raw documents or structured
systems.
- **Version graph schemas.** Changes to labels, relationship names, or primary
keys can break Text2Cypher and traversal prompts.
- **Rebuild derived indexes safely.** Rebuild AI Search indexes, graph indexes,
and reranker datasets after source schema changes.
- **Use Databricks Apps resources for credentials.** Do not hardcode graph store,
AI Search, Model Serving, or SQL credentials in app code.
- **Prefer read-only graph access.** Graph write pipelines should be separate
jobs with tests, not agent runtime actions.
- **Monitor drift.** Track entity-linking quality, orphan nodes, stale edges,
retrieval failure categories, and latency/cost by retriever.

## Failure Modes

| Symptom | Likely cause | Fix |
|---|---|---|
| Correct document found, wrong relationship answer | Graph expansion too broad or stale edge | Add edge provenance, tighter hop limits, and relationship filters |
| Graph answer is correct but uncited | Missing source_doc_id or chunk IDs on graph nodes | Store provenance on every entity and relationship |
| Text2Cypher times out | Unbounded path or missing indexes | Validate query, add limits, index common properties |
| Hybrid beats GraphRAG | Relationships are not needed | Keep AI Search hybrid; do not ship graph complexity |
| Reranker improves quality but misses SLA | Candidate set too large or model too slow | Lower k, cache, or move reranker to Model Serving with clear p95 budget |

## Promotion Bar

Ship GraphRAG only when it has evidence over baseline:

- Better groundedness or context sufficiency in MLflow evaluation.
- Better relationship/path recall on graph-specific questions.
- Acceptable latency and cost at p95.
- Clear provenance from answer to documents, entities, and relationships.
- Safe Text2Cypher validation or no generated graph query path at all.
Original file line number Diff line number Diff line change
@@ -0,0 +1,91 @@
# GraphRAG Retrieval Patterns

GraphRAG retrieval should be explicit about the entrypoint, expansion strategy,
ranking, and provenance. Keep each retriever small and measurable; combine them
through a router rather than one large prompt.

## Pattern 1: AI Search Baseline

Start with AI Search hybrid retrieval. It combines semantic similarity with
keyword matching and is the default baseline for Databricks RAG.

```python
results = w.vector_search_indexes.query_index(
index_name="catalog.schema.support_docs_index",
columns=["doc_id", "chunk_id", "text", "product", "updated_at"],
query_text=user_query,
query_type="HYBRID",
num_results=20,
filters_json='{"status": "active"}',
)
```

Use this before GraphRAG. Add metadata filters for tenant, region, product,
classification, or document status. Add a reranker when the candidate set is
large enough and latency budget allows.

## Pattern 2: Vector-to-Graph Expansion

Use AI Search to find seed chunks or entities, then expand through graph
relationships for context.

1. Query AI Search for candidate chunks/entities.
2. Extract seed entity IDs from metadata or an entity-linking table.
3. Traverse bounded graph neighborhoods in Neo4j or another graph store.
4. Return passages plus graph facts with provenance.

```cypher
MATCH (seed:Entity {entity_id: $entity_id})-[r:RELATED_TO*1..2]-(neighbor:Entity)
RETURN seed.name, type(r[0]) AS relationship, neighbor.name, neighbor.source_doc_id
LIMIT 50
```

Use when semantically similar text points to the right area, but the answer needs
relationships such as "owned by", "depends on", "contraindicated with", or
"part of".

## Pattern 3: Full-Text + Graph Traversal

Use a full-text graph index when exact entity names, SKUs, contract IDs, or
technical identifiers drive recall. Then expand from matched nodes.

```cypher
CALL db.index.fulltext.queryNodes('entityText', $query, {limit: 20})
YIELD node, score
MATCH (node)-[r:MENTIONS|CAUSES|OWNS*1..2]-(related)
RETURN node.name, score, collect(DISTINCT related.name)[0..20] AS related_entities
```

This is often better than pure vector search for canonical entity names and
controlled vocabularies.

## Pattern 4: Text2Cypher

Text2Cypher turns a user question into a read-only Cypher query. Use it only when
the graph schema is curated and the question asks for graph facts or paths.

Safety requirements:

- Provide the model a compact schema: labels, relationship types, allowed
properties, examples, and forbidden operations.
- Reject or rewrite generated Cypher containing `CREATE`, `MERGE`, `DELETE`,
`SET`, `CALL dbms`, file access, or unbounded path expansion.
- Add `LIMIT` and timeout controls.
- Return the generated Cypher in traces for review.

## Pattern 5: Router / Supervisor

Use a router when query types differ:

| Query type | Retriever |
|---|---|
| Semantic document lookup | AI Search hybrid |
| Exact entity relationship | Full-text + graph traversal |
| Path/multi-hop graph question | Text2Cypher or bounded traversal |
| SQL aggregate over governed tables | Genie or SQL tool |
| Managed document Q&A baseline | Knowledge Assistant |

Databricks Apps can host the router and tool calls. Model Serving can host an
agent endpoint when the application needs a served API. MLflow traces should tag
the chosen retriever, candidate count, latency, token count, and final answer
groundedness.
14 changes: 14 additions & 0 deletions manifest.json
Original file line number Diff line number Diff line change
Expand Up @@ -189,6 +189,20 @@
"repo_dir": "experimental",
"version": "0.0.1"
},
"databricks-graphrag": {
"description": "Design GraphRAG and knowledge-graph augmented retrieval on Databricks. Use when comparing GraphRAG vs traditional RAG, combining AI Search with graph traversal, designing entity/relationship models, evaluating hybrid retrieval quality, or integrating Neo4j/Text2Cypher patterns with Unity Catalog, Databricks Apps, Model Serving, and MLflow.",
"files": [
"SKILL.md",
"agents/openai.yaml",
"assets/databricks.png",
"assets/databricks.svg",
"references/architecture-decisions.md",
"references/evaluation-and-operations.md",
"references/retrieval-patterns.md"
],
"repo_dir": "experimental",
"version": "0.1.0"
},
"databricks-iceberg": {
"description": "Apache Iceberg tables on Databricks \u2014 Managed Iceberg tables, External Iceberg Reads (fka Uniform), Compatibility Mode, Iceberg REST Catalog (IRC), Iceberg v3, Snowflake interop, PyIceberg, OSS Spark, external engine access and credential vending. Use when creating Iceberg tables, enabling External Iceberg Reads (uniform) on Delta tables (including Streaming Tables and Materialized Views via compatibility mode), configuring external engines to read Databricks tables via Unity Catalog IRC, integrating with Snowflake catalog to read Foreign Iceberg tables",
"files": [
Expand Down
Loading