【指南】SQLite中的检索增强生成

2025年02月20日 由 alex 发表 2724 0

我们真的可以用单文件架构来实现 RAG 吗?


在本文中,我将向你介绍如何在最简单的堆栈上执行 RAG。事实上,这个“堆栈”只是带有 sqlite-vec 扩展的 SQLite 和用于嵌入和聊天模型的 OpenAI API。就我们的目的而言,只要理解 SQLite 是最简单的数据库类型就足够了:存储库中的一个文件。


SQLite-Vec

SQLite数据库的强大之处之一在于其扩展的使用。对于熟悉Python的人来说,扩展很像库。它们是用C语言编写的模块化代码,用于扩展SQLite的功能,使曾经不可能的事情变得可能。一个流行的SQLite扩展示例是全文搜索(FTS)扩展。该扩展使SQLite能够在大量文本数据中执行高效搜索。由于该扩展完全是用C语言编写的,因此我们可以在任何可以运行SQLite数据库的地方运行它,包括Rasbperry Pis和浏览器。


在本文中,我将介绍一个名为sqlite-vec的扩展。它赋予了SQLite执行向量搜索的能力。向量搜索与全文搜索相似,因为它允许在文本数据上进行高效搜索。然而,向量搜索并不是在文本中搜索确切的单词或短语,而是具有语义理解能力。换句话说,搜索“horses”将找到“equestrian”、“pony”、“Clydesdale”等匹配项。全文搜索无法做到这一点。全文搜索则无法做到这一点。


sqlite-vec利用虚拟表,就像SQLite中的大多数扩展一样。虚拟表与普通表相似,但具有额外的功能:

  • 自定义数据源:SQLite中标准表的数据存储在一个单独的数据库文件中。而对于虚拟表,数据可以存储在外部来源中,例如CSV文件或API调用。
  • 灵活的功能:虚拟表可以添加专门的索引或查询功能,并支持复杂的数据类型,如JSON或XML。
  • 与SQLite查询引擎集成:虚拟表与SQLite的标准查询语法(如SELECT、INSERT、UPDATE和DELETE选项)无缝集成。最终,是否支持这些操作取决于扩展的编写者。
  • 使用模块:虚拟表的工作后端逻辑由一个模块(用C语言或其他语言编写)实现。


创建虚拟表的典型语法如下:


CREATE VIRTUAL TABLE my_table USING my_extension_module();


这段语句 my_extension_module() 的导入部分。这指定了将为 my_table 虚拟表的后端提供动力的模块。在 sqlite-vec 中,我们将使用 vec0 模块。


代码解析

本文的代码可以在这里找到。这是一个简单的目录,其中大部分文件是 .txt 文件,我们将它们作为示例数据使用。因为我是一个物理迷,所以大部分文件都与物理有关,只有少数文件涉及其他随机领域。在本解析中,我不会展示完整代码,而是会突出显示重要的部分。以下是仓库的树状视图。请注意,my_docs.db 是 SQLite 用来管理我们所有数据的单文件数据库。


.
├── data
│   ├── cooking.txt
│   ├── gardening.txt
│   ├── general_relativity.txt
│   ├── newton.txt
│   ├── personal_finance.txt
│   ├── quantum.txt
│   ├── thermodynamics.txt
│   └── travel.txt
├── my_docs.db
├── requirements.txt
└── sqlite_rag_tutorial.py


第一步是安装必要的库。以下是我们的 requirements.txt 文件。如你所见,它只有三个库。我建议使用最新版本的 Python(本文使用的是 3.13.1)创建一个虚拟环境,然后运行 pip install -r requirements.txt 来安装这些库。


# requirements.txt
sqlite-vec==0.1.6
openai==1.63.0
python-dotenv==1.0.1


第二步是创建一个 OpenAI API 密钥(如果你还没有的话)。我们将使用 OpenAI 为文本文件生成嵌入,以便执行向量搜索。


