From Neo4j to Postgres: Building a Sentence Graph for AI with Vector Memory

Working with Knowledge Graphs can be complex, especially when it comes to scalability and performance. In a recent project, I needed to manage a sentence graph for an AI memory system. I initially opted for Neo4j, a highly popular graph database. However, management costs, the complexity of configuring it in cloud environments, and limitations in horizontal scaling pushed me to look for a more efficient alternative that integrated better with my existing infrastructure. The solution? Postgres with the pgvector extension. This tutorial walks through my migration process, the benefits I gained, and how you can do the same.

Prerequisites

Before diving in, make sure you have the following tools and knowledge:

• An active Postgres server (version 13 or later).
• The pgvector extension installed and enabled.
• Basic knowledge of SQL and vector embeddings concepts.
• Python 3.6+ with the psycopg2, numpy, and sentence-transformers libraries.
• Command-line access to the Postgres server.

Installing and Configuring pgvector

The first step is installing and configuring the pgvector extension on your Postgres server. This extension allows you to store and manipulate vectors directly within the database, making Postgres suitable for similarity search and vector memory applications.

Installing the pgvector Package

Depending on your operating system, the installation process may vary. On Debian/Ubuntu systems, you can use apt:

 sudo apt update
 sudo apt install postgresql-server-dev-all
 sudo apt install postgresql-15-vector #adatta la versione di postgresql

On RHEL-based systems (CentOS, Fedora), you can use yum or dnf:

 sudo dnf install postgresql15-server-devel #adatta la versione di postgresql

Enabling the Extension in Your Database

After installation, you need to enable the extension in the specific database you’ll use for your sentence graph. Connect to the database as a superuser and run the following SQL command:

 CREATE EXTENSION vector;

Verify that the extension was enabled successfully:

 dx

You should see vector in the list of installed extensions.

Creating the Tables for the Sentence Graph

The sentence graph will consist of two main tables: nodes and edges. The nodes table will contain the sentences and their corresponding vector embeddings, while the edges table will store the relationships between the sentences.

The nodes Table

The nodes table will have the following columns:

• id: Unique identifier for the sentence (UUID).
• text: The text of the sentence.
• embedding: The vector embedding of the sentence (using the vector type from pgvector).

Here is the SQL script to create the table:

 CREATE TABLE nodes (
  id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
  text TEXT NOT NULL,
  embedding VECTOR(384) -- Adatta la dimensione del vettore al tuo modello
 );

In this case, I used a vector dimension of 384, which is the typical output size for vector embeddings generated by models like sentence-transformers/all-MiniLM-L6-v2.

The edges Table

The edges table will have the following columns:

• source_id: Identifier for the source node (foreign key to nodes).
• target_id: Identifier for the target node (foreign key to nodes).
• relation_type: The type of relationship between the nodes (e.g., “similar”, “contradictory”, etc.).

Here is the SQL script to create the table:

 CREATE TABLE edges (
  source_id UUID NOT NULL REFERENCES nodes(id),
  target_id UUID NOT NULL REFERENCES nodes(id),
  relation_type TEXT NOT NULL,
  PRIMARY KEY (source_id, target_id)
 );

Populating the Graph with Data

Now that the tables are set up, we can start populating the graph with data. This process involves generating vector embeddings for the sentences and inserting the data into the nodes and edges tables.

Generating Vector Embeddings

