搜 索

如何建立本地ai知识库

  • 57阅读
  • 2026年01月10日
  • 0评论
首页 / AI/大数据 / 正文

起因:我受够了重复造轮子

作为一个每天和YOLO模型、FastAPI打交道的AI架构师,我发现自己陷入了一个诡异的循环:

  • 明明三个月前写过一段完美的数据预处理代码,死活想不起来放哪了
  • 每次问AI"怎么导出ONNX模型",它给的答案总是差那么点意思,因为它不知道我的项目结构
  • 积累了一堆设计文档,结果用的时候还是靠Ctrl+F大法

于是我决定搞一个本地知识库,把自己的代码片段、设计文档、踩坑记录都扔进去,让AI能够基于我的知识体系来回答问题。

选型:一场短暂的纠结

在正式开始折腾ChromaDB之前,我确实调研了一圈主流方案。简单说说为什么最后选了它:

Pinecone - 企业级向量数据库,性能杠杠的,十亿级向量毫秒响应。但是...我只是想存点自己的代码啊,又不是要给淘宝做推荐系统。更关键的是,数据要上传到云端,我那些写得乱七八糟的代码要是被人看到,社死程度堪比公开我的浏览器历史记录。Pass。

NotebookLM - Google家的产品,上传文档就能问答,还能生成播客风格的摘要,听起来很酷对吧?但它没有API,只能在网页上点点点。我一个写代码的,你让我每次都手动上传文件?而且它对代码的理解约等于我对时尚的理解——基本为零。Pass。

Weaviate / Qdrant - 都是正经的向量数据库,功能强大,部署灵活。但需要跑独立服务,配置文件比我的头发还多。我只是想要一个本地知识库,不是想运维一套分布式系统。先Pass,以后规模大了再说。

FAISS - Facebook出品,性能怪兽,学术圈的最爱。但它只是一个索引库,持久化、元数据管理、API封装全都要自己搞。我只是想用,不是想从零开始写一个向量数据库。敬而远之。

LanceDB - 新秀选手,Rust写的,性能比ChromaDB好。但生态还不够成熟,文档有点少,遇到问题只能靠自己摸索。等它再长大一点我再考虑。观望中。

最后的选择:ChromaDB。理由很简单:

pip install chromadb

没错,就这一行,完事。数据存本地,开箱即用,和LangChain配合得天衣无缝。ChromaDB已经跨越了长达数年的alpha阶段,当前可用版本为1.4.x,对于个人知识库这个量级,它完全够用了。

ChromaDB:从入门到真香

第一步:Hello World(入门)

先来个最简单的例子,感受一下它有多丝滑:

import chromadb

# 创建一个持久化的客户端,数据存在本地
client = chromadb.PersistentClient(path="./my_brain")

# 创建一个集合,就像数据库里的表
collection = client.get_or_create_collection(name="code_snippets")

# 往里面塞点东西
collection.add(
    documents=[
        "FastAPI路由示例:@app.get('/users/{user_id}')",
        "YOLO模型加载:model = YOLO('yolov8n.pt')",
        "PyTorch保存模型:torch.save(model.state_dict(), 'model.pth')"
    ],
    ids=["doc1", "doc2", "doc3"]
)

# 语义搜索!
results = collection.query(
    query_texts=["怎么加载预训练模型"],
    n_results=2
)
print(results['documents'])
# 输出会包含YOLO和PyTorch那两条,因为语义相近

就这么简单。没有配置文件,没有Docker,没有注册账号,没有填信用卡。它甚至自带了一个embedding模型(all-MiniLM-L6-v2),你什么都不用管就能跑起来。

第二步:正经地组织一下(进阶)

Hello World跑通之后,我们来搞一个正经的项目结构。毕竟以后东西多了,总不能一股脑全扔一个collection里吧。

my_knowledge_base/
├── kb/                     # 核心代码
│   ├── __init__.py
│   ├── config.py           # 配置
│   ├── processor.py        # 文档处理(分块、提取元数据)
│   ├── store.py            # ChromaDB封装
│   └── retriever.py        # 检索接口
├── data/
│   ├── raw/                # 原始文档,按类型分文件夹
│   │   ├── designs/        # 设计文档
│   │   ├── snippets/       # 代码片段
│   │   └── notes/          # 学习笔记
│   └── chroma_db/          # ChromaDB数据目录
├── ingest.py               # 导入脚本
├── query.py                # 查询脚本
└── requirements.txt

