从零拼一个最小 RAG 问答

用户问题
   ↓ ① 检索:把问题和知识库都变成向量,找最相似的几段
相关片段
   ↓ ② 增强:把片段拼进 prompt("根据以下资料回答…")
拼好的 prompt
   ↓ ③ 生成:让 LLM 基于资料作答
带依据的回答

python -m pip install numpy
export OPENAI_API_KEY="替换为支持该embedding模型的OpenAI API key"

import os
import numpy as np
from openai import OpenAI

embedding_client = OpenAI(api_key=os.environ["OPENAI_API_KEY"], timeout=30.0, max_retries=2)
chat_client = OpenAI(base_url=os.environ["LLM_BASE_URL"], api_key=os.environ["LLM_API_KEY"], timeout=30.0, max_retries=2)
MODEL = os.environ["LLM_MODEL"]

# 0) 一个迷你知识库
docs = [
    "本公司年假:入职满 1 年 5 天,满 3 年 10 天。",
    "报销流程:登录 OA 提交发票,主管审批后 3 个工作日到账。",
    "考勤:弹性工作,核心时段 10:00-16:00 需在岗。",
]

def embed(texts):
    r = embedding_client.embeddings.create(model="text-embedding-3-small", input=texts)
    return np.array([d.embedding for d in sorted(r.data, key=lambda item: item.index)])

doc_vecs = embed(docs)                       # ① 离线:知识库向量化(只算一次)

def cosine(a, b):
    denominator = np.linalg.norm(a) * np.linalg.norm(b)
    return float(a @ b / denominator) if denominator else 0.0

def answer(question, k=2):
    q_vec = embed([question])[0]             # ① 在线:问题向量化
    sims = [cosine(q_vec, dv) for dv in doc_vecs]
    top = np.argsort(sims)[::-1][:k]         # 取最相似的 k 段
    context = "\n".join(docs[i] for i in top)

    prompt = f"根据以下资料回答问题,资料没提到就说不知道。\n\n资料:\n{context}\n\n问题:{question}"  # ②
    resp = chat_client.chat.completions.create(   # ③ 生成
        model=MODEL,
        messages=[{"role": "user", "content": prompt}],
    )
    return resp.choices[0].message.content

print(answer("入职两年有几天年假?"))   # → 5 天(满 1 年 5 天,满 3 年才 10 天)