ai_pipeline/Demo.md

12 KiB
Raw Blame History

太好了 — 我把 两套完整示例代码 给你:

  1. Elasticsearch + Qdrant 双库 RAG pipelinePython 模块)

    • 作用:把文档同时写入 ES做 BM25 倒排检索)和 Qdrant做 embedding/语义检索),检索时合并两个检索结果并调用 vLLM reranker 返回最相关结果。
  2. FastAPI 混合检索服务(直接可运行 API

    • 作用:把上面的 pipeline 封装成 HTTP 接口(/ingest_text, /search 等),方便接入前端或其它服务。

两套示例以工程可用、易读为主。生产化时建议:把阻塞客户端改为 async或在线程池里并发、增加重试/限流、加入监控与认证。

下面先给出依赖 & 配置(最小),然后是两套代码。


依赖requirements

# requirements.txt
fastapi
uvicorn[standard]
httpx
qdrant-client
elasticsearch>=8.6.0
python-dotenv
numpy
python-multipart

在项目根放 .env(或使用环境变量):

# .env
VLLM_API_KEY=
VLLM_EMBED_URL=http://localhost:8000/v1/embeddings
VLLM_RERANK_URL=http://localhost:8000/v1/rerank
VLLM_CHAT_URL=http://localhost:8000/v1/chat/completions

QDRANT_HOST=localhost
QDRANT_PORT=6333

ES_HOST=http://localhost:9200
ES_INDEX=my_kb_index

EMBED_MODEL=text-embedding-3-large
RERANK_MODEL=bge-reranker-v2-m3
EMBED_DIM=1536

TOP_K=32
TOP_N=5

1) Elasticsearch + Qdrant 双库 RAG pipeline模块

把这个文件保存为 pipeline_es_qdrant.py

# pipeline_es_qdrant.py
import os
import uuid
import math
import httpx
from typing import List, Dict, Tuple
from dotenv import load_dotenv

from qdrant_client import QdrantClient
from qdrant_client.models import VectorParams, Distance

from elasticsearch import Elasticsearch

load_dotenv()

# config from env
VLLM_EMBED_URL = os.getenv("VLLM_EMBED_URL")
VLLM_RERANK_URL = os.getenv("VLLM_RERANK_URL")
EMBED_MODEL = os.getenv("EMBED_MODEL", "text-embedding-3-large")
EMBED_DIM = int(os.getenv("EMBED_DIM", "1536"))

QDRANT_HOST = os.getenv("QDRANT_HOST", "localhost")
QDRANT_PORT = int(os.getenv("QDRANT_PORT", "6333"))

ES_HOST = os.getenv("ES_HOST", "http://localhost:9200")
ES_INDEX = os.getenv("ES_INDEX", "my_kb_index")

TOP_K = int(os.getenv("TOP_K", "32"))

# clients (synchronous)
qdrant = QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT)
es = Elasticsearch(ES_HOST)

# ensure qdrant collection & es index
def ensure_qdrant_collection(name: str, dim: int = EMBED_DIM):
    try:
        qdrant.get_collection(collection_name=name)
    except Exception:
        qdrant.recreate_collection(
            collection_name=name,
            vectors_config=VectorParams(size=dim, distance=Distance.COSINE)
        )

def ensure_es_index(index_name: str):
    if not es.indices.exists(index=index_name):
        body = {
            "mappings": {
                "properties": {
                    "text": {"type": "text"},
                    "source": {"type": "keyword"},
                    "meta": {"type": "object"}
                }
            }
        }
        es.indices.create(index=index_name, body=body)

# embeddings via vLLM OpenAI-style endpoint
async def embed_texts(texts: List[str]) -> List[List[float]]:
    async with httpx.AsyncClient(timeout=60) as client:
        r = await client.post(VLLM_EMBED_URL, json={"model": EMBED_MODEL, "input": texts})
        r.raise_for_status()
        data = r.json()
    # support OpenAI-style response
    return [item["embedding"] for item in data["data"]]

