# 更新系统
apt update && apt upgrade -y

# 安装基础工具
apt install -y python3-pip git curl wget tmux htop

# 安装 Docker（推荐，简化部署）
curl -fsSL https://get.docker.com | sh
usermod -aG docker $USER

# requirements.txt
langchain==0.3.1
langchain-community==0.3.1
langchain-text-splitters==0.3.0
chromadb==0.5.0
ollama==0.3.0
unstructured==0.15.0
python-docx==1.1.0
openpyxl==3.1.0

# document_loader.py
from langchain_community.document_loaders import (
    PyPDFLoader,
    TextLoader,
    UnstructuredMarkdownLoader,
    Docx2txtLoader,
    DirectoryLoader
)
from langchain.text_splitter import RecursiveCharacterTextSplitter
import os

def load_documents(directory: str):
    """加载目录下所有支持格式的文档"""
    
    # 按文件类型加载
    pdf_loader = DirectoryLoader(
        directory, glob="**/*.pdf", loader_cls=PyPDFLoader
    )
    text_loader = DirectoryLoader(
        directory, glob="**/*.txt", loader_cls=TextLoader
    )
    md_loader = DirectoryLoader(
        directory, glob="**/*.md", loader_cls=UnstructuredMarkdownLoader
    )
    docx_loader = DirectoryLoader(
        directory, glob="**/*.docx", loader_cls=Docx2txtLoader
    )
    
    # 合并所有文档
    documents = []
    documents.extend(pdf_loader.load())
    documents.extend(text_loader.load())
    documents.extend(md_loader.load())
    documents.extend(docx_loader.load())
    
    print(f"加载了 {len(documents)} 个文档")
    return documents

def chunk_documents(documents, chunk_size=500, chunk_overlap=50):
    """将文档分割为合适大小的文本块"""
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        separators=["\n\n", "\n", "。", ".", " "]
    )
    chunks = text_splitter.split_documents(documents)
    print(f"分割为 {len(chunks)} 个文本块")
    return chunks

# embedding_store.py
import chromadb
from langchain_community.embeddings import OllamaEmbeddings
from langchain_community.vectorstores import Chroma

def create_vector_store(chunks, persist_directory="./chroma_db"):
    """创建向量数据库并存储文档块"""
    
    # 使用 Ollama 本地 Embedding
    embeddings = OllamaEmbeddings(
        model="nomic-embed-text",
        base_url="http://localhost:11434"
    )
    
    # 创建/加载 ChromaDB
    vectorstore = Chroma.from_documents(
        documents=chunks,
        embedding=embeddings,
        persist_directory=persist_directory
    )
    
    # 持久化到磁盘
    vectorstore.persist()
    print(f"向量库已保存至 {persist_directory}")
    
    return vectorstore

def load_vector_store(persist_directory="./chroma_db"):
    """从磁盘加载已有向量库"""
    embeddings = OllamaEmbeddings(
        model="nomic-embed-text",
        base_url="http://localhost:11434"
    )
    vectorstore = Chroma(
        persist_directory=persist_directory,
        embedding_function=embeddings
    )
    return vectorstore

# rag_qa.py
from langchain.chains import RetrievalQA
from langchain_community.llms import Ollama
from langchain.prompts import PromptTemplate

# 加载向量库
vectorstore = load_vector_store()

# 创建检索器
retriever = vectorstore.as_retriever(
    search_type="similarity",      # 相似度检索
    search_kwargs={"k": 5}         # 返回 Top-5 相关片段
)

# 本地 LLM
llm = Ollama(
    model="qwen2.5:7b",
    base_url="http://localhost:11434",
    temperature=0.1
)

# 自定义提示词模板
qa_prompt = PromptTemplate(
    input_variables=["context", "question"],
    template="""使用以下上下文信息回答问题。如果上下文中没有相关信息，请诚实回答"我无法从现有资料中找到答案"。

上下文：
{context}

问题：{question}

回答（请用中文回答）："""
)

