Blog

The tradeoffs between building bespoke code-based agents and the major agent frameworks.

Thanks to John Gilhuly for his contributions to this piece.

Agents are having a moment. With multiple new frameworks and fresh investment in the space, modern AI agents are overcoming shaky origins to rapidly supplant RAG as an implementation priority. So will 2024 finally be the year that autonomous AI systems that can take over writing our emails, booking flights, talking to our data, or seemingly any other task?

Maybe, but much work remains to get to that point. Any developer building an agent must not only choose foundations — which model, use case, and architecture to use — but also which framework to leverage. Do you go with the long-standing LangGraph, or the newer entrant LlamaIndex Workflows? Or do you go the traditional route and code the whole thing yourself?

This post aims to make that choice a bit easier. Over the past few weeks, I built the same agent in major frameworks to examine some of the strengths and weaknesses of each at a technical level. All of the code for each agent is available in this repo.

Background on the Agent Used for Testing

The agent used for testing includes function calling, multiple tools or skills, connections to outside resources, and shared state or memory.

The agent has the following capabilities:

1. Answering questions from a knowledge base
2. Talking to data: answering questions about telemetry data of an LLM application
3. Analyzing data: analyzing higher-level trends and patterns in retrieved telemetry data

In order to accomplish these, the agent has three starting skills: RAG with product documentation, SQL generation on a trace database, and data analysis. A simple gradio-powered interface is used for the agent UI, with the agent itself structured as a chatbot.

Code-Based Agent (No Framework)

The first option you have when developing an agent is to skip the frameworks entirely and build the agent fully yourself. When embarking on this project, this was the approach I started with.

Pure Code Architecture

The code-based agent below is made up of an OpenAI-powered router that uses function calling to select the right skill to use. After that skill completes, it returns back to the router to either call another skill or respond to the user.

The agent keeps an ongoing list of messages and responses that is passed fully into the router on each call to preserve context through cycles.

def router(messages):
    if not any(
        isinstance(message, dict) and message.get("role") == "system" for message in messages
    ):
        system_prompt = {"role": "system", "content": SYSTEM_PROMPT}
        messages.append(system_prompt)

    response = client.chat.completions.create(
        model="gpt-4o",
        messages=messages,
        tools=skill_map.get_combined_function_description_for_openai(),
    )

    messages.append(response.choices[0].message)
    tool_calls = response.choices[0].message.tool_calls
    if tool_calls:
        handle_tool_calls(tool_calls, messages)
        return router(messages)
    else:
        return response.choices[0].message.content

The skills themselves are defined in their own classes (e.g. GenerateSQLQuery) that are collectively held in a SkillMap. The router itself only interacts with the SkillMap, which it uses to load skill names, descriptions, and callable functions. This approach means that adding a new skill to the agent is as simple as writing that skill as its own class, then adding it to the list of skills in the SkillMap. The idea here is to make it easy to add new skills without disturbing the router code.

class SkillMap:
    def __init__(self):
        skills = [AnalyzeData(), GenerateSQLQuery()]

        self.skill_map = {}
        for skill in skills:
            self.skill_map[skill.get_function_name()] = (
                skill.get_function_dict(),
                skill.get_function_callable(),
            )

    def get_function_callable_by_name(self, skill_name) -> Callable:
        return self.skill_map[skill_name][1]

    def get_combined_function_description_for_openai(self):
        combined_dict = []
        for _, (function_dict, _) in self.skill_map.items():
            combined_dict.append(function_dict)
        return combined_dict

    def get_function_list(self):
        return list(self.skill_map.keys())

    def get_list_of_function_callables(self):
        return [skill[1] for skill in self.skill_map.values()]

    def get_function_description_by_name(self, skill_name):
        return str(self.skill_map[skill_name][0]["function"])

Overall, this approach is fairly straightforward to implement but comes with a few challenges.

Challenges with Pure Code Agents

The first difficulty lies in structuring the router system prompt. Often, the router in the example above insisted on generating SQL itself instead of delegating that to the right skill. If you’ve ever tried to get an LLM not to do something, you know how frustrating that experience can be; finding a working prompt took many rounds of debugging. Accounting for the different output formats from each step was also tricky. Since I opted not to use structured outputs, I had to be ready for multiple different formats from each of the LLM calls in my router and skills.

Benefits of a Pure Code Agent

A code-based approach provides a good baseline and starting point, offering a great way to learn how agents work without relying on canned agent tutorials from prevailing frameworks. Although convincing the LLM to behave can be challenging, the code structure itself is simple enough to use and might make sense for certain use cases (more in the analysis section below).

LangGraph

LangGraph is one of the longest-standing agent frameworks, first releasing in January 2024. The framework is built to address the acyclic nature of existing pipelines and chains by adopting a Pregel graph structure instead. LangGraph makes it easier to define loops in your agent by adding the concepts of nodes, edges, and conditional edges to traverse a graph. LangGraph is built on top of LangChain, and uses the objects and types from that framework.

LangGraph Architecture

The LangGraph agent looks similar to the code-based agent on paper, but the code behind it is drastically different. LangGraph still uses a “router” technically, in that it calls OpenAI with functions and uses the response to continue to a new step. However the way the program moves between skills is controlled completely differently.

tools = [generate_and_run_sql_query, data_analyzer]
model = ChatOpenAI(model="gpt-4o", temperature=0).bind_tools(tools)

def create_agent_graph():
    workflow = StateGraph(MessagesState)

    tool_node = ToolNode(tools)
    workflow.add_node("agent", call_model)
    workflow.add_node("tools", tool_node)

    workflow.add_edge(START, "agent")
    workflow.add_conditional_edges(
        "agent",
        should_continue,
    )
    workflow.add_edge("tools", "agent")

    checkpointer = MemorySaver()
    app = workflow.compile(checkpointer=checkpointer)
    return app

The graph defined here has a node for the initial OpenAI call, called “agent” above, and one for the tool handling step, called “tools.” LangGraph has a built-in object called ToolNode that takes a list of callable tools and triggers them based on a ChatMessage response, before returning to the “agent” node again.

def should_continue(state: MessagesState):
    messages = state["messages"]
    last_message = messages[-1]
    if last_message.tool_calls:
        return "tools"
    return END

def call_model(state: MessagesState):
    messages = state["messages"]
    response = model.invoke(messages)
    return {"messages": [response]}

After each call of the “agent” node (put another way: the router in the code-based agent), the should_continue edge decides whether to return the response to the user or pass on to the ToolNode to handle tool calls.

Throughout each node, the “state” stores the list of messages and responses from OpenAI, similar to the code-based agent’s approach.

Challenges with LangGraph

Most of the difficulties with LangGraph in the example stem from the need to use Langchain objects for things to flow nicely.

Challenge #1: Function Call Validation

In order to use the ToolNode object, I had to refactor most of my existing Skill code. The ToolNode takes a list of callable functions, which originally made me think I could use my existing functions, however things broke down due to my function parameters.

The skills were defined as classes with a callable member function, meaning they had “self” as their first parameter. GPT-4o was smart enough to not include the “self” parameter in the generated function call, however LangGraph read this as a validation error due to a missing parameter.

This took hours to figure out, because the error message instead marked the third parameter in the function (“args” on the data analysis skill) as the missing parameter:

pydantic.v1.error_wrappers.ValidationError: 1 validation error for data_analysis_toolSchema
args field required (type=value_error.missing)

It is worth mentioning that the error message originated from Pydantic, not from LangGraph.

I eventually bit the bullet and redefined my skills as basic methods with Langchain’s @tool decorator, and was able to get things working.

@tool
def generate_and_run_sql_query(query: str):
    """Generates and runs an SQL query based on the prompt.

    Args:
        query (str): A string containing the original user prompt.

    Returns:
        str: The result of the SQL query.
    """

Challenge #2: Debugging

As mentioned, debugging in a framework is difficult. This primarily comes down to confusing error messages and abstracted concepts that make it harder to view variables.

The abstracted concepts primarily show up when trying to debug the messages being sent around the agent. LangGraph stores these messages in state[“messages”]. Some nodes within the graph pull from these messages automatically, which can make it difficult to understand the value of messages when they are accessed by the node.

LangGraph Benefits

One of the main benefits of LangGraph is that it’s easy to work with. The graph structure code is clean and accessible. Especially if you have complex node logic, having a single view of the graph makes it easier to understand how the agent is connected together. LangGraph also makes it straightforward to convert an existing application built in LangChain.

