🇨🇳 DeepSeek国产RAG方案

# ========== DeepSeek RAG完整方案（全国产技术栈） ==========

# 步骤1：安装依赖
# pip install openai sentence-transformers qdrant-client

import os
from openai import OpenAI
from sentence_transformers import SentenceTransformer
import numpy as np
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct

# ========== 配置DeepSeek API ==========
client = OpenAI(
    api_key="sk-your-deepseek-key",  # 注册：https://platform.deepseek.com
    base_url="https://api.deepseek.com"
)

# ========== 加载中文Embedding模型 ==========
print("📥 加载bge-large-zh-v1.5模型...")
embed_model = SentenceTransformer('BAAI/bge-large-zh-v1.5')
print("✅ 模型加载完成")

# ========== 初始化Qdrant向量数据库 ==========
qdrant = QdrantClient(path="./qdrant_db")  # 本地存储

# 创建集合
collection_name = "chinese_knowledge_base"
try:
    qdrant.create_collection(
        collection_name=collection_name,
        vectors_config=VectorParams(size=1024, distance=Distance.COSINE)
    )
    print(f"✅ 创建集合: {collection_name}")
except:
    print(f"集合已存在: {collection_name}")

# ========== 知识库文档 ==========
documents = [
    "RAG是检索增强生成系统，结合信息检索和大语言模型，解决知识时效性问题。",
    "DeepSeek是国产大语言模型，95分能力，成本仅$0.27/1M输入，性价比极高。",
    "bge-large-zh-v1.5是智源AI开源的中文Embedding模型，中文检索准确率95%。",
    "Qdrant是开源向量数据库，支持高性能向量检索和过滤功能。",
    "向量数据库用于存储文档的向量表示，实现语义相似度检索。"
]

# ========== 文档入库 ==========
print("\n📚 文档向量化并入库...")
embeddings = embed_model.encode(documents, normalize_embeddings=True)

points = [
    PointStruct(
        id=i,
        vector=embeddings[i].tolist(),
        payload={"text": doc}
    )
    for i, doc in enumerate(documents)
]

qdrant.upsert(collection_name=collection_name, points=points)
print(f"✅ {len(documents)}个文档已入库")

# ========== RAG检索函数 ==========
def rag_search(query, top_k=3):
    """DeepSeek RAG检索"""
    print(f"\n🔍 查询: {query}")
    
    # 1. 查询向量化
    query_vector = embed_model.encode(query, normalize_embeddings=True)
    
    # 2. 向量检索
    search_result = qdrant.search(
        collection_name=collection_name,
        query_vector=query_vector.tolist(),
        limit=top_k
    )
    
    # 3. 提取检索文档
    retrieved_docs = [hit.payload["text"] for hit in search_result]
    print(f"✅ 检索到{len(retrieved_docs)}个相关文档")
    
    for i, hit in enumerate(search_result, 1):
        print(f"  {i}. {hit.payload['text'][:50]}... (分数: {hit.score:.3f})")
    
    # 4. 构建Prompt
    context = "\n".join([f"{i+1}. {doc}" for i, doc in enumerate(retrieved_docs)])
    prompt = f"""你是一个专业的问答助手。请基于以下参考文档回答问题。

参考文档：
{context}

用户问题：{query}

要求：
1. 仅基于参考文档回答，不要编造
2. 回答要准确、简洁
3. 如果文档中没有相关信息，请说明

回答："""
    
    # 5. 调用DeepSeek生成
    print("\n🤖 DeepSeek生成回答...")
    response = client.chat.completions.create(
        model="deepseek-chat",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.3,
        max_tokens=500
    )
    
    answer = response.choices[0].message.content
    
    # 6. 成本计算
    input_tokens = response.usage.prompt_tokens
    output_tokens = response.usage.completion_tokens
    cost = (input_tokens / 1_000_000 * 0.27) + (output_tokens / 1_000_000 * 1.1)
    
    print(f"\n💡 回答：\n{answer}")
    print(f"\n💰 成本：${cost:.6f} (输入{input_tokens} + 输出{output_tokens})")
    
    return answer

# ========== 使用示例 ==========
if __name__ == "__main__":
    # 测试查询
    rag_search("什么是RAG系统？")
    rag_search("DeepSeek有什么优势？")
    
    print("\n🎉 全国产RAG系统搭建完成！")
    print("💰 成本对比：DeepSeek比GPT-4便宜18倍！")