AI 工具开发实战2开发一个本地 RAG 知识库——丢一个文件夹进去直接问答上一篇做了一个命令行翻译工具这篇做一个更实用的本地 RAG 知识库。把 PDF、Markdown、TXT 文件丢到一个文件夹里运行一条命令就能向这些文件提问AI 会基于文件内容回答。全程本地运行数据不出本机。项目结构localrag/ ├── localrag.py # CLI 主程序 ├── indexer.py # 文档索引 ├── retriever.py # 检索引擎 ├── requirements.txt └── docs/ # 放文档的文件夹安装依赖# requirements.txt click8.1.7 openai1.50.0 python-dotenv1.0.0 sentence-transformers3.1.0 PyMuPDF1.24.0 numpy1.26.0核心实现文档索引器indexer.py# indexer.pyimportosimportjsonimporthashlibfrompathlibimportPathimportfitz# PyMuPDFclassDocumentIndexer:文档索引读取文件夹、切分文本、向量化。def__init__(self,chunk_size512,chunk_overlap128):self.chunk_sizechunk_size self.chunk_overlapchunk_overlapdefload_folder(self,folder_path:str)-list:加载文件夹中的所有文档。docs[]forfinsorted(Path(folder_path).iterdir()):iff.suffix.lower().pdf:textself._parse_pdf(str(f))eliff.suffix.lower()in(.txt,.md):textf.read_text(encodingutf-8,errorsignore)else:continueiftext.strip():docs.append({source:f.name,content:text})returndocsdef_parse_pdf(self,path):docfitz.open(path)text\n.join(p.get_text()forpindoc)doc.close()returntextdefchunk_documents(self,docs:list)-list:将文档切分为重叠的 chunk。chunks[]fordocindocs:textdoc[content]# 按段落切paragraphstext.split(\n\n)bufferforparainparagraphs:parapara.strip()ifnotpara:continueiflen(buffer)len(para)self.chunk_size:buffer(buffer\n\npara).strip()ifbufferelseparaelse:ifbuffer:chunks.append({text:buffer,source:doc[source],id:hashlib.md5(buffer.encode()).hexdigest()[:8],})bufferparaifbuffer:chunks.append({text:buffer,source:doc[source],id:hashlib.md5(buffer.encode()).hexdigest()[:8],})returnchunks检索引擎retriever.py# retriever.pyimportjsonimportosimportnumpyasnpfromsentence_transformersimportSentenceTransformerclassLocalRetriever:本地检索引擎Embedding 向量检索。def__init__(self,model_nameBAAI/bge-small-zh-v1.5):self.modelSentenceTransformer(model_name)self.chunks[]self.vectorsNoneself.index_pathNonedefindex(self,chunks:list,cache_dir:str):建立向量索引。self.chunkschunks texts[c[text]forcinchunks]self.vectorsself.model.encode(texts,normalize_embeddingsTrue)self.index_pathos.path.join(cache_dir,index.json)self._save()defsearch(self,query:str,top_k:int3)-list:检索最相关的 chunk。q_vecself.model.encode([query],normalize_embeddingsTrue)[0]scoresnp.dot(self.vectors,q_vec)# 余弦相似度top_indicesnp.argsort(scores)[-top_k:][::-1]results[]foridxintop_indices:ifscores[idx]0.3:# 相似度太低的不返回continueresults.append({text:self.chunks[idx][text],source:self.chunks[idx][source],score:float(scores[idx]),})returnresultsdef_save(self):data{chunks:self.chunks,vectors:self.vectors.tolist()ifself.vectorsisnotNoneelse[],}withopen(self.index_path,w,encodingutf-8)asf:json.dump(data,f,ensure_asciiFalse)defload(self,cache_dir:str):pathos.path.join(cache_dir,index.json)ifnotos.path.exists(path):returnFalsewithopen(path,encodingutf-8)asf:datajson.load(f)self.chunksdata[chunks]self.vectorsnp.array(data[vectors])returnTrueCLI 主程序localrag.py#!/usr/bin/env python3# localrag.pyimportosimportclickfrompathlibimportPathfromopenaiimportOpenAIfromdotenvimportload_dotenvfromindexerimportDocumentIndexerfromretrieverimportLocalRetriever load_dotenv()CACHE_DIRos.path.expanduser(~/.localrag)os.makedirs(CACHE_DIR,exist_okTrue)clientOpenAI(api_keyos.getenv(DEEPSEEK_API_KEY),base_urlhttps://api.deepseek.com/v1,)SYSTEM_PROMPT你是一个基于本地文档的 AI 助手。 请严格基于以下检索到的资料回答问题。 如果资料不足以回答请说资料中没有相关信息。 检索资料 {context}click.group()defcli():localrag - 本地 RAG 知识库cli.command()click.argument(folder,typeclick.Path(existsTrue))defindex(folder):索引一个文件夹中的所有文档。indexerDocumentIndexer(chunk_size512,chunk_overlap128)retrieverLocalRetriever()docsindexer.load_folder(folder)chunksindexer.chunk_documents(docs)retriever.index(chunks,CACHE_DIR)click.echo(f✅ 索引完成{len(docs)}个文档 →{len(chunks)}个片段)cli.command()click.argument(question)click.option(--top,-k,default3,help检索片段数)click.option(--show-sources,-s,is_flagTrue,help显示引用来源)defask(question,top,show_sources):向知识库提问。retrieverLocalRetriever()ifnotretriever.load(CACHE_DIR):click.echo(❌ 没有索引请先运行 localrag index folder)return# 检索resultsretriever.search(question,top_ktop)ifnotresults:click.echo(❌ 没有找到相关内容)return# 构建 Promptcontext\n\n.join(f[{r[source]}]{r[text][:500]}forrinresults)# 调用 LLMresponseclient.chat.completions.create(modeldeepseek-chat,messages[{role:system,content:SYSTEM_PROMPT.format(contextcontext)},{role:user,content:question},],temperature0.3,streamTrue,)# 流式输出forchunkinresponse:ifchunk.choices[0].delta.content:click.echo(chunk.choices[0].delta.content,nlFalse)click.echo()# 显示来源ifshow_sources:click.echo(f\n 引用来源)fori,rinenumerate(results):click.echo(f [{i1}]{r[source]}(相关度:{r[score]:.2f}))if__name____main__:cli()使用方式# 1. 索引文档python localrag.py index ./docs# ✅ 索引完成5 个文档 → 42 个片段# 2. 提问python localrag.py ask这个项目的架构是什么# 这个项目采用微服务架构包含 API 服务、数据库、缓存三层...# 3. 带来源引用python localrag.py ask部署流程是怎样的--show-sources# ...回答内容...# 引用来源# [1] deploy.md (相关度: 0.87)# [2] architecture.md (相关度: 0.72)性能特点全本地运行数据不出本机支持 PDF、MD、TXT 三种格式首次索引后缓存下次提问不用重新索引512 维 BGE 向量500 个 chunk 检索 50ms下一版可以加的功能增量索引新文件不需要重建全部索引Web UI 界面支持更多文件格式docx、csv切换不同 Embedding 模型总结一个本地 RAG 知识库核心就三步读文件 → 切分 chunk向量化 → 建立索引提问 → 检索 → LLM 回答把这三个功能封装成 CLI就是一个实用的生产力工具。本文是《AI 开发者工具链实战》系列的第 2 篇。上一篇命令行 AI 翻译工具下一篇AI 代码审查 CLI本文由 Zyentor智元界原创发布本文发布于 Zyentor智元界 —— AI 开发者社区原文链接https://www.zyentor.com/news/4130