Takeaway

If you use everything in the framework, LangGraph works cleanly; if you step outside of it, prepare for some debugging headaches.

LlamaIndex Workflows

Workflows is a newer entrant into the agent framework space, premiering earlier this summer. Like LangGraph, it aims to make looping agents easier to build. Workflows also has a particular focus on running asynchronously.

Some elements of Workflows seem to be in direct response to LangGraph, specifically its use of events instead of edges and conditional edges. Workflows use steps (analogous to nodes in LangGraph) to house logic, and emitted and received events to move between steps.

The structure above looks similar to the LangGraph structure, save for one addition. I added a setup step to the Workflow to prepare the agent context, more on this below. Despite the similar structure, there is very different code powering it.

Workflows Architecture

The code below defines the Workflow structure. Similar to LangGraph, this is where I prepared the state and attached the skills to the LLM object.

class AgentFlow(Workflow):
    def __init__(self, llm, timeout=300):
        super().__init__(timeout=timeout)
        self.llm = llm
        self.memory = ChatMemoryBuffer(token_limit=1000).from_defaults(llm=llm)
        self.tools = []
        for func in skill_map.get_function_list():
            self.tools.append(
                FunctionTool(
                    skill_map.get_function_callable_by_name(func),
                    metadata=ToolMetadata(
                        name=func, description=skill_map.get_function_description_by_name(func)
                    ),
                )
            )

    @step
    async def prepare_agent(self, ev: StartEvent) -> RouterInputEvent:
        user_input = ev.input
        user_msg = ChatMessage(role="user", content=user_input)
        self.memory.put(user_msg)

        chat_history = self.memory.get()
        return RouterInputEvent(input=chat_history)

This is also where I define an extra step, “prepare_agent”. This step creates a ChatMessage from the user input and adds it to the workflow memory. Splitting this out as a separate step means that we do return to it as the agent loops through steps, which avoids repeatedly adding the user message to the memory.

In the LangGraph case, I accomplished the same thing with a run_agent method that lived outside the graph. This change is mostly stylistic, however it’s cleaner in my opinion to house this logic with the Workflow and graph as we’ve done here.

With the Workflow set up, I then defined the routing code:

@step
async def router(self, ev: RouterInputEvent) -> ToolCallEvent | StopEvent:
    messages = ev.input

    if not any(
        isinstance(message, dict) and message.get("role") == "system" for message in messages
    ):
        system_prompt = ChatMessage(role="system", content=SYSTEM_PROMPT)
        messages.insert(0, system_prompt)

    with using_prompt_template(template=SYSTEM_PROMPT, version="v0.1"):
        response = await self.llm.achat_with_tools(
            model="gpt-4o",
            messages=messages,
            tools=self.tools,
        )

    self.memory.put(response.message)

    tool_calls = self.llm.get_tool_calls_from_response(response, error_on_no_tool_call=False)
    if tool_calls:
        return ToolCallEvent(tool_calls=tool_calls)
    else:
        return StopEvent(result=response.message.content)

And the tool call handling code:

@step
async def tool_call_handler(self, ev: ToolCallEvent) -> RouterInputEvent:
    tool_calls = ev.tool_calls

    for tool_call in tool_calls:
        function_name = tool_call.tool_name
        arguments = tool_call.tool_kwargs
        if "input" in arguments:
            arguments["prompt"] = arguments.pop("input")

        try:
            function_callable = skill_map.get_function_callable_by_name(function_name)
        except KeyError:
            function_result = "Error: Unknown function call"

        function_result = function_callable(arguments)
        message = ChatMessage(
            role="tool",
            content=function_result,
            additional_kwargs={"tool_call_id": tool_call.tool_id},
        )

        self.memory.put(message)

    return RouterInputEvent(input=self.memory.get())

Both of these look more similar to the code-based agent than the LangGraph agent. This is mainly because Workflows keeps the conditional routing logic in the steps as opposed to in conditional edges — lines 18–24 were a conditional edge in LangGraph, whereas now they are just part of the routing step — and the fact that LangGraph has a ToolNode object that does just about everything in the tool_call_handler method automatically.

Moving past the routing step, one thing I was very happy to see is that I could use my SkillMap and existing skills from my code-based agent with Workflows. These required no changes to work with Workflows, which made my life much easier.

Challenges with Workflows

Challenge #1: Sync vs Async

While asynchronous execution is preferable for a live agent, debugging a synchronous agent is much easier. Workflows is designed to work asynchronously, and trying to force synchronous execution was very difficult.

I initially thought I would just be able to remove the “async” method designations and switch from “achat_with_tools” to “chat_with_tools”. However, since the underlying methods within the Workflow class were also marked as asynchronous, it was necessary to redefine those in order to run synchronously. I ended up sticking to an asynchronous approach, but this didn’t make debugging more difficult.

Challenge #2: Pydantic Validation Errors

In a repeat of the woes with LangGraph, similar problems emerged around confusing Pydantic validation errors on skills. Fortunately, these were easier to address this time since Workflows was able to handle member functions just fine. I ultimately just ended up having to be more prescriptive in creating LlamaIndex FunctionTool objects for my skills:

for func in skill_map.get_function_list(): 
            self.tools.append(FunctionTool(
                skill_map.get_function_callable_by_name(func), 
                metadata=ToolMetadata(name=func, description=skill_map.get_function_description_by_name(func))))

Excerpt from AgentFlow.__init__ that builds FunctionTools

Benefits of Workflows

I had a much easier time building the Workflows agent than I did the LangGraph agent, mainly because Workflows still required me to write routing logic and tool handling code myself instead of providing built-in functions. This also meant that my Workflow agent looked extremely similar to my code-based agent.

The biggest difference came in the use of events. I used two custom events to move between steps in my agent:

class ToolCallEvent(Event):
    tool_calls: list[ToolSelection]

class RouterInputEvent(Event):
    input: list[ChatMessage]

The emitter-receiver, event-based architecture took the place of directly calling some of the methods in my agent, like the tool call handler.

If you have more complex systems with multiple steps that are triggering asynchronously and might emit multiple events, this architecture becomes very helpful to manage that cleanly.

Other benefits of Workflows include the fact that it is very lightweight and doesn’t force much structure on you (aside from the use of certain LlamaIndex objects) and that its event-based architecture provides a helpful alternative to direct function calling — especially for complex, asynchronous applications.

Comparing Frameworks

Looking across the three approaches, each one has its benefits.

The no framework approach is the simplest to implement. Because any abstractions are defined by the developer (i.e. SkillMap object in the above example), keeping various types and objects straight is easy. The readability and accessibility of the code entirely comes down to the individual developer however, and it’s easy to see how increasingly complex agents could get messy without some enforced structure.

LangGraph provides quite a bit of structure, which makes the agent very clearly defined. If a broader team is collaborating on an agent, this structure would provide a helpful way of enforcing an architecture. LangGraph also might provide a good starting point with agents for those not as familiar with the structure. There is a tradeoff, however — since LangGraph does quite a bit for you, it can lead to headaches if you don’t fully buy into the framework; the code may be very clean, but you may pay for it with more debugging.

Workflows falls somewhere in the middle. The event-based architecture might be extremely helpful for some projects, and the fact that less is required in terms of using of LlamaIndex types provides greater flexibility for those not be fully using the framework across their application.

Ultimately, the core question may just come down to “are you already using LlamaIndex or LangChain to orchestrate your application?” LangGraph and Workflows are both so entwined with their respective underlying frameworks that the additional benefits of each agent-specific framework might not cause you to switch on merit alone.

The pure code approach will likely always be an attractive option. If you have the rigor to document and enforce any abstractions created, then ensuring nothing in an external framework slows you down is easy.

Key Questions To Help In Choosing An Agent Framework

Of course, “it depends” is never a satisfying answer. These three questions should help you decide which framework to use in your next agent project.

Are you already using LlamaIndex or LangChain for significant pieces of your project?

If yes, explore that option first.

Are you familiar with common agent structures, or do you want something telling you how you should structure your agent?

If you fall into the latter group, try Workflows. If you really fall into the latter group, try LangGraph.

Has your agent been built before?

One of the framework benefits is that there are many tutorials and examples built with each. There are far fewer examples of pure code agents to build from.

Conclusion

Picking an agent framework is just one choice among many that will impact outcomes in production for generative AI systems. As always, it pays to have robust guardrails and LLM tracing in place — and to be agile as new agent frameworks, research, and models upend established techniques.

