SQLite-vec memory for RAG w/ DB migration builtin!

memory_sqlite.py

'''

    ./llama-server --port 8081 -fa -c 0 --metrics \
      --embeddings \
      -hfr nomic-ai/nomic-embed-text-v1.5-GGUF -hff nomic-embed-text-v1.5.Q4_K_M.gguf \
      --rope-freq-scale 0.75 --verbose

    EMBEDDINGS_MODEL_FILE=~/Library/Caches/llama.cpp/nomic-embed-text-v1.5.Q4_K_M.gguf \
        MEMORY_SQLITE_DB=memory_local.db \
            python -m examples.agent.tools.memory_sqlite

    EMBEDDINGS_ENDPOINT=http://localhost:8081/v1/embeddings \
        MEMORY_SQLITE_DB=memory_remote.db \
            python -m examples.agent.tools.memory_sqlite

'''
import asyncio
import os
import aiosqlite
from typing import List
from tqdm import tqdm
import time


import sqlite_lembed
import sqlite_rembed
import sqlite_vec

from .utils import deindent_code

db_path = os.environ['MEMORY_SQLITE_DB']
embeddings_endpoint = os.environ.get('EMBEDDINGS_ENDPOINT')
embeddings_api_key = os.environ.get('EMBEDDINGS_API_KEY')
embeddings_model_file = os.environ.get('EMBEDDINGS_MODEL_FILE')

class MemoryDb:
    def __init__(self, db, embed_fn):
        self.db = db
        self.embed_fn = embed_fn

    @staticmethod
    async def setup(db: aiosqlite.Connection) -> 'MemoryDb':
        
        await db.enable_load_extension(True)
        await db.load_extension(sqlite_vec.loadable_path())
        
        if embeddings_model_file:
            await db.load_extension(sqlite_lembed.loadable_path())
            local = True
            embed_fn = lambda x: f'lembed("default", {x})'
        elif embeddings_endpoint:
            await db.load_extension(sqlite_rembed.loadable_path())
            local = False
            embed_fn = lambda x: f'rembed("default", {x})'
        else:
            raise Exception('Must define EMBEDDINGS_ENDPOINT for remote embeddings or EMBEDDINGS_MODEL_FILE for local embeddings (CPU)')
        
        await db.enable_load_extension(False)

        if local:
            await db.execute('''
                INSERT INTO lembed_models(name, model) VALUES (
                    'default',lembed_model_from_file(?)
                );
            ''', (embeddings_model_file,))
        else:
            await db.execute('''
                INSERT INTO rembed_clients(name, options) VALUES (
                    'default', rembed_client_options('format', 'llamafile', 'url', ?, 'key', ?)
                );
            ''', (embeddings_endpoint, embeddings_api_key))
            
        await db.execute('''
            CREATE TABLE IF NOT EXISTS _schema_history (
                rowid INTEGER PRIMARY KEY AUTOINCREMENT,
                statement TEXT NOT NULL
            )
        ''')

        schema_history = [
            #
            # BEWARE: ONLY APPEND NEW STATEMENTS TO THE END OF THIS LIST.
            #         DO NOT MODIFY OR DELETE ANY EXISTING STATEMENTS
            #         OTHERWISE, AUTOMATIC SCHEMA MIGRATION WILL BREAK!
            #
            '''
            CREATE TABLE IF NOT EXISTS documents (
                rowid INTEGER PRIMARY KEY AUTOINCREMENT,
                content TEXT NOT NULL
            )
            ''',
            '''
            CREATE VIRTUAL TABLE IF NOT EXISTS document_embeddings USING vec0(
                embedding float[768]
            )
            ''',
            f'''
            CREATE TRIGGER IF NOT EXISTS insert_document_embedding
            AFTER INSERT ON documents
            BEGIN
                INSERT INTO document_embeddings (rowid, embedding)
                VALUES (NEW.rowid, {embed_fn('NEW.content')});
            END;
            ''',
            f'''
            CREATE TRIGGER IF NOT EXISTS update_document_embedding
            AFTER UPDATE OF content ON documents
            BEGIN
                UPDATE document_embeddings
                SET embedding = {embed_fn('NEW.content')}
                WHERE rowid = NEW.rowid;
            END;
            ''',
            '''
            CREATE TRIGGER IF NOT EXISTS delete_document_embedding
            AFTER DELETE ON documents
            BEGIN
                DELETE FROM document_embeddings
                WHERE rowid = OLD.rowid;
            END;
            ''',
        ]

        inserted_count = (await (await db.execute('SELECT max(rowid) FROM _schema_history')).fetchone() or (0,))[0] or 0
        if inserted_count > len(schema_history):
            raise Exception(f'Schema history mismatch: inserted count `{inserted_count}` > known schema count `{len(schema_history)}`')

        for (i, statement) in enumerate(map(deindent_code, schema_history)):
            if i < inserted_count:
                inserted_statement = (await (await db.execute('SELECT statement FROM _schema_history WHERE rowid = ?', (i + 1,))).fetchone() or (None,))[0]
                if inserted_statement != statement:
                    raise Exception(f'Schema history mismatch at statement {i + 1}: inserted statement `{inserted_statement}` != proposed statement `{statement}`')
            else:
                print(f'Executing schema statement {i + 1}: {statement}')
                await db.execute('INSERT INTO _schema_history (statement) VALUES (?)', (statement,))
                await db.execute(statement)

        await db.commit()

        return MemoryDb(db, embed_fn)

    async def search(self, question: str, top_n: int = 10):
        async with self.db.execute(
            f'''
            SELECT documents.rowid, documents.content, distance
            FROM (
                select rowid, distance
                from document_embeddings
                WHERE document_embeddings.embedding MATCH {self.embed_fn('?')}
                ORDER BY distance
                LIMIT ?
            )
            JOIN documents using (rowid)
            ''',
            (question, top_n)
        ) as cursor:
            results = await cursor.fetchall()
            return results

    async def memorize(self, documents: List[str]):
        await self.db.executemany(
            'INSERT INTO documents (content) VALUES (?)',
            [(doc,) for doc in documents]
        )
        await self.db.commit()

async def search_memory(question: str, top_n: int = 10):
    '''
        Search the memory for a question.
    '''

    async with aiosqlite.connect(db_path) as db:
        memory = await MemoryDb.setup(db)
        return await memory.search(question, top_n)

async def memorize(documents: List[str]):
    '''
        Memorize a set of statements / facts.
    '''

    async with aiosqlite.connect(db_path) as db:
        memory = await MemoryDb.setup(db)
        await memory.memorize(documents)

async def main():
    # time method
    async with aiosqlite.connect(db_path) as db:
        memory = await MemoryDb.setup(db)

        await memory.memorize([
            "User's name is Olivier Chafik",
            "User lives in London",
        ])
        
        
        N = 1000
        print(f"Inserting {N} embeddings...")
        start_time = time.perf_counter()
        await memory.memorize(
            [f"User is a human {i}" for i in range(N)]
        )
        end_time = time.perf_counter()
        elapsed_sec = end_time - start_time
        print(f"Insertion of {N} embeddings took {elapsed_sec} sec")

        for q in ["Where does user live?", "What is user's name?", "Who is the user?"]:
            print(f"Searching for: {q}")
            print(await memory.search(q))
            print()
            
    

if __name__ == '__main__':
    asyncio.run(main())