Create embeddings with cloud AI providers

The Cypher function ai.text.embed and procedure ai.text.embedBatch allow you to generate embeddings for one or more pieces of text through external AI providers. You need an API token for one of the supported providers (OpenAI, Vertex AI, Azure OpenAI, Amazon Bedrock).

This page assumes you have already imported the recommendations dataset and set up your environment, and shows how to generate and store embeddings for Movie nodes basing on their title and plot.

Embeddings are always generated outside of Neo4j, but stored in the Neo4j database.

Setup environment

The encoding functions are part of the Neo4j GenAI plugin.

  • On Aura instances the plugin is enabled by default, so you don’t need to take any further actions if you are using Neo4j on Aura.

  • For self-managed instances, the plugin needs to be installed. You do so by moving the neo4j-genai.jar file from /products to /plugins in your Neo4j home directory, or by starting the Docker container with the extra parameter --env NEO4J_PLUGINS='["genai"]'.
    For more information, see Configuration → Plugins.

Create embeddings for movies

The example below fetches all Movie nodes from the database, generates an embedding of the concatenation of movie title and plot, and adds that as an extra embedding property to each node.

import neo4j


URI = '<database-uri>'
AUTH = ('<username>', '<password>')
DB_NAME = '<database-name>'  # examples: 'recommendations-5.26', 'neo4j'

openAI_token = '<OpenAI API token>'


def main():
    with neo4j.GraphDatabase.driver(URI, auth=AUTH) as driver:  (1)
        driver.verify_connectivity()

        batch_size = 100
        batch_n = 1
        movies_batch = []
        with driver.session(database=DB_NAME) as session:
            # Fetch `Movie` nodes
            result = session.run('MATCH (m:Movie) RETURN m.plot AS plot, m.title AS title')
            for record in result:
                title = record.get('title')
                plot = record.get('plot')

                if title is not None and plot is not None:
                    movies_batch.append({
                        'title': title,
                        'plot': plot,
                        'to_encode': f'Title: {title}\nPlot: {plot}'  (2)
                    })

                # Import a batch; flush buffer
                if len(movies_batch) == batch_size:  (3)
                    import_batch(driver, movies_batch, batch_n)
                    movies_batch = []
                    batch_n += 1

            # Flush last batch
            import_batch(driver, movies_batch, batch_n)

        # Import complete, show counters
        records, _, _ = driver.execute_query('''
        MATCH (m:Movie WHERE m.embedding IS NOT NULL)
        RETURN count(*) AS countMoviesWithEmbeddings, size(m.embedding) AS embeddingSize
        ''', database_=DB_NAME)
        print(f"""
    Embeddings generated and attached to nodes.
    Movie nodes with embeddings: {records[0].get('countMoviesWithEmbeddings')}.
    Embedding size: {records[0].get('embeddingSize')}.
        """)


def import_batch(driver, nodes, batch_n):
    # Generate and store embeddings for Movie nodes
    driver.execute_query('''
    CALL ai.text.embedBatch($listToEncode, 'OpenAI', { (4)
      token: $token, model: 'text-embedding-3-small'
    }) YIELD index, vector  (5)
    MATCH (m:Movie {title: $movies[index].title, plot: $movies[index].plot})
    CALL db.create.setNodeVectorProperty(m, 'embedding', toFloatList(vector))  (6)
    ''', movies=nodes, listToEncode=[movie['to_encode'] for movie in nodes], token=openAI_token,
    database_=DB_NAME)
    print(f'Processed batch {batch_n}')


if __name__ == '__main__':
    main()

'''
Movie nodes with embeddings: 9083.
Embedding size: 1536.
'''
1 The driver object is the interface to interact with your Neo4j instance. For more information, see Build applications with Neo4j and Python.
2 The strings that OpenAI should encode into embeddings.
3 A number of embeddings are collected before a whole batch is submitted to the database. This avoids holding the whole dataset into memory and potential timeouts (especially relevant for larger datasets).
4 The procedure ai.text.embedBatch() submits the batch for encoding to OpenAI. OpenAI’s model text-embedding-3-small embeds text into vectors of size 1536 (i.e. lists of 1536 numbers). See GenAI providers for a list of supported providers and options. Always use the same model to generate embeddings for a dataset: pick one and stick to it for your whole project.
5 The returned index from ai.text.embedBatch allows to relate embeddings to movies, so that it’s possible to retrieve each movie node and attach its embedding to it; vector is of type VECTOR, the storage of which is available only in Enterprise Edition.
6 The import query sets a new embedding property on each node m, with the embedding vector as value. The Cypher procedure db.create.setNodeVectorProperty stores vector properties more efficiently than if they were stored as lists. As storing properties of VECTOR type is available only in Enterprise Edition, the extra call to toFloatList ensures compatibility across all editions. To set vector properties on relationships, use db.create.setRelationshipVectorProperty.

With Enterprise Edition, you can avoid calling db.create.setNodeVectorProperty and instead set the embeddings as properties directly via the Cypher clause SET. The query would thus become:

CALL ai.text.embedBatch($listToEncode, 'OpenAI', {
  token: $token, model: 'text-embedding-3-small'
}) YIELD index, vector
MATCH (m:Movie {title: $movies[index].title, plot: $movies[index].plot})
SET m.embedding = vector

Once embeddings are in the database, you can use them to compare how similar one movie is to another.