---

Choosing Between LLM Agent Frameworks was originally published in Towards Data Science on Medium, where people are continuing the conversation by highlighting and responding to this story.

Elevating RAG accuracy and performance by structuring long documents into explorable graphs and implementing graph-based agent systems

Large Language Models (LLMs) are great at traditional NLP tasks like summarization and sentiment analysis but the stronger models also demonstrate promising reasoning abilities. LLM reasoning is often understood as the ability to tackle complex problems by formulating a plan, executing it, and assessing progress at each step. Based on this evaluation, they can adapt by revising the plan or taking alternative actions. The rise of agents is becoming an increasingly compelling approach to answering complex questions in RAG applications.

In this blog post, we’ll explore the implementation of the GraphReader agent. This agent is designed to retrieve information from a structured knowledge graph that follows a predefined schema. Unlike the typical graphs you might see in presentations, this one is closer to a document or lexical graph, containing documents, their chunks, and relevant metadata in the form of atomic facts.

The image above illustrates a knowledge graph, beginning at the top with a document node labeled Joan of Arc. This document is broken down into text chunks, represented by numbered circular nodes (0, 1, 2, 3), which are connected sequentially through NEXT relationships, indicating the order in which the chunks appear in the document. Below the text chunks, the graph further breaks down into atomic facts, where specific statements about the content are represented. Finally, at the bottom level of the graph, we see the key elements, represented as circular nodes with topics like historical icons, Dane, French nation, and France. These elements act as metadata, linking the facts to the broader themes and concepts relevant to the document.

Once we have constructed the knowledge graph, we will follow the implementation provided in the GraphReader paper.

The agent exploration process involves initializing the agent with a rational plan and selecting initial nodes to start the search in a graph. The agent explores these nodes by first gathering atomic facts, then reading relevant text chunks, and updating its notebook. The agent can decide to explore more chunks, neighboring nodes, or terminate based on gathered information. When the agent decided to terminate, the answer reasoning step is executed to generate the final answer.

In this blog post, we will implement the GraphReader paper using Neo4j as the storage layer and LangChain in combination with LangGraph to define the agent and its flow.

The code is available on GitHub.

Environment Setup

You need to setup a Neo4j to follow along with the examples in this blog post. The easiest way is to start a free instance on Neo4j Aura, which offers cloud instances of Neo4j database. Alternatively, you can also setup a local instance of the Neo4j database by downloading the Neo4j Desktop application and creating a local database instance.

The following code will instantiate a LangChain wrapper to connect to Neo4j Database.

os.environ["NEO4J_URI"] = "bolt://localhost:7687"
os.environ["NEO4J_USERNAME"] = "neo4j"
os.environ["NEO4J_PASSWORD"] = "password"

graph = Neo4jGraph(refresh_schema=False)

graph.query("CREATE CONSTRAINT IF NOT EXISTS FOR (c:Chunk) REQUIRE c.id IS UNIQUE")
graph.query("CREATE CONSTRAINT IF NOT EXISTS FOR (c:AtomicFact) REQUIRE c.id IS UNIQUE")
graph.query("CREATE CONSTRAINT IF NOT EXISTS FOR (c:KeyElement) REQUIRE c.id IS UNIQUE")

Additionally, we have also added constraints for the node types we will be using. The constraints ensure faster import and retrieval performance.

Additionally, you will require an OpenAI api key that you pass in the following code:

os.environ["OPENAI_API_KEY"] = getpass.getpass("OpenAI API Key:")

Graph construction

We will be using the Joan of Arc Wikipedia page in this example. We will use LangChain built-in utility to retrieve the text.

wikipedia = WikipediaQueryRun(
    api_wrapper=WikipediaAPIWrapper(doc_content_chars_max=10000)
)
text = wikipedia.run("Joan of Arc")

As mentioned before, the GraphReader agent expects knowledge graph that contains chunks, related atomic facts, and key elements.

First, the document is split into chunks. In the paper they maintained paragraph structure while chunking. However, that is hard to do in a generic way. Therefore, we will use naive chunking here.

Next, each chunk is processed by the LLM to identify atomic facts, which are the smallest, indivisible units of information that capture core details. For instance, from the sentence “The CEO of Neo4j, which is in Sweden, is Emil Eifrem” an atomic fact could be broken down into something like “The CEO of Neo4j is Emil Eifrem.” and “Neo4j is in Sweden.” Each atomic fact is focused on one clear, standalone piece of information.

From these atomic facts, key elements are identified. For the first fact, “The CEO of Neo4j is Emil Eifrem,” the key elements would be “CEO,” “Neo4j,” and “Emil Eifrem.” For the second fact, “Neo4j is in Sweden,” the key elements would be “Neo4j” and “Sweden.” These key elements are the essential nouns and proper names that capture the core meaning of each atomic fact.

The prompt used to extract the graph are provided in the appendix of the paper.

The authors used prompt-based extraction, where you instruct the LLM what it should output and then implement a function that parses the information in a structured manner. My preference for extracting structured information is to use the with_structured_output method in LangChain, which utilizes the tools feature to extract structured information. This way, we can skip defining a custom parsing function.

Here is the prompt that we can use for extraction.

construction_system = """
You are now an intelligent assistant tasked with meticulously extracting both key elements and
atomic facts from a long text.
1. Key Elements: The essential nouns (e.g., characters, times, events, places, numbers), verbs (e.g.,
actions), and adjectives (e.g., states, feelings) that are pivotal to the text’s narrative.
2. Atomic Facts: The smallest, indivisible facts, presented as concise sentences. These include
propositions, theories, existences, concepts, and implicit elements like logic, causality, event
sequences, interpersonal relationships, timelines, etc.
Requirements:
#####
1. Ensure that all identified key elements are reflected within the corresponding atomic facts.
2. You should extract key elements and atomic facts comprehensively, especially those that are
important and potentially query-worthy and do not leave out details.
3. Whenever applicable, replace pronouns with their specific noun counterparts (e.g., change I, He,
She to actual names).
4. Ensure that the key elements and atomic facts you extract are presented in the same language as
the original text (e.g., English or Chinese).
"""

construction_human = """Use the given format to extract information from the 
following input: {input}"""

construction_prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            construction_system,
        ),
        (
            "human",
            (
                "Use the given format to extract information from the "
                "following input: {input}"
            ),
        ),
    ]
)

We have put the instruction in the system prompt, and then in the user message we provide relevant text chunks that need to be processed.

To define the desired output, we can use the Pydantic object definition.

class AtomicFact(BaseModel):
    key_elements: List[str] = Field(description="""The essential nouns (e.g., characters, times, events, places, numbers), verbs (e.g.,
actions), and adjectives (e.g., states, feelings) that are pivotal to the atomic fact's narrative.""")
    atomic_fact: str = Field(description="""The smallest, indivisible facts, presented as concise sentences. These include
propositions, theories, existences, concepts, and implicit elements like logic, causality, event
sequences, interpersonal relationships, timelines, etc.""")

class Extraction(BaseModel):
    atomic_facts: List[AtomicFact] = Field(description="List of atomic facts")

We want to extract a list of atomic facts, where each atomic fact contains a string field with the fact, and a list of present key elements. It is important to add description to each element to get the best results.

Now we can combine it all in a chain.

model = ChatOpenAI(model="gpt-4o-2024-08-06", temperature=0.1)
structured_llm = model.with_structured_output(Extraction)

construction_chain = construction_prompt | structured_llm

To put it all together, we’ll create a function that takes a single document, chunks it, extracts atomic facts and key elements, and stores the results into Neo4j.

async def process_document(text, document_name, chunk_size=2000, chunk_overlap=200):
    start = datetime.now()
    print(f"Started extraction at: {start}")
    text_splitter = TokenTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
    texts = text_splitter.split_text(text)
    print(f"Total text chunks: {len(texts)}")
    tasks = [
        asyncio.create_task(construction_chain.ainvoke({"input":chunk_text}))
        for index, chunk_text in enumerate(texts)
    ]
    results = await asyncio.gather(*tasks)
    print(f"Finished LLM extraction after: {datetime.now() - start}")
    docs = [el.dict() for el in results]
    for index, doc in enumerate(docs):
        doc['chunk_id'] = encode_md5(texts[index])
        doc['chunk_text'] = texts[index]
        doc['index'] = index
        for af in doc["atomic_facts"]:
            af["id"] = encode_md5(af["atomic_fact"])
    # Import chunks/atomic facts/key elements
    graph.query(import_query, 
            params={"data": docs, "document_name": document_name})
    # Create next relationships between chunks
    graph.query("""MATCH (c:Chunk) WHERE c.document_name = $document_name
WITH c ORDER BY c.index WITH collect(c) AS nodes
UNWIND range(0, size(nodes) -2) AS index
WITH nodes[index] AS start, nodes[index + 1] AS end
MERGE (start)-[:NEXT]->(end)
""",
           params={"document_name":document_name})
    print(f"Finished import at: {datetime.now() - start}")

