neo4j-graphrag Python Package

The neo4j-graphrag package provides several retriever implementations to fetch context from Neo4j or external vector databases:

• VectorRetriever: Similarity search using a Neo4j vector index. Returns matched node and similarity score.
• VectorCypherRetriever: Similarity search followed by a custom Cypher query to fetch additional graph context.
• HybridRetriever: Combines vector search and full-text index search in Neo4j.
• HybridCypherRetriever: Hybrid search followed by a custom Cypher traversal.
• ToolsRetriever: Uses an LLM to select and execute multiple tools (like other retrievers) based on the query.
• Text2Cypher: Translates natural language into Cypher queries to fetch structured data.
• External Vector Database Retrievers:
  • WeaviateNeo4jRetriever (for Weaviate)
  • PineconeNeo4jRetriever (for Pinecone)
  • QdrantNeo4jRetriever (for Qdrant)

All retrievers expose a .search() method.

You can guide the LLM's extraction process by providing a schema dictionary to the SimpleKGPipeline. This prevents unguided extraction and ensures the graph follows your desired structure.

Node and Relationship Types

Types can be defined as:

• Simple strings: Just the label (e.g., "Person").
• Dictionaries: For detailed control. Must include a label key. Optionally include description and properties.
  • If no properties are provided, a default "name" property is added automatically.
  • properties is a list of dicts with name and type (e.g., {"name": "age", "type": "INTEGER"}).

Patterns

Patterns define how nodes are connected using triplets in the format: (source_node_label, relationship_label, target_node_label).

Schema Parameter Modes

• schema=None or schema="EXTRACTED" (Default): The LLM automatically extracts a schema from the input text once, then uses that schema for all subsequent chunks.
• schema="FREE" or an empty schema: No schema extraction is performed; extraction is unguided.

NODE_TYPES = [
    "Person",
    {"label": "House", "description": "Family the person belongs to"},
    {"label": "Planet", "properties": [{"name": "name", "type": "STRING"}, {"name": "weather", "type": "STRING"}]},
]

RELATIONSHIP_TYPES = [
    "PARENT_OF",
    {
        "label": "HEIR_OF",
        "description": "Used for inheritor relationship between father and sons",
    },
    {"label": "RULES", "properties": [{"name": "fromYear", "type": "INTEGER"}]},
]

PATTERNS = [
    ("Person", "PARENT_OF", "Person"),
    ("Person", "HEIR_OF", "House"),
    ("House", "RULES", "Planet"),
]

kg_builder = SimpleKGPipeline(
    # ...
    schema={
        "node_types": NODE_TYPES,
        "relationship_types": RELATIONSHIP_TYPES,
        "patterns": PATTERNS,
        "additional_node_types": False,
    },
    # ...
)

The GraphPruning component cleans up the extracted graph based on your GraphSchema configuration.

Schema Options:

• constraints: Use EXISTENCE to mark properties that must be present. Nodes/relationships violating this are pruned.
• additional_properties: (bool) If True, keeps properties not in the schema. If False, removes them. (Default: False if properties are defined, True otherwise).
• additional_node_types: (bool) If False, removes nodes not in the schema.
• additional_relationship_types: (bool) If False, removes relationships not in the schema.
• additional_patterns: (bool, default True) If False, only keeps relationship patterns explicitly listed in the schema.

Pruning Rules:

• Nodes with empty labels/IDs are removed.
• Nodes/relationships missing EXISTENCE properties are removed.
• Nodes with no remaining properties after pruning are removed.
• Relationships with empty types or invalid source/target nodes are removed.

The Knowledge Graph construction pipeline is composed of several components that transform unstructured data into a structured graph. While some components are optional, the core workflow involves extracting text, splitting it, extracting entities/relations, and writing them to a graph.

Core Components:

• Data loader: Extracts text from files (e.g., PDFs).
• Text splitter: Breaks text into manageable chunks based on LLM token limits.
• Chunk embedder (optional): Computes embeddings for the text chunks.
• Schema builder: Provides a schema to ground the LLM's extraction of node and relationship types. This can be manual or automatically extracted.
• Lexical graph builder (optional): Builds a graph representing the relationship between Documents and Chunks.
• Entity and relation extractor: Identifies relevant entities and relations within the text.
• Graph pruner: Cleans the graph based on a provided schema.
• Knowledge Graph writer: Saves the identified entities and relations to a database.
• Entity resolver: Merges similar entities into a single node to reduce redundancy.

