Update app.py

app.py

@@ -5,12 +5,13 @@ import re
 from huggingface_hub import InferenceClient
 import os
 import numpy as np
+from openai import OpenAI
 
 
 
 
-st.set_page_config(layout="wide")
 
+st.set_page_config(layout="wide")
 
 
 my_initial_rag_text = f"""This is a RAG (Retrieval-Augmented Generation) chatbot application built with Streamlit that combines document context with LLM responses. Here's a breakdown of its main components:
@@ -62,18 +63,28 @@ A disclaimer at the bottom reminds users about potential LLM inaccuracies and th
 if "my_llm_model" not in st.session_state:
     # Set the default LLM model to "mistralai/Mistral-7B-Instruct-v0.3"
     st.session_state['my_llm_model'] = "mistralai/Mistral-7B-Instruct-v0.3"
+
 # Check if the SPACE_ID environment variable is not already in the session state
 if "my_space" not in st.session_state:
     st.session_state['my_space'] = os.environ.get("SPACE_ID") 
 
 # Function to update the LLM model client
 def update_llm_model():
-    if st.session_state['my_space']:
-        # Initialize the client with the model if SPACE_ID is available
-        st.session_state['client'] = InferenceClient(st.session_state['my_llm_model'])
+    if st.session_state['my_llm_model'].startswith("gemini-"):
+        # Initialize the client for gemini models. We use the OpenAI API to interact with gemini models.
+        st.session_state['client'] = OpenAI(api_key = os.getenv("GOOGLE_API_KEY"),
+                                            base_url = "https://generativelanguage.googleapis.com/v1beta/openai/")
+    elif st.session_state['my_llm_model'].startswith("gpt-"):
+        # Initialize the client for openai models
+        st.session_state['client'] = OpenAI(api_key = os.getenv("OPENAI_API_KEY"))
+                                            # ,base_url = "https://eu.api.openai.com/" # gives error
     else:
-        # Initialize the client with the model and token if SPACE_ID is not available
-        st.session_state['client'] = InferenceClient(st.session_state['my_llm_model'], token=os.getenv("HF_TOKEN"))
+        if st.session_state['my_space']:
+            # Initialize the client with the model if SPACE_ID is available
+            st.session_state['client'] = InferenceClient(st.session_state['my_llm_model'])
+        else:
+            # Initialize the client with the model and token if SPACE_ID is not available
+            st.session_state['client'] = InferenceClient(st.session_state['my_llm_model'], token=os.getenv("HF_TOKEN"))
         
 # Check if the client is not already in the session state
 if "client" not in st.session_state:
@@ -84,7 +95,6 @@ if "embeddings_model" not in st.session_state:
     # We will use the all-MiniLM-L6-v2 model for embeddings
     st.session_state['embeddings_model'] = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
 
-
 my_system_instructions = "You are a helpful assistant. Be brief and concise. Provide your answers in 100 words or less."
 
 first_message = "Hello, how can I help you today?"
@@ -93,32 +103,33 @@ def delete_chat_messages():
     for key in st.session_state.keys():
         if key != "my_rag_text" and key != "my_system_instructions":
             del st.session_state[key]
+    update_llm_model()
+            
 
 def create_sentences_rag():
-    text = st.session_state['my_rag_text']
+    with rag_status_placeholder:
+        # The pattern splits text at any of the punctuation marks .?!;: followed by one or more spaces, or at a newline character
+        pattern = r'(?<=[.?!;:])\s+|\n'
+        st.session_state['my_sentences'] = [sentence.strip() for sentence in re.split(pattern, st.session_state['my_rag_text']) if sentence.strip()]
+        with st.spinner(f"Encoding {len(st.session_state['my_sentences'])} sentences..."):
+            sentences_ids = [i for i in range(len(st.session_state['my_sentences']))]
+            # Rolling window: include partial windows at end
+            st.session_state['my_sentences_rag_ids'] = []
+            st.session_state['my_sentences_rag'] = []
+            for rolling_window_size in range(st.session_state['min_window_size'], st.session_state['max_window_size']+1):
+                for i in range(0, len(st.session_state['my_sentences'])-rolling_window_size+1):
+                    chunk = " ".join(st.session_state['my_sentences'][i:i+rolling_window_size]).strip()
+                    if chunk:   
+                        st.session_state['my_sentences_rag'].append(chunk)
+                        st.session_state['my_sentences_rag_ids'].append(sentences_ids[i:i+rolling_window_size])
+                        # print(f"*****{chunk}*****\n")
 
-    # Split only on .?!;: followed by space OR on \n
-    pattern = r'(?<=[.?!;:])\s+|\n'
-    st.session_state['my_sentences'] = [sentence.strip() for sentence in re.split(pattern, text) if sentence.strip()]
-    sentences_ids = [i for i in range(len(st.session_state['my_sentences']))]
+            st.session_state['my_embeddings'] = st.session_state['embeddings_model'].encode(st.session_state['my_sentences_rag'])
+        st.success(f"{len(st.session_state['my_sentences_rag'])} chunks have been encoded!")
 
-    # Rolling window: include partial windows at end
-    st.session_state['my_sentences_rag_ids'] = []
-    st.session_state['my_sentences_rag'] = []
-    for rolling_window_size in range(st.session_state['min_window_size'], st.session_state['max_window_size']+1):
-        for i in range(0, len(st.session_state['my_sentences'])-rolling_window_size+1):
-            chunk = " ".join(st.session_state['my_sentences'][i:i+rolling_window_size]).strip()
-            if chunk:   
-                st.session_state['my_sentences_rag'].append(chunk)
-                st.session_state['my_sentences_rag_ids'].append(sentences_ids[i:i+rolling_window_size])
-                # print(f"*****{chunk}*****\n")
 
-    print(len(st.session_state['my_sentences_rag']))
-    st.session_state['my_embeddings'] = st.session_state['embeddings_model'].encode(st.session_state['my_sentences_rag'])
 
 
-augmented_prompt = ""
-
 # Create two columns with a 1:2 ratio
 column_1, column_2 = st.columns([1, 2])
 
@@ -136,17 +147,25 @@ The user agrees to indemnify and hold harmless the developers of the Software fr
 By using the Software, you agree to the terms and conditions of the disclaimer.""")
 
     # Add a selectbox for model selection
