post-graph-rag

Production-Grade, High-Performance Knowledge Graph RAG Engine Native to PostgreSQL.

post-graph-rag seamlessly combines automated LLM-based entity & triple extraction, vector similarity search via pgvector, and graph relationship traversal directly on PostgreSQL using the post-graph graph database library.

It connects to any OpenAI-compatible API (LiteLLM, vLLM, Ollama, DeepSeek, OpenAI) for zero-shot domain-agnostic knowledge extraction, structured document metadata tracking, and context-aware answer synthesis.

🌟 Why post-graph-rag?

Traditional Vector RAG systems suffer from "chunk isolation"—they retrieve isolated text passages based purely on semantic similarity, missing higher-level relationships and cross-document entity connections.

post-graph-rag solves this by building a dual representation inside PostgreSQL:

1. Unstructured Vector Passages: Full document chunks indexed with pgvector HNSW embeddings.
2. Knowledge Graph Triples: Extracted Subject-Predicate-Object entities connected by graph edges.
3. Structured Document Metadata: Rich metadata tracking (source, category, collection, document, page, paragraph, space).
4. Application-Level Space Sub-grouping (space): Scopes indexing and vector similarity search to application-specific environments (e.g. production, sandbox, staging, user_workspace) within {realm} tenant partitions.

Relationship quality

Extracted relations are context-specific, drawn from what the text actually states:

(Zeus) --[is_king_of]--> (Olympian gods)
(Zeus) --[son_of]--> (Cronus)
(Zeus) --[married_to]--> (Hera)
(Zeus) --[defeated]--> (Titans)

Vague connectors (relates_to, associated_with, connected_to, …), self-loops and blank endpoints are rejected at extraction time. A relation that reaches the graph always says something specific about the pair it connects, and two entities merely appearing near each other never produces an edge.

If the LLM cannot produce usable structure, indexing raises ExtractionError. Placeholder edges are never invented as a fallback: once written they are indistinguishable from genuine extracted structure.

Relations the text explicitly denies are stored with the positive predicate and negated: true, rather than as an inverted predicate like did_not_have_relationship_with. Traversal and synthesis can then exclude them instead of reading them as assertions.

Entity resolution

Entities are unique per (realm, space, lower(name)), enforced by a unique index, and are additionally resolved through aliases. Extraction records every other surface form it sees, so Babbage, Charles Babbage and C. Babbage converge on one vertex; the fuller name becomes canonical and the rest become aliases. Pronouns and relative references (he, his father, the company) are rejected outright — they cannot resolve to a stable vertex.

The same entity mentioned in many documents is one vertex, which is what allows the graph to connect chunks that share no vocabulary.

Chunking and document context

index_text() chunks a document (with overlap, so relations spanning a boundary survive) and threads a DocumentContext through the chunks — title, source, and the canonical entity names found so far. Without it, every chunk after the first is extracted blind and its pronouns become junk vertices.

Bring your own splitter by passing chunker= to GraphRAG, or use index_document() directly with your own DocumentContext.

rag = GraphRAG(config, chunker=my_splitter)
await rag.index_text(long_text, metadata=DocumentMetadata(document="babbage.txt"))

Repeated relations

The same triple extracted from several chunks is one edge whose weight increments, not several edges. Weight then breaks ties when ranking relations.

🏗️ Architecture Workflow

graph TD
    subgraph INDEXING ["1. Knowledge Graph & Vector Indexing"]
        A[Document Text + Metadata] --> B[Embedding Service]
        A --> C[LLM GraphExtractor]
        
        B -->|Vectors| D[post-graph Store]
        C -->|Entities & Triples| D
        
        D --> E[(PostgreSQL + pgvector)]
        E -->|Tables| E1[documents]
        E -->|Tables| E2[entities]
        E -->|Edges| E3[relations]
        E -->|Edges| E4[doc_mentions]
    end

    subgraph RETRIEVAL ["2. Hybrid Retrieval & Synthesis"]
        Q[User Question] --> R[GraphRAG Query Engine]
        R -->|Embedding| S[pgvector Similarity Search]
        E1 & E2 -->|Top-K Passages & Entities| S
        S --> T[1-Hop Graph Relationship Traversal]
        E3 -->|Subject-Predicate-Object| T
        
        S & T --> U[LLM Answer Synthesis]
        U --> V[Final Answer + Citations + Graph Triples]
    end

