GraphRAG SDK API Reference

from graphrag_sdk import GraphRAG

GraphRAG(
    connection: FalkorDBConnection | ConnectionConfig,
    llm: LLMInterface,
    embedder: Embedder,
    ontology: Ontology | None = None,
    retrieval_strategy: RetrievalStrategy | None = None,
    embedding_dimension: int = 256,
    *,
    enable_cypher: bool = False,
    schema: Ontology | None = None,  # deprecated alias of ontology=
)

async def ingest(
    source: str | list[str] | None = None,
    *,
    text: str | None = None,
    document_id: str | None = None,
    record_loader: RecordLoaderStrategy | None = None,
    strict_mapping: bool = False,
    loader: LoaderStrategy | None = None,
    chunker: ChunkingStrategy | None = None,
    extractor: ExtractionStrategy | None = None,
    resolver: ResolutionStrategy | None = None,
    max_concurrency: int = 3,
    ctx: Context | None = None,
) -> IngestionResult | list[IngestionResult | Exception] | StructuredIngestionResult

async def retrieve(
    question: str,
    *,
    strategy: RetrievalStrategy | None = None,
    reranker: RerankingStrategy | None = None,
    ctx: Context | None = None,
) -> RetrieverResult

async def completion(
    question: str,
    *,
    history: list[ChatMessage | dict[str, str]] | None = None,
    strategy: RetrievalStrategy | None = None,
    reranker: RerankingStrategy | None = None,
    prompt_template: str | None = None,
    return_context: bool = False,
    ctx: Context | None = None,
) -> RagResult

from graphrag_sdk import ChatMessage

# Using ChatMessage objects
result = await rag.completion(
    "What happened next?",
    history=[
        ChatMessage(role="user", content="Who is Alice?"),
        ChatMessage(role="assistant", content="Alice is an engineer."),
    ],
)

# Using plain dicts
result = await rag.completion(
    "Tell me more.",
    history=[
        {"role": "user", "content": "What is Acme?"},
        {"role": "assistant", "content": "A tech company."},
    ],
)

async def query(cypher: str, params: dict[str, Any] | None = None) -> list[list[Any]]

rows = await rag.query(
    "MATCH (p:Person)-[r:RELATES]->(o:Organization) "
    "WHERE r.rel_type = 'WORKS_AT' AND o.name = $company RETURN avg(p.hr__age)",
    {"company": "Acme Corp"},
)

async def drop_table(source: str) -> Ontology

async def deduplicate_entities(
    *,
    batch_size: int = 500,
    judge: DedupJudge | None = None,
    dedup_all: bool = False,
    fuzzy: bool | None = None,                # deprecated: False is ignored, True raises
    similarity_threshold: float | None = None,  # deprecated and ignored
) -> int

async def finalize(
    *,
    judge: DedupJudge | None = None,
    dedup_all: bool = False,
) -> FinalizeResult

def retrieve_sync(question: str, **kwargs) -> RetrieverResult
def completion_sync(question: str, **kwargs) -> RagResult
def ingest_sync(source: str | list[str], **kwargs) -> IngestionResult | list[IngestionResult | Exception] | StructuredIngestionResult
def finalize_sync(*, judge=None, dedup_all=False) -> FinalizeResult
def query_sync(cypher: str, params: dict[str, Any] | None = None) -> list[list[Any]]
def drop_table_sync(source: str) -> Ontology

from graphrag_sdk import ConnectionConfig, FalkorDBConnection

@dataclass
class ConnectionConfig:
    host: str = "localhost"
    port: int = 6379
    username: str | None = None
    password: str | None = None
    graph_name: str = "knowledge_graph"
    max_connections: int = 16
    retry_count: int = 3
    retry_delay: float = 1.0
    query_timeout_ms: int = 10_000
    pool_timeout: int = 30

FalkorDBConnection(config: ConnectionConfig | None = None)

from graphrag_sdk import LLMInterface, Embedder, LLMBatchItem
from graphrag_sdk import LiteLLM, LiteLLMEmbedder, OpenRouterLLM, OpenRouterEmbedder

LLMInterface(model_name: str, model_params: dict | None = None, max_concurrency: int = 12)

@dataclass
class LLMBatchItem:
    index: int
    response: LLMResponse | None = None
    error: Exception | None = None

    @property
    def ok(self) -> bool  # True if response is not None

LiteLLM(model: str, *, api_key=None, api_base=None, api_version=None, temperature=0.0, max_tokens=None, **kwargs)

LiteLLMEmbedder(model: str, *, api_key=None, api_base=None, api_version=None, **kwargs)

