From f31f910bd4c81b077f106398dbc9a9e628477a31 Mon Sep 17 00:00:00 2001 From: Joe Martinez Date: Sat, 6 Jul 2024 16:20:18 -0500 Subject: [PATCH] refactored llm into class --- yt_fts/llm.py | 387 +++++++++++++++++++++-------------------------- yt_fts/yt_fts.py | 7 +- 2 files changed, 174 insertions(+), 220 deletions(-) diff --git a/yt_fts/llm.py b/yt_fts/llm.py index 01e34e4..3882d32 100644 --- a/yt_fts/llm.py +++ b/yt_fts/llm.py @@ -1,244 +1,199 @@ +# llm_class.py + import sys import traceback - from openai import OpenAI - from .db_utils import ( get_channel_id_from_input, get_channel_name_from_video_id, get_title_from_db ) - from .get_embeddings import get_embedding from .utils import time_to_secs from .config import get_chroma_client -def init_llm(openai_api_key: str, prompt: str, channel: str) -> None: - openai_client = OpenAI(api_key=openai_api_key) - channel_id = get_channel_id_from_input(channel) - messages = start_llm(openai_client, prompt, channel_id) - - print(messages[len(messages) - 1]["content"]) - while True: - user_input = input("> ") - if user_input == "exit": - sys.exit(0) - messages.append({"role": "user", "content": user_input}) - messages = continue_llm(openai_client, channel_id, messages) - print(messages[len(messages) - 1]["content"]) - - -def start_llm(openai_client: object, prompt: str, channel_id: str) -> list: - - try: - context = create_context( - text=prompt, - channel_id=channel_id, - openai_client=openai_client - ) - - user_str = f"Context: {context}\n\n---\n\nQuestion: {prompt}\nAnswer:" +class LLMHandler: + def __init__(self, openai_api_key: str, channel: str): + self.openai_client = OpenAI(api_key=openai_api_key) + self.channel_id = get_channel_id_from_input(channel) + self.chroma_client = get_chroma_client() - system_prompt = """ - Answer the question based on the context below, The context are - subtitles and timestamped links from videos related to the question. - In your answer, provide the link to the video where the answer can - be found. and if the question can't be answered based on the context, - say \"I don't know\" AND ONLY I don't know\n\n - """ - messages = [ - {"role": "system", "content": system_prompt}, - {"role": "user", "content": user_str}, + def init_llm(self, prompt: str): + messages = self.start_llm(prompt) + print(messages[-1]["content"]) + + while True: + user_input = input("> ") + if user_input == "exit": + sys.exit(0) + messages.append({"role": "user", "content": user_input}) + messages = self.continue_llm(messages) + print(messages[-1]["content"]) + + def start_llm(self, prompt: str) -> list: + try: + context = self.create_context(prompt) + user_str = f"Context: {context}\n\n---\n\nQuestion: {prompt}\nAnswer:" + system_prompt = """ + Answer the question based on the context below, The context are + subtitles and timestamped links from videos related to the question. + In your answer, provide the link to the video where the answer can + be found. and if the question can't be answered based on the context, + say \"I don't know\" AND ONLY I don't know\n\n + """ + messages = [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": user_str}, ] - - response_text = get_completion(openai_client, messages) - - if "i don't know" in response_text.lower(): - expanded_query = get_expand_context_query(messages, openai_client) - print(f"Expanded query: {expanded_query}") - expanded_context = create_context(expanded_query, channel_id, openai_client) - - messages.append({ - "role": "user", - "content": f"Okay here is some more context:\n---\n\n{expanded_context}\n\n---" - }) - - response_text = get_completion(openai_client, messages) - + + response_text = self.get_completion(messages) + + if "i don't know" in response_text.lower(): + expanded_query = self.get_expand_context_query(messages) + expanded_context = self.create_context(expanded_query) + messages.append({ + "role": "user", + "content": f"Okay here is some more context:\n---\n\n{expanded_context}\n\n---" + }) + response_text = self.get_completion(messages) + messages.append({ "role": "assistant", "content": response_text }) - return messages - - messages.append({ - "role": "assistant", - "content": response_text - }) - return messages - - except Exception as e: - traceback.print_exc() - print(f"Error: {e}") - sys.exit(1) - - -def continue_llm(openai_client: object, channel_id: str, messages: list) -> list: - - try: - response_text = get_completion(openai_client, messages) - - if "i don't know" in response_text.lower(): - expanded_query = get_expand_context_query(messages, openai_client) - print(f"Expanded query: {expanded_query}") - expanded_context = create_context(expanded_query, channel_id, openai_client) - messages.append({ - "role": "user", - "content": f"Okay here is some more context:\n---\n\n{expanded_context}\n\n---" - }) - - response_text = get_completion(openai_client, messages) - messages.append({ - "role": "assistant", - "content": response_text - }) - return messages - else: + + except Exception as e: + traceback.print_exc() + print(f"Error: {e}") + sys.exit(1) + + def continue_llm(self, messages: list) -> list: + try: + response_text = self.get_completion(messages) + + if "i don't know" in response_text.lower(): + expanded_query = self.get_expand_context_query(messages) + print(f"Expanded query: {expanded_query}") + expanded_context = self.create_context(expanded_query) + messages.append({ + "role": "user", + "content": f"Okay here is some more context:\n---\n\n{expanded_context}\n\n---" + }) + response_text = self.get_completion(messages) + messages.append({ "role": "assistant", "content": response_text }) return messages - - except Exception as e: - traceback.print_exc() - print(f"Error: {e}") - sys.exit(1) - - - -def create_context(text: str, channel_id: str, openai_client: object) -> str: - - chroma_client = get_chroma_client() - collection = chroma_client.get_collection(name="subEmbeddings") - - search_embedding = get_embedding(text, "text-embedding-ada-002", openai_client) - - scope_options = {} - - - scope_options = {"channel_id": channel_id} - - chroma_res = collection.query( - query_embeddings=[search_embedding], - n_results=10, - where=scope_options, - ) - - documents = chroma_res["documents"][0] - metadata = chroma_res["metadatas"][0] - distances = chroma_res["distances"][0] - - res = [] - - for i in range(len(documents)): - text = documents[i] - video_id = metadata[i]["video_id"] - start_time = metadata[i]["start_time"] - link = f"https://youtu.be/{video_id}?t={time_to_secs(start_time)}" - channel_name = get_channel_name_from_video_id(video_id) - channel_id = metadata[i]["channel_id"] - title = get_title_from_db(video_id) - - match = { - "distance": distances[i], - "channel_name": channel_name, - "channel_id": channel_id, - "video_title": title, - "subs": text, - "start_time": start_time, - "video_id": video_id, - "link": link, - } - res.append(match) - - formatted_context = format_context(res) - - return formatted_context - - -def get_expand_context_query(messages: list, openai_client) -> str: - try: - system_prompt = """ - Your task is to generate a question to input into a vector search - engine of youtube subitles to find strings that can answer the question - asked in the previous message. - """ - formatted_context = format_message_history_context(messages) - messages = [ - {"role": "system", "content": system_prompt}, - {"role": "user", "content": formatted_context}, + + except Exception as e: + traceback.print_exc() + print(f"Error: {e}") + sys.exit(1) + + def create_context(self, text: str) -> str: + collection = self.chroma_client.get_collection(name="subEmbeddings") + search_embedding = get_embedding(text, "text-embedding-ada-002", self.openai_client) + scope_options = {"channel_id": self.channel_id} + + chroma_res = collection.query( + query_embeddings=[search_embedding], + n_results=10, + where=scope_options, + ) + + documents = chroma_res["documents"][0] + metadata = chroma_res["metadatas"][0] + distances = chroma_res["distances"][0] + + res = [] + for i in range(len(documents)): + text = documents[i] + video_id = metadata[i]["video_id"] + start_time = metadata[i]["start_time"] + link = f"https://youtu.be/{video_id}?t={time_to_secs(start_time)}" + channel_name = get_channel_name_from_video_id(video_id) + channel_id = metadata[i]["channel_id"] + title = get_title_from_db(video_id) + + match = { + "distance": distances[i], + "channel_name": channel_name, + "channel_id": channel_id, + "video_title": title, + "subs": text, + "start_time": start_time, + "video_id": video_id, + "link": link, + } + res.append(match) + + return self.format_context(res) + + def get_expand_context_query(self, messages: list) -> str: + try: + system_prompt = """ + Your task is to generate a question to input into a vector search + engine of youtube subitles to find strings that can answer the question + asked in the previous message. + """ + formatted_context = self.format_message_history_context(messages) + messages = [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": formatted_context}, ] + + return self.get_completion(messages) - expand_query = get_completion(openai_client, messages) - - return expand_query - - - except Exception as e: - traceback.print_exc() - print(e) - sys.exit(1) - - -def get_completion(openai_client, messages: list) -> str: - - try: - response = openai_client.chat.completions.create( - model="gpt-4o", - messages=messages, - temperature=0, - max_tokens=2000, - top_p=1, - frequency_penalty=0, - presence_penalty=0, - stop=None, + except Exception as e: + traceback.print_exc() + print(e) + sys.exit(1) + + def get_completion(self, messages: list) -> str: + try: + response = self.openai_client.chat.completions.create( + model="gpt-4", + messages=messages, + temperature=0, + max_tokens=2000, + top_p=1, + frequency_penalty=0, + presence_penalty=0, + stop=None, ) - - return response.choices[0].message.content - - except Exception as e: - traceback.print_exc() - print(f"Error: {e}") - sys.exit(1) - - -def format_message_history_context(messages: list) -> str: - - formatted_context = "" - for message in messages: - if message["role"] == "system": - continue - role = message["role"] - content = message["content"] - formatted_context += f"{role}: {content}\n" - return formatted_context - - -def format_context(chroma_res: str) -> str: - formatted_context = "" - - for obj in chroma_res: - tmp = f""" - Video Title: {obj["video_title"]} - Text: {obj["subs"]} - Time: {obj["start_time"]} - Similarity: {obj["distance"]} - Link: {obj["link"]} - ------------------------- - """ - formatted_context += tmp - - return formatted_context + return response.choices[0].message.content + + except Exception as e: + traceback.print_exc() + print(f"Error: {e}") + sys.exit(1) + + @staticmethod + def format_message_history_context(messages: list) -> str: + formatted_context = "" + for message in messages: + if message["role"] == "system": + continue + role = message["role"] + content = message["content"] + formatted_context += f"{role}: {content}\n" + return formatted_context + + @staticmethod + def format_context(chroma_res: list) -> str: + formatted_context = "" + for obj in chroma_res: + tmp = f""" + Video Title: {obj["video_title"]} + Text: {obj["subs"]} + Time: {obj["start_time"]} + Similarity: {obj["distance"]} + Link: {obj["link"]} + ------------------------- + """ + formatted_context += tmp + return formatted_context diff --git a/yt_fts/yt_fts.py b/yt_fts/yt_fts.py index fec79a1..7ef4e1b 100644 --- a/yt_fts/yt_fts.py +++ b/yt_fts/yt_fts.py @@ -383,7 +383,7 @@ def embeddings(channel, openai_api_key, interval=10): help="OpenAI API key. If not provided, the script will attempt to read it from" " the OPENAI_API_KEY environment variable.") def llm(prompt, channel, openai_api_key=None): - from yt_fts.llm import init_llm + from yt_fts.llm import LLMHandler if openai_api_key is None: openai_api_key = os.environ.get("OPENAI_API_KEY") @@ -396,9 +396,8 @@ def llm(prompt, channel, openai_api_key=None): """) sys.exit(1) - init_llm(openai_api_key=openai_api_key, - prompt=prompt, - channel=channel) + llm_handler = LLMHandler(openai_api_key, channel) + llm_handler.init_llm(prompt) sys.exit(0)