Concepts

Understanding these concepts allow you to use GraphRAGZen more efficiently

If you are not familiar with Graph RAG yet, I found this to be a good write-up

Entity extraction

RAG relies on retrieving documents related to a query and adding those to the query to create a final prompt to be send to an LLM.

GraphRAG differs by first retrieving concepts and relations that are extracted from the documents, and adds the relevant concepts and relations to the query to create a final prompt. These concepts and relations contain denser information and can handle more complex queries.

A graph consists of entities. Entities are the nodes (concepts) and edges (relationships between the concepts). These are extracted from the documents in advance, not during query-time.

To extract a graph, documents are fed into an LLM with a prompt that asks it to:

  1. Extract the entities in the document

  2. For each entity say if it’s a node or and relationship (edge)

  3. Give the entity a name

  4. Give each entity a category (e.g. concept, location or organization)

  5. Give the entity a description

  6. Format this all in a json string

These steps are in a single prompt. GraphRAGZen does query the LLM a few times per document to check if all entities have been extracted, so multple json strings can be returned per document.

see graphragzen.entity_extraction.extract_entities.extract_raw_entities()

The output strings are parsed to an actual graph by simply loading the json strings, doing some simple checks on the contents, and using is as input to networkx.

see graphragzen.entity_extraction.extract_entities.raw_entities_to_graph()

""""
[
    {{
        "type": "node",
        "name": "WASHINGTON",
        "category": "LOCATION",
        "description": "Washington is a location where communications are being received, indicating its importance in the decision-making process."
    }},
    {{
        "type": "node",
        "name": "OPERATION: DULCE",
        "category": "MISSION",
        "description": "Operation: Dulce is described as a mission that has evolved to interact and prepare, indicating a significant shift in objectives and activities."
    }},
    {{
        "type": "node",
        "name": "THE TEAM",
        "category": "ORGANIZATION",
        "description": "The team is portrayed as a group of individuals who have transitioned from passive observers to active participants in a mission, showing a dynamic change in their role."
    }},
    {{
        "type": "edge",
        "source": "THE TEAM",
        "target": "WASHINGTON",
        "descripton": "The team receives communications from Washington, which influences their decision-making process.",
        "weight": 1.0
    }},
    {{
        "type": "edge",
        "source": "THE TEAM",
        "target": "OPERATION: DULCE",
        "descripton": "The team is directly involved in Operation: Dulce, executing its evolved objectives and activities.",
        "weight": 1.0
    }}
]"""

The prompt used for extraction can be found here Entity extraction

document size

While extracting a graph, we could feed each document one at a time into the LLM and ask it to extract the entities. However, if documents are large we run the risk of exceeding the LLM context size.

To overcome this we first split each document into smaller chunks and extract the entities from each chunk. It’s a good idea to have overlap between the chunks so no entities spanning two chunks are missed.

see graphragzen.preprocessing.preprocess.chunk_documents()

Node merging

It’s not unlikely that the LLM finds nodes in separate documents that are synonymous. For instance, it might find ‘Pierce Brosnan’ and ‘pierce_brosnan’.

GraphRAGZen can find and merge nodes that are very similar by text embedding the node name and selected features (e.g. the node description), and merging nodes that have very similar embeddings.

See graphragzen.merge.merge_nodes.merge_similar_graph_nodes()

Feature merging

When entities are extracted from the documents the same entity (node or edge) can be found multiple times in different documents. This means that the same entity will have multiple versions of it’s features (e.g. multiple descriptions).

It’s a good idea to consolidate these multiplicated features for the final graph.

GraphRAGZen can

  • Summarize a list of features

  • Return the most occuring from a list

  • Average a list of features if it’s numeric

For descriptions we can ask the LLM to summarize them, edge weights we can simply average, etc.

See graphragzen.merge.merge_features.merge_graph_features()

Clustering

There are latent structures to Graphs that can be usefull when extracting information from them.
For instance, in a graph about vehicles, part of the graph might talk about cars and another about bicycles. When we want information about cars we can safely presume to start looking at the car part.

