饮墨

# requirements: langchain chromadb openai tiktoken
from langchain.text_splitter import RecursiveCharacterTextSplitter
import chromadb
from chromadb.utils import embedding_functions
import os

# 初始化 Embedding 函数
openai_ef = embedding_functions.OpenAIEmbeddingFunction(
    api_key=os.getenv("OPENAI_API_KEY"),
    model_name="text-embedding-3-small"
)

# 初始化 ChromaDB（持久化存储）
client = chromadb.PersistentClient(path="./chroma_db")
collection = client.get_or_create_collection(
    name="knowledge_base",
    embedding_function=openai_ef,
    metadata={"hnsw:space": "cosine"}  # 余弦相似度
)

# 文档分块 — 核心参数
splitter = RecursiveCharacterTextSplitter(
    chunk_size=512,        # 每块最大512字符
    chunk_overlap=64,      # 相邻块重叠64字符，避免语义断裂
    separators=["\n\n", "\n", "。", ".", " "],  # 中英文语义边界
    length_function=len
)

def ingest_document(file_path: str, doc_id: str):
    """将文档分块后写入向量库"""
    with open(file_path, "r", encoding="utf-8") as f:
        text = f.read()

    chunks = splitter.split_text(text)
    ids = [f"{doc_id}_chunk_{i}" for i in range(len(chunks))]
    metadatas = [{"source": file_path, "chunk_index": i} for i in range(len(chunks))]

    collection.upsert(
        ids=ids,
        documents=chunks,
        metadatas=metadatas
    )
    print(f"✅ 入库完成: {file_path}, 共 {len(chunks)} 个分块")

# 批量入库
ingest_document("docs/ops-handbook.md", "ops-handbook")
ingest_document("docs/incident-reports.md", "incidents")

import cohere

co = cohere.Client(os.getenv("COHERE_API_KEY"))

def retrieve_with_rerank(query: str, top_k: int = 5, final_k: int = 3):
    """向量检索 + Rerank 双重过滤"""
    # 第一阶段：向量召回（粗筛）
    results = collection.query(
        query_texts=[query],
        n_results=top_k,
        include=["documents", "metadatas", "distances"]
    )

    documents = results["documents"][0]
    metadatas = results["metadatas"][0]

    if not documents:
        return []

    # 第二阶段：Rerank（精排）
    rerank_resp = co.rerank(
        model="rerank-v3.5",
        query=query,
        documents=documents,
        top_n=final_k
    )

    # 组装最终结果
    ranked_docs = []
    for item in rerank_resp.results:
        ranked_docs.append({
            "text": documents[item.index],
            "score": item.relevance_score,
            "source": metadatas[item.index]["source"]
        })

    return ranked_docs

from litellm import completion

SYSTEM_PROMPT = """你是一个技术知识助手。基于提供的参考文档回答问题。
规则：
1. 只基于参考文档内容回答，不要编造
2. 如果文档中没有相关信息，明确告知用户
3. 回答末尾标注引用来源"""

def rag_answer(query: str) -> str:
    """完整 RAG 问答流程"""
    # 检索相关文档
    docs = retrieve_with_rerank(query)

    if not docs:
        return "未找到相关文档，无法回答此问题。"

    # 构造上下文
    context = "\n\n---\n\n".join(
        [f"[来源: {d['source']}]\n{d['text']}" for d in docs]
    )

    # 调用 LLM
    response = completion(
        model="gpt-4o",  # 可切换为 "anthropic/claude-3-5-sonnet"
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": f"参考文档:\n{context}\n\n问题: {query}"}
        ],
        temperature=0.1  # 低温度减少幻觉
    )

    return response.choices[0].message.content

# 使用示例
answer = rag_answer("上次数据库故障的根因是什么？")
print(answer)

# 基于文件哈希判断是否需要重新入库
import hashlib

def file_hash(path):
    return hashlib.md5(open(path, "rb").read()).hexdigest()

def sync_document(file_path, doc_id):
    current_hash = file_hash(file_path)
    # 对比存储的 hash，变化时重新 ingest
    stored = collection.get(ids=[f"{doc_id}_meta"], include=["metadatas"])
    if stored["metadatas"] and stored["metadatas"][0].get("hash") == current_hash:
        return  # 未变化，跳过
    ingest_document(file_path, doc_id)
    collection.upsert(
        ids=[f"{doc_id}_meta"],
        documents=["meta"],
        metadatas=[{"hash": current_hash}]
    )