Update app.py

app.py

@@ -20,6 +20,7 @@ from langchain_core.runnables import RunnableParallel, RunnablePassthrough
 from langchain_core.documents import Document
 from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.metrics.pairwise import cosine_similarity
+from datetime import datetime
 
 huggingface_token = os.environ.get("HUGGINGFACE_TOKEN")
 
@@ -142,7 +143,7 @@ _useragent_list = [
     "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Safari/537.36",
 ]
 
-def google_search(term, num_results=5, lang="en", timeout=5, safe="active", ssl_verify=None):
+def google_search(term, num_results=20, lang="en", timeout=5, safe="active", ssl_verify=None):
     escaped_term = urllib.parse.quote_plus(term)
     start = 0
     all_results = []
@@ -221,6 +222,30 @@ def google_search(term, num_results=5, lang="en", timeout=5, safe="active", ssl_
 
     return all_results
 
+def summarize_content(content, model):
+    summary_prompt = f"""
+    Summarize the following content in a concise manner:
+    {content}
+    Summary:
+    """
+    summary = generate_chunked_response(model, summary_prompt, max_tokens=200)
+    return summary
+
+def rank_search_results(titles, summaries, model):
+    ranking_prompt = f"""
+    Rank the following search results from a financial analyst perspective. 
+    Assign a rank from 1 to {len(titles)} based on relevance, with 1 being the most relevant.
+    Return only the numeric ranks in order, separated by commas.
+
+    Titles and summaries:
+    {', '.join([f"{i+1}. Title: {title}\nSummary: {summary}" for i, (title, summary) in enumerate(zip(titles, summaries))])}
+
+    Ranks:
+    """
+    ranks_str = generate_chunked_response(model, ranking_prompt)
+    ranks = [float(rank.strip()) for rank in ranks_str.split(',')]
+    return ranks
+
 def ask_question(question, temperature, top_p, repetition_penalty, web_search):
     global conversation_history
 
@@ -235,19 +260,19 @@ def ask_question(question, temperature, top_p, repetition_penalty, web_search):
         database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
     else:
         database = None
-
+    
     if web_search:
         search_results = google_search(question)
-        web_docs = [Document(page_content=result["text"], metadata={"source": result["link"]}) for result in search_results if result["text"]]
+        model = get_model(temperature, top_p, repetition_penalty)
         
-        if database is None:
-            database = FAISS.from_documents(web_docs, embed)
-        else:
-            database.add_documents(web_docs)
+        summaries = [summarize_content(result["text"], model) for result in search_results]
+        titles = [result["title"] for result in search_results]
+        ranks = rank_search_results(titles, summaries, model)
         
-        database.save_local("faiss_database")
+        update_vector_db_with_search_results(search_results, summaries, ranks)
         
-        context_str = "\n".join([doc.page_content for doc in web_docs])
+        context_str = "\n".join([f"Title: {result['title']}\nSummary: {summary}\nRank: {rank}" 
+                                 for result, summary, rank in zip(search_results, summaries, ranks)])
         
         prompt_template = """
         Answer the question based on the following web search results:
@@ -325,6 +350,48 @@ def update_vectors(files, use_recursive_splitter):
     
     return f"Vector store updated successfully. Processed {total_chunks} chunks from {len(files)} files."
 
+def update_vector_db_with_search_results(search_results, summaries, ranks):
+    embed = get_embeddings()
+    database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True) if os.path.exists("faiss_database") else FAISS.from_documents([], embed)
+    
+    current_date = datetime.now().strftime("%Y-%m-%d")
+    
+    for result, summary, rank in zip(search_results, summaries, ranks):
+        doc = Document(
+            page_content=summary,
+            metadata={
+                "search_date": current_date,
+                "search_title": result["title"],
+                "search_content": result["text"],
+                "search_summary": summary,
+                "rank": rank
+            }
+        )
+        database.add_documents([doc])
+    
+    database.save_local("faiss_database")
+
+def export_vector_db_to_excel():
+    embed = get_embeddings()
+    database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
+    
+    documents = database.docstore._dict.values()
+    data = [{
+        "Search Date": doc.metadata["search_date"],
+        "Search Title": doc.metadata["search_title"],
+        "Search Content": doc.metadata["search_content"],
+        "Search Summary": doc.metadata["search_summary"],
+        "Rank": doc.metadata["rank"]
+    } for doc in documents]
+    
+    df = pd.DataFrame(data)
+    
+    with NamedTemporaryFile(delete=False, suffix='.xlsx') as tmp:
+        excel_path = tmp.name
+        df.to_excel(excel_path, index=False)
+    
+    return excel_path
+    
 def extract_db_to_excel():
     embed = get_embeddings()
     database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
@@ -384,6 +451,10 @@ with gr.Blocks() as demo:
     
     submit_button.click(chat, inputs=[question_input, chatbot, temperature_slider, top_p_slider, repetition_penalty_slider, web_search_checkbox], outputs=[question_input, chatbot])
     
+    export_vector_db_button = gr.Button("Export Vector DB to Excel")
+    vector_db_excel_output = gr.File(label="Download Vector DB Excel File")
+    export_vector_db_button.click(export_vector_db_to_excel, inputs=[], outputs=vector_db_excel_output)
+
     extract_button = gr.Button("Extract Database to Excel")
     excel_output = gr.File(label="Download Excel File")
     extract_button.click(extract_db_to_excel, inputs=[], outputs=excel_output)