When GraphRAGZen first creates a graph these clusters are not yet known.
We don’t want to go through the graph manually and assign each node to a topic. We rather use unsupervised clustering for this, specifically the leiden algorithm see see graphragzen.clustering.leiden.leiden()

This algorithm only assigns a number to each node, indicating to which cluster it belongs. The semantic topic of each cluster still needs to be extracted. This can be done using graphragzen.clustering.describe.describe_clusters()

Text Embedding

In order to find the part of the graph that’s relevant to a query text embeddings are used.

Any feature of a node and edge that is text can be embedded, and during querying these embeddings are compared to the query embedding. The nodes and edges coupled to the k embeddings closest to the query embedding are used to build the final context that’s inject into the prompt.

Although any feature of a node and edge that is text can be embedded, for retrieving the relevant entities often the embeddings of entity descriptions are used.

See graphragzen.text_embedding.embed.embed_graph_features() for embedding entity features.

Vector Database

After a graph is created, and before querying, the vectors of the entity feature embedding are stored in a vector database. GraphRAGZen implements local Qdrant as a vector database backend out of the box, but any backend can be used.

In order to use your own backend, simply make a class that inherits from graphragzen.text_embedding.vector_databases.VectorDatabase and implement the methods that are @abstractmethods in the VectorDatabase class (click source to see the methods you’ll need to implement)

LLM interaction

GraphRAGZen uses an LLM for the following:

  • Creating custom prompts for extracting graphs. These prompts are specific to the domain of your documents .

  • Extracting a graph from your documents.

  • Summarizing a node or edge feature that has multiple descriptions of that feature.
    (this happens when a node or edge is found multiple times in your documents. The features assigned to this entity are concatenated when creating the graph, where the entity exists only once.)

  • Creating descriptions of graph clusters.

  • Querying while adding context from the graph.

Two methods are supported to interact with an LLM:

  1. By loading a model locally in-memory.

  2. With an LLM running on a server through an openAI API compatible endpoint.

A server can be remote or deployed locally depending on your own preference.

Loading a model in-memory uses llama-cpp-python and is unlikely to use your GPU unless configured well. Thus, using locally in-memory is good for development and testing, but for production deployment it is recommended to communicate with an LLM that is properly set-up on a server.

These examples show how the different ways to initiate interaction with an LLM

from graphragzen.llm import OpenAICompatibleClient
from graphragzen.llm import Phi35MiniGGUF


# Using OpenAI's API
llm = OpenAICompatibleClient(
    api_key_env_variable = "OPENAI_API_KEY"  # the env variable, not the actual key!
    model_name="gpt-4o-mini",
    context_size = 32768,
    # Caches responses locally, saving time and OpenAI credits if the same request is made twice
    use_cache=True,
    # Append the cache to a file on disk so it can be re-used between runs.
    cache_persistent=True,
    # The file to store the persistent cache.
    persistent_cache_file="./OpenAI_persistent_cache.yaml"
)

# If you're running your own server running an LLM
llm = OpenAICompatibleClient(
    base_url = "http://localhost:8081",
    model_name = "joeAI/phi3.5:latest"
    context_size = 32768,
    # Caches responses locally, saving time and server load if the same request is made twice
    use_cache=True,
    # Append the cache to a file on disk so it can be re-used between runs.
    cache_persistent=True,
    # The file to store the persistent cache.
    persistent_cache_file="./phi35_mini_instruct_server_persistent_cache.yaml"
)

# Load model locally in-memory using llama-cpp-python
llm = Phi35MiniGGUF(
    model_storage_path="path/to/Phi-3.5-mini-instruct-Q4_K_M.gguf",
    tokenizer_URI="microsoft/Phi-3.5-mini-instruct",
    context_size=32786,
    persistent_cache_file="./phi35_mini_instruct_persistent_cache.yaml",
)

Which model to use

The latest OpenAI models give very good results, but might be costly; even with a moderate amount of documents the number of LLM calls to create a graph get large.

For choosing an open-source model, I found Phi 3.5 mini instruct to give good results in my tests. It’s a relatively small model so deployment should be easy and inference fast. When querying your knowledge graph you can switch to a larger, more capable model if desired.