To support custom endpoints with specific credential handling or different defaults, you can subclass the provider base classes: BaseAnthropicLLM, BaseOpenAILLM, or BaseGeminiLLM. These base classes provide all provider-agnostic logic (message building, schema conversion, etc.), so you only need to implement the client construction logic.

For Anthropic and OpenAI, use the neo4j_graphrag.llm.utils.split_http_client_kwargs helper to correctly route http_client arguments to the appropriate sync and async SDK clients.

from typing import Any, Optional
import anthropic
from neo4j_graphrag.llm import BaseAnthropicLLM
from neo4j_graphrag.llm.utils import split_http_client_kwargs


class MyCustomAnthropicLLM(BaseAnthropicLLM):
    """Talks to a self-hosted, Anthropic-compatible endpoint."""

    DEFAULT_ENDPOINT = "https://my-custom-endpoint.example.com"

    def __init__(
        self,
        model_name: str,
        model_params: Optional[dict[str, Any]] = None,
        **kwargs: Any,
    ):
        super().__init__(model_name=model_name, model_params=model_params, **kwargs)
        # Route an optional http_client kwarg to the matching sync/async
        # client, exactly as the built-in AnthropicLLM does.
        sync_params, async_params = split_http_client_kwargs(kwargs)
        sync_params.setdefault("base_url", self.DEFAULT_ENDPOINT)
        async_params.setdefault("base_url", self.DEFAULT_ENDPOINT)
        self.client = anthropic.Anthropic(**sync_params)
        self.async_client = anthropic.AsyncAnthropic(**async_params)


llm = MyCustomAnthropicLLM(model_name="claude-3-opus-20240229")
llm.invoke("Who is the mother of Paul Atreides?")

Structured output allows LLMs to return responses conforming to a Pydantic model or JSON schema.

V2 Interface (Recommended): For OpenAILLM, VertexAILLM, and AnthropicLLM, pass response_format (a Pydantic model or JSON schema) to the invoke() method when using the V2 interface (passing a list of LLMMessage).

V1 Interface (Legacy): For string-based input, standard JSON mode is supported via constructor parameters (model_params for OpenAI, generation_config for VertexAI). The response_format parameter is not permitted in invoke() for V1.

Provider Specifics:

• OpenAI: Pydantic models must include ConfigDict(extra="forbid") to support strict mode.
• VertexAI: Uses GenerationConfig internally. Additional parameters like temperature can be passed as kwargs to invoke().
• Anthropic: Requires anthropic>=0.77.0 and recent Claude models (Sonnet 4.5 / Opus 4.5+) for native output_config support.

from pydantic import BaseModel, ConfigDict
from neo4j_graphrag.llm import OpenAILLM
from neo4j_graphrag.types import LLMMessage

class Person(BaseModel):
    model_config = ConfigDict(extra="forbid")  # Required for OpenAI structured output
    name: str
    age: int
    occupation: str

lm = OpenAILLM(model_name="gpt-5-mini")

# V2: Pass response_format to invoke()
messages = [LLMMessage(role="user", content="Extract: John is a 30 year old engineer.")]
response = lm.invoke(messages, response_format=Person, temperature=0)
person = Person.model_validate_json(response.content)  # {"name": "John", "age": 30, ...}

# V1: Use constructor parameters for standard JSON mode
lm_v1 = OpenAILLM(
    model_name="gpt-5-mini",
    model_params={"response_format": {"type": "json_object"}, "temperature": 0}
)
response_v1 = lm_v1.invoke("Extract person in JSON format: John is 30 years old.")

The ToolsRetriever uses an LLM to intelligently select and execute one or more tools based on a user query. This is ideal for complex queries that require both vector similarity and structured Cypher lookups.

Tool Creation Patterns:

1. Converting Retrievers: Use retriever.convert_to_tool(name=..., description=...). This automatically infers parameters from the retriever's signature.
2. Custom Tools: Inherit from the Tool class and define parameters using ObjectParameter, StringParameter, etc.

Customizing Selection Logic: You can provide a system_instruction to guide the LLM on when to use specific tools.

from neo4j_graphrag.retrievers import ToolsRetriever, VectorRetriever, Text2CypherRetriever
from neo4j_graphrag.llm import OpenAILLM
from neo4j_graphrag.embeddings import OpenAIEmbeddings

