Update app.py

app.py

@@ -95,7 +95,7 @@ def transform_query(query: str) -> str:
     """
     return f'Represent this sentence for searching relevant passages: {query}'
 
-def query_hybrid_search(query: str, client: QdrantClient, collection_name: str, query_embeddings, sparse_embeddings):
+def query_hybrid_search(query: str, client: QdrantClient, collection_name: str, dense_embeddings, sparse_embeddings):
     return client.query_points(
         collection_name=collection_name,
         prefetch=[
@@ -121,7 +121,7 @@ def build_prompt_conv():
         },
         {
             'role': 'user',
-            'content': f"""Generate a short, single-sentence summary of the user's intent or topic based on their question, capturing the main focus of what they want to discuss. Do NOT cite the user.
+            'content': f"""Generate a short, single-sentence summary, in 10 tokens maximum, of the user's intent or topic based on their question, capturing the main focus of what they want to discuss. Do NOT cite the user.
             
             Question : {st.session_state.user_input}
             """
@@ -228,59 +228,74 @@ def self_knowledge(query: str):
     Question: {{ query }}
     """
 
-def main(query: str, client: QdrantClient, collection_name: str, llm, dense_model: AsyncEmbeddingEngine, sparse_model: SparseTextEmbedding, past_messages: str):
-    dense_embeddings, tokens_count = asyncio.run(embed_text(dense_model[0], transform_query(query)))
+def generate_answer(query: str, 
+                    client: QdrantClient, 
+                    collection_name: str, 
+                    llm, 
+                    dense_model: AsyncEmbeddingEngine, 
+                    sparse_model: SparseTextEmbedding, 
+                    past_messages: str, 
+                    search_strategy: str, 
+                   ):
     sparse_embeddings = list(sparse_model.query_embed(query))[0].as_object()
     
     s = time.time()
-    scored_points = query_hybrid_search(query, client, collection_name, dense_embeddings, sparse_embeddings).points
-    print(f'Score : {scored_points[0]}')   
-    docs = [(scored_point.payload['text'], scored_point.payload['metadata']) for scored_point in scored_points]
-    contents, metadatas = [list(t) for t in zip(*docs)]
-        
-    context = "\n".join(contents)
-    print(f'Context : \n + {context}')
 
-    regex = build_regex_from_schema(schema, r"[\n ]?")
-    gen_text = outlines.generate.regex(llm, regex)
+    if search_strategy == 'Exact Search':
+        scored_points = query_keywords_search(query, client, collection_name, sparse_embeddings).points
 
-    gen_choice = outlines.generate.choice(llm, choices=['Yes', 'No'])
-    prompt = route_llm(context, query)
-    action = gen_choice(prompt, max_tokens=2, sampling_params=SamplingParams(temperature=0))
-    print(f'Choice: {action}')
-    
-    if action == 'Yes':
-        filtered_metadatas = {
-            value
-            for metadata in metadatas
-            if 'url' in metadata
-            for value in [metadata['url']]
-        }
-        result_metadatas = "\n\n".join(f'{value}' for value in filtered_metadatas)
+        answer = f"{scored_points[0].payload['text']}\n\n\nSource :\n\n{scored_points[0].payload['metadata']}"
+    else:
+        regex = build_regex_from_schema(schema, r"[\n ]?")
+        gen_text = outlines.generate.regex(llm, regex)
 
-        prompt = answer_with_context(context, query)
-        answer = json.loads(gen_text(prompt, max_tokens=300, sampling_params=SamplingParams(temperature=0)))['answer']
-        answer = f"{answer}\n\n\nSource(s) :\n\n{result_metadatas}"
+        gen_choice = outlines.generate.choice(llm, choices=['Yes', 'No'])
+        prompt = route_llm(context, query)
+        action = gen_choice(prompt, max_tokens=2, sampling_params=SamplingParams(temperature=0))
+        print(f'Choice: {action}')
+    
+        if action == 'Yes':
+            dense_embeddings, tokens_count = asyncio.run(embed_text(dense_model[0], transform_query(query)))
             
-        if st.session_state.initialized == False:
-            answer = f'Documents Based :\n\n{answer}'
-    else:
-        gen_choice = outlines.generate.choice(llm, choices=['Domain-Specific Question', 'General Question'])
-        prompt = question_type_prompt(query)
-        action = gen_choice(prompt, max_tokens=3, sampling_params=SamplingParams(temperature=0))
-        print(f'Choice 2: {action}')
-        if action == 'General Question':
-            prompt = open_query_prompt(past_messages, query)
-            answer = json.loads(gen_text(prompt, max_tokens=300, sampling_params=SamplingParams(temperature=0.6, top_p=0.9, top_k=10)))['answer']
+            scored_points = query_hybrid_search(query, client, collection_name, dense_embeddings, sparse_embeddings).points
+            print(f'Score : {scored_points[0]}')   
+            docs = [(scored_point.payload['text'], scored_point.payload['metadata']) for scored_point in scored_points]
+            contents, metadatas = [list(t) for t in zip(*docs)]
+            
+            context = "\n".join(contents)
+            print(f'Context : \n + {context}')
+        
+            filtered_metadatas = {
+                value
+                for metadata in metadatas
+                if 'url' in metadata
+                for value in [metadata['url']]
+            }
+            result_metadatas = "\n\n".join(f'{value}' for value in filtered_metadatas)
+            
+            prompt = answer_with_context(context, query)
+            answer = json.loads(gen_text(prompt, max_tokens=300, sampling_params=SamplingParams(temperature=0)))['answer']
+            answer = f"{answer}\n\n\nSource(s) :\n\n{result_metadatas}"
+            
+            if search_strategy == 'Documents + LLM Search':
+                answer = f'Documents Based :\n\n{answer}'
         else:
-            print(f'GLOBAL STATE : {st.session_state.initialized}')
-            if st.session_state.initialized == True:
-                prompt = idk(query)
-                answer = json.loads(gen_text(prompt, max_tokens=128, sampling_params=SamplingParams(temperature=0.6, top_p=0.9, top_k=10)))['answer']
-            else:
-                prompt = self_knowledge(query)
+            gen_choice = outlines.generate.choice(llm, choices=['Domain-Specific Question', 'General Question'])
+            prompt = question_type_prompt(query)
+            action = gen_choice(prompt, max_tokens=3, sampling_params=SamplingParams(temperature=0))
+            print(f'Choice 2: {action}')
+            
+            if action == 'General Question':
+                prompt = open_query_prompt(past_messages, query)
                 answer = json.loads(gen_text(prompt, max_tokens=300, sampling_params=SamplingParams(temperature=0.6, top_p=0.9, top_k=10)))['answer']
-                answer = f'Internal Knowledge :\n\n{answer}'
+            else:
+                if search_strategy == 'Documents Only Search':
+                    prompt = idk(query)
+                    answer = json.loads(gen_text(prompt, max_tokens=128, sampling_params=SamplingParams(temperature=0.6, top_p=0.9, top_k=10)))['answer']
+                elif search_strategy == 'Documents + LLM Search':
+                    prompt = self_knowledge(query)
+                    answer = json.loads(gen_text(prompt, max_tokens=300, sampling_params=SamplingParams(temperature=0.6, top_p=0.9, top_k=10)))['answer']
+                    answer = f'Internal Knowledge :\n\n{answer}'
 
     torch.cuda.empty_cache()
 
@@ -577,28 +592,11 @@ def load_models_and_documents():
     container.empty()
     
     return client, collection_name, llm, model, dense_model, sparse_model, nlp, conn, cursor
-
-def on_change_documents_only():
-    if st.session_state.documents_only:
-        st.session_state.initialized = True
-        st.session_state.toggle_docs = {
-            'tooltip': 'The AI answer your questions only considering the documents provided',
-            'display': True
-        }
-    else:
-        st.session_state.initialized = False
-        st.session_state.toggle_docs = {
-            'tooltip': """The AI answer your questions considering the documents provided, and if it doesn't found the answer in them, try to find in its own internal knowledge""",
-            'display': False
-        }
     
     
 if __name__ == '__main__':
     st.set_page_config(page_title="Multipurpose AI Agent",layout="wide", initial_sidebar_state='auto')
 
