Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
21 commits
Select commit Hold shift + click to select a range
31e927e
feat: initial mult-extension loader
drr00t May 17, 2026
6631d8e
feat(ingestion): default to SentenceTokenCapChunking in ingest()/upda…
galshubeli May 14, 2026
9bfbab4
test(loader): more tests for doclin-base loaders
drr00t May 19, 2026
857f870
fix(issues): from coderabbitai review
drr00t May 19, 2026
4af2f8c
fix(conflict): need to be updated
drr00t May 19, 2026
7969df9
fix(retrieval): rank MENTIONED_IN chunks by cosine in MultiPath Path C
galshubeli May 18, 2026
4f7d37c
Merge branch 'main' into 241-issue-expand-document-loader-coverage
drr00t May 19, 2026
531d29c
test: refactor docling loader tests to use local mock enumerations in…
drr00t May 21, 2026
a5ecc4e
test: implement robust sys.modules mocking for docling in unit tests …
drr00t May 22, 2026
21c55e4
Merge branch 'main' into 241-issue-expand-document-loader-coverage
drr00t May 26, 2026
a451138
refactor: remove legacy loaders and consolidate ingestion logic into …
drr00t May 28, 2026
15407de
refactor: update docling example command path and replace HTML script…
drr00t May 28, 2026
4c0a7cf
Merge branch 'main' into 241-issue-expand-document-loader-coverage
drr00t Jun 3, 2026
a86dfd2
Merge branch 'main' into 241-issue-expand-document-loader-coverage
galshubeli Jun 3, 2026
a41c9ed
Merge branch 'main' into 241-issue-expand-document-loader-coverage
drr00t Jun 8, 2026
fc1a854
feat: add support for URL sources in DoclingLoader by bypassing local…
drr00t Jun 8, 2026
c9ff0e4
feat: update default file loader mapping to prioritize DoclingLoader …
drr00t Jun 8, 2026
fee46b1
fix: update docling loader to prioritize table markdown exports over …
drr00t Jun 8, 2026
7d0f838
chore: update project dependencies, removingo docling from all in pyp…
drr00t Jun 8, 2026
ff16881
refactor: apply consistent ruff code formatting across all tests and …
drr00t Jun 8, 2026
ab647b0
test: update docling loader mock setup to support markdown export in …
drr00t Jun 8, 2026
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
2 changes: 2 additions & 0 deletions docs/ingestion.md
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,8 @@ Steps 1-7 run **sequentially** (each depends on the previous). Steps 8-9 run **i

**How:** The `LoaderStrategy` ABC handles this. The SDK auto-detects the loader based on file extension:
- `.pdf` files use `PdfLoader`
- `.docx`, `.xlsx`, `.pptx`, `.html`, `.csv`, and URLs use `DoclingLoader` (if `graphrag-sdk[docling]` is installed)
- `.md` files use `MarkdownLoader`
- Everything else uses `TextLoader`
- If you pass `text=` directly, the loader step is skipped entirely

Expand Down
20 changes: 19 additions & 1 deletion docs/strategies.md
Original file line number Diff line number Diff line change
Expand Up @@ -64,12 +64,30 @@ loader = MarkdownLoader()
**Design Note: Markup Preservation**
For complex elements like tables, lists, and code blocks, `MarkdownLoader` intentionally outputs the raw markdown source (including pipes `|`, list dashes `-`, and code fences) rather than stripping the syntax. While this introduces minor syntax "noise", it preserves critical structural cues (such as spatial column alignment and nested indentation) that the LLM requires during the Extraction phase to accurately parse relational data.

### Built-in: DoclingLoader

A universal loader utilizing the `docling` library to parse rich document formats (PDF, DOCX, XLSX, PPTX, HTML, CSV, Markdown, URLs). Requires `pip install graphrag-sdk[docling]`.

```python
from graphrag_sdk.ingestion.loaders.docling_loader import DoclingLoader

# You can pass arbitrary docling DocumentConverter arguments
loader = DoclingLoader(
allowed_formats=["docx", "pptx"],
format_options={...}
)
```

**Design Note: Hierarchical Breadcrumbs**
`DoclingLoader` extracts deep structural context by generating "breadcrumbs" (e.g., tracking that a paragraph is inside `H1 -> H2 -> List`). These breadcrumbs are attached to each extracted text chunk as metadata, ensuring that the semantic location of the text within the original document is preserved during chunking and extraction.

