python实现类拟RAG检索给到AI生成回答!

requests==2.31.0sentence-transformers==2.2.2faiss-cpu==1.7.4numpy==1.24.3

pip install faiss-cpu

pip install faiss-gpu

set HF_ENDPOINT=https://hf-mirror.com

import osimport jsonimport requestsfrom typing import List, Dictfrom sentence_transformers import SentenceTransformerimport faissimport numpy as np# 配置DeepSeek API密钥和URLDEEPSEEK_API_KEY = "DEEPSEEK_API_KEY"#DEEPSEEK_API_KEYDEEPSEEK_API_URL = "https://api.deepseek.com/v1/chat/completions"# 初始化嵌入模型embedding_model = SentenceTransformer('paraphrase-MiniLM-L6-v2')# 模拟文档库documents = [    "DeepSeek 是一个强大的语言模型平台。",    "RAG 结合了检索和生成的优势。",    "Python 是一种广泛使用的高级编程语言。",    "向量数据库用于高效存储和检索嵌入向量。",    "Faiss 是 Facebook 开发的高效相似性搜索库。"]# 构建向量索引def build_vector_index(docs: List[str]) -> faiss.IndexFlatIP:    embeddings = embedding_model.encode(docs)    dimension = embeddings.shape[1]    index = faiss.IndexFlatIP(dimension)    faiss.normalize_L2(embeddings)    index.add(embeddings.astype('float32'))    return index# 检索相关文档def retrieve_documents(query: str, index: faiss.IndexFlatIP, docs: List[str], top_k: int = 3) -> List[str]:    query_embedding = embedding_model.encode([query])    faiss.normalize_L2(query_embedding)    distances, indices = index.search(query_embedding.astype('float32'), top_k)    return [docs[i] for i in indices[0]]# 调用DeepSeek API生成回答def generate_answer(prompt: str) -> str:    headers = {        "Authorization": f"Bearer {DEEPSEEK_API_KEY}",        "Content-Type": "application/json"    }    payload = {        "model": "deepseek-chat",        "messages": [{"role": "user", "content": prompt}],        "stream": False    }    response = requests.post(DEEPSEEK_API_URL, headers=headers, data=json.dumps(payload))    if response.status_code == 200:        result = response.json()        return result["choices"][0]["message"]["content"]    else:        raise Exception(f"请求失败：{response.status_code}, {response.text}")# 主函数实现RAG流程def rag_pipeline(query: str):    print("构建向量索引...")    index = build_vector_index(documents)    print("检索相关文档...")    retrieved_docs = retrieve_documents(query, index, documents)    context = "\n".join(retrieved_docs)    prompt = f"根据以下内容回答问题：\n{context}\n\n问题：{query}"    print("生成回答中...")    answer = generate_answer(prompt)    print("回答：", answer)if __name__ == "__main__":    user_query = input("请输入你的问题：")    rag_pipeline(user_query)