Merge pull request #1 from seanpedrick-case/dev

Added Gemini and AWS Bedrock compatibility. Gemma model. Now document redaction QA.

.dockerignore

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

.gitattributes

@@ -0,0 +1 @@
+*.zip filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -8,4 +8,7 @@
 bootstrapper.py
 build/*
 dist/*
-test/*
+test/*
+config/*
+output/*
+input/*

app.py

@@ -1,95 +1,122 @@
-# 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 torch import float16
 from llama_cpp import Llama
 from huggingface_hub import hf_hub_download
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM,  AutoModelForCausalLM
+import zipfile
+
 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 get_connection_params, 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, INPUT_FOLDER, OUTPUT_FOLDER, MAX_QUEUE_SIZE, DEFAULT_CONCURRENCY_LIMIT, MAX_FILE_SIZE, GRADIO_SERVER_PORT, ROOT_PATH, DEFAULT_EMBEDDINGS_LOCATION, EMBEDDINGS_MODEL_NAME, DEFAULT_DATA_SOURCE, HF_TOKEN, LARGE_MODEL_REPO_ID, LARGE_MODEL_GGUF_FILE, LARGE_MODEL_NAME, SMALL_MODEL_NAME, SMALL_MODEL_REPO_ID, DEFAULT_DATA_SOURCE_NAME, DEFAULT_EXAMPLES, DEFAULT_MODEL_CHOICES
+from chatfuncs.model_load import torch_device, gpu_config, cpu_config, context_length
+import chatfuncs.chatfuncs as chatf
+import chatfuncs.ingest as ing
 
 PandasDataFrame = Type[pd.DataFrame]
 
 from datetime import datetime
 today_rev = datetime.now().strftime("%Y%m%d")
 
-ensure_output_folder_exists()
+host_name = HOST_NAME
+access_logs_data_folder = ACCESS_LOGS_FOLDER
+feedback_data_folder = FEEDBACK_LOGS_FOLDER
+usage_data_folder = USAGE_LOGS_FOLDER
 
-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 + '/'
+if isinstance(DEFAULT_EXAMPLES, str): default_examples_set = eval(DEFAULT_EXAMPLES)
+if isinstance(DEFAULT_MODEL_CHOICES, str): default_model_choices = eval(DEFAULT_MODEL_CHOICES)
 
 # Disable cuda devices if necessary
 #os.environ['CUDA_VISIBLE_DEVICES'] = '-1' 
 
-#from chatfuncs.chatfuncs import *
-import chatfuncs.ingest as ing
 
 ###
 # Load preset embeddings, vectorstore, and model
 ###
 
-embeddings_name =  "BAAI/bge-base-en-v1.5" #"mixedbread-ai/mxbai-embed-xsmall-v1" 
+def load_embeddings_model(embeddings_model = EMBEDDINGS_MODEL_NAME):
 
-def load_embeddings(embeddings_name = embeddings_name):
+    embeddings_func = HuggingFaceEmbeddings(model_name=embeddings_model)
 
-    embeddings_func = HuggingFaceEmbeddings(model_name=embeddings_name)
+    #global embeddings
 
-    global embeddings
+    #embeddings = embeddings_func
 
-    embeddings = embeddings_func
+    return embeddings_func
 
-    return embeddings
+def get_faiss_store(faiss_vstore_folder:str, embeddings_model:object):
 
-def get_faiss_store(faiss_vstore_folder,embeddings):
-    import zipfile
     with zipfile.ZipFile(faiss_vstore_folder + '/' + faiss_vstore_folder + '.zip', 'r') as zip_ref:
         zip_ref.extractall(faiss_vstore_folder)
 
-    faiss_vstore = FAISS.load_local(folder_path=faiss_vstore_folder, embeddings=embeddings, allow_dangerous_deserialization=True)
+    faiss_vstore = FAISS.load_local(folder_path=faiss_vstore_folder, embeddings=embeddings_model, allow_dangerous_deserialization=True)
     os.remove(faiss_vstore_folder + "/index.faiss")
     os.remove(faiss_vstore_folder + "/index.pkl")
     
-    global vectorstore
+    #global vectorstore
 
-    vectorstore = faiss_vstore
+    #vectorstore = faiss_vstore
 
-    return vectorstore
+    return faiss_vstore #vectorstore
 
-import chatfuncs.chatfuncs as chatf
+# Load in default embeddings and embeddings model name
+embeddings_model = load_embeddings_model(EMBEDDINGS_MODEL_NAME)
+vectorstore = get_faiss_store(faiss_vstore_folder=DEFAULT_EMBEDDINGS_LOCATION,embeddings_model=embeddings_model)#globals()["embeddings"])
 
-chatf.embeddings = load_embeddings(embeddings_name)
-chatf.vectorstore = get_faiss_store(faiss_vstore_folder="faiss_embedding",embeddings=globals()["embeddings"])
+chatf.embeddings = embeddings_model
+chatf.vectorstore = vectorstore
 
+def docs_to_faiss_save(docs_out:PandasDataFrame, embeddings_model=embeddings_model):
 
-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)
 
-    # 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
+    vectorstore_func = FAISS.from_documents(documents=docs_out, embedding=embeddings_model)
 
-    if model_type == "Phi 3.5 Mini (larger, slow)":
+    chatf.vectorstore = vectorstore_func
+
+    out_message = "Document processing complete"
+
+    return out_message, vectorstore_func
+ 
+
+def create_hf_model(model_name:str, hf_token=HF_TOKEN):
+    if torch_device == "cuda":
+        if "flan" in model_name:
+            model = AutoModelForSeq2SeqLM.from_pretrained(model_name, device_map="auto")#, torch_dtype=torch.float16)
+        else:
+            if hf_token:
+                model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", token=hf_token) # , torch_dtype=float16
+            else:
+                model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto") # , torch_dtype=float16
+    else:
+        if "flan" in model_name:
+            model = AutoModelForSeq2SeqLM.from_pretrained(model_name)#, torch_dtype=torch.float16)
+        else:
+            if hf_token:
+                model = AutoModelForCausalLM.from_pretrained(model_name, token=hf_token) # , torch_dtype=float16
+            else:
+                model = AutoModelForCausalLM.from_pretrained(model_name) # , torch_dtype=float16
+
+    if hf_token:
+        tokenizer = AutoTokenizer.from_pretrained(model_name, model_max_length = context_length, token=hf_token)
+    else:
+        tokenizer = AutoTokenizer.from_pretrained(model_name, model_max_length = context_length)
+
+    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 == LARGE_MODEL_NAME:
         if torch_device == "cuda":
             gpu_config.update_gpu(gpu_layers)
             print("Loading with", gpu_config.n_gpu_layers, "model layers sent to GPU.")
@@ -99,86 +126,46 @@ def load_model(model_type, gpu_layers, gpu_config=None, cpu_config=None, torch_d
 
             print("Loading with", cpu_config.n_gpu_layers, "model layers sent to GPU.")
 
-        print(vars(gpu_config))
-        print(vars(cpu_config))
-
         try:
             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"), # "QuantFactory/Meta-Llama-3-8B-Instruct-GGUF-v2"), #"microsoft/Phi-3-mini-4k-instruct-gguf"),#"TheBloke/Mistral-7B-OpenOrca-GGUF"),
-            filename=os.environ.get("MODEL_FILE", "Phi-3.5-mini-instruct.Q4_K_M.gguf") #"Phi-3-mini-128k-instruct.Q4_K_M.gguf")  #"Meta-Llama-3-8B-Instruct-v2.Q6_K.gguf") #"Phi-3-mini-4k-instruct-q4.gguf")#"mistral-7b-openorca.Q4_K_M.gguf"),
+            repo_id=LARGE_MODEL_REPO_ID,
+            filename=LARGE_MODEL_GGUF_FILE 
         ),
         **vars(gpu_config) # change n_gpu_layers if you have more or less VRAM 
         )
         
         except Exception as e:
-            print("GPU load failed")
-            print(e)
+            print("GPU load failed", e, "loading CPU version instead")
             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"),
-            filename=os.environ.get("MODEL_FILE", "Phi-3.5-mini-instruct.Q4_K_M.gguf"), # "Phi-3-mini-128k-instruct.Q4_K_M.gguf") # , #"Meta-Llama-3-8B-Instruct-v2.Q6_K.gguf") #"Phi-3-mini-4k-instruct-q4.gguf"),#"mistral-7b-openorca.Q4_K_M.gguf"),
+            repo_id=LARGE_MODEL_REPO_ID,
+            filename=LARGE_MODEL_GGUF_FILE
         ),
         **vars(cpu_config)
         )
 
         tokenizer = []
 
-    if model_type == "Qwen 2 0.5B (small, fast)":
+    if model_type == SMALL_MODEL_NAME:
         # 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' #
+        hf_checkpoint = SMALL_MODEL_REPO_ID# 'declare-lab/flan-alpaca-large'#'declare-lab/flan-alpaca-base' # # # 'Qwen/Qwen1.5-0.5B-Chat' #
         
-        def create_hf_model(model_name):
+        model, tokenizer = create_hf_model(model_name = hf_checkpoint)
 
-            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)
-
-            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,24 +175,40 @@ 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 = SMALL_MODEL_NAME
+    load_model(model_type, 0, gpu_config, cpu_config, torch_device) # chatf.model_object, chatf.tokenizer, chatf.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, 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)
-    embeddings_state = gr.State(chatf.embeddings)#globals()["embeddings"])
-    vectorstore_state = gr.State(chatf.vectorstore)#globals()["vectorstore"]) 
+    gpu_config_state = gr.State(gpu_config)
+    cpu_config_state = gr.State(cpu_config)
+    torch_device_state = gr.State(torch_device)
 
+    # Embeddings related vars
+    embeddings_model_object_state = gr.State(embeddings_model)#globals()["embeddings"])
+    vectorstore_state = gr.State(vectorstore)#globals()["vectorstore"]) 
+    default_embeddings_store_text = gr.Textbox(value=DEFAULT_EMBEDDINGS_LOCATION, visible=False)
+
+    # Is the query relevant to the sources provided?
     relevant_query_state = gr.Checkbox(value=True, visible=False) 
 
-    model_state = gr.State() # chatf.model (gives error)
+    # Storing model objects in state doesn't seem to work, so we have to load in different models in roundabout ways
+    model_state = gr.State() # chatf.model_object (gives error)
     tokenizer_state = gr.State() # chatf.tokenizer (gives error)
 
     chat_history_state = gr.State()
     instruction_prompt_out = gr.State()
 
     session_hash_state = gr.State()
-    s3_output_folder_state = gr.State()
+    output_folder_textbox = gr.Textbox(value=OUTPUT_FOLDER, visible=False)
+    input_folder_textbox = gr.Textbox(value=INPUT_FOLDER, visible=False)
 
     session_hash_textbox = gr.Textbox(value="", visible=False)
     s3_logs_output_textbox = gr.Textbox(label="S3 logs", visible=False)
@@ -219,21 +222,18 @@ with app:
 
     gr.Markdown("<h1><center>Lightweight PDF / web page QA bot</center></h1>")        
     
-    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)
-        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)
+    gr.Markdown(f"""Chat with PDF, web page or (new) csv/Excel documents. The default is a small model ({SMALL_MODEL_NAME}), that can only answer specific questions that are answered in the text. It cannot give overall impressions of, or summarise the document. The alternative ({LARGE_MODEL_NAME}), can reason a little better, but is much slower (See Advanced settings tab).\n\nBy default '[{DEFAULT_DATA_SOURCE_NAME}]({DEFAULT_DATA_SOURCE})' 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.Row():
-        current_source = gr.Textbox(label="Current data source(s)", value="Lambeth_2030-Our_Future_Our_Lambeth.pdf", scale = 10)
+        current_source = gr.Textbox(label="Current data source(s)", value=DEFAULT_DATA_SOURCE, scale = 10)
         current_model = gr.Textbox(label="Current model", value=model_type, scale = 3)
 
     with gr.Tab("Chatbot"):
 
         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
 
@@ -245,17 +245,12 @@ with app:
         with gr.Row():
             submit = gr.Button(value="Send message", variant="primary", scale = 4)
             clear = gr.Button(value="Clear chat", variant="secondary", scale=1) 
-            stop = gr.Button(value="Stop generating", variant="secondary", scale=1)
-
-        examples_set = gr.Radio(label="Examples for the Lambeth Borough Plan",
-            #value = "What were the five pillars of the previous borough plan?",
-            choices=["What were the five pillars of the previous borough plan?",
-                "What is the vision statement for Lambeth?",
-                "What are the commitments for Lambeth?",
-                "What are the 2030 outcomes for Lambeth?"])
-        
-        current_topic = gr.Textbox(label="Feature currently disabled - Keywords related to current conversation topic.", placeholder="Keywords related to the conversation topic will appear here", visible=False)      
+            stop = gr.Button(value="Stop generating", variant="stop", scale=1)
 
+        examples_set = gr.Radio(label="Example questions",
+            choices=default_examples_set)
+        
+        current_topic = gr.Textbox(label="Feature currently disabled - Keywords related to current conversation topic.", placeholder="Keywords related to the conversation topic will appear here", visible=False)
 
     with gr.Tab("Load in a different file to chat with"):
         with gr.Accordion("PDF file", open = False):
@@ -281,75 +276,91 @@ 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=SMALL_MODEL_NAME, choices = default_model_choices)
+            in_api_key = gr.Textbox(value = "", label="Enter Gemini API key (only if using Google API models)", lines=1, type="password",interactive=True, visible=True)
             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>"
+        "<center>This app is powered by Gradio, Transformers, and Llama.cpp.</center>"
     )
 
     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_model_object_state, model_type_state, out_passages, 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, in_api_key], 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_model_object_state, model_type_state, out_passages, 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, in_api_key], 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]).\
              success(ing.text_to_docs, inputs=[ingest_text], outputs=[ingest_docs]).\
-             success(embed_faiss_save_to_zip, inputs=[ingest_docs], outputs=[ingest_embed_out, vectorstore_state, file_out_box]).\
+             success(embed_faiss_save_to_zip, inputs=[ingest_docs, output_folder_textbox, embeddings_model_object_state], outputs=[ingest_embed_out, vectorstore_state, file_out_box]).\
              success(chatf.hide_block, outputs = [examples_set])
 
     # Load in a webpage
     load_web_click = load_web.click(ing.parse_html, inputs=[in_web, in_div], outputs=[ingest_text, ingest_metadata, current_source]).\
              success(ing.html_text_to_docs, inputs=[ingest_text, ingest_metadata], outputs=[ingest_docs]).\
-             success(embed_faiss_save_to_zip, inputs=[ingest_docs], outputs=[ingest_embed_out, vectorstore_state, file_out_box]).\
+             success(embed_faiss_save_to_zip, inputs=[ingest_docs, output_folder_textbox, embeddings_model_object_state], outputs=[ingest_embed_out, vectorstore_state, file_out_box]).\
              success(chatf.hide_block, outputs = [examples_set])
     
     # Load in a csv/excel file
     load_csv_click = load_csv.click(ing.parse_csv_or_excel, inputs=[in_csv, in_text_column], outputs=[ingest_text, current_source]).\
              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(embed_faiss_save_to_zip, inputs=[ingest_docs, output_folder_textbox, embeddings_model_object_state], 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])
+    ###  
+    # Load in default model and embeddings for each user  
+    app.load(get_connection_params, inputs=None, outputs=[session_hash_state, output_folder_textbox, session_hash_textbox, input_folder_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]).\
+    success(get_faiss_store, inputs=[default_embeddings_store_text, embeddings_model_object_state], outputs=[vectorstore_state])
 
     # Log usernames and times of access to file (to know who is using the app when running on AWS)
     access_callback = gr.CSVLogger()
@@ -358,12 +369,8 @@ 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')
+    if COGNITO_AUTH == "1":
+        app.queue(max_size=int(MAX_QUEUE_SIZE), default_concurrency_limit=int(DEFAULT_CONCURRENCY_LIMIT)).launch(show_error=True, inbrowser=True, auth=authenticate_user, max_file_size=MAX_FILE_SIZE, server_port=GRADIO_SERVER_PORT, root_path=ROOT_PATH)
     else:
-        app.queue().launch(show_error=True, inbrowser=True, max_file_size='50mb')
+        app.queue(max_size=int(MAX_QUEUE_SIZE), default_concurrency_limit=int(DEFAULT_CONCURRENCY_LIMIT)).launch(show_error=True, inbrowser=True, max_file_size=MAX_FILE_SIZE, server_port=GRADIO_SERVER_PORT, root_path=ROOT_PATH)

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

@@ -5,65 +5,60 @@ from typing import Type, Dict, List, Tuple
 import time
 from itertools import compress
 import pandas as pd
-import numpy as np
-
-# Model packages
-import torch.cuda
-from threading import Thread
-from transformers import pipeline, TextIteratorStreamer
-
-# Alternative model sources
-#from dataclasses import asdict, dataclass
-
-# Langchain functions
-from langchain.prompts import PromptTemplate
-from langchain_community.vectorstores import FAISS
-from langchain_community.retrievers import SVMRetriever 
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.docstore.document import Document
-
-# For keyword extraction (not currently used)
-#import nltk
-#nltk.download('wordnet')
+import google.generativeai as ai
+import gradio as gr
+from gradio import Progress
+import boto3
+import json
 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
+# Model packages
+import torch.cuda
+from threading import Thread
+from transformers import pipeline, TextIteratorStreamer
+# Langchain functions
+from langchain.prompts import PromptTemplate
+from langchain_community.vectorstores import FAISS
+from langchain_community.retrievers import SVMRetriever 
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.docstore.document import Document
 
-#from gensim.corpora import Dictionary
-#from gensim.models import TfidfModel, OkapiBM25Model
-#from gensim.similarities import SparseMatrixSimilarity
+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, instruction_prompt_gemma
+from chatfuncs.model_load import temperature, max_new_tokens, sample, repetition_penalty, top_p, top_k, torch_device, CtransGenGenerationConfig, max_tokens
+from chatfuncs.config import GEMINI_API_KEY, AWS_DEFAULT_REGION, LARGE_MODEL_NAME, SMALL_MODEL_NAME
 
-from llama_cpp import Llama
-from huggingface_hub import hf_hub_download
+model_object = [] # Define empty list for model functions to run
+tokenizer = [] # Define empty list for model functions to run
 
-from chatfuncs.prompts import instruction_prompt_template_alpaca, instruction_prompt_mistral_orca, instruction_prompt_phi3, instruction_prompt_llama3, instruction_prompt_qwen
+# ResponseObject class for AWS Bedrock calls
+class ResponseObject:
+    def __init__(self, text, usage_metadata):
+        self.text = text
+        self.usage_metadata = usage_metadata
 
-import gradio as gr
+bedrock_runtime = boto3.client('bedrock-runtime', region_name=AWS_DEFAULT_REGION)
 
 torch.cuda.empty_cache()
 
 PandasDataFrame = Type[pd.DataFrame]
 
 embeddings = None  # global variable setup
+embeddings_model = None  # global variable setup
 vectorstore = None # global variable setup
 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
-
-model = [] # Define empty list for model functions to run
-tokenizer = [] # Define empty list for model functions to run
+source_texts = "" # Define dummy source text (full text) just to enable highlight function to load
 
 ## Highlight text constants
 hlt_chunk_size = 12
@@ -77,117 +72,53 @@ 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):
+# def docs_to_faiss_save(docs_out:PandasDataFrame, embeddings=embeddings):
 
-    print(f"> Total split documents: {len(docs_out)}")
+#     print(f"> Total split documents: {len(docs_out)}")
 
-    vectorstore_func = FAISS.from_documents(documents=docs_out, embedding=embeddings)
+#     vectorstore_func = FAISS.from_documents(documents=docs_out, embedding=embeddings)
         
-    '''  
-    #with open("vectorstore.pkl", "wb") as f:
-        #pickle.dump(vectorstore, f) 
-    ''' 
+#     '''  
+#     #with open("vectorstore.pkl", "wb") as f:
+#         #pickle.dump(vectorstore, f) 
+#     ''' 
 
-    #if Path(save_to).exists():
-    #    vectorstore_func.save_local(folder_path=save_to)
-    #else:
-    #    os.mkdir(save_to)
-    #    vectorstore_func.save_local(folder_path=save_to)
+#     #if Path(save_to).exists():
+#     #    vectorstore_func.save_local(folder_path=save_to)
+#     #else:
+#     #    os.mkdir(save_to)
+#     #    vectorstore_func.save_local(folder_path=save_to)
 
-    global vectorstore
+#     global vectorstore
 
-    vectorstore = vectorstore_func
+#     vectorstore = vectorstore_func
 
-    out_message = "Document processing complete"
+#     out_message = "Document processing complete"
 
-    #print(out_message)
-    #print(f"> Saved to: {save_to}")
+#     #print(out_message)
+#     #print(f"> Saved to: {save_to}")
 
-    return out_message
+#     return out_message
+
+# def docs_to_faiss_save(docs_out:PandasDataFrame, embeddings_model=embeddings_model):
+
+#     print(f"> Total split documents: {len(docs_out)}")
+
+#     print(docs_out)
+
+#     vectorstore_func = FAISS.from_documents(documents=docs_out, embedding=embeddings_model)
+
+#     vectorstore = vectorstore_func
+
+#     out_message = "Document processing complete"
+
+#     return out_message, vectorstore_func
 
 # Prompt functions
 
-def base_prompt_templates(model_type = "Qwen 2 0.5B (small, fast)"):    
+def base_prompt_templates(model_type:str = SMALL_MODEL_NAME):    
   
     #EXAMPLE_PROMPT = PromptTemplate(
     #    template="\nCONTENT:\n\n{page_content}\n\nSOURCE: {source}\n\n",
@@ -201,24 +132,24 @@ def base_prompt_templates(model_type = "Qwen 2 0.5B (small, fast)"):
 
 # The main prompt:  
 
-    if model_type == "Qwen 2 0.5B (small, fast)":
+    if model_type == SMALL_MODEL_NAME:
         INSTRUCTION_PROMPT=PromptTemplate(template=instruction_prompt_qwen, input_variables=['question', 'summaries'])
-    elif model_type == "Phi 3.5 Mini (larger, slow)":
+    elif model_type == LARGE_MODEL_NAME:
         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, total_output_passage_chunks_size:int=5): # , 
         
     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:
@@ -243,7 +172,7 @@ def generate_expanded_prompt(inputs: Dict[str, str], instruction_prompt, content
 
     # Only expand passages if not tabular data
     if (file_type != ".csv") & (file_type != ".xlsx"):
-        docs_keep_as_doc, doc_df = get_expanded_passages(vectorstore, docs_keep_out, width=3)    
+        docs_keep_as_doc, doc_df = get_expanded_passages(vectorstore, docs_keep_out, width=total_output_passage_chunks_size)    
 
     # Build up sources content to add to user display
     doc_df['meta_clean'] = write_out_metadata_as_string(doc_df["metadata"]) # [f"<b>{'  '.join(f'{k}: {v}' for k, v in d.items() if k != 'page_section')}</b>" for d in doc_df['metadata']]
@@ -259,29 +188,28 @@ def generate_expanded_prompt(inputs: Dict[str, str], instruction_prompt, content
     sources_docs_content_string = '<br><br>'.join(doc_df['content_meta'])#.replace("  "," ")#.strip()
     
     instruction_prompt_out = instruction_prompt.format(question=new_question_kworded, summaries=docs_content_string)
-    
-    print('Final prompt is: ')
-    print(instruction_prompt_out)
             
     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_key:str="",
+                       relevant_flag:bool=True):
+    
+    if "gemini" in model_type and not GEMINI_API_KEY and not api_key:
+        raise Exception("Gemini model selected but no API key found. Please enter an API key on the Advanced settings page.")
     
     #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":
-         print("API model choice detected")
-         if api_key:
-            print("API key detected")
-            return history, "", None, relevant_flag       
-         else:
-            return history, "", None, relevant_flag
          
     # Create instruction prompt
     instruction_prompt, content_prompt = base_prompt_templates(model_type=model_type)
@@ -291,38 +219,15 @@ def create_full_prompt(user_input, history, extracted_memory, vectorstore, embed
         relevant_flag = False
     else:
         relevant_flag = True
-
-    print("User input:", user_input)
    
     instruction_prompt_out, docs_content_string, new_question_kworded =\
                 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)
-    
-    print("Output history is:", history)
-    print("Final prompt to model is:",instruction_prompt_out)
+    history.append({"metadata":None, "options":None, "role": 'user', "content": user_input})
         
     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 +256,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 +283,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": ""}],
+            in_api_key:str=GEMINI_API_KEY,
+            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            
 ):
     #print("Model type is: ", model_type)
 
@@ -395,16 +459,18 @@ 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)": 
+    if model_type == SMALL_MODEL_NAME: 
+
         # 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 +488,7 @@ def produce_streaming_answer_chatbot(history,
             top_k=top_k
         )
 
-        #print(generate_kwargs)
-
-        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 +496,15 @@ 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
+                history[-1]['content'] = history[-1]['content'].replace('<|im_end|>','')
                 yield history
             except Exception as e:
                 print(f"Error during text generation: {e}")
@@ -450,7 +517,7 @@ def produce_streaming_answer_chatbot(history,
         print(f'Tokens per secound: {NUM_TOKENS/time_generate}')
         print(f'Time per token: {(time_generate/NUM_TOKENS)*1000}ms')
 
-    elif model_type == "Phi 3.5 Mini (larger, slow)":
+    elif model_type == LARGE_MODEL_NAME:
         #tokens = model.tokenize(full_prompt)
 
         gen_config = CtransGenGenerationConfig()
@@ -463,15 +530,17 @@ 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
+                history[-1]['content'] = history[-1]['content'].replace('<|im_end|>','')
                 yield history
             else:
                 print(f"Unexpected output structure: {out}") 
@@ -481,36 +550,80 @@ 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  
+        
+        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
+    
+    elif "gemini" in model_type:
+
+        if in_api_key: gemini_api_key = in_api_key
+        elif GEMINI_API_KEY: gemini_api_key = GEMINI_API_KEY
+        else: raise Exception("Gemini API key not found. Please enter a key on the Advanced settings page or select another 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   
         
-        # 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
+
+        print("history at end of function:", history)
 
 # Chat helper functions
 
@@ -589,9 +702,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 +798,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 +849,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 +938,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 +961,27 @@ def highlight_found_text(search_text: str, full_text: str, hlt_chunk_size:int=hl
             return text[i][0].replace("  ", " ").strip()
         else:
             return ""
-
-    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("chat_history:", chat_history)
+        
+    response_text = next(
+    (entry['content'] for entry in reversed(chat_history) if entry.get('role') == 'assistant'),
+    "")
+        
+    source_texts = extract_text_from_input(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 +1004,22 @@ 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)
+    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 +1110,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 +1119,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 +1170,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,250 @@
+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', '')
+
+HF_TOKEN = get_or_create_env_var('HF_TOKEN', '')
+
+
+# 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
+###
+
+SMALL_MODEL_NAME = get_or_create_env_var("SMALL_MODEL_NAME", "Gemma 3 1B (small, fast)") # "Qwen 2 0.5B (small, fast)"
+
+SMALL_MODEL_REPO_ID = get_or_create_env_var("SMALL_MODEL_REPO_ID", 'google/gemma-3-1b-it') #'Qwen/Qwen2-0.5B-Instruct')
+
+LARGE_MODEL_NAME = get_or_create_env_var("LARGE_MODEL_NAME", "Phi 3.5 Mini (larger, slow)")
+
+LARGE_MODEL_REPO_ID = get_or_create_env_var("LARGE_MODEL_REPO_ID", "QuantFactory/Phi-3.5-mini-instruct-GGUF") # "QuantFactory/Phi-3-mini-128k-instruct-GGUF"), # "QuantFactory/Meta-Llama-3-8B-Instruct-GGUF-v2"), #"microsoft/Phi-3-mini-4k-instruct-gguf"),#"TheBloke/Mistral-7B-OpenOrca-GGUF"),
+LARGE_MODEL_GGUF_FILE = get_or_create_env_var("LARGE_MODEL_GGUF_FILE", "Phi-3.5-mini-instruct.Q4_K_M.gguf") #"Phi-3-mini-128k-instruct.Q4_K_M.gguf")  #"Meta-Llama-3-8B-Instruct-v2.Q6_K.gguf") #"Phi-3-mini-4k-instruct-q4.gguf")#"mistral-7b-openorca.Q4_K_M.gguf"),
+
+if RUN_AWS_FUNCTIONS == "1":
+    default_model_choices = f'["{SMALL_MODEL_NAME}", "{LARGE_MODEL_NAME}", "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"]'
+else:
+    default_model_choices = f'["{SMALL_MODEL_NAME}", "{LARGE_MODEL_NAME}", "gemini-2.0-flash-001", "gemini-2.5-flash-preview-04-17", "gemini-2.5-pro-preview-03-25"]'
+
+DEFAULT_MODEL_CHOICES = get_or_create_env_var("DEFAULT_MODEL_CHOICES", default_model_choices)
+
+EMBEDDINGS_MODEL_NAME = get_or_create_env_var('EMBEDDINGS_MODEL_NAME', "BAAI/bge-base-en-v1.5") #"mixedbread-ai/mxbai-embed-xsmall-v1"
+
+DEFAULT_EMBEDDINGS_LOCATION = get_or_create_env_var('DEFAULT_EMBEDDINGS_LOCATION', "faiss_embedding")
+
+DEFAULT_DATA_SOURCE_NAME = get_or_create_env_var('DEFAULT_DATA_SOURCE_NAME', "Document redaction app documentation")
+
+DEFAULT_DATA_SOURCE = get_or_create_env_var('DEFAULT_DATA_SOURCE', "https://seanpedrick-case.github.io/doc_redaction/README.html")
+
+DEFAULT_EXAMPLES = get_or_create_env_var('DEFAULT_EXAMPLES', '[ "How can I make a custom deny list?", "How can I find page duplicates?", "How can I review and modify existing redactions?", "How can I export my review files to Adobe?"]')
+                # 
+                # ') # ["What were the five pillars of the previous borough plan?",
+                #"What is the vision statement for Lambeth?",
+                #"What are the commitments for Lambeth?",
+                #"What are the 2030 outcomes for Lambeth?"]
+
+# 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/helper_functions.py

@@ -1,6 +1,9 @@
 import os
 import gradio as gr
 import pandas as pd
+import boto3
+from botocore.exceptions import ClientError
+from chatfuncs.config import CUSTOM_HEADER_VALUE, CUSTOM_HEADER, OUTPUT_FOLDER, INPUT_FOLDER, SESSION_OUTPUT_FOLDER, AWS_USER_POOL_ID
 
 def get_or_create_env_var(var_name, default_value):
     # Get the environment variable if it exists
@@ -13,12 +16,6 @@ def get_or_create_env_var(var_name, default_value):
     
     return value
 
-# Retrieving or setting output folder
-env_var_name = 'GRADIO_OUTPUT_FOLDER'
-default_value = 'output/'
-
-output_folder = get_or_create_env_var(env_var_name, default_value)
-print(f'The value of {env_var_name} is {output_folder}')
 
 def get_file_path_with_extension(file_path):
     # First, get the basename of the file (e.g., "example.txt" from "/path/to/example.txt")
@@ -165,64 +162,129 @@ def wipe_logs(feedback_logs_loc, usage_logs_loc):
 
     
 
-async def get_connection_params(request: gr.Request):
-    base_folder = ""
-
-    if request:
-        #print("request user:", request.username)
-
-        #request_data = await request.json()  # Parse JSON body
-        #print("All request data:", request_data)
-        #context_value = request_data.get('context') 
-        #if 'context' in request_data:
-        #     print("Request context dictionary:", request_data['context'])
-
-        # print("Request headers dictionary:", request.headers)
-        # print("All host elements", request.client)           
-        # print("IP address:", request.client.host)
-        # print("Query parameters:", dict(request.query_params))
-        # To get the underlying FastAPI items you would need to use await and some fancy @ stuff for a live query: https://fastapi.tiangolo.com/vi/reference/request/
-        #print("Request dictionary to object:", request.request.body())
-        print("Session hash:", request.session_hash)
-
-        # Retrieving or setting CUSTOM_CLOUDFRONT_HEADER
-        CUSTOM_CLOUDFRONT_HEADER_var = get_or_create_env_var('CUSTOM_CLOUDFRONT_HEADER', '')
-        #print(f'The value of CUSTOM_CLOUDFRONT_HEADER is {CUSTOM_CLOUDFRONT_HEADER_var}')
-
-        # Retrieving or setting CUSTOM_CLOUDFRONT_HEADER_VALUE
-        CUSTOM_CLOUDFRONT_HEADER_VALUE_var = get_or_create_env_var('CUSTOM_CLOUDFRONT_HEADER_VALUE', '')
-        #print(f'The value of CUSTOM_CLOUDFRONT_HEADER_VALUE_var is {CUSTOM_CLOUDFRONT_HEADER_VALUE_var}')
-
-        if CUSTOM_CLOUDFRONT_HEADER_var and CUSTOM_CLOUDFRONT_HEADER_VALUE_var:
-            if CUSTOM_CLOUDFRONT_HEADER_var in request.headers:
-                supplied_cloudfront_custom_value = request.headers[CUSTOM_CLOUDFRONT_HEADER_var]
-                if supplied_cloudfront_custom_value == CUSTOM_CLOUDFRONT_HEADER_VALUE_var:
-                    print("Custom Cloudfront header found:", supplied_cloudfront_custom_value)
+# async def get_connection_params(request: gr.Request):
+#     base_folder = ""
+
+#     if request:
+#         #print("request user:", request.username)
+
+#         #request_data = await request.json()  # Parse JSON body
+#         #print("All request data:", request_data)
+#         #context_value = request_data.get('context') 
+#         #if 'context' in request_data:
+#         #     print("Request context dictionary:", request_data['context'])
+
+#         # print("Request headers dictionary:", request.headers)
+#         # print("All host elements", request.client)           
+#         # print("IP address:", request.client.host)
+#         # print("Query parameters:", dict(request.query_params))
+#         # To get the underlying FastAPI items you would need to use await and some fancy @ stuff for a live query: https://fastapi.tiangolo.com/vi/reference/request/
+#         #print("Request dictionary to object:", request.request.body())
+#         print("Session hash:", request.session_hash)
+
+#         # Retrieving or setting CUSTOM_CLOUDFRONT_HEADER
+#         CUSTOM_CLOUDFRONT_HEADER_var = get_or_create_env_var('CUSTOM_CLOUDFRONT_HEADER', '')
+#         #print(f'The value of CUSTOM_CLOUDFRONT_HEADER is {CUSTOM_CLOUDFRONT_HEADER_var}')
+
+#         # Retrieving or setting CUSTOM_CLOUDFRONT_HEADER_VALUE
+#         CUSTOM_CLOUDFRONT_HEADER_VALUE_var = get_or_create_env_var('CUSTOM_CLOUDFRONT_HEADER_VALUE', '')
+#         #print(f'The value of CUSTOM_CLOUDFRONT_HEADER_VALUE_var is {CUSTOM_CLOUDFRONT_HEADER_VALUE_var}')
+
+#         if CUSTOM_CLOUDFRONT_HEADER_var and CUSTOM_CLOUDFRONT_HEADER_VALUE_var:
+#             if CUSTOM_CLOUDFRONT_HEADER_var in request.headers:
+#                 supplied_cloudfront_custom_value = request.headers[CUSTOM_CLOUDFRONT_HEADER_var]
+#                 if supplied_cloudfront_custom_value == CUSTOM_CLOUDFRONT_HEADER_VALUE_var:
+#                     print("Custom Cloudfront header found:", supplied_cloudfront_custom_value)
+#                 else:
+#                     raise(ValueError, "Custom Cloudfront header value does not match expected value.")
+
+#         # Get output save folder from 1 - username passed in from direct Cognito login, 2 - Cognito ID header passed through a Lambda authenticator, 3 - the session hash.
+
+#         if request.username:
+#             out_session_hash = request.username
+#             base_folder = "user-files/"
+#             print("Request username found:", out_session_hash)
+
+#         elif 'x-cognito-id' in request.headers:
+#             out_session_hash = request.headers['x-cognito-id']
+#             base_folder = "user-files/"
+#             print("Cognito ID found:", out_session_hash)
+
+#         else:
+#             out_session_hash = request.session_hash
+#             base_folder = "temp-files/"
+#             # print("Cognito ID not found. Using session hash as save folder:", out_session_hash)
+
+#         output_folder = base_folder + out_session_hash + "/"
+#         #if bucket_name:
+#         #    print("S3 output folder is: " + "s3://" + bucket_name + "/" + output_folder)
+
+#         return out_session_hash, output_folder, out_session_hash
+#     else:
+#         print("No session parameters found.")
+#         return "",""
+    
+async def get_connection_params(request: gr.Request,
+                                output_folder_textbox:str=OUTPUT_FOLDER,
+                                input_folder_textbox:str=INPUT_FOLDER,
+                                session_output_folder:str=SESSION_OUTPUT_FOLDER):
+
+    #print("Session hash:", request.session_hash)
+
+    if CUSTOM_HEADER and CUSTOM_HEADER_VALUE:
+            if CUSTOM_HEADER in request.headers:
+                supplied_custom_header_value = request.headers[CUSTOM_HEADER]
+                if supplied_custom_header_value == CUSTOM_HEADER_VALUE:
+                    print("Custom header supplied and matches CUSTOM_HEADER_VALUE")
                 else:
-                    raise(ValueError, "Custom Cloudfront header value does not match expected value.")
+                    print("Custom header value does not match expected value.")
+                    raise ValueError("Custom header value does not match expected value.")
+            else:
+                print("Custom header value not found.")
+                raise ValueError("Custom header value not found.")   
+
+    # Get output save folder from 1 - username passed in from direct Cognito login, 2 - Cognito ID header passed through a Lambda authenticator, 3 - the session hash.
+
+    if request.username:
+        out_session_hash = request.username
+        #print("Request username found:", out_session_hash)
+
+    elif 'x-cognito-id' in request.headers:
+        out_session_hash = request.headers['x-cognito-id']
+        #print("Cognito ID found:", out_session_hash)
+
+    elif 'x-amzn-oidc-identity' in request.headers:
+        out_session_hash = request.headers['x-amzn-oidc-identity']
+
+        # Fetch email address using Cognito client
+        cognito_client = boto3.client('cognito-idp')
+        try:
+            response = cognito_client.admin_get_user(
+                UserPoolId=AWS_USER_POOL_ID,  # Replace with your User Pool ID
+                Username=out_session_hash
+            )
+            email = next(attr['Value'] for attr in response['UserAttributes'] if attr['Name'] == 'email')
+            #print("Email address found:", email)
+
+            out_session_hash = email
+        except ClientError as e:
+            print("Error fetching user details:", e)
+            email = None
+
+        print("Cognito ID found:", out_session_hash)
 
-        # Get output save folder from 1 - username passed in from direct Cognito login, 2 - Cognito ID header passed through a Lambda authenticator, 3 - the session hash.
-
-        if request.username:
-            out_session_hash = request.username
-            base_folder = "user-files/"
-            print("Request username found:", out_session_hash)
+    else:
+        out_session_hash = request.session_hash
 
-        elif 'x-cognito-id' in request.headers:
-            out_session_hash = request.headers['x-cognito-id']
-            base_folder = "user-files/"
-            print("Cognito ID found:", out_session_hash)
+    if session_output_folder == 'True':
+        output_folder = output_folder_textbox + out_session_hash + "/"
+        input_folder = input_folder_textbox + out_session_hash + "/"
 
-        else:
-            out_session_hash = request.session_hash
-            base_folder = "temp-files/"
-            # print("Cognito ID not found. Using session hash as save folder:", out_session_hash)
+    else:
+        output_folder = output_folder_textbox
+        input_folder = input_folder_textbox
 
-        output_folder = base_folder + out_session_hash + "/"
-        #if bucket_name:
-        #    print("S3 output folder is: " + "s3://" + bucket_name + "/" + output_folder)
+    if not os.path.exists(output_folder): os.mkdir(output_folder)
+    if not os.path.exists(input_folder): os.mkdir(input_folder)
 
-        return out_session_hash, output_folder, out_session_hash
-    else:
-        print("No session parameters found.")
-        return "",""
+    return out_session_hash, output_folder, out_session_hash, input_folder

chatfuncs/ingest.py

@@ -7,13 +7,14 @@ import requests
 import pandas as pd
 import dateutil.parser
 from typing import Type, List
-import shutil
+#import shutil
 
-from langchain_community.embeddings import HuggingFaceEmbeddings # HuggingFaceInstructEmbeddings, 
+#from langchain_community.embeddings import HuggingFaceEmbeddings # HuggingFaceInstructEmbeddings, 
 from langchain_community.vectorstores.faiss import FAISS
 #from langchain_community.vectorstores import Chroma
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.docstore.document import Document
+#from chatfuncs.config import EMBEDDINGS_MODEL_NAME
 
 from bs4 import BeautifulSoup
 from docx import Document as Doc
@@ -557,31 +558,24 @@ def docs_elements_from_csv_save(docs_path="documents.csv"):
 
 # ## Create embeddings and save faiss vector store to the path specified in `save_to`
 
-def load_embeddings(model_name = "BAAI/bge-base-en-v1.5"):
+# def load_embeddings_model(embeddings_model = EMBEDDINGS_MODEL_NAME):
 
-    #if model_name == "hkunlp/instructor-large":
-    #    embeddings_func = HuggingFaceInstructEmbeddings(model_name=model_name,
-    #    embed_instruction="Represent the paragraph for retrieval: ",
-    #    query_instruction="Represent the question for retrieving supporting documents: "
-    #    )
+#     embeddings_func = HuggingFaceEmbeddings(model_name=embeddings_model)
 
-    #else: 
-    embeddings_func = HuggingFaceEmbeddings(model_name=model_name)
+#     #global embeddings
 
-    global embeddings
+#     #embeddings = embeddings_func
 
-    embeddings = embeddings_func
+#     return embeddings_func
 
-    return embeddings_func
-
-def embed_faiss_save_to_zip(docs_out, save_to="output", model_name="BAAI/bge-base-en-v1.5"):
-    load_embeddings(model_name=model_name)
+def embed_faiss_save_to_zip(docs_out, save_folder, embeddings_model_object, save_to="faiss_embeddings", model_name="BAAI/bge-base-en-v1.5"):
+    #load_embeddings(model_name=model_name)
 
     print(f"> Total split documents: {len(docs_out)}")
 
-    vectorstore = FAISS.from_documents(documents=docs_out, embedding=embeddings)
+    vectorstore = FAISS.from_documents(documents=docs_out, embedding=embeddings_model_object)
 
-    save_to_path = Path(save_to)
+    save_to_path = Path(save_folder, save_to)
     save_to_path.mkdir(parents=True, exist_ok=True)
 
     vectorstore.save_local(folder_path=str(save_to_path))
@@ -619,20 +613,20 @@ def embed_faiss_save_to_zip(docs_out, save_to="output", model_name="BAAI/bge-bas
 
 
 
-def sim_search_local_saved_vec(query, k_val, save_to="faiss_lambeth_census_embedding"):
+# def sim_search_local_saved_vec(query, k_val, save_to="faiss_lambeth_census_embedding"):
 
-    load_embeddings()
+#     load_embeddings()
 
-    docsearch = FAISS.load_local(folder_path=save_to, embeddings=embeddings)
+#     docsearch = FAISS.load_local(folder_path=save_to, embeddings=embeddings)
 
 
-    display(Markdown(question))
+#     display(Markdown(question))
 
-    search = docsearch.similarity_search_with_score(query, k=k_val)
+#     search = docsearch.similarity_search_with_score(query, k=k_val)
 
-    for item in search:
-        print(item[0].page_content)
-        print(f"Page: {item[0].metadata['source']}")
-        print(f"Date: {item[0].metadata['date']}")
-        print(f"Score: {item[1]}")
-        print("---")
+#     for item in search:
+#         print(item[0].page_content)
+#         print(f"Page: {item[0].metadata['source']}")
+#         print(f"Date: {item[0].metadata['date']}")
+#         print(f"Score: {item[1]}")
+#         print("---")

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.
@@ -73,4 +71,9 @@ Answer the QUESTION using information from the following CONTENT. Respond with s
 CONTENT: {summaries}
 QUESTION: {question}\n
 Answer:<|im_end|>
-<|im_start|>assistant\n"""
+<|im_start|>assistant\n"""
+
+instruction_prompt_gemma = """Answer the QUESTION using information from the following CONTENT. Respond with short answers that directly answer the question.
+CONTENT: {summaries}
+QUESTION: {question}
+assistant:"""

requirements.txt

@@ -4,7 +4,7 @@ langchain-community==0.3.22
 beautifulsoup4==4.13.4
 google-generativeai==0.8.5
 pandas==2.2.3
-transformers==4.41.2
+transformers==4.51.3
 # For Windows https://github.com/abetlen/llama-cpp-python/releases/download/v0.3.2/llama_cpp_python-0.3.2-cp311-cp311-win_amd64.whl -C cmake.args="-DGGML_BLAS=ON;-DGGML_BLAS_VENDOR=OpenBLAS" 
 llama-cpp-python==0.3.2 --extra-index-url https://abetlen.github.io/llama-cpp-python/whl/cpu
 #-C cmake.args="-DGGML_BLAS=ON;-DGGML_BLAS_VENDOR=OpenBLAS"
@@ -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