Update app.py

app.py

@@ -28,6 +28,7 @@ from groq import Groq
 import os
 from mistralai import Mistral
 from dotenv import load_dotenv
+import re
 
 # Automatically get the current year
 current_year = datetime.datetime.now().year
@@ -222,7 +223,13 @@ Rephrased query:
         logger.error(f"Error rephrasing query with LLM: {e}")
         return query  # Fallback to original query if rephrasing fails
 
-def rerank_documents(query, documents, similarity_threshold=0.95, max_results=5):
+def extract_entity_domain(query):
+    # Use a simple regex pattern to extract domain names from the query
+    domain_pattern = r'\b(?:https?://)?(?:www\.)?([a-zA-Z0-9-]+(?:\.[a-zA-Z0-9-]+)+)\b'
+    matches = re.findall(domain_pattern, query)
+    return matches[0] if matches else None
+
+def rerank_documents_with_priority(query, documents, entity_domain, similarity_threshold=0.95, max_results=5):
     try:
         # Step 1: Encode the query and document summaries
         query_embedding = similarity_model.encode(query, convert_to_tensor=True)
@@ -240,8 +247,8 @@ def rerank_documents(query, documents, similarity_threshold=0.95, max_results=5)
         # Combine documents and cosine scores
         scored_documents = list(zip(documents, cosine_scores))
 
-        # Step 3: Sort documents by cosine similarity score
-        scored_documents.sort(key=lambda x: x[1], reverse=True)
+        # Step 3: Sort documents by cosine similarity score and prioritize entity domain
+        scored_documents.sort(key=lambda x: (not x[0]['is_entity_domain'], -x[1]), reverse=False)
         
         # Step 4: Filter out similar documents
         filtered_docs = []
@@ -428,7 +435,11 @@ def search_and_scrape(query, chat_history, num_results=5, max_chars=3000, time_r
             logger.info("No need to perform search based on the rephrased query.")
             return "No search needed for the provided input."
 
-        # Step 2: Perform search
+        # Step 2: Extract entity domain
+        entity_domain = extract_entity_domain(rephrased_query)
+        logger.info(f"Extracted entity domain: {entity_domain}")
+
+        # Step 3: Perform search
         # Search query parameters
         params = {
             'q': rephrased_query,
@@ -534,7 +545,7 @@ def search_and_scrape(query, chat_history, num_results=5, max_chars=3000, time_r
 
         logger.info(f"Successfully scraped {len(scraped_content)} documents.")
 
-        # Step 3: Assess relevance, summarize, and check for uniqueness
+         # Step 4: Assess relevance, summarize, and check for uniqueness
         relevant_documents = []
         unique_summaries = []
         for doc in scraped_content:
@@ -545,11 +556,14 @@ def search_and_scrape(query, chat_history, num_results=5, max_chars=3000, time_r
                 summary_text = summary.replace("Summary: ", "").strip()
                 
                 if is_content_unique(summary_text, unique_summaries):
+                    doc_domain = urlparse(doc['url']).netloc
+                    is_entity_domain = doc_domain == entity_domain
                     relevant_documents.append({
                         "title": doc['title'],
                         "url": doc['url'],
                         "summary": summary_text,
-                        "scraper": doc['scraper']
+                        "scraper": doc['scraper'],
+                        "is_entity_domain": is_entity_domain
                     })
                     unique_summaries.append(summary_text)
                 else:
@@ -559,8 +573,8 @@ def search_and_scrape(query, chat_history, num_results=5, max_chars=3000, time_r
             logger.warning("No relevant and unique documents found.")
             return "No relevant and unique financial news found for the given query."
 
-        # Step 4: Rerank documents based on similarity to query
-        reranked_docs = rerank_documents(rephrased_query, relevant_documents, similarity_threshold=0.95, max_results=num_results)
+        # Step 5: Rerank documents based on similarity to query and prioritize entity domain
+        reranked_docs = rerank_documents_with_priority(rephrased_query, relevant_documents, entity_domain, similarity_threshold=0.95, max_results=num_results)
         
         if not reranked_docs:
             logger.warning("No documents remained after reranking.")
@@ -595,7 +609,6 @@ def search_and_scrape(query, chat_history, num_results=5, max_chars=3000, time_r
         logger.error(f"Unexpected error in search_and_scrape: {e}")
         return f"An unexpected error occurred during the search and scrape process: {e}"
 
-
 def chat_function(message, history, num_results, max_chars, time_range, language, category, engines, safesearch, method, llm_temperature, model, use_pydf2):
     chat_history = "\n".join([f"{role}: {msg}" for role, msg in history])