At a high level, this code processes a document by breaking it into chunks, extracting information from each chunk using an AI model, and storing the results in a graph database. Here’s a summary:

1. It splits the document text into chunks of a specified size, allowing for some overlap. The chunk size of 2000 tokens is used by the authors in the paper.
2. For each chunk, it asynchronously sends the text to an LLM for extraction of atomic facts and key elements.
3. Each chunk and fact is given a unique identifier using an md5 encoding function.
4. The processed data is imported into a graph database, with relationships established between consecutive chunks.

We can now run this function on our Joan of Arc text.

await process_document(text, "Joan of Arc", chunk_size=500, chunk_overlap=100)

We used a smaller chunk size because it’s a small document, and we want to have a couple of chunks for demonstration purposes. If you explore the graph in Neo4j Browser, you should see a similar visualization.

At the center of the structure is the document node (blue), which branches out to chunk nodes (pink). These chunk nodes, in turn, are linked to atomic facts (orange), each of which connects to key elements (green).

Let’s examine the constructed graph a bit. We’ll start of by examining the token count distribution of atomic facts.

def num_tokens_from_string(string: str) -> int:
    """Returns the number of tokens in a text string."""
    encoding = tiktoken.encoding_for_model("gpt-4")
    num_tokens = len(encoding.encode(string))
    return num_tokens


atomic_facts = graph.query("MATCH (a:AtomicFact) RETURN a.text AS text")
df = pd.DataFrame.from_records(
    [{"tokens": num_tokens_from_string(el["text"])} for el in atomic_facts]
)

sns.histplot(df["tokens"])

Results

Atomic facts are relatively short, with the longest being only about 50 tokens. Let’s examine a couple to get a better idea.

graph.query("""MATCH (a:AtomicFact) 
RETURN a.text AS text
ORDER BY size(text) ASC LIMIT 3
UNION ALL
MATCH (a:AtomicFact) 
RETURN a.text AS text
ORDER BY size(text) DESC LIMIT 3""")

Results

Some of the shortest facts lack context. For example, the original score and screenplay don’t directly mention which. Therefore, if we processed multiple documents, these atomic facts might be less helpful. This lack of context might be solved with additional prompt engineering.

Let’s also examine the most frequent keywords.

data = graph.query("""
MATCH (a:KeyElement) 
RETURN a.id AS key, 
       count{(a)<-[:HAS_KEY_ELEMENT]-()} AS connections
ORDER BY connections DESC LIMIT 5""")
df = pd.DataFrame.from_records(data)
sns.barplot(df, x='key', y='connections')

Results

Unsurprisingly, Joan of Arc is the most mentioned keyword or element. Following are broad keywords like film, English, and France. I suspect that if we parsed many documents, broad keywords would end up having a lot of connections, which might lead to some downstream problems that aren’t dealt with in the original implementation. Another minor problem is the non-determinism of the extraction, as the results will be slight different on every run.

Additionally, the authors employ key element normalization as described in Lu et al. (2023), specifically using frequency filtering, rule, semantic, and association aggregation. In this implementation, we skipped this step.

GraphReader Agent

We’re ready to implement GraphReader, a graph-based agent system. The agent starts with a couple of predefined steps, followed by the steps in which it can traverse the graph autonomously, meaning the agent decides the following steps and how to traverse the graph.

Here is the LangGraph visualization of the agent we will implement.

The process begins with a rational planning stage, after which the agent makes an initial selection of nodes (key elements) to work with. Next, the agent checks atomic facts linked to the selected key elements. Since all these steps are predefined, they are visualized with a full line.

Depending on the outcome of the atomic fact check, the flow proceeds to either read relevant text chunks or explore the neighbors of the initial key elements in search of more relevant information. Here, the next step is conditional and based on the results of an LLM and is, therefore, visualized with a dotted line.

In the chunk check stage, the LLM reads and evaluates whether the information gathered from the current text chunk is sufficient. Based on this evaluation, the LLM has a few options. It can decide to read additional text chunks if the information seems incomplete or unclear. Alternatively, the LLM may choose to explore neighboring key elements, looking for more context or related information that the initial selection might not have captured. If, however, the LLM determines that enough relevant information has been gathered, it will proceed directly to the answer reasoning step. At this point, the LLM generates the final answer based on the collected information.

Throughout this process, the agent dynamically navigates the flow based on the outcomes of the conditional checks, making decisions on whether to repeat steps or continue forward depending on the specific situation. This provides flexibility in handling different inputs while maintaining a structured progression through the steps.

Now, we’ll go over the steps and implement them using LangGraph abstraction. You can learn more about LangGraph through LangChain’s academy course.

LangGraph state

To build a LangGraph implementation, we start by defining a state passed along the steps in the flow.

class InputState(TypedDict):
    question: str

class OutputState(TypedDict):
    answer: str
    analysis: str
    previous_actions: List[str]

class OverallState(TypedDict):
    question: str
    rational_plan: str
    notebook: str
    previous_actions: Annotated[List[str], add]
    check_atomic_facts_queue: List[str]
    check_chunks_queue: List[str]
    neighbor_check_queue: List[str]
    chosen_action: str

For more advanced use cases, multiple separate states can be used. In our implementation, we have separate input and output states, which define the input and output of the LangGraph, and a separate overall state, which is passed between steps.

By default, the state is overwritten when returned from a node. However, you can define other operations. For example, with the previous_actions we define that the state is appended or added instead of overwritten.

The agent begins by maintaining a notebook to record supporting facts, which are eventually used to derive the final answer. Other states will be explained as we go along.

Let’s move on to defining the nodes in the LangGraph.

Rational plan

In the rational plan step, the agent breaks the question into smaller steps, identifies the key information required, and creates a logical plan. The logical plan allows the agent to handle complex multi-step questions.

While the code is unavailable, all the prompts are in the appendix, so we can easily copy them.

The authors don’t explicitly state whether the prompt is provided in the system or user message. For the most part, I have decided to put the instructions as a system message.

The following code shows how to construct a chain using the above rational plan as the system message.

rational_plan_system = """As an intelligent assistant, your primary objective is to answer the question by gathering
supporting facts from a given article. To facilitate this objective, the first step is to make
a rational plan based on the question. This plan should outline the step-by-step process to
resolve the question and specify the key information required to formulate a comprehensive answer.
Example:
#####
User: Who had a longer tennis career, Danny or Alice?
Assistant: In order to answer this question, we first need to find the length of Danny’s
and Alice’s tennis careers, such as the start and retirement of their careers, and then compare the
two.
#####
Please strictly follow the above format. Let’s begin."""

rational_prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            rational_plan_system,
        ),
        (
            "human",
            (
                "{question}"
            ),
        ),
    ]
)

rational_chain = rational_prompt | model | StrOutputParser()

Now, we can use this chain to define a rational plan node. A node in LangGraph is a function that takes the state as input and updates it as output.

def rational_plan_node(state: InputState) -> OverallState:
    rational_plan = rational_chain.invoke({"question": state.get("question")})
    print("-" * 20)
    print(f"Step: rational_plan")
    print(f"Rational plan: {rational_plan}")
    return {
        "rational_plan": rational_plan,
        "previous_actions": ["rational_plan"],
    }

The function starts by invoking the LLM chain, which produces the rational plan. We do a little printing for debugging and then update the state as the function’s output. I like the simplicity of this approach.

Initial node selection

In the next step, we select the initial nodes based on the question and rational plan. The prompt is the following:

The prompt starts by giving the LLM some context about the overall agent system, followed by the task instructions. The idea is to have the LLM select the top 10 most relevant nodes and score them. The authors simply put all the key elements from the database in the prompt for an LLM to select from. However, I think that approach doesn’t really scale. Therefore, we will create and use a vector index to retrieve a list of input nodes for the prompt.