# 1. Setup components
llm = OpenAILLM(model_name="gpt-5")
embedder = OpenAIEmbeddings(model="text-embedding-3-large")

# 2. Create retrievers
vector_retriever = VectorRetriever(driver=driver, index_name="vector-index", embedder=embedder)
text2cypher_retriever = Text2CypherRetriever(driver=driver, llm=llm)

# 3. Convert to tools
vector_tool = vector_retriever.convert_to_tool(
    name="vector_search",
    description="Search for similar documents using vector similarity",
)
cypher_tool = text2cypher_retriever.convert_to_tool(
    name="cypher_search",
    description="Generate and execute Cypher queries for structured data retrieval",
)

# 4. Initialize ToolsRetriever
tools_retriever = ToolsRetriever(
    driver=driver,
    llm=llm,
    tools=[vector_tool, cypher_tool],
)

# 5. Use in GraphRAG pipeline
from neo4j_graphrag.generation import GraphRAG
rag = GraphRAG(retriever=tools_retriever, llm=llm)
response = rag.search("What movies did Tom Hanks act in and what are their plots?")
print(response.answer)

To create a custom LLM, subclass neo4j_graphrag.llm.LLMBase. LLMBase combines LLMInterface (string input) and LLMInterfaceV2 (list of LLMMessage input). You must implement both invoke and ainvoke methods, which should branch based on whether the input is a str or a List[LLMMessage].

from typing import Any, List, Optional, Type, Union
import ollama
from pydantic import BaseModel
from neo4j_graphrag.llm import LLMBase, LLMResponse
from neo4j_graphrag.message_history import MessageHistory
from neo4j_graphrag.types import LLMMessage

class MyOllamaLLM(LLMBase):

    def invoke(
        self,
        input: Union[str, List[LLMMessage]],
        message_history=None,
        system_instruction=None,
        response_format=None,
        **kwargs: Any,
    ) -> LLMResponse:
        if isinstance(input, str):
            messages = [{"role": "user", "content": input}]
        else:
            messages = list(input)
        response = ollama.chat(model=self.model_name, messages=messages)
        return LLMResponse(content=response["message"]["content"])

    async def ainvoke(
        self,
        input: Union[str, List[LLMMessage]],
        message_history=None,
        system_instruction=None,
        response_format=None,
        **kwargs: Any,
    ) -> LLMResponse:
        return self.invoke(input)  # TODO: implement with ollama.AsyncClient

# retriever = ...

lm = MyOllamaLLM("llama3:8b")

rag = GraphRAG(retriever=retriever, llm=lm)
query_text = "How do I do similarity search in Neo4j?"
response = rag.search(query_text=query_text, retriever_config={"top_k": 5})
print(response.answer)

To perform a GraphRAG query, you need to coordinate three main components: a Neo4j driver, a Retriever, and an LLM. The GraphRAG class orchestrates these components to search the graph and generate answers.

Core Components:

1. Neo4j Driver: Connects to your database.
2. Retriever: Fetches relevant context from the graph (e.g., VectorRetriever). Requires an Embedder to convert text queries into vectors.
3. LLM: Generates the final answer. The package's LLM interface is compatible with LangChain.

Note: Using OpenAILLM requires the openai Python client. Install it using: pip install "neo4j_graphrag[openai]".

from neo4j import GraphDatabase
from neo4j_graphrag.retrievers import VectorRetriever
from neo4j_graphrag.llm import OpenAILLM
from neo4j_graphrag.generation import GraphRAG
from neo4j_graphrag.embeddings import OpenAIEmbeddings

# 1. Neo4j driver
URI = "neo4j://localhost:7687"
AUTH = ("neo4j", "password")
INDEX_NAME = "index-name"

driver = GraphDatabase.driver(URI, auth=AUTH)

# 2. Retriever
# Create Embedder object, needed to convert the user question (text) to a vector
embedder = OpenAIEmbeddings(model="text-embedding-3-large")

# Initialize the retriever
retriever = VectorRetriever(driver, INDEX_NAME, embedder)

# 3. LLM
# Note: the OPENAI_API_KEY must be in the env vars
llm = OpenAILLM(model_name="gpt-5", model_params={"temperature": 0})

# Initialize the RAG pipeline
rag = GraphRAG(retriever=retriever, llm=llm)

# Query the graph
query_text = "How do I do similarity search in Neo4j?"
response = rag.search(query_text=query_text, retriever_config={"top_k": 5})
print(response.answer)