diff --git a/web/api/chat.py b/web/api/chat.py index 3483b82..23bf86f 100644 --- a/web/api/chat.py +++ b/web/api/chat.py @@ -15,22 +15,202 @@ class handler(BaseHTTPRequestHandler): post_data = self.rfile.read(content_length) data = json.loads(post_data) - self.wfile.write(chat(data['history'], data['query']).encode('utf-8')) + self.wfile.write(chat(data['query'], data['history']).encode('utf-8')) # ------------------------------- chat gpt stuff ------------------------------- +import os +import requests +from typing import List, Dict -def chat(history, query) -> str: +try: + import tiktoken +except ImportError as e: + print(e) + print("Please install tiktoken with `pip install tiktoken`") - # history = [ - # {'role': 'user', 'content': 'Open the pod bay doors'}, - # {'role': 'assistant', 'content': 'I'm sorry, Dave. I'm afraid I can't do that.'}, - # {'role': 'user', 'content': 'Who won the world series in 2020?'}, - # {'role': 'assistant', 'content': 'The Los Angeles Dodgers won the World Series in 2020.'}, - # ] - # - # query = 'Who will win the world series in 2023?' - # - # (if you want any system message, add it yourself) +import openai +try: + import config + openai.api_key = config.OPENAI_API_KEY +except ImportError: + openai.api_key = os.environ.get('OPENAI_API_KEY') - return "no, you're a " + query +from get_blocks import get_top_k_blocks, Block +# OpenAI models +EMBEDDING_MODEL = "text-embedding-ada-002" +COMPLETIONS_MODEL = "gpt-3.5-turbo" +MODERATION_ENDPOINT = "https://api.openai.com/v1/moderations" + +# OpenAI parameters +LEN_EMBEDDINGS = 1536 +MAX_TOKEN_LEN_PROMPT = 4095 # This may be 8191, unsure. +TRUNCATE_CONTEXT = 2000 + +def moderate_query(query: str) -> List[str]: + """This function uses the OpenAI Moderation API to check if a query contains any offensive language. + + Args: + query (str): The query to be checked. + + Raises: + Exception: If the API call fails. + + Returns: + List[str]: A list of categories that the query was flagged for. + """ + + headers = {"Content-Type": "application/json","Authorization": f"Bearer {OPENAI_API_KEY}"} + + data = {"input": query} + + response = requests.post(MODERATION_ENDPOINT, headers=headers, data=json.dumps(data)) + flagged_categories = [] + + if response.status_code == 200: + moderation_results = response.json() + flagged = moderation_results['results'][0]['flagged'] + categories = moderation_results['results'][0]['categories'] + + if flagged: + for category, is_flagged in categories.items(): + if is_flagged: + flagged_categories.append(category) + else: + raise Exception(f"Error: {response.status_code} {response.reason}") + + return flagged_categories + +def limit_tokens(text: str, max_tokens: int, encoding_name: str = "cl100k_base") -> str: + encoding = tiktoken.get_encoding(encoding_name) + tokens = encoding.encode(text)[:max_tokens] + return encoding.decode(tokens) + +def generate_prompt(query: str, history: List[Dict[str, str]] = [], blocks: List[Block] = [], mode: str = "standard") -> List[Dict[str, str]]: + """ + This function generates a prompt in messages format for the OpenAI ChatCompletions API. + First, it picks a system description using the mode. + Second, it adds the previous dialogue to the prompt. + Third, it adds an instruction to the prompt based on the mode. + Fourth, it adds the context from the top-k most relevant blocks from the Alignment Research Dataset to the prompt. + Fifth, it adds the user query to the prompt. + + Messages take the following format: + messages=[ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Who won the world series in 2020?"}, + {"role": "assistant", "content": "The Los Angeles Dodgers won the World Series in 2020."}, + {"role": "user", "content": "Where was it played?"} + ] + + Args: + query (str): The user query. + history (List[Dict[str, str]]): The previous dialogue. Defaults to []. + blocks (List[Block]): The top-k most relevant blocks from the Alignment Research Dataset. Defaults to []. + mode (str): The mode of the assistant. Can be "standard", etc. Defaults to "standard". + + Returns: + List[Dict[str, str]]: The prompt in messages format. + """ + # Initialize prompt + prompt = [] + + # Generate system description + if mode == "standard": + prompt.append({"role": "system", "content": "You are a helpful assistant knowledgeable about AI Alignment and Safety."}) + # elif mode == "other": + else: + raise Exception(f"Invalid mode: {mode}") + + # Add previous dialogue + for message in history: + prompt.append(message) + + # Add instruction + if mode == "standard": + instruction_prompt = "Please answer my question (after the Q:) using the provided context." + prompt.append({"role": "assistant", "content": instruction_prompt}) + # elif mode == "other": + else: + raise Exception(f"Invalid mode: {mode}") + + # Add context from top-k blocks + if blocks is None: + return "Context missing." + context_prompt = "Context:\n\n" + for i, block in enumerate(blocks): + context_prompt += f"[{i}] {block.text}\n\n" + context_prompt = context_prompt[:-2] + + try: + context_prompt = limit_tokens(context_prompt, TRUNCATE_CONTEXT) + except: # If the tokenizer does not work, just truncate the context prompt to 8000 characters + context_prompt = context_prompt[:TRUNCATE_CONTEXT*4] + prompt.append({"role": "user", "content": f"{context_prompt}"}) + + # Add user query + prompt.append({"role": "user", "content": f"Q: {query}"}) + + return prompt + +def normal_completion(prompt: List[Dict[str, str]]) -> str: + """ + This function uses the OpenAI ChatCompletions API to answer a user query. + + Args: + messages (Dict[str, str]): A dictionary containing the system prompt and user prompt, in addition to any previous dialogue. + + Returns: + str: The answer generated by the API. + + Raises: + Exception: If the API call fails. + """ + try: + return openai.ChatCompletion.create( + model=COMPLETIONS_MODEL, + messages=prompt + )["choices"][0]["message"]["content"] + except Exception as e: + print(e) + return "I'm sorry, I failed to process your query. Please try again. If the problem persists, please contact the administrator." + +def chat(query: str, history: List[Dict[str, str]] = [], k: str = 10, mode: str = "standard", HyDE: bool = False) -> str: + """ + This function uses the OpenAI ChatCompletions API to answer a user query. + It first checks if the query is offensive, and if so, raises an exception. + Then, it finds the top-k most relevant blocks from the Alignment Research Dataset and uses them as context for the ChatCompletions API. + It uses the blocks to generate a prompt for the ChatCompletions API. + Finally, it uses the ChatCompletions API to generate an answer to the user query. + + Args: + query (str): The user query. + previous_dialogue (List[Dict[str, str]]): The previous dialogue. Defaults to []. + k (str): The number of blocks to use as context. + mode (str): The mode to use for the ChatCompletions API. Defaults to "standard". + HyDE (bool): Whether to use the HyDE technique for semantic search. This makes search slower, but better. Defaults to False. + stream (bool): Whether to stream the results from the ChatCompletions API. Defaults to True. + stream_delay (float): The delay between each word in the streamed response when streaming a hard-coded response. Defaults to 0.1. + + Returns: + str: The answer to the user query. + + Raises: + Exception: If the query is offensive. + """ + + + # 1. Check if the query is offensive + flagged_categories: List[str] = moderate_query(query) + if len(flagged_categories) > 0: + return f"Your query contains offensive language. Please try again." + + # 2. Find the top-k most relevant blocks from the Alignment Research Dataset + top_k_blocks: List[Block] = get_top_k_blocks(query, k, HyDE) + + # 3. Generate a prompt for the ChatCompletions API + prompt: List[Dict[str, str]] = generate_prompt(query, history, top_k_blocks, mode) + + # 4. Use the top-k most relevant blocks as context for the ChatCompletions API, and generate an answer to the user query + completion: str = normal_completion(prompt) + return completion, top_k_blocks \ No newline at end of file diff --git a/web/api/get_blocks.py b/web/api/get_blocks.py new file mode 100644 index 0000000..bceb357 --- /dev/null +++ b/web/api/get_blocks.py @@ -0,0 +1,114 @@ +import time +import numpy as np +import json +from typing import List +import openai + +EMBEDDING_MODEL = "text-embedding-ada-002" +COMPLETIONS_MODEL = "gpt-3.5-turbo" + +import pathlib +project_path = pathlib.Path(__file__).parent +PATH_TO_DATASET_JSON = project_path / "data" / "dataset.json" # Path to the saved dataset (.json) file, containing the dataset class object. + +class Dataset: + def __init__(self, path_to_dataset: str = PATH_TO_DATASET_JSON): + self.path_to_dataset = path_to_dataset # .json + self.load_dataset() + + def load_dataset(self): # Load the dataset from the saved .json file + with open(self.path_to_dataset, 'rb') as f: + dataset_dict = json.load(f) + self.metadata = dataset_dict['metadata'] + self.embedding_strings = dataset_dict['embedding_strings'] + self.embeddings_metadata_index = dataset_dict['embeddings_metadata_index'] + self.articles_count = dataset_dict['articles_count'] + self.total_articles_count = dataset_dict['total_articles_count'] + self.total_char_count = dataset_dict['total_char_count'] + self.total_word_count = dataset_dict['total_word_count'] + self.total_sentence_count = dataset_dict['total_sentence_count'] + self.total_block_count = dataset_dict['total_block_count'] + self.sources_so_far = dataset_dict['sources_so_far'] + self.info_types = dataset_dict['info_types'] + self.embeddings = np.array(dataset_dict['embeddings']) + +class Block: + def __init__(self, title: str, author: str, date: str, url: str, tags: str, text: str): + self.title = title + self.author = author + self.date = date + self.url = url + self.tags = tags + self.text = text + +def get_embedding(text: str) -> np.ndarray: + """Get the embedding for a given text. The function will retry with exponential backoff if the API rate limit is reached, up to 4 times. + + Args: + text (str): The text to get the embedding for. + + Returns: + np.ndarray: The embedding for the given text. + """ + max_retries = 4 + max_wait_time = 10 + + for attempt in range(max_retries): + try: + result = openai.Embedding.create( + model=EMBEDDING_MODEL, + input=text + ) + return result["data"][0]["embedding"] + except openai.error.RateLimitError as e: + if attempt + 1 == max_retries: + raise e + wait_time = min(max_wait_time, (2 ** attempt)) # Exponential backoff + time.sleep(wait_time) + +def get_top_k_blocks(user_query: str, k: int = 10, HyDE: bool = False) -> List[Block]: + """Get the top k blocks that are most semantically similar to the query, using the provided dataset. + + Args: + query (str): The query to be searched for. + k (int, optional): The number of blocks to return. + HyDE (bool, optional): Whether to use HyDE or not. Defaults to False. + + Returns: + List[Block]: A list of the top k blocks that are most semantically similar to the query. + """ + # Get the dataset (in data/dataset.json) + metadataset = Dataset() + + # Get the embedding for the query. + query_embedding = get_embedding(user_query) + + # If HyDE is enabled, produce a no-context ChatCompletion to the query. + if HyDE: + messages = [ + {"role": "system", "content": "You are a knowledgeable AI Alignment assistant."}, + {"role": "user", "content": f"Do your best to answer the question/instruction, even if you don't know the correct answer or action for sure.\nQ: {user_query}"}, + ] + HyDE_completion = openai.ChatCompletion.create( + model=COMPLETIONS_MODEL, + messages=messages + )["choices"][0]["message"]["content"] + HyDe_completion_embedding = get_embedding(f"Question: {user_query}\n\nAnswer: {HyDE_completion}") + + similarity_scores = np.dot(metadataset.embeddings, HyDe_completion_embedding) + else: + similarity_scores = np.dot(metadataset.embeddings, query_embedding) + + ordered_blocks = np.argsort(similarity_scores)[::-1] # Sort the blocks by similarity score + top_k_block_indices = ordered_blocks[:k] # Get the top k indices of the blocks + top_k_metadata_indexes = [metadataset.embeddings_metadata_index[i] for i in top_k_block_indices] + + # Get the top k blocks (title, author, date, url, tags, text) + top_k_texts = [metadataset.embedding_strings[i] for i in top_k_block_indices] # Get the top k texts + top_k_metadata = [metadataset.metadata[i] for i in top_k_metadata_indexes] # Get the top k metadata (title, author, date, url, tags) + + # Combine the top k texts and metadata into a list of Block objects + top_k_metadata_and_text = [list(top_k_metadata[i]) + [top_k_texts[i]] for i in range(k)] + blocks = [Block(*block) for block in top_k_metadata_and_text] + + return blocks \ No newline at end of file