+    model_list_all = [  'mistralai/Mistral-7B-Instruct-v0.3',
+                        'Qwen/Qwen2.5-72B-Instruct', 
+                        'HuggingFaceH4/zephyr-7b-beta']
+    if os.getenv("GOOGLE_API_KEY"):
+        model_list_all.append('gemini-2.5-flash-preview-05-20')
+    if os.getenv("OPENAI_API_KEY"):
+        model_list_all.append('gpt-4.1-nano-2025-04-14')
     st.selectbox("Select the model to use:",
-                ['mistralai/Mistral-7B-Instruct-v0.3',
-                 'Qwen/Qwen2.5-72B-Instruct', 
-                'HuggingFaceH4/zephyr-7b-beta'], 
-                key="my_llm_model", on_change=update_llm_model)
+                model_list_all, 
+                key="my_llm_model", 
+                on_change=update_llm_model)
     
     # Add a text are for the system instructions
-    st.text_area(label="Please enter your system instructions here:", value=my_system_instructions, height=100, key="my_system_instructions", on_change=delete_chat_messages)
+    st.text_area(label="Please enter your system instructions here:", value=my_system_instructions, height=80, key="my_system_instructions", on_change=delete_chat_messages)
     
+    # Placeholder right after text_area
+    rag_status_placeholder = st.empty()
     # Add a text area for RAG text input