config.py - 配置集中管理

from pathlib import Path

class Config:
    BASE_DIR = Path(__file__).parent.parent
    DATA_DIR = BASE_DIR / "data"
    RAW_DIR = DATA_DIR / "raw"
    CHROMA_DIR = DATA_DIR / "chroma_db"
    
    # 分块参数,这两个数字我调了好几次才满意
    CHUNK_SIZE = 500      # 每块大约500字符
    CHUNK_OVERLAP = 50    # 块之间重叠50字符,防止语义被切断
    
    # 集合定义
    COLLECTIONS = {
        "designs": "系统设计文档",
        "snippets": "代码片段", 
        "notes": "学习笔记和踩坑记录"
    }

config = Config()

processor.py - 文档处理器

这是最容易踩坑的地方。直接把整个文件扔进去?太长了,检索效果差。切太碎?上下文丢失,AI看了也懵。

import re
from pathlib import Path
from dataclasses import dataclass
from typing import List, Dict, Any

@dataclass
class Document:
    content: str
    metadata: Dict[str, Any]
    doc_id: str

class DocumentProcessor:
    """
    文档处理器:把大文件切成小块
    
    踩坑记录:
    1. 代码文件别按字数切,要按函数/类切,不然一个函数被切两半,AI看了想打人
    2. Markdown文件可以按段落切,但要保留标题作为上下文
    3. 元数据很重要!后面过滤全靠它
    """
    
    def __init__(self, chunk_size: int = 500, chunk_overlap: int = 50):
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
    
    def process_file(self, file_path: Path) -> List[Document]:
        content = file_path.read_text(encoding='utf-8')
        suffix = file_path.suffix.lower()
        
        # 元数据:后面检索过滤全靠这些
        metadata = {
            "source": str(file_path),
            "filename": file_path.name,
            "file_type": suffix,
            "category": file_path.parent.name,
        }
        
        # 根据文件类型选择分块策略
        if suffix in ['.py', '.js', '.ts']:
            chunks = self._chunk_code(content)
        else:
            chunks = self._chunk_text(content)
        
        return [
            Document(
                content=chunk,
                metadata={**metadata, "chunk_index": i},
                doc_id=f"{file_path.stem}_{i}"
            )
            for i, chunk in enumerate(chunks)
        ]
    
    def _chunk_text(self, text: str) -> List[str]:
        """普通文本:按段落切,尽量保持语义完整"""
        paragraphs = text.split('\n\n')
        chunks = []
        current = ""
        
        for para in paragraphs:
            if len(current) + len(para) < self.chunk_size:
                current += para + "\n\n"
            else:
                if current.strip():
                    chunks.append(current.strip())
                current = para + "\n\n"
        
        if current.strip():
            chunks.append(current.strip())
        
        return chunks if chunks else [text]
    
    def _chunk_code(self, code: str) -> List[str]:
        """代码文件:按函数/类切割"""
        # 匹配Python的函数和类定义
        pattern = r'((?:def |class |async def )\w+[^\n]*\n(?:[ \t]+[^\n]*\n)*)'
        matches = re.findall(pattern, code)
        
        if matches:
            return [m.strip() for m in matches if len(m.strip()) > 20]
        
        # 如果没匹配到(可能是脚本文件),回退到普通分块
        return self._chunk_text(code)

store.py - ChromaDB封装

import chromadb
from chromadb.config import Settings
from pathlib import Path
from typing import List, Dict, Any, Optional

