Redis和GCP AI服务搭建RAG参考架构解决方案

!pip install redis "google-cloud-aiplatform==1.25.0" --upgrade --user

curl -fsSL https://packages.redis.io/gpg | sudo gpg --dearmor -o /usr/share/keyrings/redis-archive-keyring.gpg
echo "deb [signed-by=/usr/share/keyrings/redis-archive-keyring.gpg] https://packages.redis.io/deb $(lsb_release -cs) main" | sudo tee /etc/apt/sources.list.d/redis.list
sudo apt-get update  > /dev/null 2>&1
sudo apt-get install redis-stack-server  > /dev/null 2>&1
redis-stack-server --daemonize yes

import os
import redis

# Redis connection params
REDIS_HOST = os.getenv("REDIS_HOST", "redis-10610.c274.us-east-1-3.ec2.cloud.redislabs.com") #"redis-10610.c274.us-east-1-3.ec2.cloud.redislabs.com"
REDIS_PORT = os.getenv("REDIS_PORT", "10610")      #10610
REDIS_PASSWORD = os.getenv("REDIS_PASSWORD", "******")  #"pobhBJP7Psicp2gV0iqa2ZOc1WdXXXXX"

# Create Redis client
redis_client = redis.Redis(
  host=REDIS_HOST,
  port=REDIS_PORT,
  password=REDIS_PASSWORD
)

# Test connection
redis_client.ping()
# Clear Redis database (optional)
redis_client.flushdb()

from google.colab import auth
auth.authenticate_user()
print('Authenticated')
from getpass import getpass

# input your GCP project ID and region for Vertex AI
PROJECT_ID = getpass("PROJECT_ID:") #'central-beach-194106'
REGION = input("REGION:") #'us-central1'

import vertexai

vertexai.init(project=PROJECT_ID, location=REGION)

from google.cloud import bigquery

# Create bigquery client
bq = bigquery.Client(project=PROJECT_ID)

TABLE_NAME = input("Input a Big Query TABLE_NAME:") #hackernews
DATASET_ID = f"{PROJECT_ID}.google_redis_llms"

# Create dataset
dataset = bigquery.Dataset(DATASET_ID)
dataset.location = "US"
dataset = bq.create_dataset(dataset, timeout=30, exists_ok=True)

# Define table ID
TABLE_ID = f"{DATASET_ID}.{TABLE_NAME}"

# Create source table
def create_source_table(table_id: str):
    create_job = bq.query(f'''
      CREATE OR REPLACE TABLE `{table_id}` AS (
        SELECT
          title, text, time, timestamp, id
        FROM `bigquery-public-data.hacker_news.full`
        WHERE
          type ='story'
        LIMIT 1000
      )
    ''')
    res = create_job.result() # Make an API request
    table = bq.get_table(table_id)
    return table

# Create table
table = create_source_table(TABLE_ID)

# List schema
table.schema

# Query 5 sample records from BQ
query_job = bq.query(f'''
SELECT *
FROM {TABLE_ID}
LIMIT 5
''')

query_job.result().to_dataframe()

from typing import Generator, List, Any

from tenacity import retry, stop_after_attempt, wait_random_exponential
from vertexai.preview.language_models import TextEmbeddingModel

# Embedding model definition from VertexAI PaLM API
embedding_model = TextEmbeddingModel.from_pretrained("textembedding-gecko@001")
VECTOR_DIMENSIONS = 768

@retry(wait=wait_random_exponential(min=1, max=20), stop=stop_after_attempt(3))
def embed_text(text=[]):
    embeddings = embedding_model.get_embeddings(text)
    return [each.values for each in embeddings]

# Convert embeddings to bytes for Redis storage
def convert_embedding(emb: List[float]):
  return np.array(emb).astype(np.float32).tobytes()

import pandas as pd
import numpy as np

QUERY_TEMPLATE = f"""
SELECT id, title, text
FROM {TABLE_ID}
LIMIT {{limit}} OFFSET {{offset}};
"""