### Default Behavior

If no loader is specified in `ingest()`:
- `.pdf` files use `PdfLoader`
- `.docx`, `.xlsx`, `.pptx`, `.html`, `.csv` files and URLs use `DoclingLoader` (if installed)
- `.md` files use `MarkdownLoader`
- Everything else uses `TextLoader`
- Everything else uses `TextLoader`
- If `text=` is passed directly, the loader is skipped

### Writing Your Own
Expand Down
10 changes: 9 additions & 1 deletion graphrag_sdk/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ asyncio.run(main())
pip install graphrag-sdk[litellm] # OpenAI, Azure, Anthropic, 100+ models
pip install graphrag-sdk[openrouter] # OpenRouter models
pip install graphrag-sdk[pdf] # PDF ingestion
pip install graphrag-sdk[docling] # DOCX, XLSX, PPTX, HTML, CSV, URLs
pip install graphrag-sdk[all] # Everything
```

Expand All @@ -59,7 +60,14 @@ async def main():
llm=LiteLLM(model="openai/gpt-4o"),
embedder=LiteLLMEmbedder(model="openai/text-embedding-3-small"),
) as rag:
# Supported formats: PDF, DOCX, XLSX, PPTX, HTML, CSV, Markdown, URLs, TXT
await rag.ingest("report.pdf") # PDF
await rag.ingest("document.docx") # Word
await rag.ingest("spreadsheet.xlsx") # Excel
await rag.ingest("presentation.pptx") # PowerPoint
await rag.ingest("page.html") # HTML
await rag.ingest("data.csv") # CSV
await rag.ingest("notes.md") # Markdown
await rag.ingest("source_id", text="Alice works at Acme.") # Raw text
await rag.finalize() # Dedup + index

Expand Down Expand Up @@ -135,7 +143,7 @@ Every algorithmic concern is a swappable strategy behind an abstract base class:

| Concern | ABC | Built-in Options | Default |
|---------|-----|-----------------|---------|
| **Loading** | `LoaderStrategy` | `TextLoader`, `PdfLoader` | Auto-detect by extension |
| **Loading** | `LoaderStrategy` | `TextLoader`, `PdfLoader`, `DoclingLoader` (universal: DOCX/XLSX/PPTX/HTML/CSV/URL) | Auto-detect by extension |
| **Chunking** | `ChunkingStrategy` | `FixedSizeChunking`, `SentenceTokenCapChunking`, `ContextualChunking`, `CallableChunking` | `FixedSizeChunking` |
| **Extraction** | `ExtractionStrategy` | `GraphExtraction` (GLiNER2 + LLM) | `GraphExtraction` |
| **Resolution** | `ResolutionStrategy` | `ExactMatchResolution`, `DescriptionMergeResolution`, `SemanticResolution`, `LLMVerifiedResolution` | `ExactMatch` |
Expand Down
5 changes: 2 additions & 3 deletions graphrag_sdk/examples/01_quickstart.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,9 +20,8 @@
"""

import asyncio
import os

from graphrag_sdk import GraphRAG, ConnectionConfig, LiteLLM, LiteLLMEmbedder
from graphrag_sdk import ConnectionConfig, GraphRAG, LiteLLM, LiteLLMEmbedder

TEXT = (
"Alice Johnson is a software engineer at Acme Corp in London. "
Expand Down Expand Up @@ -99,7 +98,7 @@ async def main():
ChatMessage(role="assistant", content="Alice works at Acme Corp in London."),
],
)
print(f"\nQ: What is her role there?")
print("\nQ: What is her role there?")
print(f"A: {followup.answer}")


Expand Down
2 changes: 1 addition & 1 deletion graphrag_sdk/examples/02_pdf_with_schema.py
Original file line number Diff line number Diff line change
Expand Up @@ -121,7 +121,7 @@ async def main():
print(f"\nRetrieved {len(answer.retriever_result.items)} context items:")
for i, item in enumerate(answer.retriever_result.items[:5]):
score = item.score if item.score is not None else 0.0
print(f" [{i+1}] (score={score:.3f}) {item.content[:100]}...")
print(f" [{i + 1}] (score={score:.3f}) {item.content[:100]}...")

# Show graph stats
stats = await rag.get_statistics()
Expand Down
18 changes: 11 additions & 7 deletions graphrag_sdk/examples/03_custom_strategies.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
from graphrag_sdk import (
ConnectionConfig,
EntityType,
GraphExtraction,
GraphRAG,
GraphSchema,
LiteLLM,
Expand All @@ -31,8 +32,9 @@
)
from graphrag_sdk.core.context import Context
from graphrag_sdk.ingestion.chunking_strategies.fixed_size import FixedSizeChunking
from graphrag_sdk import GraphExtraction
from graphrag_sdk.ingestion.resolution_strategies.description_merge import DescriptionMergeResolution
from graphrag_sdk.ingestion.resolution_strategies.description_merge import (
DescriptionMergeResolution,
)

# Sample documents (replace with your own)
DOCUMENTS = [
Expand All @@ -42,14 +44,14 @@
"radioactivity. Born in Warsaw, Poland, she moved to Paris to study at the Sorbonne. "
"She was the first woman to win a Nobel Prize, and the only person to win Nobel Prizes "
"in two different sciences -- Physics in 1903 and Chemistry in 1911. She worked closely "
"with her husband Pierre Curie at the University of Paris."
"with her husband Pierre Curie at the University of Paris.",
),
(
"doc_2",
"Pierre Curie was a French physicist and Nobel laureate. He shared the 1903 Nobel Prize "
"in Physics with his wife Marie Curie and Henri Becquerel for their research on radiation. "
"Pierre was a professor at the University of Paris. He tragically died in 1906 in a "
"street accident in Paris. After his death, Marie took over his teaching position."
"street accident in Paris. After his death, Marie took over his teaching position.",
),
]

Expand Down Expand Up @@ -146,15 +148,17 @@ async def main():
print("\nRunning finalize()...")
finalize_result = await rag.finalize()
print(f" Deduplicated: {finalize_result.entities_deduplicated} entities")
print(f" Embedded: {finalize_result.entities_embedded} entities, "
f"{finalize_result.relationships_embedded} relationships")
print(
f" Embedded: {finalize_result.entities_embedded} entities, "
f"{finalize_result.relationships_embedded} relationships"
)

elapsed = time.time() - t0
print(f"\nTotal ingestion time: {elapsed:.1f}s")

# --- Graph statistics ---
stats = await rag.get_statistics()
print(f"\nGraph Statistics:")
print("\nGraph Statistics:")
print(f" Nodes: {stats['node_count']}")
print(f" Edges: {stats['edge_count']}")
print(f" Mentions: {stats['mention_edge_count']}")
Expand Down
1 change: 0 additions & 1 deletion graphrag_sdk/examples/04_custom_provider.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@
from graphrag_sdk import ConnectionConfig, Embedder, GraphRAG, LLMInterface
from graphrag_sdk.core.models import LLMResponse


# --- Custom LLM Provider ---


Expand Down
68 changes: 64 additions & 4 deletions graphrag_sdk/examples/05_notebook_demo.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,31 @@
"execution_count": null,
"metadata": {},
"outputs": [],
"source": "from graphrag_sdk import GraphRAG, ConnectionConfig, LiteLLM, LiteLLMEmbedder\n\n# Option A: OpenAI (default)\nllm = LiteLLM(model=\"openai/gpt-5.5\")\nembedder = LiteLLMEmbedder(model=\"openai/text-embedding-3-small\")\n\n# Option B: Azure OpenAI (uncomment and fill in)\n# import os\n# llm = LiteLLM(\n# model=\"azure/gpt-4.1\",\n# api_key=os.getenv(\"AZURE_OPENAI_API_KEY\"),\n# api_base=os.getenv(\"AZURE_OPENAI_ENDPOINT\"),\n# api_version=\"2024-12-01-preview\",\n# )\n# embedder = LiteLLMEmbedder(\n# model=\"azure/text-embedding-ada-002\",\n# api_key=os.getenv(\"AZURE_OPENAI_API_KEY\"),\n# api_base=os.getenv(\"AZURE_OPENAI_ENDPOINT\"),\n# api_version=\"2024-12-01-preview\",\n# )\n\nconn = ConnectionConfig(host=\"localhost\", graph_name=\"notebook_demo\")\nprint(\"Providers configured.\")"
"source": [
"from graphrag_sdk import ConnectionConfig, GraphRAG, LiteLLM, LiteLLMEmbedder\n",
"\n",
"# Option A: OpenAI (default)\n",
"llm = LiteLLM(model=\"openai/gpt-5.5\")\n",
"embedder = LiteLLMEmbedder(model=\"openai/text-embedding-3-small\")\n",
"\n",
"# Option B: Azure OpenAI (uncomment and fill in)\n",
"# import os\n",
"# llm = LiteLLM(\n",
"# model=\"azure/gpt-4.1\",\n",
"# api_key=os.getenv(\"AZURE_OPENAI_API_KEY\"),\n",
"# api_base=os.getenv(\"AZURE_OPENAI_ENDPOINT\"),\n",
"# api_version=\"2024-12-01-preview\",\n",
"# )\n",
"# embedder = LiteLLMEmbedder(\n",
"# model=\"azure/text-embedding-ada-002\",\n",
"# api_key=os.getenv(\"AZURE_OPENAI_API_KEY\"),\n",
"# api_base=os.getenv(\"AZURE_OPENAI_ENDPOINT\"),\n",
"# api_version=\"2024-12-01-preview\",\n",
"# )\n",
"\n",
"conn = ConnectionConfig(host=\"localhost\", graph_name=\"notebook_demo\")\n",
"print(\"Providers configured.\")"
]
},
{
"cell_type": "markdown",
Expand All @@ -49,7 +73,24 @@
"execution_count": null,
"metadata": {},
"outputs": [],
"source": "TEXT = (\n \"Alice Johnson is a software engineer at Acme Corp in London. \"\n \"She leads the backend team and reports to Bob Smith, the CTO. \"\n \"Acme Corp was founded in 2015 and specializes in cloud infrastructure. \"\n \"The company recently opened a new office in Berlin, managed by Clara Wei. \"\n \"Clara previously worked at TechStart in Munich before joining Acme Corp.\"\n)\n\nrag = GraphRAG(connection=conn, llm=llm, embedder=embedder)\n\nresult = await rag.ingest(text=TEXT, document_id=\"demo_doc\")\nprint(f\"Nodes created: {result.nodes_created}\")\nprint(f\"Relationships created: {result.relationships_created}\")\n\nfinalize_result = await rag.finalize()\nprint(f\"\\nFinalize complete: {finalize_result}\")"
"source": [
"TEXT = (\n",
" \"Alice Johnson is a software engineer at Acme Corp in London. \"\n",
" \"She leads the backend team and reports to Bob Smith, the CTO. \"\n",
" \"Acme Corp was founded in 2015 and specializes in cloud infrastructure. \"\n",
" \"The company recently opened a new office in Berlin, managed by Clara Wei. \"\n",
" \"Clara previously worked at TechStart in Munich before joining Acme Corp.\"\n",
")\n",
"\n",
"rag = GraphRAG(connection=conn, llm=llm, embedder=embedder)\n",
"\n",
"result = await rag.ingest(text=TEXT, document_id=\"demo_doc\")\n",
"print(f\"Nodes created: {result.nodes_created}\")\n",
"print(f\"Relationships created: {result.relationships_created}\")\n",
"\n",
"finalize_result = await rag.finalize()\n",
"print(f\"\\nFinalize complete: {finalize_result}\")"
]
},
{
"cell_type": "markdown",
Expand Down Expand Up @@ -92,7 +133,23 @@
"execution_count": null,
"metadata": {},
"outputs": [],
"source": "result = await rag.completion(\"Where does Alice work?\", return_context=True)\n\nprint(\"Answer:\", result.answer)\nprint(\"\\n--- Retrieved Context ---\")\nif result.retriever_result and result.retriever_result.items:\n for i, item in enumerate(result.retriever_result.items[:5], 1):\n score = item.score if item.score is not None else 0.0\n print(f\"\\n[{i}] Score: {score:.3f}\")\n print(f\" Text: {item.content[:200]}...\" if len(item.content) > 200 else f\" Text: {item.content}\")\nelse:\n print(\"No retrieval context available.\")"
"source": [
"result = await rag.completion(\"Where does Alice work?\", return_context=True)\n",
"\n",
"print(\"Answer:\", result.answer)\n",
"print(\"\\n--- Retrieved Context ---\")\n",
"if result.retriever_result and result.retriever_result.items:\n",
" for i, item in enumerate(result.retriever_result.items[:5], 1):\n",
" score = item.score if item.score is not None else 0.0\n",
" print(f\"\\n[{i}] Score: {score:.3f}\")\n",
" print(\n",
" f\" Text: {item.content[:200]}...\"\n",
" if len(item.content) > 200\n",
" else f\" Text: {item.content}\"\n",
" )\n",
"else:\n",
" print(\"No retrieval context available.\")"
]
},
{
"cell_type": "markdown",
Expand All @@ -104,7 +161,10 @@
"execution_count": null,
"metadata": {},
"outputs": [],
"source": "await rag.close()\nprint(\"Connection closed.\")"
"source": [
"await rag.close()\n",
"print(\"Connection closed.\")"
]
}
],
"metadata": {
Expand Down
20 changes: 13 additions & 7 deletions graphrag_sdk/examples/06_markdown_document_aware.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,10 +49,14 @@ def create_schema() -> GraphSchema:
"""Generic schema suitable for technical documentation."""
return GraphSchema(
entities=[
EntityType(label="Concept", description="A technical concept, feature, or abstraction"),
EntityType(label="Component", description="A software module, library, or service"),
EntityType(label="Technology", description="A programming language, framework, or tool"),
EntityType(label="Person", description="An author, contributor, or mentioned individual"),
EntityType(label="Concept", description="A technical concept, feature, or abstraction"),
EntityType(label="Component", description="A software module, library, or service"),
EntityType(
label="Technology", description="A programming language, framework, or tool"
),
EntityType(
label="Person", description="An author, contributor, or mentioned individual"
),
EntityType(label="Organization", description="A company, team, or standards body"),
],
relations=[
Expand Down Expand Up @@ -196,8 +200,10 @@ async def main():
loader=MarkdownLoader(),
chunker=StructuralChunking(max_tokens=512),
)
print(f"Done: {result.nodes_created} nodes, {result.relationships_created} edges, "
f"{result.chunks_indexed} chunks indexed")
print(
f"Done: {result.nodes_created} nodes, {result.relationships_created} edges, "
f"{result.chunks_indexed} chunks indexed"
)

# Post-ingestion: dedup entities, embed entity/relationship descriptions,
# and build full-text + vector indexes. Required before retrieval works.
Expand All @@ -214,7 +220,7 @@ async def main():
print(f"\nRetrieved {len(answer.retriever_result.items)} context items:")
for i, item in enumerate(answer.retriever_result.items[:5]):
score = item.score if item.score is not None else 0.0
print(f" [{i+1}] (score={score:.3f}) {item.content[:120]}...")
print(f" [{i + 1}] (score={score:.3f}) {item.content[:120]}...")

# Graph statistics
stats = await rag.get_statistics()
Expand Down
15 changes: 4 additions & 11 deletions graphrag_sdk/examples/07_incremental_updates.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,19 +26,12 @@
LiteLLMEmbedder,
)


V1_TEXT = (
"Alice Johnson is a software engineer at Acme Corp in London. "
"She reports to Bob Smith."
)
V1_TEXT = "Alice Johnson is a software engineer at Acme Corp in London. She reports to Bob Smith."
V2_TEXT = (
"Alice Johnson is the VP of Engineering at Acme Corp in Berlin. "
"She reports to the CEO, Carol Wei."
)
NEW_DOC = (
"Acme Corp acquired Globex Industries in 2026. "
"The combined company employs 5000 people."
)
NEW_DOC = "Acme Corp acquired Globex Industries in 2026. The combined company employs 5000 people."


async def main():
Expand All @@ -63,7 +56,7 @@ async def main():
# effectively a single Cypher lookup.
result = await rag.update(text=V1_TEXT, document_id="alice_bio")
assert result.no_op is True
print(f"No-op update detected (content hash matched).")
print("No-op update detected (content hash matched).")

# ── 3. Real update — chunks replaced, orphans cleaned up ───
# Bob Smith was only mentioned by V1; he becomes an orphan
Expand Down Expand Up @@ -107,7 +100,7 @@ async def main():

# ── 7. Verify with a query ─────────────────────────────────
answer = await rag.completion("Who acquired Globex Industries?")
print(f"\nQ: Who acquired Globex Industries?")
print("\nQ: Who acquired Globex Industries?")
print(f"A: {answer.answer}")


Expand Down
Loading