用 AI API + RAG 搭建企业知识库：从选模型到上线的完整方案

RAG检索增强生成是目前最实用的企业 AI 落地方案之一。本文从架构设计到代码实现手把手带你搭一套生产可用的知识库系统全程用 TheRouter 统一调用 Embedding Chat 模型一个 API Key 搞定所有事情。为什么需要 RAG大模型的知识是截止到训练日期的而且上下文窗口有限——你不可能把几千份公司内部文档全塞进 prompt。RAG 的思路很清晰离线阶段把文档切片 → 向量化 → 存进向量数据库在线阶段用户提问 → 把问题也向量化 → 在数据库里搜相似片段 → 把相关片段 原始问题一起喂给 LLM → 得到有据可查的回答这样既解决了模型不知道内部资料的问题也解决了塞太多 token 太贵的问题。架构概览文档(PDF/Word/MD) ↓ 文本提取 切分 ↓ Embedding 模型 (text-embedding-3-small) ↓ 向量数据库 (ChromaDB / Qdrant) ↓ ─────────── 在线检索 ─────────── 用户问题 → 向量化 → Top-K 检索 ↓ 构造 Prompt (问题 相关片段) ↓ 生成模型 (claude-sonnet-4 / gpt-4o) ↓ 返回带引用的答案技术选型Embedding 模型选openai/text-embedding-3-small理由1536 维精度够用成本极低约 $0.02/1M tokens通过 TheRouter 调用无需单独申请 OpenAI 账号比text-embedding-ada-002更新效果更好向量数据库场景推荐理由本地开发 / 小规模ChromaDBpip 安装即用零运维生产环境Qdrant支持分布式、过滤、payload 索引已有 PGpgvector不引入新组件本文用 ChromaDB 演示本地跑Qdrant 切换方案在最后给出。生成模型anthropic/claude-sonnet-4长上下文理解强适合文档类问答支持 200K contextopenai/gpt-4o速度快适合对话频繁的场景两个都通过 TheRouter 的 OpenAI 兼容接口调用切换只需改一个字符串。环境准备pipinstallopenai chromadb langchain langchain-text-splitters fastapi uvicorn python-multipart设置环境变量exportTHEROUTER_API_KEYyour-key-here核心实现1. 文档加载与切分# rag/splitter.pyfromlangchain_text_splittersimportRecursiveCharacterTextSplitterfrompathlibimportPathdefload_and_split(file_path:str)-list[dict]: 加载文档并切分成 chunks。 返回 [{text: ..., source: ..., chunk_id: 0}, ...] textPath(file_path).read_text(encodingutf-8)splitterRecursiveCharacterTextSplitter(chunk_size512,# 每个 chunk 约 512 字符chunk_overlap64,# 相邻 chunk 重叠 64 字符避免截断关键信息separators[\n\n,\n,。,,, ,],)chunkssplitter.split_text(text)source_namePath(file_path).namereturn[{text:chunk,source:source_name,chunk_id:i}fori,chunkinenumerate(chunks)]chunk_size 调优原则太小 200单片段信息不完整模型拿到的上下文不够太大 1000检索时召回的相关度被稀释而且消耗更多 tokens512 是个经验起点可以根据你的文档密度调整2. Embedding 向量存储# rag/embedder.pyimportchromadbfromopenaiimportOpenAI# 通过 TheRouter 调用 Embedding 模型clientOpenAI(api_keyyour-therouter-key,base_urlhttps://api.therouter.ai/v1,)# 初始化 ChromaDB数据持久化到本地 ./chroma_db/chromachromadb.PersistentClient(path./chroma_db)collectionchroma.get_or_create_collection(nameknowledge_base,metadata{hnsw:space:cosine},# 余弦相似度)defembed_texts(texts:list[str])-list[list[float]]:批量向量化每次最多 2048 条API 限制responseclient.embeddings.create(modelopenai/text-embedding-3-small,inputtexts,)return[item.embeddingforiteminresponse.data]defadd_documents(chunks:list[dict])-int:将文档 chunks 写入向量数据库texts[c[text]forcinchunks]embeddingsembed_texts(texts)collection.add(ids[f{c[source]}_{c[chunk_id]}forcinchunks],embeddingsembeddings,documentstexts,metadatas[{source:c[source],chunk_id:c[chunk_id]}forcinchunks],)returnlen(chunks)defindex_file(file_path:str)-int:一键索引单个文件fromrag.splitterimportload_and_split chunksload_and_split(file_path)returnadd_documents(chunks)3. 检索增强生成# rag/retriever.pyfromopenaiimportOpenAIfromrag.embedderimportclient,collection CHAT_MODELanthropic/claude-sonnet-4# 可改成 openai/gpt-4oSYSTEM_PROMPT你是企业内部知识库助手。请严格根据以下参考资料回答问题。 如果参考资料中没有相关信息请明确说明知识库中暂无相关内容不要自行编造。 回答时请注明信息来源文件名。defretrieve(query:str,top_k:int5)-list[dict]:检索最相关的 top_k 个文档片段# 把用户问题也向量化respclient.embeddings.create(modelopenai/text-embedding-3-small,input[query],)query_embeddingresp.data[0].embedding# 在向量数据库里搜相似片段resultscollection.query(query_embeddings[query_embedding],n_resultstop_k,include[documents,metadatas,distances],)chunks[]fordoc,meta,distinzip(results[documents][0],results[metadatas][0],results[distances][0],):chunks.append({text:doc,source:meta[source],score:1-dist,# 余弦距离转相似度})returnchunksdefask(query:str,top_k:int5)-dict:RAG 核心检索 生成chunksretrieve(query,top_ktop_k)# 过滤相似度太低的片段可选阈值根据场景调整relevant[cforcinchunksifc[score]0.3]ifnotrelevant:return{answer:知识库中暂无与您问题相关的内容请尝试换个问法或联系管理员补充资料。,sources:[],}# 把检索到的片段拼成上下文context_parts[]fori,cinenumerate(relevant,1):context_parts.append(f[{i}] 来源{c[source]}\n{c[text]})context\n\n---\n\n.join(context_parts)# 调用生成模型responseclient.chat.completions.create(modelCHAT_MODEL,messages[{role:system,content:SYSTEM_PROMPT},{role:user,content:f参考资料\n\n{context}\n\n问题{query},},],temperature0.3,# 知识库问答用低温度减少幻觉max_tokens1024,)return{answer:response.choices[0].message.content,sources:[c[source]forcinrelevant],chunks_used:len(relevant),}4. FastAPI 服务封装# main.pyfromfastapiimportFastAPI,UploadFile,FilefrompydanticimportBaseModelimporttempfile,osfromrag.embedderimportindex_filefromrag.retrieverimportask appFastAPI(title企业知识库 API)classQueryRequest(BaseModel):question:strtop_k:int5app.post(/upload)asyncdefupload_document(file:UploadFileFile(...)):上传文档并写入向量库withtempfile.NamedTemporaryFile(deleteFalse,suffixos.path.splitext(file.filename)[1])astmp:tmp.write(awaitfile.read())tmp_pathtmp.nametry:countindex_file(tmp_path)return{status:ok,chunks_indexed:count,filename:file.filename}finally:os.unlink(tmp_path)app.post(/query)defquery_knowledge_base(req:QueryRequest):查询知识库resultask(req.question,top_kreq.top_k)returnresultapp.get(/health)defhealth():return{status:ok}启动uvicorn main:app--host0.0.0.0--port8000--reload上传文档curl-XPOST http://localhost:8000/upload\-Ffile公司产品手册.md查询curl-XPOST http://localhost:8000/query\-HContent-Type: application/json\-d{question: 我们的退款政策是什么}切换到 Qdrant生产环境# 替换 rag/embedder.py 中的 ChromaDB 部分fromqdrant_clientimportQdrantClientfromqdrant_client.modelsimportDistance,VectorParams,PointStruct qdrantQdrantClient(urlhttp://localhost:6333)# 创建集合只需执行一次qdrant.recreate_collection(collection_nameknowledge_base,vectors_configVectorParams(size1536,distanceDistance.COSINE),)defadd_documents_qdrant(chunks:list[dict]):texts[c[text]forcinchunks]embeddingsembed_texts(texts)points[PointStruct(idhash(f{c[source]}_{c[chunk_id]})0x7FFFFFFF,vectoremb,payload{text:c[text],source:c[source]},)forc,embinzip(chunks,embeddings)]qdrant.upsert(collection_nameknowledge_base,pointspoints)Qdrant 支持按source字段过滤检索比如只搜某个部门的文档适合多租户场景。优化技巧Reranking提升检索精度Top-K 向量检索返回的未必是最有用的片段可以用 Reranker 重排# 用 cross-encoder 重排本地模型不消耗 APIfromsentence_transformersimportCrossEncoder rerankerCrossEncoder(cross-encoder/ms-marco-MiniLM-L-6-v2)defrerank(query:str,chunks:list[dict],top_n:int3)-list[dict]:pairs[(query,c[text])forcinchunks]scoresreranker.predict(pairs)rankedsorted(zip(scores,chunks),keylambdax:x[0],reverseTrue)return[cfor_,cinranked[:top_n]]先用向量检索召回 20 个再用 Reranker 精选 3-5 个喂给 LLM效果提升明显。Chunk 重叠策略对于有强上下文依赖的文档如合同、技术规范建议chunk_size800chunk_overlap150使用MarkdownTextSplitter按标题层级切分保持文档结构成本估算以中等规模企业知识库为例1000 份文档平均 10 页/份阶段操作约费用建库一次性向量化 ~500万 tokens~$0.10每次查询Embedding问题 Chat生成~$0.003月活 500 用户人均 10 次查询5000 次查询~$15/月相比自建模型或购买 SaaS 服务成本极低。小结核心流程只有三步切分 → 向量化 → 检索生成。通过 TheRouter 统一调用text-embedding-3-small和claude-sonnet-4不需要分别申请多家 API Key账单也在一个地方管理。拿到这套代码骨架后你可以根据业务需要扩展PDF 解析pdfplumber、权限控制按用户过滤文档、对话历史多轮 RAG、增量更新新文档自动索引……代码结构清晰每个模块职责单一加功能不需要动核心逻辑。