Python搭建AI知识库,50行代码让大模型读懂你的文档

pip install chromadb sentence-transformers openai

pip install chromadb sentence-transformers openai -i https://pypi.tuna.tsinghua.edu.cn/simple

# 你的知识库，换成你自己的文档内容knowledge_base = ["公司年假制度：入职满1年享5天年假，满3年享10天，满10年享15天。年假不可跨年累积。","报销流程：登录OA系统→财务→报销申请→上传发票→直属领导审批→财务打款。审批周期3个工作日。","产品定价：基础版299元/月，专业版599元/月，企业版按需报价。年付打8折。","技术栈：后端用Go语言，前端React，数据库PostgreSQL，部署在阿里云华东区。","客户常见问题：密码重置在登录页点忘记密码，支持手机号和邮箱两种方式。",]

import chromadbfrom sentence_transformers import SentenceTransformerfrom openai import OpenAI# 向量化模型（中文用这个，英文用all-MiniLM-L6-v2）embed_model = SentenceTransformer('shibing624/text2vec-base-chinese')# 向量数据库chroma_client = chromadb.Client()collection = chroma_client.create_collection("my_knowledge")# 知识库knowledge_base = [    "公司年假制度：入职满1年享5天年假，满3年享10天，满10年享15天。年假不可跨年累积。",    "报销流程：登录OA系统→财务→报销申请→上传发票→直属领导审批→财务打款。审批周期3个工作日。",    "产品定价：基础版299元/月，专业版599元/月，企业版按需报价。年付打8折。",    "技术栈：后端用Go语言，前端React，数据库PostgreSQL，部署在阿里云华东区。",    "客户常见问题：密码重置在登录页点忘记密码，支持手机号和邮箱两种方式。",]# 把文档转成向量存进去embeddings = embed_model.encode(knowledge_base).tolist()collection.add(    documents=knowledge_base,    embeddings=embeddings,    ids=[f"doc_{i}" for i in range(len(knowledge_base))])# DeepSeek API（去 platform.deepseek.com 注册，免费送额度）llm = OpenAI(    api_key="api_key",    base_url="https://api.deepseek.com")def ask(question):    # 先检索：找到最相关的文档    q_embedding = embed_model.encode([question]).tolist()    results = collection.query(query_embeddings=q_embedding, n_results=3)    context = "\n".join(results['documents'][0])    # 再回答：把文档和问题一起交给大模型    prompt = f"""你是一个专业的知识库助手。只基于以下参考资料回答问题。如果资料里没有相关信息，请直接说"这个问题超出了我的知识范围"。不要编造信息。参考资料：{context}用户问题：{question}"""    response = llm.chat.completions.create(        model="deepseek-chat",        messages=[{"role": "user", "content": prompt}]    )    return response.choices[0].message.content# 测试！print("Q: 年假有几天？")print("A:", ask("年假有几天？"))print()print("Q: 报销要走什么流程？")print("A:", ask("报销要走什么流程？"))print()print("Q: 你们公司食堂好吃吗？")print("A:", ask("你们公司食堂好吃吗？"))

python rag_demo.py

pip install pymupdfimport fitz  # pymupdfdefread_pdf(path):    doc = fitz.open(path)return"\n".join([page.get_text() for page in doc])

defread_markdown(path):withopen(path, 'r', encoding='utf-8') as f:return f.read()

defread_txt(path):withopen(path, 'r', encoding='utf-8') as f:return f.read()

defsplit_text(text, chunk_size=300, overlap=50):"""按字数切片，每段300字，前后重叠50字"""    chunks = []    start = 0while start < len(text):        end = start + chunk_size        chunks.append(text[start:end])        start = end - overlapreturn chunks

text = read_pdf("你的文档.pdf")knowledge_base = split_text(text)

import osos.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'