class KnowledgeStore:
    """
    ChromaDB的封装层
    
    为什么要封装一层?
    1. 统一接口,以后换数据库方便(虽然大概率不会换)
    2. 加一些便利方法,比如跨集合搜索
    3. 处理一些边界情况,比如空结果
    """
    
    def __init__(self, persist_dir: Path):
        self.client = chromadb.PersistentClient(
            path=str(persist_dir),
            settings=Settings(anonymized_telemetry=False)  # 关掉遥测,我不想被统计
        )
        self._collections = {}
    
    def get_collection(self, name: str, description: str = ""):
        if name not in self._collections:
            self._collections[name] = self.client.get_or_create_collection(
                name=name,
                metadata={"description": description}
            )
        return self._collections[name]
    
    def add_documents(
        self,
        collection_name: str,
        documents: List[str],
        metadatas: List[Dict],
        ids: List[str]
    ):
        collection = self.get_collection(collection_name)
        
        # ChromaDB会自动生成embedding,省心
        collection.add(
            documents=documents,
            metadatas=metadatas,
            ids=ids
        )
        print(f"✅ [{collection_name}] 新增 {len(documents)} 条文档")
    
    def search(
        self,
        collection_name: str,
        query: str,
        n_results: int = 5,
        where: Optional[Dict] = None
    ) -> List[Dict]:
        """单集合搜索"""
        collection = self.get_collection(collection_name)
        results = collection.query(
            query_texts=[query],
            n_results=n_results,
            where=where
        )
        
        # 把结果整理成更友好的格式
        return [
            {
                "content": results['documents'][0][i],
                "metadata": results['metadatas'][0][i],
                "distance": results['distances'][0][i]
            }
            for i in range(len(results['documents'][0]))
        ]
    
    def search_all(self, query: str, n_results: int = 5) -> List[Dict]:
        """跨所有集合搜索,然后按相似度排序"""
        all_results = []
        
        for coll in self.client.list_collections():
            results = self.search(coll.name, query, n_results)
            for r in results:
                r['collection'] = coll.name
                all_results.append(r)
        
        # 按距离排序(越小越相似)
        all_results.sort(key=lambda x: x['distance'])
        return all_results[:n_results]
    
    def stats(self) -> Dict[str, int]:
        """看看各个集合有多少文档"""
        return {
            coll.name: coll.count()
            for coll in self.client.list_collections()
        }

retriever.py - 上下文组装器

这是最终对接AI的地方,把检索结果格式化成AI能理解的上下文。

from typing import List, Optional
from .store import KnowledgeStore

class ContextRetriever:
    """
    把检索结果组装成AI友好的上下文
    
    这里有个小技巧:不要把所有结果都塞给AI
    1. 太长了AI也消化不了
    2. 不相关的内容会干扰AI的判断
    3. Token是要钱的(如果你用的是付费API)
    """
    
    def __init__(self, store: KnowledgeStore):
        self.store = store
    
    def get_context(
        self,
        query: str,
        max_chars: int = 8000,
        collections: Optional[List[str]] = None
    ) -> str:
        """获取格式化的上下文"""
        
        if collections:
            results = []
            for coll in collections:
                results.extend(self.store.search(coll, query, n_results=5))
            results.sort(key=lambda x: x['distance'])
        else:
            results = self.store.search_all(query, n_results=10)
        
        # 组装上下文,注意控制长度
        context_parts = []
        total_len = 0
        
        for r in results:
            part = self._format_result(r)
            if total_len + len(part) > max_chars:
                break
            context_parts.append(part)
            total_len += len(part)
        
        return "\n\n---\n\n".join(context_parts)
    
    def _format_result(self, result: dict) -> str:
        meta = result['metadata']
        return f"**[{meta.get('category', '')}] {meta.get('filename', '')}**\n\n{result['content']}"
    
    def export_for_ai(self, query: str) -> dict:
        """导出为可以直接喂给AI的格式"""
        context = self.get_context(query)
        
        return {
            "system_prompt": "你是一个有帮助的AI助手。请基于以下知识库内容回答问题,如果知识库中没有相关信息,请诚实说明。",
            "context": f"<knowledge_base>\n{context}\n</knowledge_base>",
            "user_query": query
        }

第三步:跑起来!

ingest.py - 导入文档

#!/usr/bin/env python3
"""
文档导入脚本
用法:python ingest.py
"""
from pathlib import Path
from kb.config import config
from kb.processor import DocumentProcessor
from kb.store import KnowledgeStore

def main():
    processor = DocumentProcessor(
        chunk_size=config.CHUNK_SIZE,
        chunk_overlap=config.CHUNK_OVERLAP
    )
    store = KnowledgeStore(config.CHROMA_DIR)
    
    print("🚀 开始导入文档...\n")
    
    for category_dir in config.RAW_DIR.iterdir():
        if not category_dir.is_dir():
            continue
        
        collection_name = category_dir.name
        print(f"📁 处理 [{collection_name}]")
        
        all_docs = []
        for file_path in category_dir.rglob('*'):
            if file_path.suffix in ['.md', '.txt', '.py', '.js', '.ts']:
                docs = processor.process_file(file_path)
                all_docs.extend(docs)
                print(f"   - {file_path.name} -> {len(docs)} 块")
        
        if all_docs:
            store.add_documents(
                collection_name=collection_name,
                documents=[d.content for d in all_docs],
                metadatas=[d.metadata for d in all_docs],
                ids=[d.doc_id for d in all_docs]
            )
    
    print("\n📊 导入完成!统计:")
    for name, count in store.stats().items():
        print(f"   {name}: {count} 条")

