🔒 Ollama完全本地化RAG

# ========== Ollama本地RAG方案（100%本地，0云端） ==========

# 步骤1：安装Ollama
# macOS/Linux: curl -fsSL https://ollama.com/install.sh | sh
# Windows: 下载安装包 https://ollama.com/download

# 步骤2：下载本地模型
# ollama pull qwen2.5:7b  # 或 qwen2.5:14b（更强）

# 步骤3：安装Python依赖
# pip install langchain-community chromadb sentence-transformers

from langchain_community.llms import Ollama
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate

# ========== 配置本地Embedding ==========
print("📥 加载本地Embedding模型...")
embeddings = HuggingFaceEmbeddings(
    model_name='BAAI/bge-large-zh-v1.5',
    model_kwargs={'device': 'cpu'},  # 或 'cuda'
    encode_kwargs={'normalize_embeddings': True}
)
print("✅ Embedding模型加载完成")

# ========== 配置本地LLM（Ollama） ==========
print("🤖 连接本地Ollama...")
llm = Ollama(
    model="qwen2.5:7b",  # 或 qwen2.5:14b、llama3.1:8b
    temperature=0.3
)
print("✅ Ollama连接成功")

# ========== 文档加载与分块 ==========
documents_text = [
    "RAG是检索增强生成系统，结合信息检索和大语言模型。",
    "Ollama是本地LLM运行工具，支持Llama、Qwen等多个开源模型。",
    "Qwen2.5是阿里云开源的中文大模型，7B版本可在16GB内存运行。",
    "bge-large-zh-v1.5是智源AI开源的中文Embedding模型。",
    "完全本地化部署可确保数据隐私和安全合规。"
]

from langchain.schema import Document
documents = [Document(page_content=text) for text in documents_text]

# 文档分块
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=500,
    chunk_overlap=50,
    separators=["\n\n", "\n", "。", "！", "？"]
)
splits = text_splitter.split_documents(documents)

# ========== 创建本地向量存储 ==========
print("\n📚 创建本地向量数据库...")
vectorstore = Chroma.from_documents(
    documents=splits,
    embedding=embeddings,
    persist_directory="./chroma_local"  # 本地持久化
)
print("✅ 向量数据库创建完成")

# ========== 创建检索器 ==========
retriever = vectorstore.as_retriever(
    search_type="similarity",
    search_kwargs={"k": 3}
)

# ========== 自定义Prompt ==========
prompt_template = """你是一个专业的问答助手。请基于以下参考文档回答问题。

参考文档：
{context}

问题：{question}

要求：
1. 仅基于参考文档回答
2. 回答要准确、简洁、专业
3. 使用中文回答

回答："""

PROMPT = PromptTemplate(
    template=prompt_template,
    input_variables=["context", "question"]
)

# ========== 创建QA链 ==========
qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type="stuff",
    retriever=retriever,
    return_source_documents=True,
    chain_type_kwargs={"prompt": PROMPT}
)

print("✅ QA链创建完成\n")

# ========== 查询函数 ==========
def local_rag_query(question):
    """完全本地化RAG查询"""
    print(f"❓ 问题：{question}\n")
    
    result = qa_chain({"query": question})
    
    print(f"💡 回答：{result['result']}\n")
    print(f"📚 参考文档：")
    for i, doc in enumerate(result['source_documents'], 1):
        print(f"  {i}. {doc.page_content[:60]}...")
    
    return result

# ========== 使用示例 ==========
if __name__ == "__main__":
    # 查询测试
    local_rag_query("什么是RAG系统？")
    local_rag_query("Ollama有什么优势？")
    
    print("\n🎉 完全本地化RAG系统运行成功！")
    print("🔒 数据安全：所有数据都在本地，未发送到云端")
    print("💰 成本：0云端费用，仅需本地硬件成本")