Build a secure RAG with knowledge base encryption

# -*- coding: utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.
import os
import re
import json
import torch
import logging
from flask import Flask, request
from modelscope import AutoModelForCausalLM, AutoTokenizer
from rai_sam.engine.client.content import SamContentClient

logging.basicConfig(level=logging.DEBUG)
log: logging.Logger = logging.getLogger(__name__)

app = Flask(__name__)

# Used to encrypt and decrypt the knowledge base. Must be consistent with the key used for knowledge base file encryption. Retrieved from environment variables here.
# SAM_KEY_ID: Knowledge base key ID (4 to 48 bytes), SAM_KEY_SECRET: Key (4 to 128 bytes).
sam_key_sets = {
    os.getenv("SAM_KEY_ID"): os.getenv("SAM_KEY_SECRET"),
}

# Configure the absolute path to the large language model file.
model_name = "/mnt/workspace/Qwen/Qwen2___5-3B-Instruct"

# Pre-defined context start and end identifier.
rai_context_start = "<|rai_sam_encrypted_context_start|>"
rai_context_end = "<|rai_sam_encrypted_context_end|>"

device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="auto",
    device_map=device
)
tokenizer = AutoTokenizer.from_pretrained(model_name)

rai_sam_client = SamContentClient()

def rai_sam_decrypt_content(key_ids: list[str], contents: list[str]) -> list[str]:
    client = SamContentClient()

    if len(set(key_ids)) == 1:
       combine = True
    else:
       combine = False

    log.info("combine: %d", combine)

    if combine == True:
       key_id = key_ids[0]
       key_secret = sam_key_sets.get(key_id)
       if key_secret == None:
           raise RuntimeError("No sam key secret found")

       dec_contents = client.SamClientDecryptContents(key_id, key_secret, contents)
       if dec_contents == None:
           log.error("Failed to decrypt contents")
           return None
    else:
       dec_contents = []
       for i in range(len(key_ids)):
           key_id = key_ids[i]
           key_secret = sam_key_sets.get(key_id)
           if key_secret == None:
               raise RuntimeError("No sam key secret found")

           content = contents[i]
           dec_content = client.SamClientDecryptContents(key_id, key_secret, [content])
           if dec_content == None:
               log.error("Failed to decrypt content")
               return None

           dec_contents.append(dec_content[0])

    return dec_contents

def generate_prompt_plaintext(in_prompt: str) -> str:
    start_pos = in_prompt.find(rai_context_start)
    if start_pos == -1:
        log.info("The input prompt is plaintext")
        return in_prompt

    log.debug("rai_context_start pos: %d", start_pos)

    end_pos = in_prompt.rfind(rai_context_end)
    if end_pos == -1:
        log.error("Not find context end tag: %s", rai_context_end)
        return None

    log.debug("rai_context_end pos: %d", end_pos)

    # Get context content in the in_prompt.
    context = in_prompt[start_pos + len(rai_context_start):end_pos]
    context_json = json.loads(context)
    log.debug("context_json: %s", context_json)

    contents = context_json["contents"]
    log.debug("contents: %s", contents)

    key_ids = context_json["key_ids"]
    log.debug("key_ids: %s", key_ids)

    if len(contents) != len(key_ids):
        raise RuntimeError("the length of contents and key_ids is not euqal")

    dec_contents = rai_sam_decrypt_content(key_ids, contents)
    log.debug("dec_contents: %s", dec_contents)
 
    context = "\n\n".join(
        [content for content in dec_contents]
    )

    out_prompt = in_prompt[:start_pos] + context + in_prompt[end_pos + len(rai_context_end):]

    return out_prompt

def generate_model_response(
        prompt: str,
        max_new_tokens: int = 512,
        temperature: float = 1.0,
        top_k: int = 50,
        top_p: float = 1.0) -> str:

    model_inputs = tokenizer(prompt, return_tensors="pt").to(device)

    generated_ids = model.generate(
        **model_inputs,
        max_new_tokens=max_new_tokens,
        temperature=temperature,
        top_k=top_k,
        top_p=top_p
    )

    generated_ids = [
        output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
    ]

    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
    log.debug("response: %s", response)

    tag = "Assistant:"
    pos = response.find(tag)
    if pos != -1:
        print("pos: %d", pos)
        response = response[pos + len(tag):]

    return response.strip()

@app.route('/model', methods=['POST'])
def process_model_generate():
    log.debug('process model generate start.')

    body = request.json
    log.debug("request body: %s", body)

    if "prompt" in body:
        in_prompt = body['prompt']
    else:
        raise RuntimeError("No prompt found")

    if "max_new_tokens" in body:
        max_new_tokens = body['max_new_tokens']
    else:
        max_new_tokens = 512

    if "temperature" in body:
        temperature = body['temperature']
    else:
        temperature = 0.95

    if "top_k" in body:
        top_k = body['top_k']
    else:
        top_k = 50

    if "top_p" in body:
        top_p = body['top_p']
    else:
        top_p = 1.0

    log.debug("prompt: %s", in_prompt)
    log.debug("max_new_tokens: %d", max_new_tokens)
    log.debug("temperature: %f", temperature)
    log.debug("top_k: %d", top_k)
    log.debug("top_p: %f", top_p)

    prompt = generate_prompt_plaintext(in_prompt)
    if prompt == None:
        log.error("Failed to generate prompt plaitext")
        raise RuntimeError("generate prompt fail")

    log.debug("generated prompt: %s", prompt)

    model_response = generate_model_response(prompt,
                         max_new_tokens, temperature, top_k, top_p)
    if model_response == None:
        log.error("Failed to generate model response")
        raise RuntimeError("generate model response fail")

    response = {
        "response": model_response
    }

    return response

@app.route('/v1/chat/completions', methods=['POST'])
def process_model_chat_completions():
    log.debug('process model chat completions start.')

    body = request.json
    log.debug("request json: %s", body)
 
    if "messages" in body:
        messages = body['messages']
    else:
        raise RuntimeError("No messages found")

    if "temperature" in body:
        temperature = body['temperature']
    else:
        temperature = 0.95

    log.debug("temperature: %s", temperature)

    in_prompt = str(messages[0])
    prompt = generate_prompt_plaintext(in_prompt)
    if prompt == None:
        raise RuntimeError("generate prompt fail")

    log.debug("generated prompt: %s", prompt)

    model_response = generate_model_response(prompt, temperature=temperature)
    if model_response == None:
        raise RuntimeError("generate model response fail")

    message = {
        'role': 'assistant',
        'content': model_response
    }

    content = {
        'message': message
    }

    response = {
        'choices': [content]
    }

    return response

if __name__ == "__main__":
    app.run(host = '0.0.0.0', port = '8000', debug=True)