kw test

src/chatbot.py

@@ -70,7 +70,8 @@ def keyword_search(query, db=db, embeddings=embeddings):
     # Format vector store query results into dataframe
     #print(results[0][0].metadata.keys())
     
-    df_res = pd.DataFrame(columns=['Speech Content','Date','Relevance']) # Add Date/Party/Politician
+    df_res = pd.DataFrame(columns=['Speech Content','Date', 'Party', 'Relevance']) # Add Date/Party/Politician
+    i = 0
     for doc in results:
         speech_content = doc[0].page_content
         speech_date = doc[0].metadata["date"]
@@ -80,6 +81,9 @@ def keyword_search(query, db=db, embeddings=embeddings):
                                                           'Date': [speech_date],
                                                           'Party': [party],
                                                           'Relevance': [score]})], ignore_index=True)
+        i = i + 1
+        if i > 2:
+            break
     
     df_res.sort_values('Relevance', inplace=True, ascending=False)
     return df_res