Retrieval Augmented Generation in SQLite

This two-part series explores using SQLite for machine learning. The previous article discussed SQLite's growing role in production-ready web applications. This article focuses on implementing retrieval-augmented generation (RAG) using SQLite.

For custom web applications with generative AI, visit losangelesaiapps.com

The code is available here.

Traditional RAG implementation often involves:

1. Searching for tutorials on RAG.
2. Selecting a popular framework (LangChain, LlamaIndex).
3. Choosing a cloud vector database (Pinecone, Weaviate).
4. Integrating these components.

While effective, this approach can be overly complex, especially for beginners. This article demonstrates a simpler method using SQLite with the sqlite-vec extension and the OpenAI API. Part 1 of this series provides a detailed overview of SQLite's capabilities. For this article, it suffices to understand SQLite's simplicity as a single-file database.

This approach eliminates the need for cloud vector databases and bulky frameworks.

SQLite-Vec: Extending SQLite's Power

SQLite's strength lies in its extensibility. Extensions, similar to Python libraries, add functionality written in C. A prime example is the Full-Text Search (FTS) extension. sqlite-vec adds vector search capabilities, enabling semantic understanding beyond keyword matching. Searching for "horses" might return "equestrian" or "pony."

sqlite-vec uses virtual tables, offering:

• Custom Data Sources: Data can reside outside the database file (e.g., CSV, API).
• Flexible Functionality: Supports specialized indexing and complex data types.
• Seamless Integration: Integrates with standard SQLite query syntax.
• Modules: Backend logic is implemented in a separate module.

Virtual tables are created using:

CREATE VIRTUAL TABLE my_table USING my_extension_module();

my_extension_module() specifies the module (here, vec0 from sqlite-vec).

Code Walkthrough

The code (repo link) uses .txt files as sample data (mostly physics-related). my_docs.db is the SQLite database file.

1. Installation: requirements.txt lists the necessary libraries (sqlite-vec, openai, python-dotenv). Create a virtual environment and run pip install -r requirements.txt.

2. OpenAI API Key: Obtain an OpenAI API key.

3. Loading the Extension: The Python code loads the sqlite-vec extension and creates a virtual table:

CREATE VIRTUAL TABLE my_table USING my_extension_module();

The documents table stores embeddings (embedding), filenames (file_name), and content (content). denotes auxiliary fields.

4. Embedding and Insertion: The code iterates through .txt files, generates embeddings using the OpenAI API, and inserts them into the database:

db.enable_load_extension(True)
sqlite_vec.load(db)
db.enable_load_extension(False)

db.execute('''
    CREATE VIRTUAL TABLE documents USING vec0(
        embedding float[1536],
        +file_name TEXT,
        +content TEXT
    )
''')

5. RAG Query: A KNN query retrieves similar documents based on embedding similarity:

# ... (OpenAI embedding function) ...

for file_name in os.listdir("data"):
    # ... (Open file, get content, get embedding) ...
    db.execute(
        'INSERT INTO documents (embedding, file_name, content) VALUES (?, ?, ?)',
        (serialize_float32(embedding), file_name, content)
    )
db.commit()

The results are then used as context for an OpenAI chat completion to answer the query.

Conclusion

sqlite-vec simplifies RAG significantly. It eliminates the need for complex frameworks and cloud services, making it cost-effective and easy to iterate upon. While scaling might eventually require a more robust database, sqlite-vec offers a compelling solution for smaller projects. The extension supports multiple programming languages.

The above is the detailed content of Retrieval Augmented Generation in SQLite. For more information, please follow other related articles on the PHP Chinese website!