OpenRouterLLM(model: str, *, api_key=None, temperature=0.0, max_tokens=None, extra_headers=None)

OpenRouterEmbedder(model: str, *, api_key=None, extra_headers=None)

from graphrag_sdk import GraphNode, GraphRelationship, GraphData
from graphrag_sdk import TextChunk, TextChunks
from graphrag_sdk import DocumentInfo, DocumentOutput
from graphrag_sdk import IngestionResult, RagResult
from graphrag_sdk import RetrieverResult, RetrieverResultItem
from graphrag_sdk import ResolutionResult, SearchType

class GraphNode(DataModel):
    id: str                                              # Unique identifier
    label: str                                           # Node label (Person, Place, etc.)
    properties: dict[str, Any] = {}                      # Key-value properties
    embedding_properties: dict[str, list[float]] | None = None

class GraphRelationship(DataModel):
    start_node_id: str
    end_node_id: str
    type: str                                            # Edge type (RELATES for all extracted rels)
    properties: dict[str, Any] = {}                      # Includes rel_type, fact, keywords, etc.
    embedding_properties: dict[str, list[float]] | None = None

class GraphData(DataModel):
    nodes: list[GraphNode] = []
    relationships: list[GraphRelationship] = []

class TextChunk(DataModel):
    text: str
    index: int
    metadata: dict[str, Any] = {}
    uid: str                                             # Auto-generated UUID

class TextChunks(DataModel):
    chunks: list[TextChunk] = []

class DocumentInfo(DataModel):
    path: str | None = None
    uid: str                                             # Auto-generated UUID
    metadata: dict[str, Any] = {}

class DocumentOutput(DataModel):
    text: str
    document_info: DocumentInfo = DocumentInfo()

class IngestionResult(DataModel):
    document_info: DocumentInfo = DocumentInfo()
    nodes_created: int = 0
    relationships_created: int = 0
    chunks_indexed: int = 0
    metadata: dict[str, Any] = {}

class RagResult(DataModel):
    answer: str
    retriever_result: RetrieverResult | None = None      # Populated when return_context=True
    metadata: dict[str, Any] = {}                        # Contains model, num_context_items, strategy

class RetrieverResult(DataModel):
    items: list[RetrieverResultItem] = []
    metadata: dict[str, Any] = {}

class RetrieverResultItem(DataModel):
    content: str
    metadata: dict[str, Any] = {}
    score: float | None = None

class ResolutionResult(DataModel):
    nodes: list[GraphNode] = []
    relationships: list[GraphRelationship] = []
    merged_count: int = 0

from graphrag_sdk import ChatMessage

class ChatMessage(DataModel):
    role: Literal["system", "user", "assistant"]
    content: str

    def to_dict(self) -> dict[str, str]  # {"role": ..., "content": ...}

class LLMResponse(DataModel):
    content: str
    tool_calls: list[dict[str, Any]] | None = None

class SearchType(str, Enum):
    VECTOR = "vector"
    FULLTEXT = "fulltext"
    HYBRID = "hybrid"

class ExtractedEntity(DataModel):
    name: str
    type: str
    description: str = ""
    source_chunk_ids: list[str] = []

class ExtractedRelation(DataModel):
    source: str
    target: str
    type: str
    keywords: str = ""
    description: str = ""
    weight: float = 1.0
    source_chunk_ids: list[str] = []

class EntityMention(DataModel):
    chunk_id: str
    entity_id: str

class ExtractionOutput(DataModel):
    entities: list[ExtractedEntity] = []
    relations: list[ExtractedRelation] = []
    mentions: list[EntityMention] = []

def compute_entity_id(name: str, entity_type: str = "") -> str

from graphrag_sdk import GraphSchema, EntityType, RelationType

class EntityType(DataModel):
    label: str                            # e.g. "Person"
    description: str | None = None        # Helps LLM understand what to extract
    properties: list[PropertyType] = []   # Optional property definitions

class RelationType(DataModel):
    label: str                            # e.g. "WORKS_AT"
    description: str | None = None
    patterns: list[tuple[str, str]] = []  # Allowed (source_label, target_label) pairs

class PropertyType(DataModel):
    name: str
    type: str = "STRING"                  # STRING, INTEGER, FLOAT, BOOLEAN, DATE, LIST
    description: str | None = None
    required: bool = False

class GraphSchema(DataModel):
    entities: list[EntityType] = []
    relations: list[RelationType] = []

class LoaderStrategy(ABC):
    @abstractmethod
    async def load(self, source: str, ctx: Context) -> DocumentOutput: ...