# Ingest: add chunks to ES (for BM25) and Qdrant (for embeddings)
async def ingest_chunks(kb: str, chunks: List[Dict]):
    """
    chunks: list of {"id": optional, "text": str, "meta": dict}
    Writes to ES (document) and Qdrant (vector)
    """
    ensure_es_index(ES_INDEX)
    ensure_qdrant_collection(kb, dim=EMBED_DIM)

    texts = [c["text"] for c in chunks]
    embeddings = await embed_texts(texts)

    # bulk index to ES
    es_actions = []
    for c, emb in zip(chunks, embeddings):
        doc_id = c.get("id") or str(uuid.uuid4())
        es.index(index=ES_INDEX, id=doc_id, document={"text": c["text"], "source": c.get("meta", {}).get("source"), "meta": c.get("meta", {})})
    # upsert to qdrant
    points = []
    for c, emb in zip(chunks, embeddings):
        pid = c.get("id") or str(uuid.uuid4())
        points.append({"id": pid, "vector": emb, "payload": {"text": c["text"], **(c.get("meta") or {})}})
    qdrant.upsert(collection_name=kb, points=points)
    return {"ok": True, "ingested": len(points)}

# Search: BM25 via ES
def es_search(query: str, top_k: int = 10) -> List[Dict]:
    resp = es.search(index=ES_INDEX, body={"query": {"match": {"text": {"query": query}}}, "size": top_k})
    hits = []
    for h in resp["hits"]["hits"]:
        hits.append({"id": h["_id"], "score": h["_score"], "text": h["_source"]["text"], "meta": h["_source"].get("meta")})
    return hits

# Qdrant search
def qdrant_search(kb: str, q_emb: List[float], top_k: int = TOP_K) -> List[Dict]:
    hits = qdrant.search(collection_name=kb, query_vector=q_emb, limit=top_k)
    out = []
    for h in hits:
        payload = h.payload or {}
        out.append({"id": h.id, "score": getattr(h, "score", None), "text": payload.get("text"), "meta": payload})
    return out

# Merge results strategy:
# - gather ES top_k and Qdrant top_k
# - deduplicate by id and produce candidate list
def merge_candidates(es_hits: List[Dict], q_hits: List[Dict], weight_es: float = 1.0, weight_q: float = 1.0) -> List[Dict]:
    # map by id with combined score (normalized)
    candidates = {}
    # normalize ES scores to 0..1 by dividing by max (if present)
    max_es = max((h["score"] for h in es_hits), default=1.0)
    max_q = max((h["score"] or 1.0 for h in q_hits), default=1.0)
    for h in es_hits:
        sid = h["id"]
        s = (h["score"] or 0.0) / max_es
        candidates.setdefault(sid, {"id": sid, "text": h["text"], "meta": h.get("meta", {}), "es_score": s, "q_score": 0.0})
        candidates[sid]["es_score"] = s
    for h in q_hits:
        sid = h["id"]
        s = (h["score"] or 0.0) / max_q
        candidates.setdefault(sid, {"id": sid, "text": h["text"], "meta": h.get("meta", {}), "es_score": 0.0, "q_score": s})
        candidates[sid]["q_score"] = s
    # compute hybrid score
    for sid, v in candidates.items():
        v["hybrid_score"] = weight_es * v["es_score"] + weight_q * v["q_score"]
    # sort by hybrid_score desc
    return sorted(candidates.values(), key=lambda x: x["hybrid_score"], reverse=True)

# Rerank via vLLM reranker endpoint (OpenAI-style)
async def rerank_with_vllm(query: str, docs: List[str], model: str = None) -> List[int]:
    model = model or os.getenv("RERANK_MODEL")
    async with httpx.AsyncClient(timeout=60) as client:
        r = await client.post(VLLM_RERANK_URL, json={"model": model, "query": query, "documents": docs})
        r.raise_for_status()
        data = r.json()
    # expect data["results"] = [{"index":i,"score":...}, ...]
    order = [item["index"] for item in sorted(data["results"], key=lambda x: -x["score"])]
    return order

# Full pipeline: query -> es + qdrant -> merge -> rerank -> return top_n
async def hybrid_search(kb: str, query: str, top_k_es: int = 8, top_k_q: int = 24, top_n: int = 5) -> Dict:
    # 1 get ES hits
    es_hits = es_search(query, top_k_es)
    # 2 embed query and qdrant search
    q_emb = (await embed_texts([query]))[0]
    q_hits = qdrant_search(kb, q_emb, top_k_q)
    # 3 merge candidates
    candidates = merge_candidates(es_hits, q_hits, weight_es=1.0, weight_q=1.0)
    # 4 rerank top M by calling reranker
    M = min(len(candidates), 50)
    docs = [c["text"] for c in candidates[:M]]
    if docs:
        order = await rerank_with_vllm(query, docs)
        ordered = [candidates[:M][i] for i in order][:top_n]
    else:
        ordered = candidates[:top_n]
    return {"query": query, "results": ordered}

