knowledge_base.py

import os
import sys
import glob
import requests
import chromadb
from chromadb.utils import embedding_functions
from openai import OpenAI

sys.path.insert(0, os.path.dirname(__file__))

BASE_DIR = os.path.dirname(__file__)
DB_DIR = os.path.join(BASE_DIR, "chroma_db")
DOCS_DIR = os.path.join(BASE_DIR, "docs")

API_KEY = os.environ.get("DEEPSEEK_API_KEY")
SILICON_KEY = os.environ.get("SILICONFLOW_API_KEY")
if not API_KEY or not SILICON_KEY:
    root_env = os.path.join(BASE_DIR, "..", ".env")
    if os.path.exists(root_env):
        with open(root_env, encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if line and not line.startswith("#"):
                    k, v = line.split("=", 1)
                    os.environ.setdefault(k, v)
    API_KEY = os.environ.get("DEEPSEEK_API_KEY")
    SILICON_KEY = os.environ.get("SILICONFLOW_API_KEY")

deepseek = OpenAI(api_key=API_KEY, base_url="https://api.deepseek.com")

class BGE(embedding_functions.EmbeddingFunction):
    def __call__(self, texts):
        resp = requests.post(
            "https://api.siliconflow.cn/v1/embeddings",
            headers={"Authorization": f"Bearer {SILICON_KEY}"},
            json={"model": "BAAI/bge-large-zh-v1.5", "input": texts, "encoding_format": "float"}
        )
        return [d["embedding"] for d in resp.json()["data"]]

chroma = chromadb.PersistentClient(path=DB_DIR)
collection = chroma.get_or_create_collection(
    name="agent_docs",
    embedding_function=BGE()
)


def build():
    if collection.count() > 0:
        return
    for filepath in glob.glob(os.path.join(DOCS_DIR, "*")):
        filename = os.path.basename(filepath)
        with open(filepath, encoding="utf-8") as f:
            text = f.read()
        chunks = text.split("\n\n")
        for i, chunk in enumerate(chunks):
            chunk = chunk.strip()
            if not chunk:
                continue
            collection.add(
                documents=[chunk],
                metadatas=[{"source": filename}],
                ids=[f"{filename}_{i}"]
            )
    print(f"[知识库] 构建完成，共 {collection.count()} 条")


def rerank(query, documents, top_n=3):
    if len(documents) <= top_n:
        return documents

    resp = requests.post(
        "https://api.siliconflow.cn/v1/rerank",
        headers={"Authorization": f"Bearer {SILICON_KEY}"},
        json={
            "model": "BAAI/bge-reranker-v2-m3",
            "query": query,
            "documents": documents,
        }
    )
    results = resp.json()["results"]
    sorted_results = sorted(results, key=lambda x: x["relevance_score"], reverse=True)
    return [documents[r["index"]] for r in sorted_results[:top_n]]


def ask(query):
    results = collection.query(query_texts=[query], n_results=10)
    contexts = results["documents"][0]
    contexts = rerank(query, contexts, top_n=3)
    joined = "\n---\n".join(contexts)

    prompt = f"""你是一个 Agent 开发学习助手。请严格根据以下知识库资料回答用户问题。
如果资料中没有相关信息，请诚实说"这部分资料暂未收录"，然后给出你的理解。

【知识库资料】
{joined}

【用户问题】
{query}"""

    resp = deepseek.chat.completions.create(
        model="deepseek-chat",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.3,
    )
    return resp.choices[0].message.content


build()