-    st.text_area(label="Please enter your RAG text here:", value=my_initial_rag_text, height=500, key="my_rag_text", on_change=delete_chat_messages)
+    st.text_area(label="Please enter your RAG text here:", value=my_initial_rag_text, height=200, key="my_rag_text", on_change=delete_chat_messages)
 
     # Add a slider for minimum window size
     st.slider("Minimum window size in original sentences", min_value=1, max_value=20, value=5, step=1, key="min_window_size", on_change=create_sentences_rag)
@@ -155,7 +174,7 @@ By using the Software, you agree to the terms and conditions of the disclaimer."
     st.slider("Maximum window size in original sentences", min_value=1, max_value=20, value=10, step=1, key="max_window_size", on_change=create_sentences_rag)
     
     # Add a slider for the similarity threshold
-    st.slider("Similarity threshold", min_value=0.0, max_value=1.0, value=0.3, step=0.01, key="my_similarity_threshold")
+    st.slider("Similarity threshold", min_value=0.0, max_value=1.0, value=0.2, step=0.01, key="my_similarity_threshold")
 
     # Add a slider for the number of sentences to keep
     st.slider("Number of original chunks to keep", min_value=1, max_value=50, value=20, step=1, key="nof_keep_sentences")
@@ -182,6 +201,8 @@ if "my_sentences_rag" not in st.session_state:
     create_sentences_rag()
 
 
+
+
 with column_2:
     # Create a container for the messages with a specified height
     messages_container = st.container(height=500)
@@ -201,10 +222,13 @@ with column_2:
     # Check if there is a new prompt from the user
     if prompt := st.chat_input("you may ask here your questions"):
         
+        # Split the prompt into words
         split_prompt = prompt.split(" ")
         all_sub_prompts = []
+        # Generate sub-prompts based on the specified range
         for jj in range(st.session_state['nof_min_sub_prompts'], st.session_state['nof_max_sub_prompts']+1):
             for ii in range(len(split_prompt)):
+                # Create sub-prompt by joining words
                 i_split = " ".join(split_prompt[ii:ii+jj]).strip()
                 if i_split:
                     all_sub_prompts.append(i_split)
@@ -225,7 +249,7 @@ with column_2:
         max_similarity = 0
         # for irag in range(st.session_state['nof_keep_sentences']):
         irag = 0
-        while len(set(sorted_indices_sentences))<st.session_state['nof_keep_sentences']:
+        while len(set(sorted_indices_sentences))<st.session_state['nof_keep_sentences'] and irag<len(sorted_indices_rag):
             sorted_indices_sentences.extend(st.session_state['my_sentences_rag_ids'][sorted_indices_rag[irag]])
             max_similarity = max(max_similarity, similarities_to_question[sorted_indices_rag[irag]])
             with bottom_col1:
@@ -237,44 +261,58 @@ with column_2:
 
         sorted_indices_sentences = sorted(list(set(sorted_indices_sentences)))
 
-        # Construct the augmented prompt with the similar sentences
-        if max_similarity > st.session_state['my_similarity_threshold']:
-            augmented_prompt = "This is my context:" + "\n\n" + 20*"-" + "\n\n" 
-            augmented_prompt += "\n".join([st.session_state['my_sentences'][idx] for idx in sorted_indices_sentences])
-            augmented_prompt += "\n\n" + 20*"-" + "\n\n" + "If the above context is not relevant to the prompt, ignore the context and reply based only on the prompt."
-            augmented_prompt += "\n\n" + 20*"-" + "\n\n" + "If the above context is relevant to the prompt, reply based on the context and the prompt."
-            augmented_prompt += "\n\n" + 20*"-" + "\n\n" + "The prompt is:"
-            augmented_prompt += "\n\n" + f"\n\n{prompt}"
-        else:
-            augmented_prompt = prompt
-        with bottom_col2:
-            # Display the augmented prompt used for generating the response
-            st.write("Augmented prompt:")
-            st.json({"max_similarity": max_similarity, "augmented_prompt": augmented_prompt}, expanded=False)
-        
         # Display the user's prompt in the chat container with a specific avatar
         messages_container.chat_message("user", avatar=":material/psychology_alt:").markdown(prompt) 
