在Pinecone中实现基于用户ID的向量检索过滤

from pinecone import Pinecone, Index
from langchain_openai import OpenAIEmbeddings
from langchain_pinecone import PineconeVectorStore
import os

# 初始化Pinecone和嵌入模型
pinecone_api_key = os.getenv("PINECONE_API_KEY")
pinecone_env = os.getenv("PINECONE_ENVIRONMENT") # 例如 "us-west-2"
index_name = os.getenv("PINECONE_INDEX")
openai_api_key = os.getenv("OPENAI_API_KEY")

# 确保Pinecone索引已存在
pc = Pinecone(api_key=pinecone_api_key)
if index_name not in pc.list_indexes():
    pc.create_index(name=index_name, dimension=1536, metric='cosine') # 假设使用OpenAI embeddings，维度为1536

embeddings_model = OpenAIEmbeddings(openai_api_key=openai_api_key)

# 示例数据：包含用户ID的文档块
documents_with_user_id = [
    {"text": "这是用户123上传的关于Python编程的文档片段。", "user_id": 123},
    {"text": "用户456的机器学习笔记。", "user_id": 456},
    {"text": "用户123的另一个关于数据科学的文档。", "user_id": 123},
    {"text": "一个公共文档，可能没有user_id或user_id为0。", "user_id": 0}
]

# 将文档块嵌入并上传到Pinecone
# 在LangChain中，通常通过`PineconeVectorStore.from_documents`或`add_texts`方法实现
# 这里为了清晰展示元数据添加，我们模拟其内部逻辑
index = pc.Index(index_name)
vectorstore = PineconeVectorStore(index=index, embedding=embeddings_model)

# 实际上传时，您会通过LangChain的API来完成，例如：
# texts = [doc["text"] for doc in documents_with_user_id]
# metadatas = [{"user_id": doc["user_id"]} for doc in documents_with_user_id]
# vectorstore.add_texts(texts=texts, metadatas=metadatas)

# 假设已经有向量和对应的元数据准备好，直接进行upsert
# 实际场景中，LangChain会为您处理嵌入和格式化
# 这是一个更接近Pinecone原生API的upsert示例，以便理解元数据结构
# from pinecone import Index
# index = Index(index_name)
# for i, doc in enumerate(documents_with_user_id):
#     vec_id = f"doc_{doc['user_id']}_{i}"
#     # 实际这里会调用embeddings_model.embed_query(doc['text'])获取向量
#     # 假设我们已经有了一个模拟的向量
#     mock_vector = [0.1] * 1536 # 替换为真实的嵌入向量
#     index.upsert(vectors=[
#         {"id": vec_id, "values": mock_vector, "metadata": {"user_id": doc["user_id"], "text": doc["text"]}}
#     ])
# print("Vectors with user_id metadata uploaded successfully.")

# 为了本教程的后续部分，我们假设向量已经带有正确的user_id元数据

import os
from flask import Flask, request, jsonify, session
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_pinecone import PineconeVectorStore
from langchain.chains import ConversationalRetrievalChain
from langchain.memory import ConversationBufferWindowMemory
from langchain_core.prompts import PromptTemplate
from pinecone import Pinecone, Index # 导入Pinecone客户端

app = Flask(__name__)
app.secret_key = os.getenv("FLASK_SECRET_KEY", "super-secret-key") # 设置一个安全的密钥

# 初始化环境变量
openai_api_key = os.getenv("OPENAI_API_KEY")
pinecone_api_key = os.getenv("PINECONE_API_KEY")
pinecone_env = os.getenv("PINECONE_ENVIRONMENT")
index_name = os.getenv("PINECONE_INDEX")
text_field = "text" # 假设您的文本内容存储在元数据的'text'字段中

# 初始化Pinecone客户端和索引
pc = Pinecone(api_key=pinecone_api_key, environment=pinecone_env)
pinecone_index = pc.Index(index_name)

# 初始化嵌入模型
embeddings = OpenAIEmbeddings(openai_api_key=openai_api_key)

# 辅助函数（如果需要，您可以根据实际情况实现）
def get_bot_temperature(user_id):
    # 根据user_id获取bot温度，这里仅为示例
    return 0.7

def get_custom_prompt(user_id):
    # 根据user_id获取自定义提示，这里仅为示例
    return "你是一个友好的AI助手，请根据提供的上下文回答问题。"

@app.route('/<int:user_id>/chat', methods=['POST'])
def chat(user_id):
    user_message = request.form.get('message')

    # 从session加载会话历史，注意key的动态性
    conversation_history_key = f'conversation_history_{user_id}'
    conversation_history = session.get(conversation_history_key, [])

    # 创建Pinecone向量存储对象
    vectorstore = PineconeVectorStore(
        index=pinecone_index, 
        embedding=embeddings, 
        text_key=text_field
    )

    bot_temperature = get_bot_temperature(user_id)
    custom_prompt = get_custom_prompt(user_id)

    # 初始化LLM
    llm = ChatOpenAI(
        openai_api_key=openai_api_key,
        model_name='gpt-3.5-turbo',
        temperature=bot_temperature
    )

    # 定义提示模板
    prompt_template = f"""
        {custom_prompt}

        CONTEXT: {{context}}

        QUESTION: {{question}}"""

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

    # 创建会话记忆
    memory = ConversationBufferWindowMemory(memory_key="chat_history", return_messages=True, k=8)

    # 核心：配置带有元数据过滤的检索器
    # 通过search_kwargs参数传递Pinecone的过滤条件
    # {"user_id": {"$eq": user_id}} 表示元数据字段user_id的值等于当前user_id
    retriever = vectorstore.as_retriever(
        search_kwargs={"filter": {"user_id": {"$eq": user_id}}}
    )

    # 创建会话检索链
    conversation_chain = ConversationalRetrievalChain.from_llm(
            llm=llm,
            retriever=retriever, # 使用带有过滤条件的检索器
            memory=memory,
            combine_docs_chain_kwargs={"prompt": TEST_PROMPT},
        )

    # 处理用户输入并获取响应
    response = conversation_chain.run({'question': user_message})

    # 保存用户消息和机器人响应到session
    conversation_history.append({'input': user_message, 'output': response})
    session[conversation_history_key] = conversation_history

    return jsonify(response=response)

if __name__ == '__main__':
    # 确保设置了环境变量
    if not all([openai_api_key, pinecone_api_key, pinecone_env, index_name, app.secret_key]):
        print("请设置所有必要的环境变量：OPENAI_API_KEY, PINECONE_API_KEY, PINECONE_ENVIRONMENT, PINECONE_INDEX, FLASK_SECRET_KEY")
        exit(1)
    app.run(debug=True)