neo4j_vector = Neo4jVector.from_existing_graph(
    embedding=embeddings,
    index_name="keyelements",
    node_label="KeyElement",
    text_node_properties=["id"],
    embedding_node_property="embedding",
    retrieval_query="RETURN node.id AS text, score, {} AS metadata"
)

def get_potential_nodes(question: str) -> List[str]:
    data = neo4j_vector.similarity_search(question, k=50)
    return [el.page_content for el in data]

The from_existing_graph method pulls the defined text_node_properties from the graph and calculates embeddings where they are missing. Here, we simply embed the id property of KeyElement nodes.

Now let’s define the chain. We’ll first copy the prompt.

initial_node_system = """
As an intelligent assistant, your primary objective is to answer questions based on information
contained within a text. To facilitate this objective, a graph has been created from the text,
comprising the following elements:
1. Text Chunks: Chunks of the original text.
2. Atomic Facts: Smallest, indivisible truths extracted from text chunks.
3. Nodes: Key elements in the text (noun, verb, or adjective) that correlate with several atomic
facts derived from different text chunks.
Your current task is to check a list of nodes, with the objective of selecting the most relevant initial nodes from the graph to efficiently answer the question. You are given the question, the
rational plan, and a list of node key elements. These initial nodes are crucial because they are the
starting point for searching for relevant information.
Requirements:
#####
1. Once you have selected a starting node, assess its relevance to the potential answer by assigning
a score between 0 and 100. A score of 100 implies a high likelihood of relevance to the answer,
whereas a score of 0 suggests minimal relevance.
2. Present each chosen starting node in a separate line, accompanied by its relevance score. Format
each line as follows: Node: [Key Element of Node], Score: [Relevance Score].
3. Please select at least 10 starting nodes, ensuring they are non-repetitive and diverse.
4. In the user’s input, each line constitutes a node. When selecting the starting node, please make
your choice from those provided, and refrain from fabricating your own. The nodes you output
must correspond exactly to the nodes given by the user, with identical wording.
Finally, I emphasize again that you need to select the starting node from the given Nodes, and
it must be consistent with the words of the node you selected. Please strictly follow the above
format. Let’s begin.
"""

initial_node_prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            initial_node_system,
        ),
        (
            "human",
            (
                """Question: {question}
Plan: {rational_plan}
Nodes: {nodes}"""
            ),
        ),
    ]
)

Again, we put most of the instructions as the system message. Since we have multiple inputs, we can define them in the human message. However, we need a more structured output this time. Instead of writing a parsing function that takes in text and outputs a JSON, we can simply use the use_structured_outputmethod to define the desired output structure.

class Node(BaseModel):
    key_element: str = Field(description="""Key element or name of a relevant node""")
    score: int = Field(description="""Relevance to the potential answer by assigning
a score between 0 and 100. A score of 100 implies a high likelihood of relevance to the answer,
whereas a score of 0 suggests minimal relevance.""")

class InitialNodes(BaseModel):
    initial_nodes: List[Node] = Field(description="List of relevant nodes to the question and plan")

initial_nodes_chain = initial_node_prompt | model.with_structured_output(InitialNodes)

We want to output a list of nodes containing the key element and the score. We can easily define the output using a Pydantic model. Additionally, it is vital to add descriptions to each of the field, so we can guide the LLM as much as possible.

The last thing in this step is to define the node as a function.

def initial_node_selection(state: OverallState) -> OverallState:
    potential_nodes = get_potential_nodes(state.get("question"))
    initial_nodes = initial_nodes_chain.invoke(
        {
            "question": state.get("question"),
            "rational_plan": state.get("rational_plan"),
            "nodes": potential_nodes,
        }
    )
    # paper uses 5 initial nodes
    check_atomic_facts_queue = [
        el.key_element
        for el in sorted(
            initial_nodes.initial_nodes,
            key=lambda node: node.score,
            reverse=True,
        )
    ][:5]
    return {
        "check_atomic_facts_queue": check_atomic_facts_queue,
        "previous_actions": ["initial_node_selection"],
    }

In the initial node selection, we start by getting a list of potential nodes using the vector similarity search based on the input. An option is to use rational plan instead. The LLM is prompted to output the 10 most relevant nodes. However, the authors say that we should use only 5 initial nodes. Therefore, we simply order the nodes by their score and take the top 5 ones. We then update the check_atomic_facts_queue with the selected initial key elements.

Atomic fact check

In this step, we take the initial key elements and inspect the linked atomic facts. The prompt is:

All prompts start by giving the LLM some context, followed by task instructions. The LLM is instructed to read the atomic facts and decide whether to read the linked text chunks or if the atomic facts are irrelevant, search for more information by exploring the neighbors. The last bit of the prompt is the output instructions. We will use the structured output method again to avoid manually parsing and structuring the output.

Since chains are very similar in their implementation, different only by prompts, we’ll avoid showing every definition in this blog post. However, we’ll look at the LangGraph node definitions to better understand the flow.

def atomic_fact_check(state: OverallState) -> OverallState:
    atomic_facts = get_atomic_facts(state.get("check_atomic_facts_queue"))
    print("-" * 20)
    print(f"Step: atomic_fact_check")
    print(
        f"Reading atomic facts about: {state.get('check_atomic_facts_queue')}"
    )
    atomic_facts_results = atomic_fact_chain.invoke(
        {
            "question": state.get("question"),
            "rational_plan": state.get("rational_plan"),
            "notebook": state.get("notebook"),
            "previous_actions": state.get("previous_actions"),
            "atomic_facts": atomic_facts,
        }
    )

    notebook = atomic_facts_results.updated_notebook
    print(
        f"Rational for next action after atomic check: {atomic_facts_results.rational_next_action}"
    )
    chosen_action = parse_function(atomic_facts_results.chosen_action)
    print(f"Chosen action: {chosen_action}")
    response = {
        "notebook": notebook,
        "chosen_action": chosen_action.get("function_name"),
        "check_atomic_facts_queue": [],
        "previous_actions": [
            f"atomic_fact_check({state.get('check_atomic_facts_queue')})"
        ],
    }
    if chosen_action.get("function_name") == "stop_and_read_neighbor":
        neighbors = get_neighbors_by_key_element(
            state.get("check_atomic_facts_queue")
        )
        response["neighbor_check_queue"] = neighbors
    elif chosen_action.get("function_name") == "read_chunk":
        response["check_chunks_queue"] = chosen_action.get("arguments")[0]
    return response

The atomic fact check node starts by invoking the LLM to evaluate the atomic facts of the selected nodes. Since we are using the use_structured_output we can parse the updated notebook and the chosen action output in a straightforward manner. If the selected action is to get additional information by inspecting the neighbors, we use a function to find those neighbors and append them to the check_atomic_facts_queue. Otherwise, we append the selected chunks to the check_chunks_queue. We update the overall state by updating the notebook, queues, and the chosen action.

Text chunk check

As you might imagine by the name of the LangGraph node, in this step, the LLM reads the selected text chunk and decides the best next step based on the provided information. The prompt is the following:

The LLM is instructed to read the text chunk and decide on the best approach. My gut feeling is that sometimes relevant information is at the start or the end of a text chunk, and parts of the information might be missing due to the chunking process. Therefore, the authors decided to give the LLM the option to read a previous or next chunk. If the LLM decides it has enough information, it can hop on to the final step. Otherwise, it has the option to search for more details using the search_morefunction.

Again, we’ll just look at the LangGraph node function.

def chunk_check(state: OverallState) -> OverallState:
    check_chunks_queue = state.get("check_chunks_queue")
    chunk_id = check_chunks_queue.pop()
    print("-" * 20)
    print(f"Step: read chunk({chunk_id})")

    chunks_text = get_chunk(chunk_id)
    read_chunk_results = chunk_read_chain.invoke(
        {
            "question": state.get("question"),
            "rational_plan": state.get("rational_plan"),
            "notebook": state.get("notebook"),
            "previous_actions": state.get("previous_actions"),
            "chunk": chunks_text,
        }
    )

    notebook = read_chunk_results.updated_notebook
    print(
        f"Rational for next action after reading chunks: {read_chunk_results.rational_next_move}"
    )
    chosen_action = parse_function(read_chunk_results.chosen_action)
    print(f"Chosen action: {chosen_action}")
    response = {
        "notebook": notebook,
        "chosen_action": chosen_action.get("function_name"),
        "previous_actions": [f"read_chunks({chunk_id})"],
    }
    if chosen_action.get("function_name") == "read_subsequent_chunk":
        subsequent_id = get_subsequent_chunk_id(chunk_id)
        check_chunks_queue.append(subsequent_id)
    elif chosen_action.get("function_name") == "read_previous_chunk":
        previous_id = get_previous_chunk_id(chunk_id)
        check_chunks_queue.append(previous_id)
    elif chosen_action.get("function_name") == "search_more":
        # Go over to next chunk
        # Else explore neighbors
        if not check_chunks_queue:
            response["chosen_action"] = "search_neighbor"
            # Get neighbors/use vector similarity
            print(f"Neighbor rational: {read_chunk_results.rational_next_move}")
            neighbors = get_potential_nodes(
                read_chunk_results.rational_next_move
            )
            response["neighbor_check_queue"] = neighbors

    response["check_chunks_queue"] = check_chunks_queue
    return response

