Bare-bones implementation of similarity search in DuckDB using Expression API

duckdb-RAG.py

import duckdb
from duckdb import (
    ConstantExpression,
    ColumnExpression,
    FunctionExpression,
    StarExpression
)
conn = duckdb.connect(database='embed.ddb')

text="Wooden floor"
from openai import OpenAI
client = OpenAI()
embedding = client.embeddings.create(input=[text], model="text-embedding-3-large").data[0].embedding

rms = conn.table("rms_embeddings_01")
list_cosine_similarity = FunctionExpression('list_cosine_similarity', ColumnExpression('embedding'), ConstantExpression(embedding))
top10 = (rms.
    select(*[StarExpression(),list_cosine_similarity.alias("similarity")]).
    order("similarity desc").
    limit(10).
    select(StarExpression(exclude=["similarity", "embedding"])).
    fetchdf()
)
conn.close()