-    if 'initialized' not in st.session_state:
-        st.session_state.initialized = True
-
     client, collection_name, llm, model, dense_model, sparse_model, nlp, conn, cursor = load_models_and_documents()
 
     styles = {
@@ -752,6 +750,10 @@ if __name__ == '__main__':
         if 'id_chat' not in st.session_state:
             st.session_state.id_chat = 'New Conversation'
 
+        if 'search_strategy' not in st.session_state:
+            st.session_state.search_strategy = 'Documents Only Search'
+            st.session_state.tooltip = 'The AI answer your questions only considering the documents provided'
+
         def options_list(conversations: Dict[str, list]):
             if st.session_state.id_chat == 'New Conversation':
                 return [st.session_state.id_chat] + list(conversations.keys())
@@ -760,10 +762,9 @@ if __name__ == '__main__':
         
         with st.sidebar:
             st.session_state.id_chat = st.selectbox(
-                label='Choose a conversation',
+                label='Choose a Conversation',
                 options=options_list(conversations),
                 index=0,
-                placeholder='_',
                 key='chat_id'
             )
 
@@ -775,7 +776,6 @@ if __name__ == '__main__':
                     packed_bytes = msgpack.packb(conversations, use_bin_type=True)
                     fp.write(packed_bytes)
                 st.session_state.chat_id = 'New Conversation'
-                    
             
             st.button(
                 'Delete Conversation', 
@@ -785,6 +785,25 @@ if __name__ == '__main__':
                 args=(conversations_path, conversations)
             )
 
+            st.divider()
+
+            def tooltip_change():
+                if st.session_state.search_id == 'Exact Search':
+                    st.session_state.tooltip = 'Search the exact definition'
+                elif st.session_state.search_id == 'Documents Only Search':
+                    st.session_state.tooltip = 'The AI answer your questions only considering the documents provided'
+                elif st.session_state.search_id == 'Documents + LLM Search':
+                    st.session_state.tooltip = 'The AI answer your questions considering the documents provided, and if it doesn't found the answer in them, try to find in its own internal knowledge'
+
+            st.session_state.search_strategy = st.radio(
+                label='Choose a Search Strategy',
+                options=['Exact Search', 'Documents Only Search', 'Documents + LLM Search'],
+                index=1,
+                on_change=tooltip_change,
+                key='search_id',
+                help=st.session_state.tooltip
+            )
+
         def generate_conv_title(llm):
             if st.session_state.chat_id == 'New Conversation':
                 output = llm.chat(
@@ -815,7 +834,7 @@ if __name__ == '__main__':
             st.chat_message("user").markdown(prompt)
             st.session_state.messages.append({"role": "user", "content": prompt})
             
-            ai_response = main(prompt, client, collection_name, model, dense_model, sparse_model, "\n".join([f'{msg["role"]}: {msg["content"]}' for msg in st.session_state.messages]))
+            ai_response = generate_answer(prompt, client, collection_name, model, dense_model, sparse_model, "\n".join([f'{msg["role"]}: {msg["content"]}' for msg in st.session_state.messages]), st.session_state.search_strategy)
             with st.chat_message("assistant"):
                 message_placeholder = st.empty()
                 full_response = ""
@@ -843,7 +862,7 @@ if __name__ == '__main__':
                     'tooltip': 'The AI answer your questions only considering the documents provided',
                     'display': True
                 }
-                
+
             st.toggle(
                 label="""Enable 'Documents-Only' Mode""", 
                 value=st.session_state.toggle_docs['display'],