利用元数据在Pinecone中实现用户ID过滤的教程

from pinecone import Pinecone, Index
from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import Pinecone as LangchainPinecone
import os

# 初始化Pinecone客户端和嵌入模型
api_key = os.getenv("PINECONE_API_KEY")
environment = os.getenv("PINECONE_ENVIRONMENT")
index_name = os.getenv("PINECONE_INDEX")

pinecone_client = Pinecone(api_key=api_key, environment=environment)
embeddings = OpenAIEmbeddings(openai_api_key=os.getenv("OPENAI_API_KEY"))

# 假设这是您要嵌入的文档和对应的用户ID
documents_for_user1 = [
    ("This is a document for user 1.", {"source": "user_document", "user_id": 1}),
    ("Another piece of text from user 1.", {"source": "user_document", "user_id": 1})
]
documents_for_user2 = [
    ("User 2's specific information.", {"source": "user_document", "user_id": 2}),
    ("A different document for user 2.", {"source": "user_document", "user_id": 2})
]

# 获取或创建Pinecone索引
if index_name not in pinecone_client.list_indexes():
    pinecone_client.create_index(
        name=index_name,
        dimension=embeddings.client.dimensions, # 确保维度匹配您的嵌入模型
        metric='cosine'
    )
pinecone_index = pinecone_client.Index(index_name)

# 批量嵌入并上传向量，包含user_id元数据
def upsert_vectors_with_metadata(index: Index, texts_and_metadatas: list, embeddings_model, batch_size=32):
    for i in range(0, len(texts_and_metadatas), batch_size):
        batch = texts_and_metadatas[i:i+batch_size]
        texts = [item[0] for item in batch]
        metadatas = [item[1] for item in batch]

        # 生成嵌入
        embeds = embeddings_model.embed_documents(texts)

        # 准备upsert数据
        # Pinecone的upsert方法需要 (id, vector, metadata) 格式
        # 这里我们简化处理，假设id是递增的
        vectors_to_upsert = []
        for j, (text, metadata) in enumerate(batch):
            # 实际应用中，id应该是一个唯一且持久的标识符
            vector_id = f"doc_{metadata['user_id']}_{i+j}" 
            vectors_to_upsert.append((vector_id, embeds[j], metadata))

        index.upsert(vectors=vectors_to_upsert)
    print(f"Upserted {len(texts_and_metadatas)} vectors to index '{index_name}'.")

# 示例调用
# upsert_vectors_with_metadata(pinecone_index, documents_for_user1, embeddings)
# upsert_vectors_with_metadata(pinecone_index, documents_for_user2, embeddings)

from flask import Flask, request, jsonify, session
import os
from langchain_openai import ChatOpenAI
from langchain.memory import ConversationBufferWindowMemory
from langchain.chains import ConversationalRetrievalChain
from langchain_core.prompts import PromptTemplate
from langchain_community.vectorstores import Pinecone as LangchainPinecone
from pinecone import Pinecone, Index

app = Flask(__name__)
app.secret_key = os.getenv("FLASK_SECRET_KEY", "supersecretkey") # 用于会话管理

# 初始化Pinecone客户端和嵌入模型
pinecone_api_key = os.getenv("PINECONE_API_KEY")
pinecone_environment = os.getenv("PINECONE_ENVIRONMENT")
openai_api_key = os.getenv("OPENAI_API_KEY")
index_name = os.getenv("PINECONE_INDEX")
text_field = "text" # 假设您的文本内容存储在元数据的'text'字段中

pinecone_client = Pinecone(api_api_key=pinecone_api_key, environment=pinecone_environment)
embeddings = OpenAIEmbeddings(openai_api_key=openai_api_key)

# 获取Pinecone索引实例
# 确保索引已经存在并包含数据
pinecone_index_instance = pinecone_client.Index(index_name)

# 使用Langchain的Pinecone集成创建vectorstore
vectorstore = LangchainPinecone(
    index=pinecone_index_instance, 
    embedding=embeddings, 
    text_key=text_field # 指定存储原始文本的元数据字段
)

# 假设这些函数用于获取用户特定的配置
def get_bot_temperature(user_id):
    # 根据user_id返回不同的温度，或默认值
    return 0.7

def get_custom_prompt(user_id):
    # 根据user_id返回不同的自定义提示，或默认值
    return "You are a helpful AI assistant. Answer the question based on the provided context."

@app.route('/<int:user_id>/chat', methods=['POST'])
def chat(user_id):
    user_message = request.form.get('message')
    if not user_message:
        return jsonify({"error": "Message is required"}), 400

    # 从会话中加载对话历史
    # 注意：为了每个用户隔离，会话键应包含user_id
    conversation_history_key = f'conversation_history_{user_id}'
    conversation_history = session.get(conversation_history_key, [])

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

    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)

    # 关键部分：在as_retriever中应用filter
    # Pinecone的过滤语法是字典形式，这里使用'$eq'操作符表示“等于”
    retriever = vectorstore.as_retriever(
        search_kwargs={
            'filter': {'user_id': user_id} # 精确匹配当前用户的user_id
        }
    )

    conversation_chain = ConversationalRetrievalChain.from_llm(
        llm=llm,
        retriever=retriever, # 使用带有过滤器的retriever
        memory=memory,
        combine_docs_chain_kwargs={"prompt": TEST_PROMPT},
    )

    response = conversation_chain.run({'question': user_message})

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

    return jsonify(response=response)

# if __name__ == '__main__':
#     # 仅用于开发测试，生产环境应使用WSGI服务器
#     app.run(debug=True)