Redis 作为聊天完成的上下文存储

! pip install -q redis openai python-dotenv 'openai[datalib]'

OPENAI_API_KEY=your_key

from openai import OpenAI
import os
from dotenv import load_dotenv

load_dotenv()
oai_client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

def get_completion(prompt, model="gpt-3.5-turbo"):
    messages = [{"role": "user", "content": prompt}]
    response = oai_client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=0,
    )
    return response.choices[0].message.content

! docker compose up -d

directory = './assets/'
model = 'text-embedding-3-small'
i = 1

for file in os.listdir(directory):
    with open(os.path.join(directory, file), 'r') as f:
        content = f.read()
        # Create the embedding using the new client-based method
        response = oai_client.embeddings.create(
            model=model,
            input=[content]
        )
        # Access the embedding from the response object
        vector = response.data[0].embedding
        
        # Store the content and vector using your JSON client
        client.json().set(f'doc:{i}', '$', {'content': content, 'vector': vector})
    i += 1

from redis.commands.search.query import Query
import numpy as np

response = oai_client.embeddings.create(
    input=[prompt],
    model=model
)
# Extract the embedding vector from the response
embedding_vector = response.data[0].embedding

# Convert the embedding to a numpy array of type float32 and then to bytes
vec = np.array(embedding_vector, dtype=np.float32).tobytes()

# Build and execute the Redis query
q = Query('*=>[KNN 1 @vector $query_vec AS vector_score]') \
    .sort_by('vector_score') \
    .return_fields('content') \
    .dialect(2)
params = {"query_vec": vec}

context = client.ft('idx').search(q, query_params=params).docs[0].content
print(context)