📦 Installation

Install post-graph-rag via pip or uv:

pip install post-graph-rag

Or using uv:

uv add post-graph-rag

PostgreSQL Requirements

pgvector is required, not optional. Without it the vertex tables are created without embedding columns and every similarity search silently returns nothing. initialize() raises SchemaError if it is missing rather than degrading.

# macOS
brew install pgvector

# Debian/Ubuntu (match your server version)
sudo apt install postgresql-17-pgvector

CREATE EXTENSION IF NOT EXISTS vector;

🚀 Quick Start

1. Basic Indexing & Querying

import asyncio
from post_graph_rag import GraphRAG, RAGConfig, DocumentMetadata

async def main():
    # 1. Configure GraphRAG engine
    config = RAGConfig(
        api_base="http://localhost:4000/v1",       # OpenAI-compatible router endpoint
        api_key=os.environ["OPENAI_API_KEY"],     # Never hardcode credentials
        model="DeepSeek-V3.2",                    # LLM model for extraction & synthesis
        embedding_model="text-embedding-3-small", # Embedding model
        embedding_dim=1536,                       # Must match the model's output width
        db_uri="postgresql://user:password@localhost:5432/postgres",
        realm="enterprise_kb",
        schema_per_realm=True                     # Give each tenant its own schema
    )

    rag = GraphRAG(config)
    
    # 2. Connect & initialize PostgreSQL graph schema
    await rag.initialize()

    # 3. Index unstructured documents
    doc_text = (
        "Zeus is the king of the Olympian gods, ruling sky and thunder from Mount Olympus. "
        "He is the son of Cronus and Rhea, and married to Hera. "
        "Zeus defeated the Titans in the Titanomachy to establish his rule."
    )
    
    result = await rag.index_document(doc_text, metadata={"source": "greek_mythology.txt"})
    print(f"Indexed document {result['document_id']}: Extracted {result['entities_extracted']} entities.")

    # 4. Perform Hybrid RAG Query
    response = await rag.query("Who are the parents of Zeus and what did he defeat?")
    
    print("\n=== SYNTHESIZED ANSWER ===")
    print(response["answer"])

    print("\n=== RETRIEVED GRAPH TRIPLES ===")
    for triple in response["retrieved_graph_triples"]:
        print(f"  - {triple}")

    # 5. Clean up
    await rag.close()

if __name__ == "__main__":
    asyncio.run(main())

📋 Document Metadata (DocumentMetadata)

post-graph-rag includes structured document metadata tracking via the DocumentMetadata model:

from post_graph_rag import DocumentMetadata

metadata = DocumentMetadata(
    source="https://mythology.org/zeus.html",  # Document origin (URL, filepath, API)
    category="greek_mythology",               # Document category/topic
    collection="olympian_deities",             # Collection namespace
    document="zeus_overview.pdf",              # Title or filename
    page=1,                                    # 1-based page number
    paragraph=2,                               # 1-based paragraph index
    space="production",                        # Sub-grouping within the realm
    extra={"author": "Homer", "year": -700}    # Custom metadata key-value pairs
)

await rag.index_document(chunk_text, metadata=metadata)

Design Rationale: Optional vs. Required

• All metadata fields are optional with default None. This allows seamless indexing of raw strings, short code snippets, webhooks, or unformatted text, while offering rich structural provenance tracking when indexing multi-page PDFs or categorized enterprise documents.

⚙️ Configuration Reference (RAGConfig)

RAGConfig can be configured explicitly or automatically loaded from environment variables:

Environment variables are read when a RAGConfig is constructed, not at import time.

schema_per_realm

Off by default for backwards compatibility, but recommended for new deployments. With it off, every realm shares one physical set of tables filtered by a realm column — so the first realm to create entities fixes the embedding column width for all of them, and a second realm with a different embedding_dim cannot work.

allow_embedding_fallback

Off by default. Fallback vectors are not comparable with API embeddings, so mixing them into the same table corrupts retrieval rather than degrading it. With it off, an embedding failure raises EmbeddingError.