# sqlite_rag_tutorial.py
import sqlite3
from sqlite_vec import serialize_float32
import sqlite_vec
import os
from openai import OpenAI
from dotenv import load_dotenv
# Set up OpenAI client
client = OpenAI(api_key=os.getenv('OPENAI_API_KEY'))


第三步是将 sqlite-vec 扩展加载到 SQLite 中。在本文的示例中,我们将使用 Python 和 SQL。在加载扩展后立即禁用加载扩展的能力是一个良好的安全实践。


# Path to the database file
db_path = 'my_docs.db'
# Delete the database file if it exists
db = sqlite3.connect(db_path)
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
    )
''')


documents 是一个具有三列的虚拟表:

  • sample_embedding:1536维浮点数,用于存储我们样本文档的嵌入。
  • file_name:文本字段,用于存储我们存储在数据库中的每个文件的名称。注意,此列和后面的列前面有一个 + 符号。这表示它们是辅助字段。在以前的 sqlite-vec 中,虚拟表只能存储嵌入数据。然而,最近推出了一项更新,允许我们在表中添加我们实际上不需要嵌入的字段。在这种情况下,我们在与嵌入相同的表中添加了文件的内容和名称。这样,我们可以轻松地查看哪些嵌入对应于哪些内容,同时省去了额外表和 JOIN 语句的需要。
  • content:文本字段,用于存储每个文件的内容。


现在,我们已经在 SQLite 数据库中设置了虚拟表,可以开始将文本文件转换为嵌入并将其存储在表中:


# Function to get embeddings using the OpenAI API
def get_openai_embedding(text):
    response = client.embeddings.create(
        model="text-embedding-3-small",
        input=text
    )
    return response.data[0].embedding

# Iterate over .txt files in the /data directory
for file_name in os.listdir("data"):
    file_path = os.path.join("data", file_name)
    with open(file_path, 'r', encoding='utf-8') as file:
        content = file.read()
        # Generate embedding for the content
        embedding = get_openai_embedding(content)
        if embedding:
            # Insert file content and embedding into the vec0 table
            db.execute(
                'INSERT INTO documents (embedding, file_name, content) VALUES (?, ?, ?)',
                (serialize_float32(embedding), file_name, content)
# Commit changes
db.commit()


我们基本上遍历每个 .txt 文件,将每个文件的内容嵌入,然后使用 INSERT INTO 语句将嵌入、文件名和内容插入到 documents 虚拟表中。最后的提交语句确保更改被持久化。请注意,我们在这里使用了 sqlite-vec 库中的 serialize_float32。SQLite 本身没有内置的向量类型,因此它将向量存储为二进制大对象(BLOB)以节省空间并允许快速操作。在内部,它使用 Python 的 struct.pack() 函数,该函数将 Python 数据转换为 C 风格的二进制表示。


最后,为了执行检索增强生成(RAG),你使用以下代码执行 K-最近邻(KNN 风格)操作。这是向量搜索的核心。


# Perform a sample KNN query
query_text = "What is general relativity?"
query_embedding = get_openai_embedding(query_text)
if query_embedding:
    rows = db.execute(
        """
        SELECT
            file_name,
            content,
            distance
        FROM documents
        WHERE embedding MATCH ?
        ORDER BY distance
        LIMIT 3
        """,
        [serialize_float32(query_embedding)]
    ).fetchall()
    print("Top 3 most similar documents:")
    top_contexts = []
    for row in rows:
        print(row)
        top_contexts.append(row[1])  # Append the 'content' column


我们首先接收用户的查询,在这个例子中是“什么是广义相对论?”,并使用与之前相同的嵌入模型对该查询进行嵌入。然后,我们执行一个 SQL 操作。让我们来分解一下:

  • SELECT 语句表示检索的数据将有三列:file_name、content 和 distance。前两列我们已经提到过。distance 将在 SQL 操作期间计算,稍后会详细说明。
  • FROM 语句确保你从 documents 表中提取数据。
  • WHERE embedding MATCH ? 语句在你的数据库中的所有向量和查询向量之间执行相似性搜索。返回的数据将包括一个 distance 列。这个距离只是一个浮点数,用于衡量查询向量和数据库向量之间的相似性。数字越大,向量越接近。sqlite-vec 提供了几种计算这种相似性的选项。
  • ORDER BY distance 确保按相似性降序(高 -> 低)对检索到的向量进行排序。
  • LIMIT 3 确保我们只获取与查询嵌入向量最接近的前三个文档。你可以调整这个数字,以查看检索更多或更少向量对你的结果有何影响。


对于我们的查询“什么是广义相对论?”,检索到了以下文档。结果相当不错!


Top 3 most similar documents:
('general_relativity.txt', 'Einstein’s theory of general relativity redefined our understanding of gravity. Instead of viewing gravity as a force acting at a distance, it interprets it as the curvature of spacetime around massive objects. Light passing near a massive star bends slightly, galaxies deflect beams traveling millions of light-years, and clocks tick at different rates depending on their gravitational potential. This groundbreaking theory led to predictions like gravitational lensing and black holes, phenomena later confirmed by observational evidence, and it continues to guide our understanding of the cosmos.', 0.8316285610198975)
('newton.txt', 'In classical mechanics, Newton’s laws of motion form the foundation of how we understand the movement of objects. Newton’s first law, often called the law of inertia, states that an object at rest remains at rest and an object in motion continues in motion unless acted upon by an external force. This concept extends into more complex physics problems, where analyzing net forces on objects allows us to predict their future trajectories and behaviors. Over time, applying Newton’s laws has enabled engineers and scientists to design safer vehicles, more efficient machines, and even guide spacecraft through intricate gravitational fields.', 1.2036118507385254)
('quantum.txt', 'Quantum mechanics revolutionized our understanding of the microscopic world. Unlike classical particles, quantum entities such as electrons can exhibit both wave-like and particle-like behaviors. Phenomena like quantum superposition suggest that particles can exist in multiple states at once, and the act of measurement often “collapses” these states into one observed outcome. This strange and counterintuitive theory underpins modern technologies like semiconductors and lasers, and it provides a conceptual framework for emerging fields like quantum computing and cryptography.', 1.251380205154419)


然后,我们可以将这三个文档的内容填入模型中,让它尝试回答我们的问题。


# Prepare the context for the query
context = "\n\n".join(top_contexts)
system_message = "You are a helpful assistant. Use the following context to answer the query."
# Send query and context to OpenAI
try:
    completion = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": system_message},
            {"role": "user", "content": f"Context: {context}\n\nQuery: {query_text}"}
        ]
    )
    print("Response:")
    print(completion.choices[0].message.content)
except Exception as e:
    print(f"Error generating response: {e}")


以下是模型的回答:


General relativity is a theory developed by Albert Einstein that redefines 
our understanding of gravity. Instead of viewing gravity as a force acting 
at a distance, general relativity interprets it as the curvature of spacetime 
caused by the presence of mass. According to this theory, massive objects 
like stars and planets create a distortion in the fabric of spacetime, and 
this curvature affects the motion of other objects, including light, which 
can bend when it passes near a massive body.


这个回答忠实地遵循了我们提供给模型的文档内容。4o-mini,干得好!


结论

sqlite-vec 是一个由 Mozilla Builders Accelerator 计划赞助的项目,因此它有着相当重要的支持。这是一个维护良好的库,定期会有更新发布。他们甚至添加了按元数据过滤的功能!

该扩展还为多种流行的编程语言提供了绑定,包括 Ruby、Go、Rust 等。


11

文章来源:https://medium.com/data-science-collective/retrieval-augmented-generation-in-sqlite-36302a28fe51
欢迎关注ATYUN官方公众号
商务合作及内容投稿请联系邮箱:bd@atyun.com
评论 登录
热门职位
Maluuba
20000~40000/月
Cisco
25000~30000/月 深圳市
PilotAILabs
30000~60000/年 深圳市
写评论取消
回复取消