Cognee + Graphiti: Integrating Temporal-Aware Graphs

At cognee, we provide innovative solutions for knowledge representation and retrieval. Cognee is a robust framework for creating, querying, and leveraging graphs to extract meaningful insights. While many tools in the market offer specialized features - like Graphiti’s temporal-aware graphs - we stand apart by integrating with such tools, allowing users to bring their own logic and extend cognee’s core capabilities.

Our recent integration with Graphiti is a prime example of this adaptability. It demonstrates how Cognee enables users to enhance their workflows by adding custom functionality, rather than relying on disparate memory systems from multiple providers

Graphiti Integration: Adding Temporal Awareness to Cognee

Graphiti specializes in building temporal-aware graphs, a feature that tracks and represents the progression of entities and relationships over time. Through this integration, cognee has not only incorporated Graphiti’s strengths but also showcased its ability to adapt external tools into its ecosystem.

How Graphiti is Integrated into Cognee

To enable seamless interaction with their graph data, we’ve developed a new Pydantic object to store and manage Graphiti nodes and edges in cognee’s vector collections. Cognee provides a DataPoint class that is a Pydantic object that can be stored in a graph and vector store with any graph/vector store adapter we support. This means adding a new tool, means just extending a class. The new nodes encapsulate textual attributes such as name, summary, and content.

• Name: An identifying label.
• Summary: A concise overview.
• Content: The main body of data generated by Graphiti.

The integration brings:

1. Feature Enrichment: Cognee enhances Graphiti’s nodes with system-specific features, tailoring the data for our platform.
2. Storage and Indexing: Textual data from Graphiti nodes is stored in our database, and vector collections are created for efficient retrieval of both nodes and edges.
3. Graph Transformation: Cognee adapts Graphiti’s temporal graphs into its modular framework, creating a seamless pipeline for handling and analyzing data.

This process enables users to use Graphiti’s temporal-awareness features while maintaining the robustness and scalability of cognee’s core framework.

Here’s a snippet showcasing the transformation and indexing process and it can be found in our Github repo.

from cognee.tasks.temporal_awareness.graphiti_model import GraphitiNode
from cognee.infrastructure.databases.vector import get_vector_engine
from cognee.infrastructure.databases.graph import get_graph_engine

async def index_and_transform_graphiti_nodes_and_edges():
vector_engine = get_vector_engine()
graph_engine = await get_graph_engine()

# Add unique identifiers to nodes and edges
await graph_engine.query("""MATCH (n) SET n.id = n.uuid RETURN n""")
await graph_engine.query("""MATCH (source)-[r]->(target)
                         SET r.source_node_id = source.id,
                         r.target_node_id = target.id
                         RETURN r""")

nodes_data, edges_data = await graph_engine.get_model_independent_graph_data()

# Transform Graphiti nodes into cognee-compatible objects
for node_data in nodes_data[0]["nodes"]:
    graphiti_node = GraphitiNode(
        **{key: node_data[key] for key in ("content", "name", "summary") if key in node_data},
        id=node_data.get("uuid"),
    )

    # Index node data for efficient retrieval
    for field_name in graphiti_node._metadata["index_fields"]:
        index_name = f"{GraphitiNode.__tablename__}.{field_name}"
        await vector_engine.create_vector_index(GraphitiNode.__tablename__, field_name)
        indexed_data_point = graphiti_node.model_copy()
        indexed_data_point._metadata["index_fields"] = [field_name]
        await vector_engine.index_data_points(GraphitiNode.__tablename__, field_name, [indexed_data_point])

Benefits of the Integration

The integration showcases the flexibility and scalability of cognee’s architecture. Key advantages include:

• Bring Your Own Logic: Users can easily integrate external tools, models, or custom logic into cognee without compromising its core functionality.
• Enhanced Functionality: By building on cognee’s foundation, users can add specialized features like temporal awareness or any other domain-specific capabilities.
• Unified Workflow: Instead of piecing together fragmented systems, users can rely on cognee as a unified platform that adapts to their needs.

Example Use Case

Imagine a user querying: “When was Kamala Harris in office?”

Here’s how cognee processes the query using the integration.

from cognee.modules.retrieval.brute_force_triplet_search import brute_force_triplet_search
from cognee.tasks.completion.graph_query_completion import retrieved_edges_to_string
from cognee.infrastructure.llm.prompts import render_prompt, read_query_prompt
from cognee.infrastructure.llm.get_llm_client import get_llm_client

async def process_query_with_graphiti():
query = "When was Kamala Harris in office?"

# Perform triplet search in the vector collections
triplets = await brute_force_triplet_search(
    query=query,
    top_k=3,
    collections=["graphitinode_content", "graphitinode_name", "graphitinode_summary"],
)

# Prepare data for LLM-based response
args = {
    "question": query,
    "context": retrieved_edges_to_string(triplets),
}
user_prompt = render_prompt("graph_context_for_question.txt", args)
system_prompt = read_query_prompt("answer_simple_question_restricted.txt")

# Use LLM to generate a structured answer
llm_client = get_llm_client()
computed_answer = await llm_client.acreate_structured_output(
    text_input=user_prompt,
    system_prompt=system_prompt,
    response_model=str,
)
print(computed_answer)

# Output: 
# “Kamala Harris served as the Attorney General of California from [specific years].”

Try it yourself in this notebook.

Cognee is flexible, so what?

The integration of Graphiti into cognee represents a leap forward in our mission to enable adaptable, contextually aware data retrieval. Instead of imposing rigid workflows, cognee enables users to integrate their preferred methods, tools or libraries.

Stay tuned as we continue to release a few more interesting integrations along the way. Whether you're exploring integrations like Graphiti or building custom solutions, the cognee community is here to support you.

Join Our Community

Love what we're building? Try cognee out, check out tutorial on GitHub. ⭐