We start by popping a chunk ID from the queue and retrieving its text from the graph. Using the retrieved text and additional information from the overall state of the LangGraph system, we invoke the LLM chain. If the LLM decides it wants to read previous or subsequent chunks, we append their IDs to the queue. On the other hand, if the LLM chooses to search for more information, we have two options. If there are any other chunks to read in the queue, we move to reading them. Otherwise, we can use the vector search to get more relevant key elements and repeat the process by reading their atomic facts and so on.

The paper is slightly dubious about the search_more function. On the one hand, it states that the search_more function can only read other chunks in the queue. On the other hand, in their example in the appendix, the function clearly explores the neighbors.

To clarify, I emailed the authors, and they confirmed that the search_morefunction first tries to go through additional chunks in the queue. If none are present, it moves on to exploring the neighbors. Since how to explore the neighbors isn’t explicitly defined, we again use the vector similarity search to find potential nodes.

Neighbor selection

When the LLM decides to explore the neighbors, we have helper functions to find potential key elements to explore. However, we don’t explore all of them. Instead, an LLM decides which of them is worth exploring, if any. The prompt is the following:

Based on the provided potential neighbors, the LLM can decide which to explore. If none are worth exploring, it can decide to terminate the flow and move on to the answer reasoning step.

The code is:

def neighbor_select(state: OverallState) -> OverallState:
    print("-" * 20)
    print(f"Step: neighbor select")
    print(f"Possible candidates: {state.get('neighbor_check_queue')}")
    neighbor_select_results = neighbor_select_chain.invoke(
        {
            "question": state.get("question"),
            "rational_plan": state.get("rational_plan"),
            "notebook": state.get("notebook"),
            "nodes": state.get("neighbor_check_queue"),
            "previous_actions": state.get("previous_actions"),
        }
    )
    print(
        f"Rational for next action after selecting neighbor: {neighbor_select_results.rational_next_move}"
    )
    chosen_action = parse_function(neighbor_select_results.chosen_action)
    print(f"Chosen action: {chosen_action}")
    # Empty neighbor select queue
    response = {
        "chosen_action": chosen_action.get("function_name"),
        "neighbor_check_queue": [],
        "previous_actions": [
            f"neighbor_select({chosen_action.get('arguments', [''])[0] if chosen_action.get('arguments', ['']) else ''})"
        ],
    }
    if chosen_action.get("function_name") == "read_neighbor_node":
        response["check_atomic_facts_queue"] = [
            chosen_action.get("arguments")[0]
        ]
    return response

Here, we execute the LLM chain and parse results. If the chosen action is to explore any neighbors, we add them to the check_atomic_facts_queue .

Answer reasoning

The last step in our flow is to ask the LLM to construct the final answer based on the collected information in the notebook. The prompt is:

This node implementation is fairly straightforward as you can see by the code:

def answer_reasoning(state: OverallState) -> OutputState:
    print("-" * 20)
    print("Step: Answer Reasoning")
    final_answer = answer_reasoning_chain.invoke(
        {"question": state.get("question"), "notebook": state.get("notebook")}
    )
    return {
        "answer": final_answer.final_answer,
        "analysis": final_answer.analyze,
        "previous_actions": ["answer_reasoning"],
    }

We simply input the original question and the notebook with the collected information to the chain and ask it to formulate the final answer and provide the explanation in the analysis part.

LangGraph flow definition

The only thing left is to define the LangGraph flow and how it should traverse between the nodes. I am quite fond of the simple approach the LangChain team has chosen.

langgraph = StateGraph(OverallState, input=InputState, output=OutputState)
langgraph.add_node(rational_plan_node)
langgraph.add_node(initial_node_selection)
langgraph.add_node(atomic_fact_check)
langgraph.add_node(chunk_check)
langgraph.add_node(answer_reasoning)
langgraph.add_node(neighbor_select)

langgraph.add_edge(START, "rational_plan_node")
langgraph.add_edge("rational_plan_node", "initial_node_selection")
langgraph.add_edge("initial_node_selection", "atomic_fact_check")
langgraph.add_conditional_edges(
    "atomic_fact_check",
    atomic_fact_condition,
)
langgraph.add_conditional_edges(
    "chunk_check",
    chunk_condition,
)
langgraph.add_conditional_edges(
    "neighbor_select",
    neighbor_condition,
)
langgraph.add_edge("answer_reasoning", END)

langgraph = langgraph.compile()

We begin by defining the state graph object, where we can define the information passed along in the LangGraph. Each node is simply added with the add_node method. Normal edges, where one step always follows the other, can be added with a add_edge method. On the other hand, if the traversals is dependent on previous actions, we can use the add_conditional_edge and pass in the function that selects the next node. For example, the atomic_fact_condition looks like this:

def atomic_fact_condition(
    state: OverallState,
) -> Literal["neighbor_select", "chunk_check"]:
    if state.get("chosen_action") == "stop_and_read_neighbor":
        return "neighbor_select"
    elif state.get("chosen_action") == "read_chunk":
        return "chunk_check"

As you can see, it’s about as simple as it gets to define the conditional edge.

Evaluation

Finally we can test our implementation on a couple of questions. Let’s begin with a simple one.

langgraph.invoke({"question":"Did Joan of Arc lose any battles?"})

Results

The agent begins by forming a rational plan to identify the battles Joan of Arc participated in during her military career and to determine whether any were lost. After setting this plan, it moves to an atomic fact check about key battles such as the Siege of Orléans, the Siege of Paris, and La Charité. Rather than expanding the graph, the agent directly confirms the facts it needs. It reads text chunks that provide further details on Joan of Arc’s unsuccessful campaigns, particularly the failed Siege of Paris and La Charité. Since this information answers the question about whether Joan lost any battles, the agent stops here without expanding its exploration further. The process concludes with a final answer, confirming that Joan did indeed lose some battles, notably at Paris and La Charité, based on the evidence gathered.

Let’s now throw it a curveball.

langgraph.invoke({"question":"What is the weather in Spain?"})

Results

After the rational plan, the agent selected the initial key elements to explore. However, the issue is that none of these key elements exists in the database, and the LLM simply hallucinated them. Maybe some prompt engineering could solve hallucinations, but I haven’t tried. One thing to note is that it’s not that terrible, as these key elements don’t exist in the database, so we can’t pull any relevant information. Since the agent didn’t get any relevant data, it searched for more information. However, none of the neighbors are relevant either, so the process is stopped, letting the user know that the information is unavailable.

Now let’s try a multi-hop question.

langgraph.invoke(
  {"question":"Did Joan of Arc visit any cities in early life where she won battles later?"})

Results

It’s a bit too much to copy the whole flow, so I copied only the answer part. The flow for this questions is quite non-deterministic and very dependent on the model being used. It’s kind of funny, but as I was testing the newer the model, the worse it performed. So the GPT-4 was the best (also used in this example), followed by GPT-4-turbo, and the last place goes to GPT-4o.

Summary

I’m very excited about GraphReader and similar approaches, specifically because I think such an approach to (Graph)RAG can be pretty generic and applied to any domain. Additionally, you can avoid the whole graph modeling part as the graph schema is static, allowing the graph agent to traverse it using predefined functions.

We discussed some issues with this implementation along the way. For example, the graph construction on many documents might result in broad key elements ending up as supernodes, and sometimes, the atomic facts don’t contain the full context.

The retriever part is super reliant on extracted and selected key elements. In the original implementation, they put all the key elements in the prompt to choose from. However, I doubt that that approach scales well. Perhaps we also need an additional function to allow the agent to search for more information in other ways than just to explore the neighbor key elements.