def query_bigquery_batches(
    max_rows: int,
    rows_per_batch: int,
    start_batch: int = 0
) -> Generator[pd.DataFrame, Any, None]:
    # Generate batches from a table in big query
    for offset in range(start_batch, max_rows, rows_per_batch):
        query = QUERY_TEMPLATE.format(limit=rows_per_batch, offset=offset)
        query_job = bq.query(query)
        rows = query_job.result()
        df = rows.to_dataframe()
        # Join title and text fields
        df["content"] = df.apply(lambda r: "Title: " + r.title + ". Content: " + r.text, axis=1)
        yield df

import math
from tqdm.auto import tqdm


# Redis key helper function
def redis_key(key_prefix: str, id: str) -> str:
  return f"{key_prefix}:{id}"

# Process a single dataset record
def process_record(record: dict) -> dict:
  return {
      'id': record['id'],
      'embedding': record['embedding'],
      'text': record['text'],
      'title': record['title']
  }

# Load batch of data into Redis as HASH objects
def load_redis_batch(
    redis_client: redis.Redis,
    dataset: list,
    key_prefix: str = "doc",
    id_column: str = "id",
):
    pipe = redis_client.pipeline()
    for i, record in enumerate(tqdm(dataset)):
        record = process_record(record)
        key = redis_key(key_prefix, record[id_column])
        pipe.hset(key, mapping=record)
    pipe.execute()

# Run the entire process
def create_embeddings_bigquery_redis(redis_client):
    # Create generator from BigQuery
    max_rows = 1000
    rows_per_batch = 100
    bq_content_query = query_bigquery_batches(max_rows, rows_per_batch)

    for batch in tqdm(bq_content_query):
      # Split batch into smaller chunks for embedding generation
      batch_splits = np.array_split(batch, math.ceil(rows_per_batch/5))
      # Create embeddings
      batch["embedding"] = [
          convert_embedding(embedding)
          for split in batch_splits
          for embedding in embed_text(split.content)
      ]
      # Write batch to Redis
      batch = batch.to_dict("records")
      load_redis_batch(redis_client, batch)

create_embeddings_bigquery_redis(redis_client)
# Validate how many records are stored in Redis
redis_client.dbsize()

from redis.commands.search.field import (
    NumericField,
    TagField,
    TextField,
    VectorField,
)
from redis.commands.search.indexDefinition import IndexDefinition, IndexType
from redis.commands.search.query import Query


INDEX_NAME = "google:idx"
PREFIX = "doc:"
VECTOR_FIELD_NAME = "embedding"

# Store vectors in redis and create index
# 对hash中的 embedding字段建向量索引FLAT,维度768,相似度测量方法:cosine
def create_redis_index(
    redis_client: redis.Redis,
    vector_field_name: str = VECTOR_FIELD_NAME,
    index_name: str = INDEX_NAME,
    prefix: list = [PREFIX],
    dim: int = VECTOR_DIMENSIONS #768
  ):

    # Construct index
    try:
        redis_client.ft(index_name).info()
        print("Existing index found. Dropping and recreating the index", flush=True)
        redis_client.ft(index_name).dropindex(delete_documents=False)
    except:
        print("Creating new index", flush=True)

    # Create new index
    redis_client.ft(index_name).create_index(
        (
            VectorField(
                vector_field_name, "FLAT",
                {
                    "TYPE": "FLOAT32",
                    "DIM": dim,
                    "DISTANCE_METRIC": "COSINE",
                }
            )
        ),
        definition=IndexDefinition(prefix=prefix, index_type=IndexType.HASH)
    )

# Create index
create_redis_index(redis_client)
# Inspect index attributes
redis_client.ft(INDEX_NAME).info()

# Retreive single HASH from Redis
key = redis_client.keys()[1]
print(key)
redis_client.hgetall(key)

# 1. Create query vector
query = "What is the best computer operating system for software dev?"
query_vector = embed_text([query])[0]

# Our query has been converted to a list of floats (this is a truncated view)
query_vector[:10]


# Helper method to perform KNN similarity search in Redis
def similarity_search(query: str, k: int, return_fields: tuple, index_name: str = INDEX_NAME) -> list:
    # create embedding from query text
    query_vector = embed_text([query])[0]
    # create redis query object
    redis_query = (
        Query(f"*=>[KNN {k} @{VECTOR_FIELD_NAME} $embedding AS score]")
            .sort_by("score")
            .return_fields(*return_fields)
            .paging(0, k)
            .dialect(2)
    )
    # execute the search
    results = redis_client.ft(index_name).search(
        redis_query, query_params={"embedding": convert_embedding(query_vector)}
    )
    return pd.DataFrame([t.__dict__ for t in results.docs ]).drop(columns=["payload"])

  
