使用Timescale/pgai构建基于FastAPI的语义搜索与RAG应用教程

最新推荐文章于 2025-06-08 09:19:38 发布

霍美予Mabel

最新推荐文章于 2025-06-08 09:19:38 发布

阅读量350

点赞数 5

CC 4.0 BY-SA版权

本文链接：https://blog.youkuaiyun.com/gitblog_00342/article/details/148508529

使用Timescale/pgai构建基于FastAPI的语义搜索与RAG应用教程

pgai Helper functions for AI workflows 项目地址: https://gitcode.com/gh_mirrors/pg/pgai

项目概述

Timescale/pgai是一个强大的PostgreSQL扩展，它集成了AI功能，使开发者能够直接在数据库中执行向量搜索和检索增强生成(RAG)等AI任务。本教程将详细介绍如何使用pgai结合FastAPI框架构建一个完整的语义搜索和RAG应用。

环境准备

在开始之前，需要准备以下环境：

安装PostgreSQL数据库并配置好连接
获取OpenAI API密钥
安装Python 3.7+环境
安装FastAPI和相关依赖

应用架构

这个FastAPI应用主要包含以下几个核心功能模块：

数据库初始化与pgai扩展安装
数据加载与向量化处理
语义搜索功能
检索增强生成(RAG)功能

详细实现步骤

1. 初始化pgai扩展

在应用启动时，我们需要安装pgai扩展并启动向量化工作器：

# 安装pgai扩展
pgai.install(DB_URL)

# 启动向量化工作器
worker = Worker(DB_URL)
task = asyncio.create_task(worker.run())

这段代码会在PostgreSQL中创建ai模式，并安装所有必要的数据库对象。向量化工作器负责在后台异步处理向量嵌入的创建和更新。

2. 创建数据表并加载示例数据

我们创建了一个articles表来存储文章数据：

async def create_articles_table():
    async with pool.connection() as conn:
        async with conn.cursor() as cur:
            await cur.execute("""
                CREATE TABLE IF NOT EXISTS articles (
                    id SERIAL PRIMARY KEY,
                    url TEXT NOT NULL,
                    title TEXT NOT NULL,
                    text TEXT NOT NULL
                )
            """)
        await conn.commit()

然后从公开数据集加载了10篇英文文章：

async def load_articles():
    num_articles = 10
    max_text_length = 1000
    
    dataset = load_dataset("public_dataset", split=f"train", streaming=True)
    
    async with pool.connection() as conn:
        async with conn.cursor() as cur:
            for article in dataset.take(num_articles):
                await cur.execute(
                    "INSERT INTO articles (url, title, text) VALUES (%s, %s, %s)",
                    (article['url'], article['title'], article['text'][:max_text_length])
                )
            await conn.commit()

3. 配置向量化器

为了使文本数据支持语义搜索，我们需要为text列创建向量嵌入：

async def create_vectorizer():
    vectorizer_statement = CreateVectorizer(
        source="articles",
        target_table='articles_embedding_storage',
        loading=LoadingColumnConfig(column_name='text'),
        embedding=EmbeddingOllamaConfig(model='all-minilm', dimensions=384, base_url="http://localhost:11434")
    ).to_sql()
    
    try:
        async with pool.connection() as conn:
            async with conn.cursor() as cur:
                await cur.execute(vectorizer_statement)
            await conn.commit()
    except Exception as e:
        if "already exists" not in str(e):
            raise e

这段代码创建了一个向量化器，它会自动为articles表的text列生成向量嵌入，并将结果存储在articles_embedding_storage表中。我们使用了Ollama的all-minilm模型来生成384维的向量。

4. 检查向量化进度

向量化是一个异步过程，我们可以通过查询vectorizer_status视图来检查进度：

@app.get("/vectorizer_status")
async def vectorizer_status():
    async with pool.connection() as conn:
        async with conn.cursor(row_factory=dict_row) as cur:
            await cur.execute("SELECT * FROM ai.vectorizer_status")
            return await cur.fetchall()

当pending_items列显示为0时，表示所有嵌入向量都已创建完成。

5. 实现语义搜索

语义搜索的核心是_find_relevant_chunks函数：

async def _find_relevant_chunks(client: ollama.AsyncClient, query: str, limit: int = 2) -> List[ArticleSearchResult]:
    response = await client.embed(model="all-minilm", input=query)
    embedding = np.array(response.embeddings[0])
    
    async with pool.connection() as conn:
        async with conn.cursor(row_factory=class_row(ArticleSearchResult)) as cur:
            await cur.execute("""
                SELECT a.id, a.url, a.title, a.text, a.chunk, a.embedding <=> %s as distance
                FROM articles_embedding a
                ORDER BY distance
                LIMIT %s
            """, (embedding, limit))
            
            return await cur.fetchall()

这个函数首先将查询文本转换为向量，然后使用pgvector的<=>操作符计算余弦距离，找出与查询最相关的文本块。

6. 实现RAG功能

RAG功能结合了语义搜索和LLM生成能力：

@app.get("/rag")
async def rag(query: str) -> Optional[str]:
    client = ollama.AsyncClient(host="http://localhost:11434")
    
    chunks = await _find_relevant_chunks(client, query)
    context = "\n\n".join(f"{article.title}:\n{article.text}" for article, _ in chunks)
    
    prompt = f"""Question: {query}

Please use the following context to provide an accurate response:

{context}

Answer:"""
            
    response = await client.generate(
        model='tinyllama',
        prompt=prompt,
        stream=False
    )
    
    return response['response']

这个端点首先通过语义搜索找到相关上下文，然后将上下文和问题一起发送给LLM生成回答。