# 创建 RAG 链
rag_chain = RetrievalQA.from_chain_type(
    llm=llm,
    retriever=retriever,
    return_source_documents=True,
    chain_type_kwargs={"prompt": qa_prompt}
)

# 执行问答
def ask_question(question: str):
    result = rag_chain.invoke({"query": question})
    return {
        "answer": result["result"],
        "sources": [doc.page_content[:200] for doc in result["source_documents"][:3]]
    }

# 混合检索：向量 + 关键词
retriever = vectorstore.as_retriever(
    search_type="similarity_score_threshold",
    search_kwargs={"k": 10, "score_threshold": 0.7}
)

# 查询扩展（Query Expansion）
from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers.document_compressors import LLMChainExtractor

compressor = LLMChainExtractor.from_llm(llm)
compression_retriever = ContextualCompressionRetriever(
    base_retriever=retriever,
    document_compressor=compressor
)

# 重排序（Rerank）提升质量
# 使用 Cross-Encoder 对检索结果重新排序

# main.py
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List, Optional
import uvicorn

app = FastAPI(title="RAG Knowledge Base API", version="1.0.0")

class QuestionRequest(BaseModel):
    question: str
    k: int = 5
    include_sources: bool = True

class QuestionResponse(BaseModel):
    answer: str
    sources: Optional[List[str]] = None
    confidence: float

@app.post("/ask", response_model=QuestionResponse)
async def ask_question(req: QuestionRequest):
    try:
        result = ask_question(req.question)
        return QuestionResponse(
            answer=result["answer"],
            sources=result["sources"] if req.include_sources else None
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/stats")
async def get_stats():
    """获取知识库统计信息"""
    vectorstore = load_vector_store()
    collection = vectorstore._collection
    count = collection.count()
    return {"document_count": count, "collection": "knowledge_base"}

@app.post("/ingest")
async def ingest_documents(directory: str):
    """重新加载文档并更新知识库"""
    documents = load_documents(directory)
    chunks = chunk_documents(documents)
    vectorstore = create_vector_store(chunks)
    return {"status": "success", "documents": len(documents), "chunks": len(chunks)}

if __name__ == "__main__":
    uvicorn.run("main:app", host="0.0.0.0", port=8000)

# 启动 Ollama（如果尚未运行）
ollama serve &

# 拉取必要模型
ollama pull qwen2.5:7b
ollama pull nomic-embed-text

# 启动 API 服务
uvicorn main:app --host 0.0.0.0 --port 8000 --reload

# 或使用 Docker Compose 一键部署
docker-compose up -d

# monitoring.py
import time
from datetime import datetime

class RAGMetrics:
    def __init__(self):
        self.query_count = 0
        self.total_latency = 0.0
        self.error_count = 0
    
    def record_query(self, latency: float, success: bool = True):
        self.query_count += 1
        self.total_latency += latency
        if not success:
            self.error_count += 1
    
    def get_stats(self):
        avg_latency = self.total_latency / self.query_count if self.query_count > 0 else 0
        return {
            "total_queries": self.query_count,
            "avg_latency_ms": round(avg_latency * 1000, 2),
            "error_rate": round(self.error_count / self.query_count * 100, 2) if self.query_count > 0 else 0,
            "uptime": datetime.now().isoformat()
        }

metrics = RAGMetrics()

# docker-compose.yml
version: '3.8'

services:
  rag-api:
    build: .
    ports:
      - "8000:8000"
    volumes:
      - ./docs:/app/docs
      - ./chroma_db:/app/chroma_db
    environment:
      - OLLAMA_HOST=http://ollama:11434
    depends_on:
      - ollama

  ollama:
    image: ollama/ollama:latest
    ports:
      - "11434:11434"
    volumes:
      - ./ollama_models:/root/.ollama
    deploy:
      resources:
        limits:
          memory: 8G

  chromadb:
    image: chromadb/chroma:latest
    ports:
      - "8001:8000"
    volumes:
      - ./chroma_db:/chroma/chroma