Lastly, the agent system is highly dependent on the performance of the LLM. Based on my testing, the best model from OpenAI is the original GPT-4, which is funny as it’s the oldest. I haven’t tested the o1, though.

All in all, I am excited to explore more of these document graphs implementations, where metadata is extracted from text chunk and used to navigate the information better. Let me know if you have any ideas how to improve this implementation or have any other you like.

As always, the code is available on GitHub.

---

Implementing GraphReader with Neo4j and LangGraph was originally published in Towards Data Science on Medium, where people are continuing the conversation by highlighting and responding to this story.

Document datasets already have structure. Take advantage of it.

There are layered challenges in building retrieval-augmented generation (RAG) applications. Document retrieval, a huge part of the RAG workflow, is itself a complex set of steps that can be approached in different ways depending on the use case.

It is difficult for RAG systems to find the best set of documents relevant to a nuanced input prompt, especially when relying entirely on vector search to find the best candidates. Yet often our documents themselves are telling us where we should look for more information on a given topic — via citations, cross-references, footnotes, hyperlinks, etc. In this article, we’ll show how a new data model — linked documents — unlocks performance improvements by enabling us to parse and preserve these direct references to other texts, making them available for simultaneous retrieval — regardless of whether they were overlooked by vector search.

AI captures complexity, but not structure

When answering complex or nuanced questions requiring supporting details from disparate documents, RAG systems often struggle to locate all of the relevant documents needed for a well-informed and complete response. Yet we keep relying almost exclusively on text embeddings and vector similarity to locate and retrieve relevant documents.

One often-understated fact: there is a lot of document information lost during the process of parsing, chunking, and embedding text. Document structure — including section hierarchy, headings, footnotes, cross-references, citations, and hyperlinks — are almost entirely lost in a typical text-to-vector workflow, unless we take specific action to preserve them. When the structure and metadata are telling us what other documents are directly related to what we are reading, why shouldn’t we preserve this information?

In particular, links and references are ignored in a typical chunking and embedding process, which means they can’t be used by the AI to help answer queries. But, links and references are valuable pieces of information that often point to more useful documents and text — why wouldn’t we want to check those target documents at query time, in case they’re useful?

Parsing and following links and references programmatically is not difficult, and in this article we present a simple yet powerful implementation designed for RAG systems. We show how to use document linking to preserve known connections between document chunks, connections which typical vector embedding and retrieval might fail to make.

Document connections get lost in vector space

Documents in a vector store are essentially pieces of knowledge embedded into a high-dimensional vector space. These vectors are essentially the internal “language” of LLMs — given an LLM and all of its internal parameter values, including previous context and state, a vector is the starting point from which a model generates text. So, all of the vectors in a vector store are embedded documents that an LLM might use to generate a response, and, similarly, we embed prompts into vectors that we then use to search for nearest neighbors in semantic vector space. These nearest neighbors correspond to documents that are likely to contain information that can address the prompt.

In a vector store, the closeness of vectors indicates document similarity in a semantic sense, but where there is no real concept of connectedness beyond similarity. However, documents that are close to each other (and typically retrieved together) can be viewed as a type of connection between those pieces of knowledge, forming an implicit knowledge graph where each chunk of text is connected to its nearest neighbors. A graph built in this sense would not be static or rigid like most knowledge graphs; it would change as new documents are added or search parameters adjusted. So it is not a perfect comparison, but this implicit graph can be helpful as a conceptual framework that is useful for thinking about how document retrieval works within RAG systems.

In terms of real-world knowledge — in contrast to vector representations — semantic similarity is just one of many ways that pieces of text may be related. Even before computers and digital representations of data, we’ve been connecting knowledge for centuries: glossaries, indexes, catalogs, tables-of-contents, dictionaries, and cross-references are all ways to connect pieces of knowledge with each other. Implementing these in software is quite simple, but they typically haven’t been included in vector stores, RAG systems, and other gen AI applications. Our documents are telling us what other knowledge is important and relevant; we just need to give our knowledge stores the capability to understand and follow the connections.

Connect knowledge with document linking

We’ve developed document linking for cases in which our documents are telling us what other knowledge is relevant, but our vector store isn’t capturing that and the document retrieval process is falling short. Document linking is a straightforward yet potent method for representing directed connections between documents. It encapsulates all the traditional ways we navigate and discover knowledge, whether through a table of contents, glossary, keyword — and of course the easiest for a programmatic parser to follow: hyperlinks. This concept of linking documents allows for relationships that can be asymmetric or tagged with qualitative metadata for filtering or other purposes. Links are not only easy to conceptualize and work with but also scale efficiently to large, dynamic datasets, supporting robust and efficient retrieval.

The data model for links

As a data type, document links are quite simple. Link information is stored alongside document vectors as metadata. That means that retrieving a given document automatically retrieves information about the links that lead from and to the given document. Outbound links point to more information that’s likely to be useful in the context of the document, inbound links show which other documents may be supported by the given document, and bi-directional (or undirected) links can represent other types of connections. Links can also be tagged with further metadata that provides qualitative information that can be used for link or document filtering, ranking, and graph traversal algorithms.

As described in more detail in the article “Scaling Knowledge Graphs by Eliminating Edges,” rather than storing every link individually, as in typical graph database implementations, our efficient and scalable implementation uses link types and link groups as intermediate data types that greatly reduce storage and compute needs during graph traversal. This implementation has a big advantage when, for example, two groups of documents are closely related.

Let’s say that we have a group of documents on the topic of the City of Seattle (call it Group A) and we have another group of documents that mention Seattle (Group B). We would like to make sure that documents mentioning Seattle can find all of the documents about the City of Seattle, and so we would like to link them. We could create a link from all of the documents in Group B to all of the documents in Group A, but unless the two groups are small, this is a lot of edges! The way we handle this is to create one link type object representing the keyword “Seattle” (kw:seattle), and then creating directed links from the documents in Group B to this kw:seattle object as well as links from the kw:seattle object to the documents in Group A. This results in far fewer links to store with each document — there is only one link each — and no information is lost.

Using document links during retrieval

The main goal of the retrieval process in RAG systems is to find a set of documents that is sufficient to answer a given query. Standard vector search and retrieval finds documents that are most “relevant” to the query in a semantic sense, but might miss some supporting documents if their overall content doesn’t closely match the content of the query.

For example, let’s say we have a large document set that includes the documents related to Seattle as described above. We have the following prompt about the Space Needle, a prominent landmark in Seattle:

“What is close to the Space Needle?”

A vector search starting with this prompt would retrieve documents mentioning the Space Needle directly, because that is the most prominent feature of the prompt text from a semantic content perspective. Documents mentioning the Space Needle are likely to mention its location in Seattle as well. Without using any document linking, a RAG system would have to try to answer the prompt using mainly documents mentioning the Space Needle, without any guarantee that other helpful documents that don’t mention the Space Needle directly would also be retrieved and used.

Below, we construct a practical example (with code!) based on this Space Needle dataset and query. Keep reading to understand how a RAG system might miss helpful documents when links are not used, and then “find” helpful documents again by simply following link information contained within the original documents themselves.

Document links in action

In order to illustrate how document linking works, and how it can make connections between documents and knowledge that might be missed otherwise, let’s look at a simple example.

We’ll start with two related documents containing some text from Wikipedia pages: one document from the page for the Space Needle, and one for the neighborhood where the Space Needle is located, Lower Queen Anne. The Space Needle document has an HTML link to the Lower Queen Anne document, but not the other way around. The document on the Space needle begins as follows:

'url': 'https://en.wikipedia.org/wiki/Space_Needle'

The Space Needle is an observation tower in Seattle, Washington, 
United States. Considered to be an icon of the city, it has been 
designated a Seattle landmark. Located in the Lower Queen Anne 
neighborhood, it was built in the Seattle Center for the 1962 
World's Fair, which drew over 2.3 million visitors...

In addition to these two documents derived from real, informative sources, we have also added four very short, uninformative documents — two that mention the Space Needle and two that don’t. These documents (and their fake URLs) are designed to be irrelevant or uninformative documents, such as social media posts that are merely commenting on the Space Needle and Seattle, such as:

“The Space Needle is TALL.”

and

“Queen Anne was a person.”

The full document set is included in the Colab notebook. They are HTML documents that we then process using BeautifulSoup4 as well as the HtmlLinkExtractor from LangChain, adding those links back to the Document objects with the add_links function specifically so we can make use of them in theGraphVectorStore , a relatively new addition to the LangChain codebase, contributed by my colleagues at DataStax. All of this is open-source.