-        # Append the augmented prompt to the chat messages in the session state
-        st.session_state['my_chat_messages'].append({"role": "user", "content": augmented_prompt})
+        
         # Create an empty container for the streaming response from the assistant
         with messages_container.chat_message("ai", avatar=":material/robot_2:"):
             response_placeholder = st.empty()
-            response = ""
-            # Stream the response from the assistant and update the placeholder
-            for chunk in st.session_state['client'].chat.completions.create(messages=st.session_state['my_chat_messages'], stream=True, max_tokens=512):
-                if chunk.choices[0].delta.content:
-                    response += chunk.choices[0].delta.content
-                    # Use markdown to update the response placeholder with the streamed content
-                    response_placeholder.markdown(response)
-
-        # Remove the last message from the chat messages in the session state
-        st.session_state['my_chat_messages'].pop()
+            if max_similarity > st.session_state['my_similarity_threshold']:
+                # Construct the augmented prompt with the similar sentences
+                augmented_prompt = "This is my context:" + "\n\n" + 20*"-" + "\n\n" 
+                augmented_prompt += "\n".join([st.session_state['my_sentences'][idx] for idx in sorted_indices_sentences])
+                augmented_prompt += "\n\n" + 20*"-" + "\n\n" + "If the above context is not relevant to the prompt, ignore the context and reply based only on the prompt."
+                augmented_prompt += "\n\n" + 20*"-" + "\n\n" + "If the above context is relevant to the prompt, reply based on the context and the prompt."
+                augmented_prompt += "\n\n" + 20*"-" + "\n\n" + "The prompt is:"
+                augmented_prompt += "\n\n" + f"\n\n{prompt}"
+                # Append the augmented prompt to the chat messages in the session state
+                st.session_state['my_chat_messages'].append({"role": "user", "content": augmented_prompt})
+                # Stream the response from the assistant and update the placeholder
+                response = ""
+                for chunk in st.session_state['client'].chat.completions.create(messages = st.session_state['my_chat_messages'], 
+                                                                                model = st.session_state['my_llm_model'],
+                                                                                stream = True, 
+                                                                                max_tokens = 1024):
+                    if chunk.choices[0].delta.content:
+                        response += chunk.choices[0].delta.content
+                        # Use markdown to update the response placeholder with the streamed content
+                        response_placeholder.markdown(response)
+                # Remove the last message from the chat messages in the session state
+                st.session_state['my_chat_messages'].pop()
+            else:
+                augmented_prompt = ""
+                response = f"I do not have enough information to reply. The maximum similarity found in the context is: {100*max_similarity:.2f}%."
+                response_placeholder.markdown(response)
+
         # Append the user's original prompt to the chat messages in the session state
         st.session_state['my_chat_messages'].append({"role": "user", "content": prompt})
         # Append the assistant's response to the chat messages in the session state
         st.session_state['my_chat_messages'].append({"role": "assistant", "content": response})
 
+
+        if len(st.session_state['my_chat_messages'])>10:
+            # Keep the first message which is the system instructions, remove the 2nd and 3rd messages which are the first user and assistant messages
+            st.session_state['my_chat_messages'] = st.session_state['my_chat_messages'][:1] + st.session_state['my_chat_messages'][3:]
+    
+            
+
         with bottom_col2:
+            # Display the augmented prompt used for generating the response
+            st.write("Augmented prompt:")
+            st.json({"max_similarity": max_similarity, "augmented_prompt": augmented_prompt}, expanded=False)
+
             # Display the chat messages history
             st.write("Messages History All:")
-            st.json(st.session_state['my_chat_messages'], expanded=False)    
+            st.json(st.session_state['my_chat_messages'], expanded=False)
+
+