Python Snippet

[Gaspi]

Dear group,

As promised, here is a commented version of the code snippet I showed last Friday.
It requires some setup to work: environment variables with credentials, a database filled with embeddings to search.
Feel free to get some inspiration from it but remember that your own code should be better organised than this.

#########################################################################################
#                                        Embedding                                      #
#########################################################################################

import requests, os
import numpy as np

# Set the environment variable LLM_TOKEN either
# - using a secret in your service's configuration,
# - or by manually running this command before running python: export LLM_TOKEN=sk-c....
llm_token = os.getenv("LLM_TOKEN")

# This should probably be an environment variable
api_url = "https://llm.lab.sspcloud.fr"


def get_embedding(prompt: str):
    """
    Calls the ollama embedding API to compute a vector of length 1024 representing [prompt].
    Returns the vector as a numpy.array.
    """
    headers = {
        "Authorization": f"Bearer {llm_token}",
        "Content-Type": "application/json"
    }
    data = {
        "model": "bge-m3:latest",
        "input": prompt
    }
    response = requests.post(api_url+'/ollama/api/embed', headers=headers, json=data)
    return np.array(response.json()['embeddings'][0])



#########################################################################################
#                                        Database                                       #
#########################################################################################

import psycopg
from pgvector.psycopg import register_vector

# Set the following env. variables for this to work: PGUSER, PGPASSWORD, PGHOST, PGPORT, PGDATABASE
conn = psycopg.connect(dbname="vector", autocommit=True)
register_vector(conn)


def get_similar_entries(embedding):
    """
    Returns the 5 entries from the database with the embedding closest to the given [embedding].
    """
    results = conn.execute("""
        SELECT
            fields_of_interest,
            embedding <-> %s as dst
        FROM my_table_with_embeddings
        ORDER BY dst
        LIMIT 5
        """, (embedding,))
    return results.fetchall()



#########################################################################################
#                                          Test                                         #
#########################################################################################

prompt = "A red creature card that flies and can sacrifice to do damage"
embedding = get_embedding(prompt)
for entry in get_similar_entries(embedding):
    print(entry)