Each document is processed as follows:

from langchain_core.documents import Document
from langchain_core.graph_vectorstores.links import add_links
from langchain_community.graph_vectorstores.extractors.html_link_extractor import HtmlInput, HtmlLinkExtractor

soup_doc = BeautifulSoup(html_doc, 'html.parser')
doc = Document(
        page_content=soup_doc.get_text(),
        metadata={"source": url}
    )
doc.metadata['content_id'] = url  # the ID for Links to point to this document
html_link_extractor = HtmlLinkExtractor()add_links(doc, html_link_extractor.extract_one(HtmlInput(soup_doc, url)))

Using `cassio`, we initialize the GraphVectorStore as below:

from langchain_openai import OpenAIEmbeddings
from langchain_community.graph_vectorstores.cassandra import CassandraGraphVectorStore

# Create a GraphVectorStore, combining Vector nodes and Graph edges.
EMBEDDING = 'text-embedding-3-small'
gvstore = CassandraGraphVectorStore(OpenAIEmbeddings(model=EMBEDDING))

We set up the LLM and other helpers for the RAG chain in the standard way — see the notebook for details. Note that, while almost everything used here is open-source, in the notebook we are using two SaaS products, OpenAI and DataStax’s Astra — LLM and vector data store, respectively — both of which have free usage tiers. See the LangChain documentation for alternatives.

Running RAG With default settings

We can run the RAG system end-to-end using a graph retriever with depth=0— which means no graph traversal at all — and other default parameters as below:

retriever = gvstore.as_retriever(
    search_kwargs={
        "depth": 0,  # depth of graph traversal; 0 is no traversal at all
    }
)

This gives an output such as:

Question:
 What is close to the Space Needle? 

Retrieved documents:
['https://TheSpaceNeedleisGreat',
 'https://TheSpaceNeedleisTALL',
 'https://en.wikipedia.org/wiki/Space_Needle',
 'https://SeattleIsOutWest',
 'https://en.wikipedia.org/wiki/Lower_Queen_Anne,_Seattle',
 'https://QueenAnneWasAPerson']

LLM response:
('The Space Needle is close to several locations in the Lower Queen Anne '
 'neighborhood, including Climate Pledge Arena, the Exhibition Hall, McCaw '
 'Hall, Cornish Playhouse, Bagley Wright Theater, the studios for KEXP radio, '
 'SIFF Cinema Uptown, and On the Boards.')

Of course in realistic scenarios, a RAG system would not retrieve the full document set, as we are doing here.

A more realistic scenario: We can’t retrieve all documents

Retrieving all documents for each query is impractical or even impossible in some cases. It also defeats the purpose of using vector search in the first place. For all realistic scenarios, only a small fraction of documents can be retrieved for each query, which is why it is so important to get the most relevant and helpful documents near the top of the list.

To make things a little more realistic for our example with our tiny dataset, let’s change the settings of the retriever so that k=3, meaning that a maximum of three documents are returned by each vector search. This means that three of the six total documents — the least similar or relevant according to vector similarity — will be left out of the returned document set. We can change the settings of the retriever like this:

retriever = gvstore.as_retriever(
    search_kwargs={
        "depth": 0,  # depth of graph traversal; 0 is no traversal at all
        "k": 3       # number of docs returned by initial vector search---not including graph Links
    }
)

Querying the system with these settings gives the output:

Question:
 What is close to the Space Needle? 

Retrieved documents:
['https://TheSpaceNeedleisGreat',
 'https://TheSpaceNeedleisTALL',
 'https://en.wikipedia.org/wiki/Space_Needle']

LLM response:
('The context does not provide specific information about what is close to the '
 'Space Needle. It only mentions that it is located in the Lower Queen Anne '
 'neighborhood and built for the Seattle Center for the 1962 World's Fair.')

We can see that this final response is much less informative than the previous one, now that we have access to only half of the document set, instead of having all six documents available for response generation.

There are some important points to note here.

1. One document that was left out was the document on Lower Queen Anne, which is the only document that describes some significant places in the neighborhood where the Space Needle is located.
2. The Lower Queen Anne document does not specifically mention the Space Needle, whereas three other documents do. So it makes sense that the initial query “What is close to the Space Needle?” returns those three.
3. The main document about the Space Needle has an HTML link directly to Lower Queen Anne, and any curious human would probably click on that link to learn about the area.
4. Without any sense of linking or graph traversal, this RAG system retrieves the most semantically similar documents — including two uninformative ones — and misses the one article that has the most information for answering the query.

Now, let’s look at how document linking affects results.

Following links finds the missing information

A simple change to our retriever setup — setting depth=1 — enables the retriever to follow any document links from the documents that are initially retrieved by vector search. (For reference, note that setting depth=2 would not only follow links in the initial document set, but would also follow the next set of links in the resulting document set — but we won’t go that far yet.)

We change the retriever depth parameter like this:

retriever = gvstore.as_retriever(
    search_kwargs={
        "depth": 1,  # depth of graph traversal; 0 is no traversal at all
        "k": 3       # number of docs returned by initial vector search---not including graph Links
    }
)

which gives the following output:

Question:
 What is close to the Space Needle? 

Retrieved documents:
['https://TheSpaceNeedleisGreat',
 'https://TheSpaceNeedleisTALL',
 'https://en.wikipedia.org/wiki/Space_Needle',
 'https://en.wikipedia.org/wiki/Lower_Queen_Anne,_Seattle']

LLM response:
('The Space Needle is located in the Lower Queen Anne neighborhood, which '
 'includes Climate Pledge Arena, Exhibition Hall, McCaw Hall, Cornish '
 'Playhouse, Bagley Wright Theater, the studios for KEXP radio, a three-screen '
 'movie theater (SIFF Cinema Uptown), and On the Boards, a center for '
 'avant-garde theater and music.')

We can see that the first k documents retrieved by vector search are the same three as before, but setting depth=1 instructed the system to follow links from those three documents and include those linked documents as well. So, the direct link from the Space Needle document to Lower Queen Anne included that document as well, giving the LLM access to the neighborhood information that it needed to answer the query properly.

Document linking can improve RAG apps

This hybrid approach of vector and graph retrieval can significantly enhance the context relevance and diversity of results in RAG applications. It can lead to fewer hallucinations and higher-quality outcomes by ensuring that the system retrieves the most contextually appropriate and diverse content.

Beyond improving the quality of responses of RAG systems, document linking has some advantages for implementation in a production system. Some beneficial properties of include:

1. Lossless — The original content remains intact within the nodes, ensuring that no information is discarded during the graph creation process. This preserves the integrity of the data, reducing the need for frequent re-indexing as needs evolve and leveraging the LLM’s strength in extracting answers from contextual clues.
2. Hands-off — This method does not require expert intervention to refine knowledge extraction. Instead, adding some edge extraction capabilities based on keywords, hyperlinks, or other document properties to the existing vector-search pipeline allows for the automatic addition of links.
3. Scalable — The graph creation process involves straightforward operations on the content without necessitating the use of an LLM to generate the knowledge graph.

Performance benchmarks and a more detailed analysis of scaling document linking is included in the article mentioned earlier.

As always, there are some limitations. If your document set truly doesn’t have links or other structure, the strategies presented here won’t accomplish much. Also, while building and traversing graph connections can be powerful, it also adds complexity to the retrieval process that might be challenging to debug and optimize — primarily if traversing the graph to depths of 2 or greater.

Overall, incorporating document linking into RAG systems combines the strengths of traditional, deterministic software methodologies, graph algorithms, and modern AI techniques. By explicitly defining links between documents, we enhance the AI’s ability to navigate knowledge as a human researcher might, improving not only retrieval accuracy but also the contextual depth of responses. This approach creates more robust, capable systems that align with the complex ways humans seek and use knowledge.

Getting started with document linking

Complete code from this article can be found in this Colab notebook. And, check out this introductory blog post by my colleague at DataStax, or see the documentation for GraphVectorStore in LangChain for detailed API information and how to use document linking to enhance your RAG applications and push the boundaries of what your knowledge systems can achieve.

by Brian Godsey, Ph.D. (LinkedIn)— mathematician, data scientist and engineer // works on AI products at DataStax // Wrote the book Think Like a Data Scientist

---

Your Documents Are Trying to Tell You What’s Relevant: Better RAG Using Links was originally published in Towards Data Science on Medium, where people are continuing the conversation by highlighting and responding to this story.