diff --git a/apps/chatbot-kickstarter/chatbot.py b/apps/chatbot-kickstarter/chatbot.py index 204dc29d..46c49e4b 100644 --- a/apps/chatbot-kickstarter/chatbot.py +++ b/apps/chatbot-kickstarter/chatbot.py @@ -2,24 +2,25 @@ import openai from termcolor import colored import streamlit as st -from database import get_redis_connection,get_redis_results +from database import get_redis_connection, get_redis_results -from config import CHAT_MODEL,COMPLETIONS_MODEL, INDEX_NAME +from config import CHAT_MODEL, COMPLETIONS_MODEL, INDEX_NAME redis_client = get_redis_connection() # A basic class to create a message as a dict for chat class Message: - - def __init__(self,role,content): - + def __init__(self, role,content): self.role = role self.content = content def message(self): - - return {"role": self.role,"content": self.content} + return { + "role": self.role, + "content": self.content + } + # New Assistant class to add a vector database call to its responses class RetrievalAssistant: @@ -28,7 +29,6 @@ class RetrievalAssistant: self.conversation_history = [] def _get_assistant_response(self, prompt): - try: completion = openai.ChatCompletion.create( model=CHAT_MODEL, @@ -36,34 +36,57 @@ class RetrievalAssistant: temperature=0.1 ) - response_message = Message(completion['choices'][0]['message']['role'],completion['choices'][0]['message']['content']) + response_message = Message( + completion['choices'][0]['message']['role'], + completion['choices'][0]['message']['content'] + ) return response_message.message() except Exception as e: - + return f'Request failed with exception {e}' # The function to retrieve Redis search results + def _get_search_results(self,prompt): latest_question = prompt - search_content = get_redis_results(redis_client,latest_question,INDEX_NAME)['result'][0] + search_content = get_redis_results( + redis_client,latest_question, + INDEX_NAME + )['result'][0] + return search_content - def ask_assistant(self, next_user_prompt): [self.conversation_history.append(x) for x in next_user_prompt] assistant_response = self._get_assistant_response(self.conversation_history) # Answer normally unless the trigger sequence is used "searching_for_answers" if 'searching for answers' in assistant_response['content'].lower(): - question_extract = openai.Completion.create(model=COMPLETIONS_MODEL,prompt=f"Extract the user's latest question and the year for that question from this conversation: {self.conversation_history}. Extract it as a sentence stating the Question and Year") + question_extract = openai.Completion.create( + model = COMPLETIONS_MODEL, + prompt=f''' + Extract the user's latest question and the year for that question from this + conversation: {self.conversation_history}. Extract it as a sentence stating the Question and Year" + ''' + ) search_result = self._get_search_results(question_extract['choices'][0]['text']) # We insert an extra system prompt here to give fresh context to the Chatbot on how to use the Redis results # In this instance we add it to the conversation history, but in production it may be better to hide - self.conversation_history.insert(-1,{"role": 'system',"content": f"Answer the user's question using this content: {search_result}. If you cannot answer the question, say 'Sorry, I don't know the answer to this one'"}) + self.conversation_history.insert( + -1,{ + "role": 'system', + "content": f''' + Answer the user's question using this content: {search_result}. + If you cannot answer the question, say 'Sorry, I don't know the answer to this one' + ''' + } + ) - assistant_response = self._get_assistant_response(self.conversation_history) + assistant_response = self._get_assistant_response( + self.conversation_history + ) self.conversation_history.append(assistant_response) return assistant_response @@ -71,14 +94,18 @@ class RetrievalAssistant: self.conversation_history.append(assistant_response) return assistant_response + def pretty_print_conversation_history( + self, + colorize_assistant_replies=True): - def pretty_print_conversation_history(self, colorize_assistant_replies=True): for entry in self.conversation_history: - if entry['role'] == 'system': + if entry['role']=='system': pass else: prefix = entry['role'] content = entry['content'] - output = colored(prefix +':\n' + content, 'green') if colorize_assistant_replies and entry['role'] == 'assistant' else prefix +':\n' + content - #prefix = entry['role'] - print(output) \ No newline at end of file + if colorize_assistant_replies and entry['role'] == 'assistant': + output = colored(f"{prefix}:\n{content}, green") + else: + output = colored(f"{prefix}:\n{content}") + print(output)