Though Phi 3.5 mini instruct worked well in my tests, the domain of your documents might show different results. I would advice to extract entities from a small set of your documents, check if the extraction makes sense, and try a different model if it doens’t. Pay attention that not just quality nodes are extracted, but also a good amount of edges.

Phi 3.5 mini instruct Q4 K M

Gemma 2 2B it Q4 M

Gemma 2 9B it Q4 XS

Async LLM calls

When interacting with an LLM on a server, aync LLM calls can significantly speed up the graph generation process.

Functions that make LLM calls have a async_llm_calls parameter that, when set to True, will call the LLM async.

The follow functions support this feature:

Of the LLM classes, only the OpenAICompatibleClient has async implemented, since it’s the only class that interacts with an LLM on a server.

When calling it directly for text completion, i.e. llm(‘input text to complete’), the class checks if it is called in an async context and calls the server async or sync accordingly.

For chat functionality there is an async version, see graphragzen.llm.openAI_API_client.OpenAICompatibleClient.a_run_chat()

Implementing your own local LLM class

If the server does not have an OpenAI API compatible endpoint, or you want to load an LLM locally without using llama-cpp-python, you can implement a custom LLM class.

GraphRAGZen expects certain methods when calling an LLM. The abstract base class LLM defines the required methods using @abstractmethod; you should inherit from this class when implementing your own LLM implementation.

See graphragzen.llm.base_llm.LLM (click source)

Structured Output

GraphRAGZen relies heavily on valid json strings to be produced by the LLM to extract information from the provided documents.

The default prompts asks the LLM to produce a json string, but the chances of this being a valid json string can be significantly increased by forcing the LLM to adhere it’s output to a schema

JSON Schema

Many inference solutions allow a schema to be submitted with the prompt to the LLM. The engine than tries to restrict the token generation such that the output adheres to the schema.

See here for examples using the openAI API

In GraphRAGZen both locally loaded LLM’s and the API client can be provided with a pydantic class as an output structure for the LLM to adhere to.

Note

The output structure should NOT be passed as an instance of the pydantic model, just the reference.

Correct = LLM(“some text”, output_structure=MyPydanticModel)

Wrong = LLM(“some text”, output_structure=MyPydanticModel())

For initializing interaction with an LLM see LLM interaction

Extract raw entities using a custom output structure: (By default graphragzen.entity_extraction.extract_raw_entities() already uses a build-in output structure see graphragzen.entity_extraction.llm_output_structures.ExtractedEntities)

from typing import List

from graphragzen import entity_extraction
from pydantic import BaseModel


class ExtractedNode(BaseModel):
    type: str = "node"
    name: str
    category: str
    description: str
    relevance: float
    source_sentence: str

class ExtractedEdge(BaseModel):
    type: str = "edge"
    source: str
    target: str
    description: str
    weight: float
    source_sentence: str


class ExtractedEntities(BaseModel):
    extracted_nodes: List[ExtractedNode]
    extracted_edges: List[ExtractedEdge]


raw_entities = entity_extraction.extract_raw_entities(
    chunked_documents, llm, output_structure: ExtractedEntities
)

Prompt tuning

GraphRAGZen comes with default prompts to extract entities from documents, summarize features, etc.

Although these can be used out-of-the-box to extract entities from documents and create a graph, they are not tuned to the documents.
Higher quality graphs could be obtained by making these prompts more relevant to the domain of the documents.

Your can provided you own custom prompts, but GraphRAGZen does come with functions to create these prompts from a sample of your documents.

These functions rely on their own prompts that ask the LLM to look at the documents and:

  • Create a domain for the documents

  • Create a persona that is an expert in the create domain

  • Define which entity types are present in the documents (e.g. person, location, school of thought)

  • Create some document->entities extracted examples

All of this information is then merged into a prompt that can be used to extract entities.

A similar method is used to create a description summarization prompt.

The default prompts: Default Prompts

The prompts used to create new prompts: Prompt tuning Prompts

See this example on how to create prompts specific to the domain of your documents (click on source) graphragzen.examples.autotune_custom_prompts.create_custom_prompts()