class ChunkingStrategy(ABC):
    @abstractmethod
    async def chunk(self, text: str, ctx: Context) -> TextChunks: ...

class ExtractionStrategy(ABC):
    @abstractmethod
    async def extract(self, chunks: TextChunks, schema: GraphSchema, ctx: Context) -> GraphData: ...

class ResolutionStrategy(ABC):
    @abstractmethod
    async def resolve(self, graph_data: GraphData, ctx: Context) -> ResolutionResult: ...

from graphrag_sdk import IngestionPipeline

IngestionPipeline(
    loader: LoaderStrategy,
    chunker: ChunkingStrategy,
    extractor: ExtractionStrategy,
    resolver: ResolutionStrategy,
    graph_store: GraphStore,
    vector_store: VectorStore,
    schema: GraphSchema | None = None,
    embedder: Embedder | None = None,
)

from graphrag_sdk import (
    Ontology, TableMapping, Column, Link, MappingError,
    CsvRecordLoader, RecordLoaderStrategy, RecordBatch, TextLoader,
)

ontology = Ontology(
    entities=[Entity(label="Person"), Entity(label="Organization")],
    tables=[
        TableMapping(
            source="employees.csv",
            label="Person",
            key="employee_id",
            name="full_name",
            properties={"age": Column("age", "INTEGER"), "title": "title"},
            links=[Link("WORKS_AT", to="Organization", by="org_id")],
        ),
    ],
)
rag = GraphRAG(connection=conn, llm=llm, embedder=embedder, ontology=ontology, enable_cypher=True)

result = await rag.ingest("data/employees.csv")     # StructuredIngestionResult
await rag.ingest("data/quarterly_report.pdf")       # prose, as before
report = await rag.finalize()                        # joins the two halves

@dataclass
class TableMapping:
    source: str                                  # file name; also the property signature
    label: str                                   # the label each row becomes
    key: str | None = None                       # column identifying the row; defaults to name
    name: str | None = None                      # column holding the display name
    properties: dict[str, Column | str] = field(default_factory=dict)  # property_name -> Column; a str means STRING
    links: list[Link] = field(default_factory=list)                    # columns that point at other entities
    standalone: bool = False                     # a new label that deliberately relates to nothing
    description: str | None = None

@dataclass
class Column:
    name: str                     # the column in the source record
    type: str = "STRING"          # one of COLUMN_TYPES
    description: str | None = None

@dataclass
class Link:
    type: str                                    # relationship type, e.g. "WORKS_AT"
    to: str                                      # label of the target
    by: str                                      # column holding the target's key
    name: str | None = None                      # column holding the target's display name, if denormalised
    properties: dict[str, Column | str] = {}     # columns written onto the edge
    description: str | None = None

class RecordLoaderStrategy(ABC):
    @abstractmethod
    async def load_records(self, source: str, ctx: Context) -> RecordBatch: ...

@dataclass
class RecordBatch:
    open_records: Callable[[], Iterator[dict[str, Any]]]   # re-openable; the file is streamed twice
    columns: list[str]
    document_info: DocumentInfo
    record_count: int | None = None
    metadata: dict[str, Any] = field(default_factory=dict)

class CsvRecordLoader(RecordLoaderStrategy):
    def __init__(self, *, delimiter: str | None = None, encoding: str = "utf-8-sig", document_id: str | None = None): ...

class RetrievalStrategy(ABC):
    def __init__(self, graph_store=None, vector_store=None): ...

    async def search(self, query, ctx=None, **kwargs) -> RetrieverResult: ...

    def set_ontology(self, ontology) -> None: ...

    @abstractmethod
    async def _execute(self, query, ctx, **kwargs) -> RawSearchResult: ...

LocalRetrieval(graph_store, vector_store, embedder, top_k=5, include_entities=True)

MultiPathRetrieval(
    graph_store, vector_store, embedder, llm,
    *,
    chunk_top_k=15,
    max_entities=30,
    max_relationships=20,
    rel_top_k=15,
    keyword_limit=10,
)

class RerankingStrategy(ABC):
    @abstractmethod
    async def rerank(self, query, result: RetrieverResult, ctx: Context) -> RetrieverResult: ...

CosineReranker(embedder: Embedder, top_k: int = 15)

from graphrag_sdk import GraphStore, VectorStore

GraphStore(connection: FalkorDBConnection)

VectorStore(connection, embedder=None, index_name="chunk_embeddings", embedding_dimension=256, similarity_function="cosine")

from graphrag_sdk import Context

Context(tenant_id: str = "default", latency_budget_ms: float = 60000.0)

from graphrag_sdk import GraphRAGError