if __name__ == "__main__":
    main()

query.py - 查询测试

#!/usr/bin/env python3
"""
查询测试脚本
用法:python query.py "你的问题"
"""
import sys
from kb.config import config
from kb.store import KnowledgeStore
from kb.retriever import ContextRetriever

def main():
    if len(sys.argv) < 2:
        print("用法:python query.py \"你的问题\"")
        sys.exit(1)
    
    query = sys.argv[1]
    
    store = KnowledgeStore(config.CHROMA_DIR)
    retriever = ContextRetriever(store)
    
    print(f"🔍 查询:{query}\n")
    print("=" * 50)
    
    # 获取上下文
    context = retriever.get_context(query)
    print(context)
    
    print("\n" + "=" * 50)
    print("\n💡 可以把上面的内容作为上下文喂给AI了!")

if __name__ == "__main__":
    main()

第四步:真香时刻

跑了一段时间之后,我发现这东西是真的好用:

# 导入我的所有文档
$ python ingest.py
🚀 开始导入文档...

📁 处理 [designs]
   - payment_system.md -> 12 块
   - yolo_pipeline.md -> 8 块
📁 处理 [snippets]
   - fastapi_utils.py -> 15 块
   - yolo_training.py -> 23 块
📁 处理 [notes]
   - autodl_gpu_setup.md -> 5 块

📊 导入完成!统计:
   designs: 20 条
   snippets: 38 条
   notes: 5 条

# 查询一下
$ python query.py "怎么在AutoDL上训练YOLO"

🔍 查询:怎么在AutoDL上训练YOLO

==================================================
**[notes] autodl_gpu_setup.md**

AutoDL RTX 4090 环境配置踩坑记录:
1. 先更新pip:pip install --upgrade pip
2. 安装ultralytics:pip install ultralytics
3. 数据集放到/root/autodl-tmp/,这个目录不限速
...

---

**[snippets] yolo_training.py**

def train_yolo(data_yaml: str, epochs: int = 100):
    """YOLO训练函数,已在4090上验证过"""
    from ultralytics import YOLO
    model = YOLO('yolov8n.pt')
    model.train(data=data_yaml, epochs=epochs, imgsz=640)
...
==================================================

💡 可以把上面的内容作为上下文喂给AI了!

再也不用翻历史记录了,问啥都能找到之前的笔记和代码。

一些踩坑经验

  1. 分块大小很重要:太大了检索不精准,太小了丢失上下文。我试下来500字符比较合适。
  2. 代码文件别按字数切:一定要按函数/类切,不然一个函数被切成两半,AI看了也没法用。
  3. 元数据要设计好:文件名、类型、分类这些一定要记录,后面过滤全靠它们。
  4. 定期清理:过时的文档记得删掉,不然AI会给你返回三年前的错误代码。
  5. Embedding模型可以升级:ChromaDB默认的模型够用,但如果你的文档主要是中文,可以换成BAAI/bge-small-zh-v1.5,效果更好。

后续可以折腾的方向

  • 加个Web界面:用FastAPI + React搞一个,查询更方便
  • 接入IDE:写个VSCode插件,直接在编辑器里查知识库
  • 自动同步:用watchdog监听文件变化,自动更新知识库
  • 更好的Embedding:接入OpenAI的text-embedding-3-small,效果更好(但要花钱)

最后

折腾了一圈下来,ChromaDB确实是个人知识库的最佳选择。简单、免费、够用。虽然它还是alpha版本,但对于我这个量级的使用场景,完全没问题。

等哪天我的知识库膨胀到百万级别,再考虑升级到Qdrant或者Pinecone吧。

不过那应该是很久以后的事了——毕竟我首先得写出那么多代码来。


写于 2026年1月,MacBook M4 Air 旁边的咖啡已经凉了

评论区
暂无评论
avatar