📖 API Reference

GraphRAG

The main orchestrator class for indexing and querying.

• await initialize(): Connects to PostgreSQL and creates the graph tables (documents, entities, relations, doc_mentions). Raises SchemaError if pgvector is unavailable or an existing table's embedding width disagrees with embedding_dim.
• await index_document(text, metadata=None, space=None) -> Dict[str, Any]: Embeds the chunk, extracts entities/triples via the LLM, resolves entities by name, and writes vertices, relations and doc_mentions edges. Raises rather than writing placeholder structure if extraction fails. Returns counts plus document_id and metadata.
• await query(question, param=None, top_k=None): Retrieves and synthesizes an answer. Returns question, answer, mode, keywords, retrieved_documents, retrieved_entities, retrieved_graph_triples, references. With QueryParam(stream=True) returns an async iterator of content chunks instead.
• await query_data(question, param=None) -> Dict[str, Any]: Structured retrieval with no synthesis — returns entities, relationships, chunks, references.
• await close(): Closes database connection pools.

QueryParam

• mode: one of mix, local, global, hybrid, naive, bypass. An unknown mode raises ValueError.
• top_k, max_total_tokens, max_entity_tokens, max_relation_tokens, response_type
• stream: return an async iterator of tokens instead of a dict.
• only_need_context: return retrieval output without calling the LLM.
• space: restrict retrieval to one space; __all__ queries across all spaces.
• conversation_history, hl_keywords, ll_keywords: supply keywords to skip extraction.

Errors

All inherit from RAGError, so failures surface instead of degrading into irrelevant results:

• SchemaError — pgvector missing, or embedding width mismatch.
• EmbeddingError — embedding request failed, or returned the wrong width.
• LLMError — completion or streaming call failed.
• ExtractionError — the LLM returned no usable entities or triples.

DocumentMetadata

Data container for structured document metadata.

• source: Optional[str]: Document URL, path, or origin.
• category: Optional[str]: Document category or domain.
• collection: Optional[str]: Document collection or folder.
• document: Optional[str]: File title or filename.
• page: Optional[int]: 1-based page number.
• paragraph: Optional[int]: 1-based paragraph index.
• space: Optional[str]: Sub-grouping space to index into.
• extra: Dict[str, Any]: Custom user metadata.
• to_dict() -> Dict[str, Any]: Serializes non-None fields to dictionary representation.
• from_dict(data: Dict[str, Any]) -> DocumentMetadata: Deserializes dictionary data.

RAGGraphStore

Database layer wrapping post-graph.

• add_document(text, embedding, metadata, space=None): Inserts a document vertex into the documents table.
• upsert_entity(name, entity_type, description, embedding, space=None): Upserts by canonical name within (realm, space), so an entity mentioned in many documents is a single vertex. A bare Concept stub never overwrites a richer type or description.
• find_entity_by_name(name, space=None): Resolve an entity vertex by name.
• add_relation(from_entity, to_entity, relation_type, description, space=None, embedding=None): Directed relation edge.
• add_doc_mention(doc_vertex, entity_vertex, space=None): Links a chunk to an entity it mentions.
• search_similar_entities(query_vec, top_k, space=None) / search_similar_documents(...): pgvector HNSW similarity search.
• search_similar_relations(query_vec, top_k, space=None): Semantic search over relation edges. Returns [] unless embed_relations is enabled.
• get_neighbors(entity_id, space=None): 1-hop outgoing relations, scoped to space.
• get_all_relations(limit, space=None): Relations with their endpoint vertices.

🗄️ PostgreSQL Database Schema

post-graph-rag automatically provisions and manages the following graph schema in PostgreSQL powered by post-graph:

🧪 Testing

pip install -e ".[test]"
createdb postgres && psql -d postgres -c "CREATE EXTENSION IF NOT EXISTS vector;"
POSTGRES_TEST_URI="postgresql://localhost:5432/postgres" pytest

DB-backed tests create a disposable schema per test realm and drop it afterwards. They skip automatically when PostgreSQL with pgvector is not reachable.

📄 License

This project is licensed under the MIT License - see the LICENSE file for details.

Developed by Chandan Rajah (chandan.rajah@gmail.com).