说明与要点:

  • ES 做 BM25match 查询),Qdrant 做向量召回;合并时把两边的分数 normalize 后加权得到 hybrid_score,然后再交给大型 rerankervLLM精排。
  • embed_texts 使用了异步 httpx 调用。Qdrant/ES 操作是同步的(如果你想更高并发,把 ES/Qdrant 客户端换成 async 或把调用放进线程池)。
  • merge_candidates 是基础合并逻辑可替换为更复杂的策略交叉乘积、IDF 规则、source bias 等)。
  • reranker 接口需由你在 vLLM 那端实现:接收 query + documents,返回每个 document 的 score 和 index。

2) FastAPI 版本的混合检索 API封装上面 pipeline

下面示例展示一个小服务,提供 /ingest_text/search 两个接口。把它保存为 app_fastapi.py(或在你的 FastAPI 项目中合并)。

# app_fastapi.py
import os
import asyncio
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List, Optional
from dotenv import load_dotenv

load_dotenv()

from pipeline_es_qdrant import ingest_chunks, hybrid_search  # import above module

app = FastAPI(title="Hybrid RAG API")

class IngestReq(BaseModel):
    kb: str
    chunks: List[dict]  # each { "text": "...", "meta": {...}, "id": optional }

class SearchReq(BaseModel):
    kb: str
    query: str
    top_n: Optional[int] = 5

@app.post("/ingest_text")
async def ingest_text(req: IngestReq):
    try:
        res = await ingest_chunks(req.kb, req.chunks)
        return res
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/search")
async def search(req: SearchReq):
    try:
        res = await hybrid_search(req.kb, req.query, top_n=req.top_n)
        return res
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/healthz")
async def healthz():
    return {"status": "ok"}

如何运行:

uvicorn app_fastapi:app --reload --port 8000

示例流程:

  1. 导入文档(最简单方式:把文章切成 chunks然后调用 /ingest_text 上传):
curl -X POST "http://localhost:8000/ingest_text" -H "Content-Type: application/json" -d '{
  "kb":"kb_demo",
  "chunks":[{"text":"文档1 内容 ...","meta":{"source":"doc1"}},{"text":"文档2 ..."}]
}'
  1. 检索:
curl -X POST "http://localhost:8000/search" -H "Content-Type: application/json" -d '{
  "kb":"kb_demo",
  "query":"如何使用产品A",
  "top_n":5
}'

生产部署 & 性能建议(实践要点)

  1. 并发与阻塞:上面示例中 ES/Qdrant 客户端为同步实现(阻塞)。在 FastAPI 高并发场景下建议:

    • 使用 AsyncQdrantClientqdrant-client[httpx])与 AsyncElasticsearchelasticsearch 8+ 支持 async via aiohttp/httpx)或
    • 将阻塞调用放到线程池:await asyncio.get_event_loop().run_in_executor(None, sync_call, args)
  2. 批量与分片大规模写入时应使用批量ES bulk API、Qdrant batch upsert并控制批大小例如每批 64/128

  3. Hybrid 调整

    • 合并评分的权重 (weight_es, weight_q) 可基于离线验证MRR、P@k调整。
    • Reranker 只对 top-M 做精排(例如 M=50避免太多计算成本。
  4. 安全与限流:在公开 API 上加认证API Key、请求频率限制和输入长度限制避免过大上下文请求

  5. 监控记录每个阶段耗时embed、es_search、qdrant_search、rerank、llm用于优化和 SLO。


还可以做的扩展(我可以帮你继续做)

  • 把 ES/Qdrant 客户端都改为 异步实现(完整 async 版本)并替换到代码中。
  • 增加 文档 loaderdocx/pdf + 自动切分器 并在 /ingest_file 中直接上传文件。
  • 提供 示例 docker-compose 把 Qdrant、Elasticsearch、vLLM或模拟服务和你的 API 一起编排。
  • 实现更复杂的 混合融合策略(比如先 union再基于 metadata过滤再 rerank

你想先要哪个扩展(例如:异步版本文件上传 + 自动切分docker-compose)?我可以立刻把对应代码补上。