We will use the sentence-transformers library to generate the vector embeddings. Here is a Python code example:

 from sentence_transformers import SentenceTransformer
 import psycopg2
 import numpy as np

 # Connessione al database Postgres
 conn = psycopg2.connect(
  host="localhost",
  database="your_database",
  user="your_user",
  password="your_password"
 )
 cur = conn.cursor()

 # Inizializzazione del modello SentenceTransformer
 model = SentenceTransformer('all-MiniLM-L6-v2')

 # Lista di frasi da inserire nel grafo
 sentences = [
  "Il sole è una stella.",
  "La terra è piatta.",
  "La terra ruota attorno al sole.",
  "I gatti sono mammiferi.",
  "I cani sono rettili."
 ]

 # Generazione dei vector embeddings
 embeddings = model.encode(sentences)

 # Inserimento dei dati nella tabella nodes
 for i, sentence in enumerate(sentences):
  embedding = embeddings[i].tolist()
  cur.execute(
  """
  INSERT INTO nodes (text, embedding)
  VALUES (%s, %s)
  """,
  (sentence, np.array(embedding))
  )

 # Commit delle modifiche e chiusura della connessione
 conn.commit()
 cur.close()
 conn.close()

 print("Dati inseriti correttamente.")

This script connects to the Postgres database, initializes the sentence-transformers model, generates the vector embeddings for the sentences, and inserts the data into the nodes table.

Creating Relationships Between Sentences

After inserting the sentences, you need to create the relationships between them in the edges table. This can be done manually or by using a similarity search algorithm to automatically identify the most similar sentences and create the relationships.

Here is an example of how to manually create a relationship:

 import psycopg2

 # Connessione al database Postgres
 conn = psycopg2.connect(
  host="localhost",
  database="your_database",
  user="your_user",
  password="your_password"
 )
 cur = conn.cursor()

 # ID dei nodi da collegare
 source_id = "..."  # Inserisci l'UUID del nodo di partenza
 target_id = "..."  # Inserisci l'UUID del nodo di arrivo
 relation_type = "simile"

 # Inserimento della relazione nella tabella edges
 cur.execute(
  """
  INSERT INTO edges (source_id, target_id, relation_type)
  VALUES (%s, %s, %s)
  """,
  (source_id, target_id, relation_type)
 )

 # Commit delle modifiche e chiusura della connessione
 conn.commit()
 cur.close()
 conn.close()

 print("Relazione creata correttamente.")

Querying the Graph

Once the graph is populated with data, you can execute queries to extract information and relationships. pgvector offers several functions to perform similarity searches and find the nodes closest to a given query.

Similarity Search

To find the sentences most similar to a given query, you can use the <-> operator from pgvector, which calculates the cosine distance between vector embeddings.

Here is an example SQL query:

 SELECT id, text, embedding <-> (SELECT embedding FROM nodes WHERE id = '...') AS distance
 FROM nodes
 ORDER BY distance ASC
 LIMIT 5;

This query selects the 5 nodes most similar to the node with the specified ID, ordering the results by cosine distance.

Finding Relationships

To find relationships between nodes, you can use standard SQL queries with JOIN clauses.

Here is an example query to find all nodes connected to a given node:

 SELECT n.id, n.text, e.relation_type
 FROM nodes n
 INNER JOIN edges e ON n.id = e.target_id
 WHERE e.source_id = '...';

This query selects all nodes connected to the node with the specified ID, along with the relationship type.

Common Errors and Troubleshooting

• Error: could not find function vector_in: Ensure that the vector extension has been properly installed and enabled in the database. Verify this by running dx from psql.
• Slow similarity search performance: Create an index on the embedding column using the hnsw algorithm. This will significantly speed up similarity search queries. Example: CREATE INDEX ON nodes USING hnsw (embedding vector_cosine_ops);
• Vector dimension error: Make sure the vector dimension specified during the creation of the nodes table matches the dimension of the vector embeddings generated by your model. If you are using SentenceTransformers’ all-MiniLM-L6-v2, the dimension is 384.

Conclusion

In this tutorial, we explored how to replace Neo4j with Postgres and pgvector to implement a sentence graph for an AI memory system. We covered installing and configuring pgvector, creating the graph tables, populating the graph with data, and running queries to extract information and relationships.

Migrating to Postgres with pgvector offers several advantages, including greater scalability, better integration with existing infrastructure, and reduced operational costs. Additionally, pgvector delivers excellent performance for similarity search queries, making it an ideal choice for vector memory applications.

If you have any questions or comments, drop them below. Have you ever considered using Postgres as a graph database? What challenges did you face?