test party filter

Home.py

@@ -39,10 +39,10 @@ with gr.Blocks() as App:
                         file = gr.File(file_types=[".xlsx", ".csv", ".json"], visible=False)
     
             #Keyword Search on click
-            def search(keyword, n, party): #ToDo: Include party
+            def search(keyword, n, party): #ToDo: Include party and timedate
                 return {
                     output_col: gr.Column(visible=True),
-                    results_df: keyword_search(query=keyword, n=n),
+                    results_df: keyword_search(query=keyword, n=n, party_filter=party),
                 }
 
             search_btn.click(

src/chatbot.py

@@ -1,12 +1,10 @@
-from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_community.llms.huggingface_hub import HuggingFaceHub
 from langchain_community.embeddings import HuggingFaceEmbeddings
 
 from src.vectordatabase import RAG, get_vectorstore
 import pandas as pd
-import os
-#from dotenv import load_dotenv, find_dotenv
+from dotenv import load_dotenv, find_dotenv
 
 #Load environmental variables from .env-file
 #load_dotenv(find_dotenv())
@@ -63,22 +61,33 @@ def chatbot(message, history, db=db, llm=llm, prompt=prompt2):
     return response  
 
 # Retrieve speech contents based on keywords
-def keyword_search(query,n=10, db=db, embeddings=embeddings):
+def keyword_search(query,n=10, db=db, embeddings=embeddings, method='ss', party_filter = ''):
     query_embedding = embeddings.embed_query(query)
-    results =  db.max_marginal_relevance_search_with_score_by_vector(query_embedding, k = n)
-    # Format vector store query results into dataframe
-    #print(results[0][0].metadata.keys())
-    
-    df_res = pd.DataFrame(columns=['Speech Content','Date', 'Party', 'Relevance']) # Add Date/Party/Politician
-    for doc in results:
-        speech_content = doc[0].page_content
-        speech_date = doc[0].metadata["date"]
-        party = doc[0].metadata["party"]
-        score = round(doc[1], ndigits=2) # Relevance based on relevance search
-        df_res = pd.concat([df_res, pd.DataFrame({'Speech Content': [speech_content],
-                                                          'Date': [speech_date],
-                                                          'Party': [party],
-                                                          'Relevance': [score]})], ignore_index=True)
-        
-    df_res.sort_values('Relevance', inplace=True, ascending=False)
+    if method == 'mmr':
+        df_res = pd.DataFrame(columns=['Speech Content','Date', 'Party', 'Relevance']) # Add Date/Party/Politician 
+        results =  db.max_marginal_relevance_search_with_score_by_vector(query_embedding, k = n, fetch_k = n + 10) #Add filter
+        for doc in results:
+            speech_content = doc[0].page_content
+            speech_date = doc[0].metadata["date"]
+            party = doc[0].metadata["party"]
+            score = round(doc[1], ndigits=2) # Relevance based on relevance search
+            df_res = pd.concat([df_res, pd.DataFrame({'Speech Content': [speech_content],
+                                                            'Date': [speech_date],
+                                                            'Party': [party],
+                                                            'Relevance': [score]})], ignore_index=True)
+        df_res.sort_values('Relevance', inplace=True, ascending=True)
+    else:
+        df_res = pd.DataFrame(columns=['Speech Content','Date', 'Party']) # Add Date/Party/Politician #Add filter
+        results = db.similarity_search_by_vector(query_embedding, k = n, filter={"party": party_filter})
+        for doc in results:
+            party = doc.metadata["party"]
+            #Filter by party input
+            #if party != party_filter or party_filter == '':
+            #    continue
+            speech_content = doc.page_content
+            speech_date = doc.metadata["date"]
+            
+            df_res = pd.concat([df_res, pd.DataFrame({'Speech Content': [speech_content],
+                                                            'Date': [speech_date],
+                                                            'Party': [party]})], ignore_index=True)
     return df_res