How to customize cognee and add your logic

Difficulty: Medium

Overview

cognee let’s you define your own custom logic for graph enrichment and retrieval

Let’s create custom cognee tasks

Adding data to our graph and connecting it to LLM can help LLM reason better You move from large amounts to text to structured data. See before and after bellow

Here is what we start with:

And here is how it ends:

git clone https://github.com/topoteretes/cognee.git

cd cognee

poetry install

Let’s stop for a moment to understand what the script does?

It:

1. Reads Data from JSON Files

2. The script opens two files—companies.json and people.json—that contain information about companies, their departments, and people working in those departments.

3. Creates Connected Data Objects (DataPoints)

   “DataPoints” are just special classes (like Person, Department, Company) that Cognee uses to understand how the data is linked. Each object type (e.g., Person, Department, Company) defines what information it holds (like “name”) and how it relates to other objects (people belong to departments, departments belong to companies, etc.). Stores Everything in a Cognee Knowledge Graph

   A “knowledge graph” is like a network of information. It links all the objects together: Each Company node links to its departments, Each Department node links to its employees, and so on. By saving data in a graph, questions like “Who works for GreenFuture Solutions?” become easier to answer because Cognee can “walk” the connections and gather the right information.

4. Indexes and Visualizes the Graph

   After the data is loaded, the script runs a process called “indexing.” This makes it easier and faster to search relationships in the graph. It then creates an HTML file to show you a visual representation of the network of companies, departments, and employees.

5. Answers Questions About the Data

6. Finally, the script uses a search function to answer a query like “Who works for GreenFuture Solutions?” This works because Cognee knows about the connections in your graph.

python simple_example.py

import webbrowser
import os
from cognee.api.v1.visualize.visualize import visualize_graph
html = await visualize_graph()
home_dir = os.path.expanduser("~")
html_file = os.path.join(home_dir, "graph_visualization.html")
display(html_file)
webbrowser.open(f"file://{html_file}")

Join the Conversation!

Have questions? Join our community now to connect with professionals, share insights, and get your questions answered!