Initial compatibility tested for use with Gemini and AWS Bedrock APIs

.dockerignore

@@ -0,0 +1,12 @@
+*.pyc
+*.ipynb
+*.pdf
+*.spec
+*.toc
+*.csv
+*.bin
+bootstrapper.py
+build/*
+dist/*
+test/*
+config/*

.gitignore

@@ -8,4 +8,5 @@
 bootstrapper.py
 build/*
 dist/*
-test/*
+test/*
+config/*

app.py

@@ -1,44 +1,35 @@
-# Load in packages
-
 import os
-import socket
-
 from typing import Type
-from langchain_huggingface.embeddings import HuggingFaceEmbeddings#, HuggingFaceInstructEmbeddings
+from langchain_huggingface.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import FAISS
 import gradio as gr
 import pandas as pd
 
-from transformers import AutoTokenizer
-import torch
-
-from llama_cpp import Llama
-from huggingface_hub import hf_hub_download
 from chatfuncs.ingest import embed_faiss_save_to_zip
-from chatfuncs.helper_functions import get_or_create_env_var
 
-from chatfuncs.helper_functions import ensure_output_folder_exists, get_connection_params, output_folder, get_or_create_env_var, reveal_feedback_buttons, wipe_logs
+from chatfuncs.helper_functions import ensure_output_folder_exists, get_connection_params, output_folder, reveal_feedback_buttons, wipe_logs
 from chatfuncs.aws_functions import upload_file_to_s3
-#from chatfuncs.llm_api_call import llm_query
 from chatfuncs.auth import authenticate_user
+from chatfuncs.config import FEEDBACK_LOGS_FOLDER, ACCESS_LOGS_FOLDER, USAGE_LOGS_FOLDER, HOST_NAME, COGNITO_AUTH
+
+from llama_cpp import Llama
+from huggingface_hub import hf_hub_download
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM,  AutoModelForCausalLM
+import os
 
 PandasDataFrame = Type[pd.DataFrame]
 
 from datetime import datetime
 today_rev = datetime.now().strftime("%Y%m%d")
 
-ensure_output_folder_exists()
-
-host_name = socket.gethostname()
-
-access_logs_data_folder = 'logs/' + today_rev + '/' + host_name + '/'
-feedback_data_folder = 'feedback/' + today_rev + '/' + host_name + '/'
-usage_data_folder = 'usage/' + today_rev + '/' + host_name + '/'
+host_name = HOST_NAME
+access_logs_data_folder = ACCESS_LOGS_FOLDER
+feedback_data_folder = FEEDBACK_LOGS_FOLDER
+usage_data_folder = USAGE_LOGS_FOLDER
 
 # Disable cuda devices if necessary
 #os.environ['CUDA_VISIBLE_DEVICES'] = '-1' 
 
-#from chatfuncs.chatfuncs import *
 import chatfuncs.ingest as ing
 
 ###
@@ -73,21 +64,44 @@ def get_faiss_store(faiss_vstore_folder,embeddings):
     return vectorstore
 
 import chatfuncs.chatfuncs as chatf
+from chatfuncs.model_load import torch_device, gpu_config, cpu_config, context_length
 
 chatf.embeddings = load_embeddings(embeddings_name)
 chatf.vectorstore = get_faiss_store(faiss_vstore_folder="faiss_embedding",embeddings=globals()["embeddings"])
 
+def docs_to_faiss_save(docs_out:PandasDataFrame, embeddings=embeddings):
 
-def load_model(model_type, gpu_layers, gpu_config=None, cpu_config=None, torch_device=None):
-    print("Loading model")
+    print(f"> Total split documents: {len(docs_out)}")
+
+    print(docs_out)
+
+    vectorstore_func = FAISS.from_documents(documents=docs_out, embedding=embeddings)
+
+    chatf.vectorstore = vectorstore_func
+
+    out_message = "Document processing complete"
+
+    return out_message, vectorstore_func
+ # Gradio chat
+
+def create_hf_model(model_name:str):
+    if torch_device == "cuda":
+        if "flan" in model_name:
+            model = AutoModelForSeq2SeqLM.from_pretrained(model_name, device_map="auto")#, torch_dtype=torch.float16)
+        else:
+            model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")#, torch_dtype=torch.float16)
+    else:
+        if "flan" in model_name:
+            model = AutoModelForSeq2SeqLM.from_pretrained(model_name)#, torch_dtype=torch.float16)
+        else: 
+            model = AutoModelForCausalLM.from_pretrained(model_name)#, trust_remote_code=True)#, torch_dtype=torch.float16)
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name, model_max_length = context_length)
 
-    # Default values inside the function
-    if gpu_config is None:
-        gpu_config = chatf.gpu_config
-    if cpu_config is None:
-        cpu_config = chatf.cpu_config
-    if torch_device is None:
-        torch_device = chatf.torch_device
+    return model, tokenizer
+
+def load_model(model_type:str, gpu_layers:int, gpu_config:dict=gpu_config, cpu_config:dict=cpu_config, torch_device:str=torch_device):
+    print("Loading model")
 
     if model_type == "Phi 3.5 Mini (larger, slow)":
         if torch_device == "cuda":
@@ -112,8 +126,7 @@ def load_model(model_type, gpu_layers, gpu_config=None, cpu_config=None, torch_d
         )
         
         except Exception as e:
-            print("GPU load failed")
-            print(e)
+            print("GPU load failed", e)
             model = Llama(
             model_path=hf_hub_download(
             repo_id=os.environ.get("REPO_ID", "QuantFactory/Phi-3.5-mini-instruct-GGUF"), #"QuantFactory/Phi-3-mini-128k-instruct-GGUF"), #, "microsoft/Phi-3-mini-4k-instruct-gguf"),#"QuantFactory/Meta-Llama-3-8B-Instruct-GGUF-v2"), #"microsoft/Phi-3-mini-4k-instruct-gguf"),#"TheBloke/Mistral-7B-OpenOrca-GGUF"),
@@ -128,57 +141,21 @@ def load_model(model_type, gpu_layers, gpu_config=None, cpu_config=None, torch_d
         # Huggingface chat model
         hf_checkpoint = 'Qwen/Qwen2-0.5B-Instruct'# 'declare-lab/flan-alpaca-large'#'declare-lab/flan-alpaca-base' # # # 'Qwen/Qwen1.5-0.5B-Chat' #
         
-        def create_hf_model(model_name):
-
-            from transformers import AutoModelForSeq2SeqLM,  AutoModelForCausalLM
-            
-            if torch_device == "cuda":
-                if "flan" in model_name:
-                    model = AutoModelForSeq2SeqLM.from_pretrained(model_name, device_map="auto")#, torch_dtype=torch.float16)
-                else:
-                    model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")#, torch_dtype=torch.float16)
-            else:
-                if "flan" in model_name:
-                    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)#, torch_dtype=torch.float16)
-                else: 
-                    model = AutoModelForCausalLM.from_pretrained(model_name)#, trust_remote_code=True)#, torch_dtype=torch.float16)
-
-            tokenizer = AutoTokenizer.from_pretrained(model_name, model_max_length = chatf.context_length)
+        model, tokenizer = create_hf_model(model_name = hf_checkpoint)
 
-            return model, tokenizer, model_type
-
-        model, tokenizer, model_type = create_hf_model(model_name = hf_checkpoint)
+    else:
+        model = model_type
+        tokenizer = ""
 
-    chatf.model = model
+    chatf.model_object = model
     chatf.tokenizer = tokenizer
     chatf.model_type = model_type
 
     load_confirmation = "Finished loading model: " + model_type
 
     print(load_confirmation)
-    return model_type, load_confirmation, model_type
 
-# Both models are loaded on app initialisation so that users don't have to wait for the models to be downloaded
-#model_type = "Phi 3.5 Mini (larger, slow)"
-#load_model(model_type, chatf.gpu_layers, chatf.gpu_config, chatf.cpu_config, chatf.torch_device)
-
-model_type = "Qwen 2 0.5B (small, fast)"
-load_model(model_type, 0, chatf.gpu_config, chatf.cpu_config, chatf.torch_device)
-
-def docs_to_faiss_save(docs_out:PandasDataFrame, embeddings=embeddings):
-
-    print(f"> Total split documents: {len(docs_out)}")
-
-    print(docs_out)
-
-    vectorstore_func = FAISS.from_documents(documents=docs_out, embedding=embeddings)
-
-    chatf.vectorstore = vectorstore_func
-
-    out_message = "Document processing complete"
-
-    return out_message, vectorstore_func
- # Gradio chat
+    return model_type, load_confirmation, model_type#model, tokenizer, model_type
 
 
 ###
@@ -188,17 +165,29 @@ def docs_to_faiss_save(docs_out:PandasDataFrame, embeddings=embeddings):
 app = gr.Blocks(theme = gr.themes.Base(), fill_width=True)#css=".gradio-container {background-color: black}")
 
 with app:
+    model_type = "Qwen 2 0.5B (small, fast)"
+    load_model(model_type, 0, gpu_config, cpu_config, torch_device) # chatf.model_object, chatf.tokenizer, chatf.model_type = 
+
+    print("chatf.model_object:", chatf.model_object)
+
+    # Both models are loaded on app initialisation so that users don't have to wait for the models to be downloaded
+    #model_type = "Phi 3.5 Mini (larger, slow)"
+    #load_model(model_type, gpu_layers, gpu_config, cpu_config, torch_device)
+
     ingest_text = gr.State()
     ingest_metadata = gr.State()
     ingest_docs = gr.State()
 
     model_type_state = gr.State(model_type)
+    gpu_config_state = gr.State(gpu_config)
+    cpu_config_state = gr.State(cpu_config)
+    torch_device_state = gr.State(torch_device)
     embeddings_state = gr.State(chatf.embeddings)#globals()["embeddings"])
     vectorstore_state = gr.State(chatf.vectorstore)#globals()["vectorstore"]) 
 
     relevant_query_state = gr.Checkbox(value=True, visible=False) 
 
-    model_state = gr.State() # chatf.model (gives error)
+    model_state = gr.State() # chatf.model_object (gives error)
     tokenizer_state = gr.State() # chatf.tokenizer (gives error)
 
     chat_history_state = gr.State()
@@ -222,7 +211,7 @@ with app:
     gr.Markdown("Chat with PDF, web page or (new) csv/Excel documents. The default is a small model (Qwen 2 0.5B), that can only answer specific questions that are answered in the text. It cannot give overall impressions of, or summarise the document. The alternative (Phi 3.5 Mini (larger, slow)), can reason a little better, but is much slower (See Advanced tab).\n\nBy default the Lambeth Borough Plan '[Lambeth 2030 : Our Future, Our Lambeth](https://www.lambeth.gov.uk/better-fairer-lambeth/projects/lambeth-2030-our-future-our-lambeth)' is loaded. If you want to talk about another document or web page, please select from the second tab. If switching topic, please click the 'Clear chat' button.\n\nCaution: This is a public app. Please ensure that the document you upload is not sensitive is any way as other users may see it! Also, please note that LLM chatbots may give incomplete or incorrect information, so please use with care.")
 
     with gr.Accordion(label="Use Gemini or AWS Claude model", open=False, visible=False):
-        api_model_choice = gr.Dropdown(value = "None", choices = ["gemini-1.5-flash-002", "gemini-1.5-pro-002", "anthropic.claude-3-haiku-20240307-v1:0", "anthropic.claude-3-sonnet-20240229-v1:0", "None"], label="LLM model to use", multiselect=False, interactive=True, visible=False)
+        api_model_choice = gr.Dropdown(value = "None", choices = ["gemini-2.0-flash-001", "gemini-2.5-flash-preview-04-17", "gemini-2.5-pro-preview-03-25", "anthropic.claude-3-haiku-20240307-v1:0", "anthropic.claude-3-sonnet-20240229-v1:0", "None"], label="LLM model to use", multiselect=False, interactive=True, visible=False)
         in_api_key = gr.Textbox(value = "", label="Enter Gemini API key (only if using Google API models)", lines=1, type="password",interactive=True, visible=False)
 
     with gr.Row():
@@ -233,7 +222,7 @@ with app:
 
         with gr.Row():
             #chat_height = 500
-            chatbot = gr.Chatbot(avatar_images=('user.jfif', 'bot.jpg'), scale = 1, resizable=True, show_copy_all_button=True, show_copy_button=True, show_share_button=True, type='tuples') # , height=chat_height
+            chatbot = gr.Chatbot(value=None, avatar_images=('user.jfif', 'bot.jpg'), scale = 1, resizable=True, show_copy_all_button=True, show_copy_button=True, show_share_button=True, type='messages') # , height=chat_height
             with gr.Accordion("Open this tab to see the source paragraphs used to generate the answer", open = True):
                 sources = gr.HTML(value = "Source paragraphs with the most relevant text will appear here") # , height=chat_height
 
@@ -281,13 +270,12 @@ with app:
         out_passages = gr.Slider(minimum=1, value = 2, maximum=10, step=1, label="Choose number of passages to retrieve from the document. Numbers greater than 2 may lead to increased hallucinations or input text being truncated.")
         temp_slide = gr.Slider(minimum=0.1, value = 0.5, maximum=1, step=0.1, label="Choose temperature setting for response generation.")
         with gr.Row():
-            model_choice = gr.Radio(label="Choose a chat model", value="Qwen 2 0.5B (small, fast)", choices = ["Qwen 2 0.5B (small, fast)", "Phi 3.5 Mini (larger, slow)"])
+            model_choice = gr.Radio(label="Choose a chat model", value="Qwen 2 0.5B (small, fast)", choices = ["Qwen 2 0.5B (small, fast)", "Phi 3.5 Mini (larger, slow)", "gemini-2.0-flash-001", "gemini-2.5-flash-preview-04-17", "gemini-2.5-pro-preview-03-25", "anthropic.claude-3-haiku-20240307-v1:0", "anthropic.claude-3-sonnet-20240229-v1:0"])
             change_model_button = gr.Button(value="Load model", scale=0)
         with gr.Accordion("Choose number of model layers to send to GPU (WARNING: please don't modify unless you are sure you have a GPU).", open = False):
             gpu_layer_choice = gr.Slider(label="Choose number of model layers to send to GPU.", value=0, minimum=0, maximum=100, step = 1, visible=True)
             
-        load_text = gr.Text(label="Load status")
-        
+        load_text = gr.Text(label="Load status")        
 
     gr.HTML(
         "<center>This app is based on the models Qwen 2 0.5B and Phi 3.5 Mini. It powered by Gradio, Transformers, and Llama.cpp.</a></center>"
@@ -295,11 +283,41 @@ with app:
 
     examples_set.change(fn=chatf.update_message, inputs=[examples_set], outputs=[message])
 
-    change_model_button.click(fn=chatf.turn_off_interactivity, inputs=[message, chatbot], outputs=[message, chatbot], queue=False).\
-    success(fn=load_model, inputs=[model_choice, gpu_layer_choice], outputs = [model_type_state, load_text, current_model]).\
-    success(lambda: chatf.restore_interactivity(), None, [message], queue=False).\
-    success(chatf.clear_chat, inputs=[chat_history_state, sources, message, current_topic], outputs=[chat_history_state, sources, message, current_topic]).\
-    success(lambda: None, None, chatbot, queue=False)
+
+    ###
+    # CHAT PAGE
+    ###
+
+    # Click to send message
+    response_click = submit.click(chatf.create_full_prompt, inputs=[message, chat_history_state, current_topic, vectorstore_state, embeddings_state, model_type_state, out_passages, api_model_choice, in_api_key], outputs=[chat_history_state, sources, instruction_prompt_out, relevant_query_state], queue=False, api_name="retrieval").\
+                success(chatf.turn_off_interactivity, inputs=None, outputs=[message, submit], queue=False).\
+                success(chatf.produce_streaming_answer_chatbot, inputs=[chatbot, instruction_prompt_out, model_type_state, temp_slide, relevant_query_state, chat_history_state], outputs=chatbot)
+    response_click.success(chatf.highlight_found_text, [chatbot, sources], [sources]).\
+                success(chatf.add_inputs_answer_to_history,[message, chatbot, current_topic], [chat_history_state, current_topic]).\
+                success(lambda: chatf.restore_interactivity(), None, [message, submit], queue=False)
+
+    # Press enter to send message
+    response_enter = message.submit(chatf.create_full_prompt, inputs=[message, chat_history_state, current_topic, vectorstore_state, embeddings_state, model_type_state, out_passages, api_model_choice, in_api_key], outputs=[chat_history_state, sources, instruction_prompt_out, relevant_query_state], queue=False).\
+                success(chatf.turn_off_interactivity, inputs=None, outputs=[message, submit], queue=False).\
+                success(chatf.produce_streaming_answer_chatbot, [chatbot, instruction_prompt_out, model_type_state, temp_slide, relevant_query_state, chat_history_state], chatbot)
+    response_enter.success(chatf.highlight_found_text, [chatbot, sources], [sources]).\
+                success(chatf.add_inputs_answer_to_history,[message, chatbot, current_topic], [chat_history_state, current_topic]).\
+                success(lambda: chatf.restore_interactivity(), None, [message, submit], queue=False)
+    
+    # Stop box
+    stop.click(fn=None, inputs=None, outputs=None, cancels=[response_click, response_enter])
+
+    # Clear box
+    clear.click(chatf.clear_chat, inputs=[chat_history_state, sources, message, current_topic], outputs=[chat_history_state, sources, message, current_topic])
+    clear.click(lambda: None, None, chatbot, queue=False)
+
+    # Thumbs up or thumbs down voting function
+    chatbot.like(chatf.vote, [chat_history_state, instruction_prompt_out, model_type_state], None)
+    
+
+    ###
+    # LOAD NEW DATA PAGE
+    ###
 
     # Load in a pdf
     load_pdf_click = load_pdf.click(ing.parse_file, inputs=[in_pdf], outputs=[ingest_text, current_source]).\
@@ -318,38 +336,23 @@ with app:
              success(ing.csv_excel_text_to_docs, inputs=[ingest_text, in_text_column], outputs=[ingest_docs]).\
              success(embed_faiss_save_to_zip, inputs=[ingest_docs], outputs=[ingest_embed_out, vectorstore_state, file_out_box]).\
              success(chatf.hide_block, outputs = [examples_set])
+   
 
-    # Load in a webpage
-
-    # Click/enter to send message action
-    response_click = submit.click(chatf.create_full_prompt, inputs=[message, chat_history_state, current_topic, vectorstore_state, embeddings_state, model_type_state, out_passages, api_model_choice, in_api_key], outputs=[chat_history_state, sources, instruction_prompt_out, relevant_query_state], queue=False, api_name="retrieval").\
-                success(chatf.turn_off_interactivity, inputs=[message, chatbot], outputs=[message, chatbot], queue=False).\
-                success(chatf.produce_streaming_answer_chatbot, inputs=[chatbot, instruction_prompt_out, model_type_state, temp_slide, relevant_query_state], outputs=chatbot)
-    response_click.success(chatf.highlight_found_text, [chatbot, sources], [sources]).\
-                success(chatf.add_inputs_answer_to_history,[message, chatbot, current_topic], [chat_history_state, current_topic]).\
-                success(lambda: chatf.restore_interactivity(), None, [message], queue=False)
-
-    response_enter = message.submit(chatf.create_full_prompt, inputs=[message, chat_history_state, current_topic, vectorstore_state, embeddings_state, model_type_state, out_passages, api_model_choice, in_api_key], outputs=[chat_history_state, sources, instruction_prompt_out, relevant_query_state], queue=False).\
-                success(chatf.turn_off_interactivity, inputs=[message, chatbot], outputs=[message, chatbot], queue=False).\
-                success(chatf.produce_streaming_answer_chatbot, [chatbot, instruction_prompt_out, model_type_state, temp_slide, relevant_query_state], chatbot)    
-    response_enter.success(chatf.highlight_found_text, [chatbot, sources], [sources]).\
-                success(chatf.add_inputs_answer_to_history,[message, chatbot, current_topic], [chat_history_state, current_topic]).\
-                success(lambda: chatf.restore_interactivity(), None, [message], queue=False)
-    
-    # Stop box
-    stop.click(fn=None, inputs=None, outputs=None, cancels=[response_click, response_enter])
-
-    # Clear box
-    clear.click(chatf.clear_chat, inputs=[chat_history_state, sources, message, current_topic], outputs=[chat_history_state, sources, message, current_topic])
-    clear.click(lambda: None, None, chatbot, queue=False)
+    ###
+    # LOAD MODEL PAGE
+    ###
 
-    # Thumbs up or thumbs down voting function
-    chatbot.like(chatf.vote, [chat_history_state, instruction_prompt_out, model_type_state], None)
+    change_model_button.click(fn=chatf.turn_off_interactivity, inputs=None, outputs=[message, submit], queue=False).\
+    success(fn=load_model, inputs=[model_choice, gpu_layer_choice], outputs = [model_type_state, load_text, current_model]).\
+    success(lambda: chatf.restore_interactivity(), None, [message, submit], queue=False).\
+    success(chatf.clear_chat, inputs=[chat_history_state, sources, message, current_topic], outputs=[chat_history_state, sources, message, current_topic]).\
+    success(lambda: None, None, chatbot, queue=False)
 
     ###
     # LOGGING AND ON APP LOAD FUNCTIONS
     ###    
-    app.load(get_connection_params, inputs=None, outputs=[session_hash_state, s3_output_folder_state, session_hash_textbox])
+    app.load(get_connection_params, inputs=None, outputs=[session_hash_state, s3_output_folder_state, session_hash_textbox]).\
+    success(load_model, inputs=[model_type_state, gpu_layer_choice, gpu_config_state, cpu_config_state, torch_device_state], outputs=[model_type_state, load_text, current_model])
 
     # Log usernames and times of access to file (to know who is using the app when running on AWS)
     access_callback = gr.CSVLogger()
@@ -358,10 +361,6 @@ with app:
     session_hash_textbox.change(lambda *args: access_callback.flag(list(args)), [session_hash_textbox], None, preprocess=False).\
     success(fn = upload_file_to_s3, inputs=[access_logs_state, access_s3_logs_loc_state], outputs=[s3_logs_output_textbox])
 
-# Launch the Gradio app
-COGNITO_AUTH = get_or_create_env_var('COGNITO_AUTH', '0')
-print(f'The value of COGNITO_AUTH is {COGNITO_AUTH}')
-
 if __name__ == "__main__":
     if os.environ['COGNITO_AUTH'] == "1":
         app.queue().launch(show_error=True, auth=authenticate_user, max_file_size='50mb')

chatfuncs/auth.py

@@ -1,14 +1,22 @@
-
+#import os
 import boto3
-from chatfuncs.helper_functions import get_or_create_env_var
-
-client_id = get_or_create_env_var('AWS_CLIENT_ID', '') # This client id is borrowed from async gradio app client
-print(f'The value of AWS_CLIENT_ID is {client_id}')
+#import gradio as gr
+import hmac
+import hashlib
+import base64
+from chatfuncs.config import AWS_CLIENT_ID, AWS_CLIENT_SECRET, AWS_USER_POOL_ID, AWS_REGION
 
-user_pool_id = get_or_create_env_var('AWS_USER_POOL_ID', '')
-print(f'The value of AWS_USER_POOL_ID is {user_pool_id}')
+def calculate_secret_hash(client_id:str, client_secret:str, username:str):
+    message = username + client_id
+    dig = hmac.new(
+        str(client_secret).encode('utf-8'),
+        msg=str(message).encode('utf-8'),
+        digestmod=hashlib.sha256
+    ).digest()
+    secret_hash = base64.b64encode(dig).decode()
+    return secret_hash
 
-def authenticate_user(username, password, user_pool_id=user_pool_id, client_id=client_id):
+def authenticate_user(username:str, password:str, user_pool_id:str=AWS_USER_POOL_ID, client_id:str=AWS_CLIENT_ID, client_secret:str=AWS_CLIENT_SECRET):
     """Authenticates a user against an AWS Cognito user pool.
 
     Args:
@@ -16,22 +24,39 @@ def authenticate_user(username, password, user_pool_id=user_pool_id, client_id=c
         client_id (str): The ID of the Cognito user pool client.
         username (str): The username of the user.
         password (str): The password of the user.
+        client_secret (str): The client secret of the app client
 
     Returns:
         bool: True if the user is authenticated, False otherwise.
     """
 
-    client = boto3.client('cognito-idp')  # Cognito Identity Provider client
+    client = boto3.client('cognito-idp', region_name=AWS_REGION)  # Cognito Identity Provider client
+
+    # Compute the secret hash
+    secret_hash = calculate_secret_hash(client_id, client_secret, username)
 
     try:
-        response = client.initiate_auth(
+
+        if client_secret == '':
+            response = client.initiate_auth(
+                AuthFlow='USER_PASSWORD_AUTH',
+                AuthParameters={
+                    'USERNAME': username,
+                    'PASSWORD': password,
+                },
+                ClientId=client_id
+            )
+
+        else:
+            response = client.initiate_auth(
             AuthFlow='USER_PASSWORD_AUTH',
             AuthParameters={
                 'USERNAME': username,
                 'PASSWORD': password,
+                'SECRET_HASH': secret_hash
             },
             ClientId=client_id
-        )
+            )
 
         # If successful, you'll receive an AuthenticationResult in the response
         if response.get('AuthenticationResult'):
@@ -44,5 +69,7 @@ def authenticate_user(username, password, user_pool_id=user_pool_id, client_id=c
     except client.exceptions.UserNotFoundException:
         return False
     except Exception as e:
-        print(f"An error occurred: {e}")
-        return False
+        out_message = f"An error occurred: {e}"
+        print(out_message)
+        raise Exception(out_message)
+        return False 

chatfuncs/chatfuncs.py

@@ -6,11 +6,16 @@ import time
 from itertools import compress
 import pandas as pd
 import numpy as np
+import google.generativeai as ai
+from gradio import Progress
+import boto3
+import json
 
 # Model packages
 import torch.cuda
 from threading import Thread
 from transformers import pipeline, TextIteratorStreamer
+from langchain_huggingface import HuggingFaceEmbeddings
 
 # Alternative model sources
 #from dataclasses import asdict, dataclass
@@ -22,31 +27,37 @@ from langchain_community.retrievers import SVMRetriever
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.docstore.document import Document
 
+from chatfuncs.config import GEMINI_API_KEY, AWS_DEFAULT_REGION
+
+model_object = [] # Define empty list for model functions to run
+tokenizer = [] # Define empty list for model functions to run
+
+from chatfuncs.model_load import temperature, max_new_tokens, sample, repetition_penalty, top_p, top_k, torch_device, CtransGenGenerationConfig, max_tokens
+
+# ResponseObject class for AWS Bedrock calls
+class ResponseObject:
+    def __init__(self, text, usage_metadata):
+        self.text = text
+        self.usage_metadata = usage_metadata
+
+bedrock_runtime = boto3.client('bedrock-runtime', region_name=AWS_DEFAULT_REGION)
+
 # For keyword extraction (not currently used)
 #import nltk
 #nltk.download('wordnet')
 from nltk.corpus import stopwords
 from nltk.tokenize import RegexpTokenizer
 from nltk.stem import WordNetLemmatizer
-#from nltk.stem.snowball import SnowballStemmer
 from keybert import KeyBERT
 
 # For Name Entity Recognition model
 #from span_marker import SpanMarkerModel # Not currently used
 
-
 # For BM25 retrieval
 import bm25s
 import Stemmer
 
-#from gensim.corpora import Dictionary
-#from gensim.models import TfidfModel, OkapiBM25Model
-#from gensim.similarities import SparseMatrixSimilarity
-
-from llama_cpp import Llama
-from huggingface_hub import hf_hub_download
-
-from chatfuncs.prompts import instruction_prompt_template_alpaca, instruction_prompt_mistral_orca, instruction_prompt_phi3, instruction_prompt_llama3, instruction_prompt_qwen
+from chatfuncs.prompts import instruction_prompt_template_alpaca, instruction_prompt_mistral_orca, instruction_prompt_phi3, instruction_prompt_llama3, instruction_prompt_qwen, instruction_prompt_template_orca
 
 import gradio as gr
 
@@ -60,10 +71,8 @@ model_type = None # global variable setup
 
 max_memory_length = 0 # How long should the memory of the conversation last?
 
-full_text = "" # Define dummy source text (full text) just to enable highlight function to load
+source_texts = "" # Define dummy source text (full text) just to enable highlight function to load
 
-model = [] # Define empty list for model functions to run
-tokenizer = [] # Define empty list for model functions to run
 
 ## Highlight text constants
 hlt_chunk_size = 12
@@ -77,84 +86,6 @@ ner_model = []#SpanMarkerModel.from_pretrained("tomaarsen/span-marker-mbert-base
 # Used to pull out keywords from chat history to add to user queries behind the scenes
 kw_model = pipeline("feature-extraction", model="sentence-transformers/all-MiniLM-L6-v2")
 
-# Currently set gpu_layers to 0 even with cuda due to persistent bugs in implementation with cuda
-if torch.cuda.is_available():
-    torch_device = "cuda"
-    gpu_layers = 100
-else: 
-    torch_device =  "cpu"
-    gpu_layers = 0
-
-print("Running on device:", torch_device)
-threads = 8 #torch.get_num_threads()
-print("CPU threads:", threads)
-
-# Qwen 2 0.5B (small, fast) Model parameters
-temperature: float = 0.1
-top_k: int = 3
-top_p: float = 1
-repetition_penalty: float = 1.15
-flan_alpaca_repetition_penalty: float = 1.3
-last_n_tokens: int = 64
-max_new_tokens: int = 1024
-seed: int = 42
-reset: bool = False
-stream: bool = True
-threads: int = threads
-batch_size:int = 256
-context_length:int = 2048
-sample = True
-
-
-class CtransInitConfig_gpu:
-    def __init__(self,
-                 last_n_tokens=last_n_tokens,
-                 seed=seed,
-                 n_threads=threads,
-                 n_batch=batch_size,
-                 n_ctx=4096,
-                 n_gpu_layers=gpu_layers):
-
-        self.last_n_tokens = last_n_tokens
-        self.seed = seed
-        self.n_threads = n_threads
-        self.n_batch = n_batch
-        self.n_ctx = n_ctx
-        self.n_gpu_layers = n_gpu_layers
-        # self.stop: list[str] = field(default_factory=lambda: [stop_string])
-
-    def update_gpu(self, new_value):
-        self.n_gpu_layers = new_value
-
-class CtransInitConfig_cpu(CtransInitConfig_gpu):
-    def __init__(self):
-        super().__init__()
-        self.n_gpu_layers = 0
-
-gpu_config = CtransInitConfig_gpu()
-cpu_config = CtransInitConfig_cpu()
-
-
-class CtransGenGenerationConfig:
-    def __init__(self, temperature=temperature,
-                 top_k=top_k,
-                 top_p=top_p,
-                 repeat_penalty=repetition_penalty,
-                 seed=seed,
-                 stream=stream,
-                 max_tokens=max_new_tokens
-                 ):
-        self.temperature = temperature
-        self.top_k = top_k
-        self.top_p = top_p
-        self.repeat_penalty = repeat_penalty
-        self.seed = seed
-        self.max_tokens=max_tokens
-        self.stream = stream
-
-    def update_temp(self, new_value):
-        self.temperature = new_value
-
 # Vectorstore funcs
 
 def docs_to_faiss_save(docs_out:PandasDataFrame, embeddings=embeddings):
@@ -187,7 +118,7 @@ def docs_to_faiss_save(docs_out:PandasDataFrame, embeddings=embeddings):
 
 # Prompt functions
 
-def base_prompt_templates(model_type = "Qwen 2 0.5B (small, fast)"):    
+def base_prompt_templates(model_type:str = "Qwen 2 0.5B (small, fast)"):    
   
     #EXAMPLE_PROMPT = PromptTemplate(
     #    template="\nCONTENT:\n\n{page_content}\n\nSOURCE: {source}\n\n",
@@ -205,20 +136,20 @@ def base_prompt_templates(model_type = "Qwen 2 0.5B (small, fast)"):
         INSTRUCTION_PROMPT=PromptTemplate(template=instruction_prompt_qwen, input_variables=['question', 'summaries'])
     elif model_type == "Phi 3.5 Mini (larger, slow)":
         INSTRUCTION_PROMPT=PromptTemplate(template=instruction_prompt_phi3, input_variables=['question', 'summaries'])
+    else:
+        INSTRUCTION_PROMPT=PromptTemplate(template=instruction_prompt_template_orca, input_variables=['question', 'summaries'])
+        
 
     return INSTRUCTION_PROMPT, CONTENT_PROMPT
 
-def write_out_metadata_as_string(metadata_in):
+def write_out_metadata_as_string(metadata_in:str):
     metadata_string = [f"{'  '.join(f'{k}: {v}' for k, v in d.items() if k != 'page_section')}" for d in metadata_in] # ['metadata']
     return metadata_string
 
-def generate_expanded_prompt(inputs: Dict[str, str], instruction_prompt, content_prompt, extracted_memory, vectorstore, embeddings, relevant_flag = True, out_passages = 2): # , 
+def generate_expanded_prompt(inputs: Dict[str, str], instruction_prompt:str, content_prompt:str, extracted_memory:list, vectorstore:object, embeddings:object, relevant_flag:bool = True, out_passages:int = 2): # , 
         
     question =  inputs["question"]
     chat_history = inputs["chat_history"]
-
-    print("relevant_flag in generate_expanded_prompt:", relevant_flag)
-    
     
     if relevant_flag == True:
         new_question_kworded = adapt_q_from_chat_history(question, chat_history, extracted_memory) # new_question_keywords, 
@@ -234,8 +165,6 @@ def generate_expanded_prompt(inputs: Dict[str, str], instruction_prompt, content
         return sorry_prompt, "No relevant sources found.", new_question_kworded
     
     # Expand the found passages to the neighbouring context
-    print("Doc_df columns:", doc_df.columns)
-
     if 'meta_url' in doc_df.columns:
         file_type = determine_file_type(doc_df['meta_url'][0])
     else:
@@ -265,14 +194,22 @@ def generate_expanded_prompt(inputs: Dict[str, str], instruction_prompt, content
             
     return instruction_prompt_out, sources_docs_content_string, new_question_kworded
 
-def create_full_prompt(user_input, history, extracted_memory, vectorstore, embeddings, model_type, out_passages, api_model_choice=None, api_key=None, relevant_flag = True):
+def create_full_prompt(user_input:str,
+                       history:list[dict],
+                       extracted_memory:str,
+                       vectorstore:object,
+                       embeddings:object,
+                       model_type:str,
+                       out_passages:list[str],
+                       api_model_choice=None,
+                       api_key=None,
+                       relevant_flag = True):
     
     #if chain_agent is None:
     #    history.append((user_input, "Please click the button to submit the Huggingface API key before using the chatbot (top right)"))
     #    return history, history, "", ""
     print("\n==== date/time: " + str(datetime.datetime.now()) + " ====")
-    
-    
+        
     history = history or []
 
     if api_model_choice and api_model_choice != "None":
@@ -298,31 +235,13 @@ def create_full_prompt(user_input, history, extracted_memory, vectorstore, embed
                 generate_expanded_prompt({"question": user_input, "chat_history": history}, #vectorstore,
                                     instruction_prompt, content_prompt, extracted_memory, vectorstore, embeddings, relevant_flag, out_passages)
   
-    history.append(user_input)
+    history.append({"metadata":None, "options":None, "role": 'user', "content": user_input})
     
     print("Output history is:", history)
     print("Final prompt to model is:",instruction_prompt_out)
         
     return history, docs_content_string, instruction_prompt_out, relevant_flag
 
-# Chat functions
-import boto3
-import json
-from chatfuncs.helper_functions import get_or_create_env_var
-
-# ResponseObject class for AWS Bedrock calls
-class ResponseObject:
-        def __init__(self, text, usage_metadata):
-            self.text = text
-            self.usage_metadata = usage_metadata
-
-max_tokens = 4096
-
-AWS_DEFAULT_REGION = get_or_create_env_var('AWS_DEFAULT_REGION', 'eu-west-2')
-print(f'The value of AWS_DEFAULT_REGION is {AWS_DEFAULT_REGION}')
-
-bedrock_runtime = boto3.client('bedrock-runtime', region_name=AWS_DEFAULT_REGION)
-
 def call_aws_claude(prompt: str, system_prompt: str, temperature: float, max_tokens: int, model_choice: str) -> ResponseObject:
     """
     This function sends a request to AWS Claude with the following parameters:
@@ -351,6 +270,8 @@ def call_aws_claude(prompt: str, system_prompt: str, temperature: float, max_tok
         ],
     }
 
+    print("prompt_config:", prompt_config)
+
     body = json.dumps(prompt_config)
 
     modelId = model_choice
@@ -376,16 +297,173 @@ def call_aws_claude(prompt: str, system_prompt: str, temperature: float, max_tok
     
     return response
 
-def produce_streaming_answer_chatbot(history,
-            full_prompt,
-            model_type,
-            temperature=temperature,
-            relevant_query_bool=True,
-            max_new_tokens=max_new_tokens,
-            sample=sample,
-            repetition_penalty=repetition_penalty,
-            top_p=top_p,
-            top_k=top_k
+def construct_gemini_generative_model(in_api_key: str, temperature: float, model_choice: str, system_prompt: str, max_tokens: int) -> Tuple[object, dict]:
+    """
+    Constructs a GenerativeModel for Gemini API calls.
+
+    Parameters:
+    - in_api_key (str): The API key for authentication.
+    - temperature (float): The temperature parameter for the model, controlling the randomness of the output.
+    - model_choice (str): The choice of model to use for generation.
+    - system_prompt (str): The system prompt to guide the generation.
+    - max_tokens (int): The maximum number of tokens to generate.
+
+    Returns:
+    - Tuple[object, dict]: A tuple containing the constructed GenerativeModel and its configuration.
+    """
+    # Construct a GenerativeModel
+    try:
+        if in_api_key:
+            #print("Getting API key from textbox")
+            api_key = in_api_key
+            ai.configure(api_key=api_key)
+        elif "GOOGLE_API_KEY" in os.environ:
+            #print("Searching for API key in environmental variables")
+            api_key = os.environ["GOOGLE_API_KEY"]
+            ai.configure(api_key=api_key)
+        else:
+            print("No API key foound")
+            raise gr.Error("No API key found.")
+    except Exception as e:
+        print(e)
+    
+    config = ai.GenerationConfig(temperature=temperature, max_output_tokens=max_tokens)
+
+    print("model_choice:", model_choice)
+
+    #model = ai.GenerativeModel.from_cached_content(cached_content=cache, generation_config=config)
+    model = ai.GenerativeModel(model_name=model_choice, system_instruction=system_prompt, generation_config=config)
+    
+    return model, config
+
+# Function to send a request and update history
+def send_request(prompt: str, conversation_history: List[dict], model: object, config: dict, model_choice: str, system_prompt: str, temperature: float, progress=Progress(track_tqdm=True)) -> Tuple[str, List[dict]]:
+    """
+    This function sends a request to a language model with the given prompt, conversation history, model configuration, model choice, system prompt, and temperature.
+    It constructs the full prompt by appending the new user prompt to the conversation history, generates a response from the model, and updates the conversation history with the new prompt and response.
+    If the model choice is specific to AWS Claude, it calls the `call_aws_claude` function; otherwise, it uses the `model.generate_content` method.
+    The function returns the response text and the updated conversation history.
+    """
+    # Constructing the full prompt from the conversation history
+    full_prompt = "Conversation history:\n"
+    
+    for entry in conversation_history:
+        role = entry['role'].capitalize()  # Assuming the history is stored with 'role' and 'content'
+        message = ' '.join(entry['parts'])  # Combining all parts of the message
+        full_prompt += f"{role}: {message}\n"
+    
+    # Adding the new user prompt
+    full_prompt += f"\nUser: {prompt}"
+
+    # Print the full prompt for debugging purposes
+    #print("full_prompt:", full_prompt)
+
+    # Generate the model's response
+    if "gemini" in model_choice:
+        try:
+            response = model.generate_content(contents=full_prompt, generation_config=config)
+        except Exception as e:
+            # If fails, try again after 10 seconds in case there is a throttle limit
+            print(e)
+            try:
+                print("Calling Gemini model")
+                out_message = "API limit hit - waiting 30 seconds to retry."
+                print(out_message)
+                progress(0.5, desc=out_message)
+                time.sleep(30)
+                response = model.generate_content(contents=full_prompt, generation_config=config)
+            except Exception as e:
+                print(e)
+                return "", conversation_history
+    elif "claude" in model_choice:
+        try:
+            print("Calling AWS Claude model")
+            print("prompt:", prompt)
+            print("system_prompt:", system_prompt)
+            response = call_aws_claude(prompt, system_prompt, temperature, max_tokens, model_choice)
+        except Exception as e:
+            # If fails, try again after x seconds in case there is a throttle limit
+            print(e)
+            try:
+                out_message = "API limit hit - waiting 30 seconds to retry."
+                print(out_message)
+                progress(0.5, desc=out_message)
+                time.sleep(30)
+                response = call_aws_claude(prompt, system_prompt, temperature, max_tokens, model_choice)
+            
+            except Exception as e:
+                print(e)
+                return "", conversation_history
+    else:
+        raise Exception("Model not found")
+
+    # Update the conversation history with the new prompt and response
+    conversation_history.append({"metadata":None, "options":None, "role": 'user', 'parts': [prompt]})
+    conversation_history.append({"metadata":None, "options":None, "role": "assistant", 'parts': [response.text]})
+    
+    # Print the updated conversation history
+    #print("conversation_history:", conversation_history)
+    
+    return response, conversation_history
+
+def process_requests(prompts: List[str], system_prompt_with_table: str, conversation_history: List[dict], whole_conversation: List[str], whole_conversation_metadata: List[str], model: object, config: dict, model_choice: str, temperature: float, batch_no:int = 1, master:bool = False) -> Tuple[List[ResponseObject], List[dict], List[str], List[str]]:
+    """
+    Processes a list of prompts by sending them to the model, appending the responses to the conversation history, and updating the whole conversation and metadata.
+
+    Args:
+        prompts (List[str]): A list of prompts to be processed.
+        system_prompt_with_table (str): The system prompt including a table.
+        conversation_history (List[dict]): The history of the conversation.
+        whole_conversation (List[str]): The complete conversation including prompts and responses.
+        whole_conversation_metadata (List[str]): Metadata about the whole conversation.
+        model (object): The model to use for processing the prompts.
+        config (dict): Configuration for the model.
+        model_choice (str): The choice of model to use.        
+        temperature (float): The temperature parameter for the model.
+        batch_no (int): Batch number of the large language model request.
+        master (bool): Is this request for the master table.
+
+    Returns:
+        Tuple[List[ResponseObject], List[dict], List[str], List[str]]: A tuple containing the list of responses, the updated conversation history, the updated whole conversation, and the updated whole conversation metadata.
+    """
+    responses = []
+    #for prompt in prompts:
+
+    response, conversation_history = send_request(prompts[0], conversation_history, model=model, config=config, model_choice=model_choice, system_prompt=system_prompt_with_table, temperature=temperature)
+    
+    print(response.text)
+    #"Okay, I'm ready. What source are we discussing, and what's your question about it? Please provide as much context as possible so I can give you the best answer."]
+    print(response.usage_metadata)
+    responses.append(response)
+
+    # Create conversation txt object
+    whole_conversation.append(prompts[0])
+    whole_conversation.append(response.text)
+
+    # Create conversation metadata
+    if master == False:
+        whole_conversation_metadata.append(f"Query batch {batch_no} prompt {len(responses)} metadata:")
+    else:
+        whole_conversation_metadata.append(f"Query summary metadata:")
+
+    whole_conversation_metadata.append(str(response.usage_metadata))
+
+    return responses, conversation_history, whole_conversation, whole_conversation_metadata
+
+def produce_streaming_answer_chatbot(
+            history:list,
+            full_prompt:str,
+            model_type:str,
+            temperature:float=temperature,
+            relevant_query_bool:bool=True,
+            chat_history:list[dict]=[{"metadata":None, "options":None, "role": 'user', "content": ""}],
+            max_new_tokens:int=max_new_tokens,
+            sample:bool=sample,
+            repetition_penalty:float=repetition_penalty,
+            top_p:float=top_p,
+            top_k:float=top_k,
+            max_tokens:int=max_tokens,
+            in_api_key:str=GEMINI_API_KEY
 ):
     #print("Model type is: ", model_type)
 
@@ -395,16 +473,19 @@ def produce_streaming_answer_chatbot(history,
 
     #    return history
 
-    
+    history = chat_history
+
+    print("history at start of streaming function:", history)
 
     if relevant_query_bool == False:
-        out_message = [("","No relevant query found. Please retry your question")]
-        history.append(out_message)
+        history.append({"metadata":None, "options":None, "role": "assistant", "content": 'No relevant query found. Please retry your question'})
 
         yield history
         return
 
     if model_type == "Qwen 2 0.5B (small, fast)": 
+
+        print("tokenizer:", tokenizer)
         # Get the model and tokenizer, and tokenize the user text.
         model_inputs = tokenizer(text=full_prompt, return_tensors="pt", return_attention_mask=False).to(torch_device)
         
@@ -422,9 +503,9 @@ def produce_streaming_answer_chatbot(history,
             top_k=top_k
         )
 
-        #print(generate_kwargs)
+        print("model_object:", model_object)
 
-        t = Thread(target=model.generate, kwargs=generate_kwargs)
+        t = Thread(target=model_object.generate, kwargs=generate_kwargs)
         t.start()
 
         # Pull the generated text from the streamer, and update the model output.
@@ -432,12 +513,14 @@ def produce_streaming_answer_chatbot(history,
         NUM_TOKENS=0
         print('-'*4+'Start Generation'+'-'*4)
 
-        history[-1][1] = ""
+        history.append({"metadata":None, "options":None, "role": "assistant", "content": ""})
+
         for new_text in streamer:
             try:
-                if new_text == None: new_text = ""
-                history[-1][1] += new_text
-                NUM_TOKENS+=1
+                if new_text is None:
+                    new_text = ""
+                history[-1]['content'] += new_text
+                NUM_TOKENS += 1
                 yield history
             except Exception as e:
                 print(f"Error during text generation: {e}")
@@ -463,14 +546,15 @@ def produce_streaming_answer_chatbot(history,
         NUM_TOKENS=0
         print('-'*4+'Start Generation'+'-'*4)
 
-        output = model(
+        output = model_object(
         full_prompt, **vars(gen_config))
 
-        history[-1][1] = ""
+        history.append({"metadata":None, "options":None, "role": "assistant", "content": ""})
+
         for out in output:
 
             if "choices" in out and len(out["choices"]) > 0 and "text" in out["choices"][0]:
-                history[-1][1] += out["choices"][0]["text"]
+                history[-1]['content'] += out["choices"][0]["text"]
                 NUM_TOKENS+=1
                 yield history
             else:
@@ -481,36 +565,75 @@ def produce_streaming_answer_chatbot(history,
         print('-'*4+'End Generation'+'-'*4)
         print(f'Num of generated tokens: {NUM_TOKENS}')
         print(f'Time for complete generation: {time_generate}s')
-        print(f'Tokens per secound: {NUM_TOKENS/time_generate}')
+        print(f'Tokens per second: {NUM_TOKENS/time_generate}')
         print(f'Time per token: {(time_generate/NUM_TOKENS)*1000}ms')
 
-    elif model_type == "anthropic.claude-3-haiku-20240307-v1:0" or model_type == "anthropic.claude-3-sonnet-20240229-v1:0":
+    elif "claude" in model_type:
         system_prompt = "You are answering questions from the user based on source material. Respond with short, factually correct answers."
 
-        try:
-            print("Calling AWS Claude model")
-            response = call_aws_claude(full_prompt, system_prompt, temperature, max_tokens, model_type)
-        except Exception as e:
-            # If fails, try again after 10 seconds in case there is a throttle limit
-            print(e)
-            try:
-                out_message = "API limit hit - waiting 30 seconds to retry."
-                print(out_message)
+        print("full_prompt:", full_prompt)
+
+        if isinstance(full_prompt, str):
+            full_prompt = [full_prompt]
+
+        model = model_type
+        config = {}
+
+        responses, summary_conversation_history, whole_summary_conversation, whole_conversation_metadata = process_requests(full_prompt, system_prompt, conversation_history=[], whole_conversation=[], whole_conversation_metadata=[], model=model, config = config, model_choice = model_type, temperature = temperature)
+        
+        if isinstance(responses[-1], ResponseObject):
+            response_texts = [resp.text for resp in responses]
+        elif "choices" in responses[-1]:
+            response_texts = [resp["choices"][0]['text'] for resp in responses]
+        else:
+            response_texts = [resp.text for resp in responses]
+
+        latest_response_text = response_texts[-1]
 
-                time.sleep(30)
-                response = call_aws_claude(full_prompt, system_prompt, temperature, max_tokens, model_type)
-            
-            except Exception as e:
-                print(e)
-                return "", history
         # Update the conversation history with the new prompt and response
-        history.append({'role': 'user', 'parts': [full_prompt]})
-        history.append({'role': 'assistant', 'parts': [response.text]})
+        clean_text = re.sub(r'[\n\t\r]', ' ', latest_response_text)  # Replace newlines, tabs, and carriage returns with a space
+        clean_response_text = re.sub(r'[^\x20-\x7E]', '', clean_text).strip()  # Remove all non-ASCII printable characters  
         
-        # Print the updated conversation history
-        #print("conversation_history:", conversation_history)
+        history.append({"metadata":None, "options":None, "role": "assistant", "content": ''})
         
-        return response, history
+        for char in clean_response_text:
+            time.sleep(0.005)
+            history[-1]['content'] += char
+            yield history
+    
+    elif "gemini" in model_type:
+        print("Using Gemini model:", model_type)
+        print("full_prompt:", full_prompt)
+
+        if isinstance(full_prompt, str):
+            full_prompt = [full_prompt]
+
+        system_prompt = "You are answering questions from the user based on source material. Respond with short, factually correct answers."
+
+        model, config = construct_gemini_generative_model(GEMINI_API_KEY, temperature, model_type, system_prompt, max_tokens)
+
+        responses, summary_conversation_history, whole_summary_conversation, whole_conversation_metadata = process_requests(full_prompt, system_prompt, conversation_history=[], whole_conversation=[], whole_conversation_metadata=[], model=model, config = config, model_choice = model_type, temperature = temperature)
+
+        if isinstance(responses[-1], ResponseObject):
+            response_texts = [resp.text for resp in responses]
+        elif "choices" in responses[-1]:
+            response_texts = [resp["choices"][0]['text'] for resp in responses]
+        else:
+            response_texts = [resp.text for resp in responses]
+
+        latest_response_text = response_texts[-1]
+
+        clean_text = re.sub(r'[\n\t\r]', ' ', latest_response_text)  # Replace newlines, tabs, and carriage returns with a space
+        clean_response_text = re.sub(r'[^\x20-\x7E]', '', clean_text).strip()  # Remove all non-ASCII printable characters   
+        
+        history.append({"metadata":None, "options":None, "role": "assistant", "content": ''})
+        
+        for char in clean_response_text:
+            time.sleep(0.005)
+            history[-1]['content'] += char
+            yield history
+
+        print("history at end of function:", history)
 
 # Chat helper functions
 
@@ -589,9 +712,6 @@ def hybrid_retrieval(new_question_kworded, vectorstore, embeddings, k_val, out_p
 
             docs = vectorstore.similarity_search_with_score(new_question_kworded, k=k_val)
 
-            print("Docs from similarity search:")
-            print(docs)
-
             # Keep only documents with a certain score
             docs_len = [len(x[0].page_content) for x in docs]
             docs_scores = [x[1] for x in docs]
@@ -688,12 +808,8 @@ def hybrid_retrieval(new_question_kworded, vectorstore, embeddings, k_val, out_p
             # 3rd level check on retrieved docs with SVM retriever
             # Check the type of the embeddings object
             embeddings_type = type(embeddings)
-            print("Type of embeddings object:", embeddings_type)
-
 
-            print("embeddings:", embeddings)
 
-            from langchain_huggingface import HuggingFaceEmbeddings
             #hf_embeddings = HuggingFaceEmbeddings(**embeddings)
             hf_embeddings = embeddings
             
@@ -743,10 +859,6 @@ def hybrid_retrieval(new_question_kworded, vectorstore, embeddings, k_val, out_p
             # Make df of best options
             doc_df = create_doc_df(docs_keep_out)
 
-            print("doc_df:",doc_df)
-            print("docs_keep_as_doc:",docs_keep_as_doc)
-            print("docs_keep_out:", docs_keep_out)
-
             return docs_keep_as_doc, doc_df, docs_keep_out
 
 def get_expanded_passages(vectorstore, docs, width):
@@ -836,16 +948,16 @@ def get_expanded_passages(vectorstore, docs, width):
 
     return expanded_docs, doc_df
 
-def highlight_found_text(search_text: str, full_text: str, hlt_chunk_size:int=hlt_chunk_size, hlt_strat:List=hlt_strat, hlt_overlap:int=hlt_overlap) -> str:
+def highlight_found_text(chat_history: list[dict], source_texts: list[dict], hlt_chunk_size:int=hlt_chunk_size, hlt_strat:List=hlt_strat, hlt_overlap:int=hlt_overlap) -> str:
     """
-    Highlights occurrences of search_text within full_text.
+    Highlights occurrences of chat_history within source_texts.
     
     Parameters:
-    - search_text (str): The text to be searched for within full_text.
-    - full_text (str): The text within which search_text occurrences will be highlighted.
+    - chat_history (str): The text to be searched for within source_texts.
+    - source_texts (str): The text within which chat_history occurrences will be highlighted.
     
     Returns:
-    - str: A string with occurrences of search_text highlighted.
+    - str: A string with occurrences of chat_history highlighted.
     
     Example:
     >>> highlight_found_text("world", "Hello, world! This is a test. Another world awaits.")
@@ -859,32 +971,31 @@ def highlight_found_text(search_text: str, full_text: str, hlt_chunk_size:int=hl
             return text[i][0].replace("  ", " ").strip()
         else:
             return ""
+        
+    print("chat_history:", chat_history)
+        
+    response_text = next(
+    (entry['content'] for entry in reversed(chat_history) if entry.get('role') == 'assistant'),
+    "")
 
-    def extract_search_text_from_input(text):
-        if isinstance(text, str):
-            return text.replace("  ", " ").strip()
-        elif isinstance(text, list):
-            return text[-1][1].replace("  ", " ").strip()
-        else:
-            return ""
-
-    full_text = extract_text_from_input(full_text)
-    search_text = extract_search_text_from_input(search_text)
-
+    print("response_text:", response_text)
+        
+    source_texts = extract_text_from_input(source_texts)
 
+    print("source_texts:", source_texts)
 
     text_splitter = RecursiveCharacterTextSplitter(
         chunk_size=hlt_chunk_size,
         separators=hlt_strat,
         chunk_overlap=hlt_overlap,
     )
-    sections = text_splitter.split_text(search_text)
+    sections = text_splitter.split_text(response_text)
 
     found_positions = {}
     for x in sections:
         text_start_pos = 0
         while text_start_pos != -1:
-            text_start_pos = full_text.find(x, text_start_pos)
+            text_start_pos = source_texts.find(x, text_start_pos)
             if text_start_pos != -1:
                 found_positions[text_start_pos] = text_start_pos + len(x)
                 text_start_pos += 1
@@ -907,20 +1018,24 @@ def highlight_found_text(search_text: str, full_text: str, hlt_chunk_size:int=hl
     prev_end = 0
     for start, end in combined_positions:
         if end-start > 15: # Only combine if there is a significant amount of matched text. Avoids picking up single words like 'and' etc.
-            pos_tokens.append(full_text[prev_end:start])
-            pos_tokens.append('<mark style="color:black;">' + full_text[start:end] + '</mark>')
+            pos_tokens.append(source_texts[prev_end:start])
+            pos_tokens.append('<mark style="color:black;">' + source_texts[start:end] + '</mark>')
             prev_end = end
-    pos_tokens.append(full_text[prev_end:])
+    pos_tokens.append(source_texts[prev_end:])
+
+    out_pos_tokens = "".join(pos_tokens)
 
-    return "".join(pos_tokens)
+    print("out_pos_tokens:", out_pos_tokens)
+
+    return out_pos_tokens
 
 
 # # Chat history functions
 
 def clear_chat(chat_history_state, sources, chat_message, current_topic):
-    chat_history_state = []
+    chat_history_state = None
     sources = ''
-    chat_message = ''
+    chat_message = None
     current_topic = ''
 
     return chat_history_state, sources, chat_message, current_topic
@@ -1011,8 +1126,7 @@ def remove_q_stopwords(question): # Remove stopwords from question. Not used at
     for word in tokens_without_sw:
         if word not in ordered_tokens:
             ordered_tokens.add(word)
-            result.append(word)
-     
+            result.append(word)   
 
 
     new_question_keywords = ' '.join(result)
@@ -1021,9 +1135,6 @@ def remove_q_stopwords(question): # Remove stopwords from question. Not used at
 def remove_q_ner_extractor(question):
     
     predict_out = ner_model.predict(question)
-
-
-
     predict_tokens = [' '.join(v for k, v in d.items() if k == 'span') for d in predict_out]
 
     # Remove duplicate words while preserving order
@@ -1075,11 +1186,11 @@ def keybert_keywords(text, n, kw_model):
     return keywords_list
     
 # Gradio functions
-def turn_off_interactivity(user_message, history):
-        return gr.update(value="", interactive=False), history + [[user_message, None]]
+def turn_off_interactivity():
+        return gr.Textbox(interactive=False), gr.Button(interactive=False)
 
 def restore_interactivity():
-        return gr.update(interactive=True)
+        return gr.Textbox(interactive=True), gr.Button(interactive=True)
 
 def update_message(dropdown_value):
         return gr.Textbox(value=dropdown_value)

chatfuncs/config.py

@@ -0,0 +1,217 @@
+import os
+import tempfile
+import socket
+import logging
+from datetime import datetime
+from dotenv import load_dotenv
+
+today_rev = datetime.now().strftime("%Y%m%d")
+HOST_NAME = socket.gethostname()
+
+# Set or retrieve configuration variables for the redaction app
+
+def get_or_create_env_var(var_name:str, default_value:str, print_val:bool=False):
+    '''
+    Get an environmental variable, and set it to a default value if it doesn't exist
+    '''
+    # Get the environment variable if it exists
+    value = os.environ.get(var_name)
+    
+    # If it doesn't exist, set the environment variable to the default value
+    if value is None:
+        os.environ[var_name] = default_value
+        value = default_value
+
+    if print_val == True:
+        print(f'The value of {var_name} is {value}')
+    
+    return value
+
+def ensure_folder_exists(output_folder:str):
+    """Checks if the specified folder exists, creates it if not."""   
+
+    if not os.path.exists(output_folder):
+        # Create the folder if it doesn't exist
+        os.makedirs(output_folder, exist_ok=True)
+        print(f"Created the {output_folder} folder.")
+    else:
+        print(f"The {output_folder} folder already exists.")
+
+def add_folder_to_path(folder_path: str):
+    '''
+    Check if a folder exists on your system. If so, get the absolute path and then add it to the system Path variable if it doesn't already exist. Function is only relevant for locally-created executable files based on this app (when using pyinstaller it creates a _internal folder that contains tesseract and poppler. These need to be added to the system path to enable the app to run)
+    '''
+
+    if os.path.exists(folder_path) and os.path.isdir(folder_path):
+        print(folder_path, "folder exists.")
+
+        # Resolve relative path to absolute path
+        absolute_path = os.path.abspath(folder_path)
+
+        current_path = os.environ['PATH']
+        if absolute_path not in current_path.split(os.pathsep):
+            full_path_extension = absolute_path + os.pathsep + current_path
+            os.environ['PATH'] = full_path_extension
+            #print(f"Updated PATH with: ", full_path_extension)
+        else:
+            print(f"Directory {folder_path} already exists in PATH.")
+    else:
+        print(f"Folder not found at {folder_path} - not added to PATH")
+
+ensure_folder_exists("config/")
+
+# If you have an aws_config env file in the config folder, you can load in app variables this way, e.g. 'config/app_config.env'
+APP_CONFIG_PATH = get_or_create_env_var('APP_CONFIG_PATH', 'config/app_config.env') # e.g. config/app_config.env
+
+if APP_CONFIG_PATH:
+    if os.path.exists(APP_CONFIG_PATH):
+        print(f"Loading app variables from config file {APP_CONFIG_PATH}")
+        load_dotenv(APP_CONFIG_PATH)
+    else: print("App config file not found at location:", APP_CONFIG_PATH)
+
+# Report logging to console?
+LOGGING = get_or_create_env_var('LOGGING', 'False')
+
+if LOGGING == 'True':
+    # Configure logging
+    logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+
+###
+# AWS CONFIG
+###
+
+# If you have an aws_config env file in the config folder, you can load in AWS keys this way, e.g. 'env/aws_config.env'
+AWS_CONFIG_PATH = get_or_create_env_var('AWS_CONFIG_PATH', '') # e.g. config/aws_config.env
+
+if AWS_CONFIG_PATH:
+    if os.path.exists(AWS_CONFIG_PATH):
+        print(f"Loading AWS variables from config file {AWS_CONFIG_PATH}")
+        load_dotenv(AWS_CONFIG_PATH)
+    else: print("AWS config file not found at location:", AWS_CONFIG_PATH)
+
+RUN_AWS_FUNCTIONS = get_or_create_env_var("RUN_AWS_FUNCTIONS", "0")
+
+AWS_REGION = get_or_create_env_var('AWS_REGION', '')
+
+AWS_DEFAULT_REGION = get_or_create_env_var('AWS_DEFAULT_REGION', '')
+
+AWS_CLIENT_ID = get_or_create_env_var('AWS_CLIENT_ID', '')
+
+AWS_CLIENT_SECRET = get_or_create_env_var('AWS_CLIENT_SECRET', '')
+
+AWS_USER_POOL_ID = get_or_create_env_var('AWS_USER_POOL_ID', '')
+
+AWS_ACCESS_KEY = get_or_create_env_var('AWS_ACCESS_KEY', '')
+if AWS_ACCESS_KEY: print(f'AWS_ACCESS_KEY found in environment variables')
+
+AWS_SECRET_KEY = get_or_create_env_var('AWS_SECRET_KEY', '')
+if AWS_SECRET_KEY: print(f'AWS_SECRET_KEY found in environment variables')
+
+DOCUMENT_REDACTION_BUCKET = get_or_create_env_var('DOCUMENT_REDACTION_BUCKET', '')
+
+# Custom headers e.g. if routing traffic through Cloudfront
+# Retrieving or setting CUSTOM_HEADER
+CUSTOM_HEADER = get_or_create_env_var('CUSTOM_HEADER', '')
+#if CUSTOM_HEADER: print(f'CUSTOM_HEADER found')
+
+# Retrieving or setting CUSTOM_HEADER_VALUE
+CUSTOM_HEADER_VALUE = get_or_create_env_var('CUSTOM_HEADER_VALUE', '')
+#if CUSTOM_HEADER_VALUE: print(f'CUSTOM_HEADER_VALUE found')
+
+###
+# File I/O config
+###
+SESSION_OUTPUT_FOLDER = get_or_create_env_var('SESSION_OUTPUT_FOLDER', 'False') # i.e. do you want your input and output folders saved within a subfolder based on session hash value within output/input folders 
+
+OUTPUT_FOLDER = get_or_create_env_var('GRADIO_OUTPUT_FOLDER', 'output/') # 'output/'
+INPUT_FOLDER = get_or_create_env_var('GRADIO_INPUT_FOLDER', 'input/') # 'input/'
+
+ensure_folder_exists(OUTPUT_FOLDER)
+ensure_folder_exists(INPUT_FOLDER)
+
+# Allow for files to be saved in a temporary folder for increased security in some instances
+if OUTPUT_FOLDER == "TEMP" or INPUT_FOLDER == "TEMP": 
+    # Create a temporary directory
+    with tempfile.TemporaryDirectory() as temp_dir:
+        print(f'Temporary directory created at: {temp_dir}')
+
+        if OUTPUT_FOLDER == "TEMP": OUTPUT_FOLDER = temp_dir + "/"
+        if INPUT_FOLDER == "TEMP": INPUT_FOLDER = temp_dir + "/"
+
+# By default, logs are put into a subfolder of today's date and the host name of the instance running the app. This is to avoid at all possible the possibility of log files from one instance overwriting the logs of another instance on S3. If running the app on one system always, or just locally, it is not necessary to make the log folders so specific.
+# Another way to address this issue would be to write logs to another type of storage, e.g. database such as dynamodb. I may look into this in future.
+
+USE_LOG_SUBFOLDERS = get_or_create_env_var('USE_LOG_SUBFOLDERS', 'True')
+
+if USE_LOG_SUBFOLDERS == "True":
+    day_log_subfolder = today_rev + '/'
+    host_name_subfolder = HOST_NAME + '/'
+    full_log_subfolder = day_log_subfolder + host_name_subfolder
+else:
+    full_log_subfolder = ""
+
+FEEDBACK_LOGS_FOLDER = get_or_create_env_var('FEEDBACK_LOGS_FOLDER', 'feedback/' + full_log_subfolder)
+ACCESS_LOGS_FOLDER = get_or_create_env_var('ACCESS_LOGS_FOLDER', 'logs/' + full_log_subfolder)
+USAGE_LOGS_FOLDER = get_or_create_env_var('USAGE_LOGS_FOLDER', 'usage/' + full_log_subfolder)
+
+ensure_folder_exists(FEEDBACK_LOGS_FOLDER)
+ensure_folder_exists(ACCESS_LOGS_FOLDER)
+ensure_folder_exists(USAGE_LOGS_FOLDER)
+
+# Should the redacted file name be included in the logs? In some instances, the names of the files themselves could be sensitive, and should not be disclosed beyond the app. So, by default this is false.
+DISPLAY_FILE_NAMES_IN_LOGS = get_or_create_env_var('DISPLAY_FILE_NAMES_IN_LOGS', 'False')
+
+###
+# RUN CONFIG
+GEMINI_API_KEY = get_or_create_env_var('GEMINI_API_KEY', '')
+
+
+# Number of pages to loop through before breaking the function and restarting from the last finished page (not currently activated).
+PAGE_BREAK_VALUE = get_or_create_env_var('PAGE_BREAK_VALUE', '99999')
+
+MAX_TIME_VALUE = get_or_create_env_var('MAX_TIME_VALUE', '999999')
+
+###
+# APP RUN CONFIG
+###
+
+# Get some environment variables and Launch the Gradio app
+COGNITO_AUTH = get_or_create_env_var('COGNITO_AUTH', '0')
+
+RUN_DIRECT_MODE = get_or_create_env_var('RUN_DIRECT_MODE', '0')
+
+MAX_QUEUE_SIZE = int(get_or_create_env_var('MAX_QUEUE_SIZE', '5'))
+
+MAX_FILE_SIZE = get_or_create_env_var('MAX_FILE_SIZE', '250mb')
+
+GRADIO_SERVER_PORT = int(get_or_create_env_var('GRADIO_SERVER_PORT', '7860'))
+
+ROOT_PATH = get_or_create_env_var('ROOT_PATH', '')
+
+DEFAULT_CONCURRENCY_LIMIT = get_or_create_env_var('DEFAULT_CONCURRENCY_LIMIT', '3')
+
+GET_DEFAULT_ALLOW_LIST = get_or_create_env_var('GET_DEFAULT_ALLOW_LIST', 'False')
+
+ALLOW_LIST_PATH = get_or_create_env_var('ALLOW_LIST_PATH', '') # config/default_allow_list.csv
+
+S3_ALLOW_LIST_PATH = get_or_create_env_var('S3_ALLOW_LIST_PATH', '') # default_allow_list.csv # This is a path within the DOCUMENT_REDACTION_BUCKET
+
+if ALLOW_LIST_PATH: OUTPUT_ALLOW_LIST_PATH = ALLOW_LIST_PATH
+else: OUTPUT_ALLOW_LIST_PATH = 'config/default_allow_list.csv'
+
+SHOW_COSTS = get_or_create_env_var('SHOW_COSTS', 'False')
+
+GET_COST_CODES = get_or_create_env_var('GET_COST_CODES', 'False')
+
+DEFAULT_COST_CODE = get_or_create_env_var('DEFAULT_COST_CODE', '')
+
+COST_CODES_PATH = get_or_create_env_var('COST_CODES_PATH', '') # 'config/COST_CENTRES.csv' # file should be a csv file with a single table in it that has two columns with a header. First column should contain cost codes, second column should contain a name or description for the cost code
+
+S3_COST_CODES_PATH = get_or_create_env_var('S3_COST_CODES_PATH', '') # COST_CENTRES.csv # This is a path within the DOCUMENT_REDACTION_BUCKET
+
+if COST_CODES_PATH: OUTPUT_COST_CODES_PATH = COST_CODES_PATH
+else: OUTPUT_COST_CODES_PATH = 'config/COST_CENTRES.csv'
+
+ENFORCE_COST_CODES = get_or_create_env_var('ENFORCE_COST_CODES', 'False') # If you have cost codes listed, is it compulsory to choose one before redacting?
+
+if ENFORCE_COST_CODES == 'True': GET_COST_CODES = 'True'

chatfuncs/model_load.py

@@ -0,0 +1,82 @@
+import torch
+
+# Currently set gpu_layers to 0 even with cuda due to persistent bugs in implementation with cuda
+if torch.cuda.is_available():
+    torch_device = "cuda"
+    gpu_layers = 100
+else: 
+    torch_device =  "cpu"
+    gpu_layers = 0
+
+print("Running on device:", torch_device)
+threads = 8 #torch.get_num_threads()
+print("CPU threads:", threads)
+
+# Qwen 2 0.5B (small, fast) Model parameters
+temperature: float = 0.1
+top_k: int = 3
+top_p: float = 1
+repetition_penalty: float = 1.15
+flan_alpaca_repetition_penalty: float = 1.3
+last_n_tokens: int = 64
+max_new_tokens: int = 1024
+seed: int = 42
+reset: bool = False
+stream: bool = True
+threads: int = threads
+batch_size:int = 256
+context_length:int = 2048
+sample = True
+
+# Bedrock parameters
+max_tokens = 4096
+
+
+class CtransInitConfig_gpu:
+    def __init__(self,
+                 last_n_tokens=last_n_tokens,
+                 seed=seed,
+                 n_threads=threads,
+                 n_batch=batch_size,
+                 n_ctx=max_tokens,
+                 n_gpu_layers=gpu_layers):
+
+        self.last_n_tokens = last_n_tokens
+        self.seed = seed
+        self.n_threads = n_threads
+        self.n_batch = n_batch
+        self.n_ctx = n_ctx
+        self.n_gpu_layers = n_gpu_layers
+        # self.stop: list[str] = field(default_factory=lambda: [stop_string])
+
+    def update_gpu(self, new_value):
+        self.n_gpu_layers = new_value
+
+class CtransInitConfig_cpu(CtransInitConfig_gpu):
+    def __init__(self):
+        super().__init__()
+        self.n_gpu_layers = 0
+
+gpu_config = CtransInitConfig_gpu()
+cpu_config = CtransInitConfig_cpu()
+
+
+class CtransGenGenerationConfig:
+    def __init__(self, temperature=temperature,
+                 top_k=top_k,
+                 top_p=top_p,
+                 repeat_penalty=repetition_penalty,
+                 seed=seed,
+                 stream=stream,
+                 max_tokens=max_new_tokens
+                 ):
+        self.temperature = temperature
+        self.top_k = top_k
+        self.top_p = top_p
+        self.repeat_penalty = repeat_penalty
+        self.seed = seed
+        self.max_tokens=max_tokens
+        self.stream = stream
+
+    def update_temp(self, new_value):
+        self.temperature = new_value

chatfuncs/prompts.py

@@ -23,8 +23,7 @@ QUESTION - {question}
 """
 
 
-instruction_prompt_template_orca = """
-### System:
+instruction_prompt_template_orca = """### System:
 You are an AI assistant that follows instruction extremely well. Help as much as you can.
 ### User:
 Answer the QUESTION with a short response using information from the following CONTENT.
@@ -33,8 +32,7 @@ CONTENT: {summaries}
 
 ### Response:"""
 
-instruction_prompt_template_orca_quote = """
-### System:
+instruction_prompt_template_orca_quote = """### System:
 You are an AI assistant that follows instruction extremely well. Help as much as you can.
 ### User:
 Quote text from the CONTENT to answer the QUESTION below.

requirements.txt

@@ -22,3 +22,4 @@ PyStemmer==2.2.0.3
 scipy==1.15.2
 numpy==1.26.4
 boto3==1.38.0
+python-dotenv==1.1.0

requirements_gpu.txt

@@ -20,4 +20,5 @@ bm25s==0.2.12
 PyStemmer==2.2.0.3
 scipy==1.15.2
 numpy==1.26.4
-boto3==1.38.0
+boto3==1.38.0
+python-dotenv==1.1.0