# 2. Perform vector similarity search with given query
results = similarity_search(query, k=5, return_fields=("score", "title", "text"))


# 3. Review and return the results
display(results)

from vertexai.preview.language_models import TextGenerationModel

# Define generation model
generation_model = TextGenerationModel.from_pretrained("text-bison@001")

response = generation_model.predict(prompt="What is a large language model?")

print("Example response:\n", response.text)

def create_prompt(prompt_template: str, **kwargs) -> str:
  return prompt_template.format(**kwargs)

def rag(query: str, prompt: str, verbose: bool = True) -> str:
    """
    Simple pipeline for performing retrieval augmented generation with
    Google Vertex PaLM API and Redis Enterprise.
    """
    # Perform a vector similarity search in Redis
    if verbose:
        print("Pulling relevant data sources from Redis", flush=True)
    relevant_sources = similarity_search(query, k=3, return_fields=("text",))
    if verbose:
        print("Relevant sources found!", flush=True)
    # Combine the relevant sources and inject into the prompt
    sources_text = "-" + "\n-".join([source for source in relevant_sources.text.values])
    full_prompt = create_prompt(
        prompt_template=prompt,
        sources=sources_text,
        query=query
      )
    if verbose:
        print("\nFull prompt:\n\n", full_prompt, flush=True)
    # Perform text generation to get a response from PaLM API
    response = generation_model.predict(prompt=full_prompt)
    return response.text

PROMPT = """You are a helpful virtual technology and IT assistant. Use the hacker news posts below as relevant context and sources to help answer the user question. Don't blindly make things up.

SOURCES:
{sources}

QUESTION:
{query}?

ANSWER:"""

query = "What are the best operating systems for software development?"
response = rag(query=query, prompt=PROMPT)
print(response)

query = "What are some commonly reported problems with MacBooks and MacOS for software dev purposes?"
response = rag(query=query, prompt=PROMPT)
print(response)

# Some boiler plate helper methods
import hashlib

def hash_input(prefix: str, _input: str):
    return prefix + hashlib.sha256(_input.encode("utf-8")).hexdigest()

def standard_check(key: str):
  # function to perform a standard cache check
    res = redis_client.hgetall(key)
    if res:
      return res[b'response'].decode('utf-8')

def cache_response(query: str, response: str):
    key = hash_input("llmcache:", query)
    redis_client.hset(key, mapping={"prompt": query, "response": response})

# LLM Cache wrapper / decorator function
def standard_llmcache(llm_callable):
    def wrapper(*args, **kwargs):
        # Check LLM Cache first
        key = hash_input("llmcache:", *args, **kwargs)
        response = standard_check(key)
        # Check if we have a cached response we can use
        if response:
            return response
        # Otherwise execute the llm callable here
        response = llm_callable(*args, **kwargs)
        cache_response(query, response)
        return response

    return wrapper
  
# Define a function that invokes the PaLM API wrapped with a cache check
@standard_llmcache
def ask_palm(query: str):
  prompt = PROMPT
  response = rag(query, prompt, verbose=False)
  return response

query = "What are the best operating systems for software development?"

ask_palm(query)

%%time # for ipynb

ask_palm(query)

import json

def add_message(prompt: str, response: str):
    msg = {
        "prompt": prompt,
        "response": response
    }
    redis_client.lpush("chat-history", json.dumps(msg))

def get_messages(k: int = 5):
    return [json.loads(msg) for msg in redis_client.lrange("chat-history", 0, k)]
  

query = "Do you have any advice for getting started in the tech field as a software dev?"
response = rag(query, PROMPT, verbose=False)
print(response)
add_message(query, response)

query = "What if I am still in college, any tips there?"
response = rag(query, PROMPT, verbose=False)
print(response)
add_message(query, response)

get_messages()

# Clean up bigquery
bq.delete_table(TABLE_ID, not_found_ok=True)

bq.delete_dataset(
    DATASET_ID, delete_contents=True, not_found_ok=True
)

# Clean up redis
#!redis-stack-server stop
redis_client.flushdb();
redis_client.close();