Update app.py

app.py

@@ -26,7 +26,7 @@ from concurrent.futures import ThreadPoolExecutor, as_completed
 import tiktoken
 
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
-logger = logging.getLogger(name)
+logger = logging.getLogger(__name__)
 
 HF_API_KEY = os.environ.get("HF_API_KEY")
 if not HF_API_KEY:
@@ -37,25 +37,20 @@ client = InferenceClient(provider="hf-inference", api_key=HF_API_KEY)
 MAIN_LLM_MODEL = "mistralai/Mistral-Nemo-Instruct-2407"
 REASONING_LLM_MODEL = "deepseek-ai/DeepSeek-R1-Distill-Qwen-32B"
 CRITIC_LLM_MODEL = "Qwen/QwQ-32B-Preview"
-SPECIALIST_MODELS = {
-    "medical": "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B",
-    "scientific": "mistralai/Mistral-7B-Instruct-v0.2",
-    "financial": "mistralai/Mistral-7B-Instruct-v0.3",
-    "legal": "Qwen/Qwen2.5-Coder-32B-Instruct"
-}
-ENSEMBLE_MODELS = [MAIN_LLM_MODEL, REASONING_LLM_MODEL, CRITIC_LLM_MODEL] + list(SPECIALIST_MODELS.values())
-
-MAX_ITERATIONS = 100
-TIMEOUT = 300
-RETRY_DELAY = 15
-NUM_RESULTS = 50
-SIMILARITY_THRESHOLD = 0.12
-MAX_CONTEXT_ITEMS = 100
-MAX_HISTORY_ITEMS = 25
-MAX_FULL_TEXT_LENGTH = 50000
+ENSEMBLE_MODELS = [MAIN_LLM_MODEL, REASONING_LLM_MODEL, CRITIC_LLM_MODEL]  # Keep, but expand upon.
+
+MAX_ITERATIONS = 40  # Increased for deeper research.
+TIMEOUT = 180 # Longer timeout for larger models / complex tasks.
+RETRY_DELAY = 10 # longer delay
+NUM_RESULTS = 20
+SIMILARITY_THRESHOLD = 0.15
+MAX_CONTEXT_ITEMS = 50  # Increased context window.
+MAX_HISTORY_ITEMS = 12
+MAX_FULL_TEXT_LENGTH = 20000 # larger document size
 FAISS_INDEX_PATH = "research_index.faiss"
 RESEARCH_DATA_PATH = "research_data.pkl"
-PAPER_SUMMARIES_PATH = "paper_summaries.pkl"
+PAPER_SUMMARIES_PATH = "paper_summaries.pkl" #New path for storing paper summary
+
 
 try:
     main_similarity_model = SentenceTransformer('sentence-transformers/all-mpnet-base-v2')
@@ -67,9 +62,8 @@ try:
         index = faiss.read_index(FAISS_INDEX_PATH)
         logger.info(f"Loaded FAISS index from {FAISS_INDEX_PATH}")
     else:
-        index = faiss.IndexFlatIP(embedding_dim)
+        index = faiss.IndexFlatIP(embedding_dim)  # Use IndexFlatIP for inner product (cosine similarity).
         logger.info("Created a new FAISS index.")
-
 except Exception as e:
     logger.error(f"Failed to load models or initialize FAISS: {e}")
     raise
@@ -95,8 +89,8 @@ def load_research_data():
         try:
             with open(RESEARCH_DATA_PATH, "rb") as f:
                 data = pickle.load(f)
-            logger.info(f"Loaded research data from {RESEARCH_DATA_PATH}")
-            return data
+                logger.info(f"Loaded research data from {RESEARCH_DATA_PATH}")
+                return data
         except Exception as e:
             logger.error(f"Error loading research data: {e}")
             return {}
@@ -117,8 +111,8 @@ def load_paper_summaries() -> Dict[str, str]:
         try:
             with open(PAPER_SUMMARIES_PATH, "rb") as f:
                 data = pickle.load(f)
-            logger.info(f"Loaded paper summaries from {PAPER_SUMMARIES_PATH}")
-            return data
+                logger.info(f"Loaded paper summaries from {PAPER_SUMMARIES_PATH}")
+                return data
         except Exception as e:
             logger.error(f"Error loading paper summaries: {e}")
             return {}
@@ -126,7 +120,8 @@ def load_paper_summaries() -> Dict[str, str]:
         logger.info("No existing paper summaries found.")
         return {}
 
-def hf_inference(model_name, prompt, max_tokens=2000, retries=5, stream=False):
+
+def hf_inference(model_name, prompt, max_tokens=2000, retries=5, stream=False):  # Added stream parameter
     for attempt in range(retries):
         try:
             messages = [{"role": "user", "content": prompt}]
@@ -134,24 +129,25 @@ def hf_inference(model_name, prompt, max_tokens=2000, retries=5, stream=False):
                 model=model_name,
                 messages=messages,
                 max_tokens=max_tokens,
-                stream=stream
+                stream=stream  # Pass the stream parameter
             )
             if stream:
-                return response_generator
+                return response_generator # Return the generator directly
             else:
-                response = next(response_generator)
+                # If not streaming, get the full response
+                response = next(response_generator) # Consume the first chunk to get complete object
                 return {"generated_text": response.choices[0].message.content}
         except Exception as e:
             if attempt == retries - 1:
                 logger.error(f"Request failed after {retries} retries: {e}")
                 return {"error": f"Request failed after {retries} retries: {e}"}
             time.sleep(RETRY_DELAY * (1 + attempt))
-            return {"error": "Request failed after multiple retries."}
+    return {"error": "Request failed after multiple retries."}
 
-def ensemble_inference(prompt, models=ENSEMBLE_MODELS, max_tokens=1500, stream=False):
+def ensemble_inference(prompt, models=ENSEMBLE_MODELS, max_tokens=1500, stream=False): #Added stream
     results = []
 
-    if stream:
+    if stream: # If streaming, return a generator that yields from each model
         def generate_responses():
             with ThreadPoolExecutor(max_workers=len(models)) as executor:
                 futures = {executor.submit(hf_inference, model, prompt, max_tokens, stream=True): model for model in models}
@@ -159,14 +155,14 @@ def ensemble_inference(prompt, models=ENSEMBLE_MODELS, max_tokens=1500, stream=F
                 for future in as_completed(futures):
                     model = future_to_model[future]
                     try:
-                        for chunk in future.result():
-                            yield {"model": model, "text": chunk.choices[0].delta.content}
+                        for chunk in future.result():  # Iterate through chunks
+                            yield {"model": model, "text": chunk.choices[0].delta.content} #yield the content of the chunk
                     except Exception as e:
                         logger.error(f"Error with model {model}: {e}")
                         yield {"model": model, "text": f"Error: {e}"}
-        return generate_responses()
+        return generate_responses() # return the generator
 
-    else:
+    else: #Non-streaming behavior
         with ThreadPoolExecutor(max_workers=len(models)) as executor:
             future_to_model = {executor.submit(hf_inference, model, prompt, max_tokens, stream=False): model for model in models}
             for future in as_completed(future_to_model):
@@ -188,14 +184,14 @@ def ensemble_inference(prompt, models=ENSEMBLE_MODELS, max_tokens=1500, stream=F
         for result in results:
             synthesis_prompt += f"Expert {results.index(result) + 1} ({result['model'].split('/')[-1]}):\n{result['text']}\n\n"
 
-        synthesis = hf_inference(MAIN_LLM_MODEL, synthesis_prompt)
+        synthesis = hf_inference(MAIN_LLM_MODEL, synthesis_prompt) # Use a consistent model for final synthesis
         if "generated_text" in synthesis:
             return synthesis
         else:
-            return {"generated_text": max(results, key=lambda x: len(x["text"]))["text"]}
+            return {"generated_text": max(results, key=lambda x: len(x["text"]))["text"]} # Fallback
 
 def tool_search_web(query: str, num_results: int = NUM_RESULTS, safesearch: str = "moderate",
-                    time_filter: Optional[str] = None, region: str = "wt-wt", language: str = "en-us") -> list:
+                   time_filter: Optional[str] = None, region: str = "wt-wt", language: str = "en-us") -> list:
     try:
         with DDGS() as ddgs:
             kwargs = {
@@ -293,7 +289,7 @@ def tool_search_wikipedia(query: str, max_results: int = 3) -> list:
     except Exception as e:
         logger.error(f"Wikipedia search error: {e}")
         return []
-
+    
 def tool_search_scholar(query: str, max_results: int = 5) -> list:
     try:
         search_query = scholarly.search_pubs(query)
@@ -330,7 +326,7 @@ def extract_article_content(url: str) -> str:
         return ""
 
 def tool_reason(prompt: str, search_results: list, reasoning_context: list = [],
-              critique: str = "", focus_areas: list = []) -> str:
+               critique: str = "", focus_areas: list = []) -> str:
     if not search_results:
         return "No search results to reason about."
 
@@ -342,7 +338,7 @@ def tool_reason(prompt: str, search_results: list, reasoning_context: list = [],
 
     results_by_source = {}
     for i, result in enumerate(search_results):
-        source = result.get('source', 'Web Search')
+        source = result.get('source', 'Web Search') # Default to 'Web Search'
         if source not in results_by_source:
             results_by_source[source] = []
         results_by_source[source].append((i, result))
@@ -358,7 +354,7 @@ def tool_reason(prompt: str, search_results: list, reasoning_context: list = [],
             reasoning_input += "\n"
 
     if reasoning_context:
-        recent_context = reasoning_context[-MAX_HISTORY_ITEMS:]
+        recent_context = reasoning_context[-MAX_HISTORY_ITEMS:]  # Limit history
         reasoning_input += "\nPrevious Reasoning Context:\n" + "\n".join(recent_context)
 
     if critique:
@@ -366,7 +362,7 @@ def tool_reason(prompt: str, search_results: list, reasoning_context: list = [],
 
     reasoning_input += "\nProvide a thorough, nuanced analysis that builds upon previous reasoning if applicable. Consider multiple perspectives, potential contradictions in the search results, and the reliability of different sources.  Address any specific critiques."
 
-    reasoning_output = ensemble_inference(reasoning_input)
+    reasoning_output = ensemble_inference(reasoning_input) # Use ensemble for high-quality reasoning.
 
     if isinstance(reasoning_output, dict) and "generated_text" in reasoning_output:
         return reasoning_output["generated_text"].strip()
@@ -380,7 +376,7 @@ def tool_summarize(insights: list, prompt: str, contradictions: list = []) -> st
 
     summarization_input = f"Synthesize the following insights into a cohesive and comprehensive summary regarding: '{prompt}'\n\n"
 
-    max_tokens = 12000
+    max_tokens = 12000  # Increased token limit
     selected_insights = []
     token_count = get_token_count(summarization_input) + get_token_count("\n\n".join(contradictions))
 
@@ -408,7 +404,7 @@ def tool_summarize(insights: list, prompt: str, contradictions: list = []) -> st
         return "Could not generate a summary due to an error."
 
 def tool_generate_search_query(prompt: str, previous_queries: list = [],
-                               failed_queries: list = [], focus_areas: list = []) -> str:
+                              failed_queries: list = [], focus_areas: list = []) -> str:
     query_gen_input = f"Generate an effective search query for the following prompt: {prompt}\n"
 
     if previous_queries:
@@ -433,7 +429,7 @@ def tool_generate_search_query(prompt: str, previous_queries: list = [],
     return ""
 
 def tool_critique_reasoning(reasoning_output: str, prompt: str,
-                            previous_critiques: list = []) -> str:
+                           previous_critiques: list = []) -> str:
     critique_input = f"Critically evaluate the following reasoning output in relation to the prompt:\n\nPrompt: {prompt}\n\nReasoning: {reasoning_output}\n\n"
 
     if previous_critiques:
@@ -441,7 +437,7 @@ def tool_critique_reasoning(reasoning_output: str, prompt: str,
 
     critique_input += "Identify any flaws, biases, logical fallacies, unsupported claims, or areas for improvement. Be specific and constructive. Suggest concrete ways to enhance the reasoning. Also evaluate the strength of evidence and whether conclusions are proportionate to the available information."
 
-    critique_output = hf_inference(CRITIC_LLM_MODEL, critique_input)
+    critique_output = hf_inference(CRITIC_LLM_MODEL, critique_input) # Use specialized critique model.
 
     if isinstance(critique_output, dict) and "generated_text" in critique_output:
         return critique_output["generated_text"].strip()
@@ -453,7 +449,7 @@ def tool_identify_contradictions(insights: list) -> list:
     if len(insights) < 2:
         return []
 
-    max_tokens = 12000
+    max_tokens = 12000  # Increased token limit for potentially more contradictions
     selected_insights = []
     token_count = 0
 
@@ -468,12 +464,13 @@ def tool_identify_contradictions(insights: list) -> list:
     contradiction_input = "Identify specific contradictions in these insights:\n\n" + "\n\n".join(selected_insights)
     contradiction_input += "\n\nList each contradiction as a separate numbered point. For each contradiction, cite the specific claims that are in tension and evaluate which claim is better supported. If no contradictions exist, respond with 'No contradictions found.'"
 
-    contradiction_output = hf_inference(CRITIC_LLM_MODEL, contradiction_input)
+    contradiction_output = hf_inference(CRITIC_LLM_MODEL, contradiction_input)  # Use critique model
 
     if isinstance(contradiction_output, dict) and "generated_text" in contradiction_output:
         result = contradiction_output["generated_text"].strip()
         if result == "No contradictions found.":
             return []
+        # More robust contradiction extraction, handles multi-sentence contradictions
         contradictions = re.findall(r'\d+\.\s+(.*?)(?=\d+\.|$)', result, re.DOTALL)
         return [c.strip() for c in contradictions if c.strip()]
 
@@ -481,7 +478,7 @@ def tool_identify_contradictions(insights: list) -> list:
     return []
 
 def tool_identify_focus_areas(prompt: str, insights: list = [],
-                               failed_areas: list = []) -> list:
+                             failed_areas: list = []) -> list:
     focus_input = f"Based on this research prompt: '{prompt}'\n\n"
 
     if insights:
@@ -493,10 +490,11 @@ def tool_identify_focus_areas(prompt: str, insights: list = [],
 
     focus_input += "Identify 3-5 specific aspects that should be investigated further to get a complete understanding. Be precise and prioritize underexplored areas. For each suggested area, briefly explain why it's important to investigate."
 
-    focus_output = hf_inference(MAIN_LLM_MODEL, focus_input)
+    focus_output = hf_inference(MAIN_LLM_MODEL, focus_input)  # Consistent model
 
     if isinstance(focus_output, dict) and "generated_text" in focus_output:
         result = focus_output["generated_text"].strip()
+        # More robust extraction, handles different list formats
         areas = re.findall(r'(?:^|\n)(?:\d+\.|\*|\-)\s*(.*?)(?=(?:\n(?:\d+\.|\*|\-|$))|$)', result)
         return [area.strip() for area in areas if area.strip()][:5]
 
@@ -509,7 +507,7 @@ def add_to_faiss_index(text: str):
     if embedding_np.shape[1] != embedding_dim:
         logger.error(f"Embedding dimension mismatch: expected {embedding_dim}, got {embedding_np.shape[1]}")
         return
-    faiss.normalize_L2(embedding_np)
+    faiss.normalize_L2(embedding_np)  # Normalize for cosine similarity.
     index.add(embedding_np)
 
 def search_faiss_index(query: str, top_k: int = 5) -> List[str]:
@@ -534,7 +532,7 @@ def filter_results(search_results, prompt, previous_snippets=None):
         for result in search_results:
             combined_text = result['title'] + " " + result['snippet']
 
-            if result['snippet'] in seen_snippets:
+            if result['snippet'] in seen_snippets:  # Prevent exact duplicates
                 continue
 
             result_embedding = main_similarity_model.encode(combined_text, convert_to_tensor=True)
@@ -543,16 +541,16 @@ def filter_results(search_results, prompt, previous_snippets=None):
             if cosine_score >= SIMILARITY_THRESHOLD:
                 result['relevance_score'] = cosine_score
                 filtered_results.append(result)
-                seen_snippets.add(result['snippet'])
+                seen_snippets.add(result['snippet']) # Add snippets after filtering
                 add_to_faiss_index(result['snippet'])
 
 
-        filtered_results.sort(key=lambda x: x.get('relevance_score', 0), reverse=True)
+        filtered_results.sort(key=lambda x: x.get('relevance_score', 0), reverse=True)  # Sort by relevance.
         return filtered_results
 
     except Exception as e:
         logger.error(f"Error during filtering: {e}")
-        return search_results
+        return search_results  # Return original results on error.
 
 def tool_extract_key_entities(prompt: str) -> list:
     entity_input = f"Extract the key entities (people, organizations, concepts, technologies, events, time periods, locations, etc.) from this research prompt that should be investigated individually:\n\n{prompt}\n\nList the 5-7 most important entities, one per line, with a brief explanation (2-3 sentences) of why each is central to the research question."
@@ -562,7 +560,7 @@ def tool_extract_key_entities(prompt: str) -> list:
     if isinstance(entity_output, dict) and "generated_text" in entity_output:
         result = entity_output["generated_text"].strip()
         entities = [e.strip() for e in result.split('\n') if e.strip()]
-        return entities[:7]
+        return entities[:7] # Limit to top 7 entities
 
     logger.error(f"Failed to extract key entities: {entity_output}")
     return []
@@ -575,11 +573,11 @@ def tool_meta_analyze(entity_insights: Dict[str, list], prompt: str) -> str:
 
     for entity, insights in entity_insights.items():
         if insights:
-            meta_input += f"\n--- {entity} ---\n" + insights[-1] + "\n"
+            meta_input += f"\n--- {entity} ---\n" + insights[-1] + "\n"  # Most recent insight for each entity
 
     meta_input += "\nProvide a high-level synthesis that identifies:\n1. Common themes across entities\n2. Important differences and contradictions\n3. How these entities interact or influence each other\n4. The broader implications for the original research question\n5. A systems-level understanding of how these elements fit together"
 
-    meta_output = ensemble_inference(meta_input)
+    meta_output = ensemble_inference(meta_input)  # Ensemble for meta-analysis
 
     if isinstance(meta_output, dict) and "generated_text" in meta_output:
         return meta_output["generated_text"].strip()
@@ -604,7 +602,7 @@ def tool_draft_research_plan(prompt: str, entities: list, focus_areas: list = []
     plan_input += "5. Potential challenges and how to address them\n"
     plan_input += "6. Criteria for evaluating the quality of findings"
 
-    plan_output = hf_inference(REASONING_LLM_MODEL, plan_input)
+    plan_output = hf_inference(REASONING_LLM_MODEL, plan_input) # Use reasoning model
 
     if isinstance(plan_output, dict) and "generated_text" in plan_output:
         return plan_output["generated_text"].strip()
@@ -616,40 +614,27 @@ def tool_extract_article(url: str) -> str:
     extracted_text = extract_article_content(url)
     return extracted_text if extracted_text else f"Could not extract content from {url}"
 
+# New tool for summarizing a single paper
 def tool_summarize_paper(paper_text: str) -> str:
-    summarization_prompt = f"""Summarize this academic paper, focusing on the following:
-
-Main Research Question(s): What questions does the paper address?
-
-Methodology: Briefly describe the methods used (e.g., experiments, surveys, simulations, theoretical analysis).
+  summarization_prompt = f"""Summarize this academic paper, focusing on the following:
 
-Key Findings: What are the most important results or conclusions?
+1.  **Main Research Question(s):** What questions does the paper address?
+2.  **Methodology:** Briefly describe the methods used (e.g., experiments, surveys, simulations, theoretical analysis).
+3.  **Key Findings:** What are the most important results or conclusions?
+4.  **Limitations:**  What are the acknowledged limitations of the study?
+5.  **Implications:** What are the broader implications of the findings, according to the authors?
 
-Limitations:  What are the acknowledged limitations of the study?
-
-Implications: What are the broader implications of the findings, according to the authors?
 Paper Text:
 {paper_text[:MAX_FULL_TEXT_LENGTH]}
-"""
-    summary = hf_inference(REASONING_LLM_MODEL, summarization_prompt, max_tokens=500)
-
-    if isinstance(summary, dict) and "generated_text" in summary:
-        return summary["generated_text"].strip()
-    else:
-        logger.error(f"Failed to generate summary: {summary}")
-        return "Could not generate a summary due to an error."
+"""  # Truncate if necessary
+  summary = hf_inference(REASONING_LLM_MODEL, summarization_prompt, max_tokens=500)
 
-def tool_search_patents(query: str, max_results: int = 10) -> list:
-    pass
+  if isinstance(summary, dict) and "generated_text" in summary:
+      return summary["generated_text"].strip()
+  else:
+      logger.error(f"Failed to generate summary: {summary}")
+      return "Could not generate a summary due to an error."
 
-def tool_search_clinical_trials(query: str, max_results: int = 10) -> list:
-    pass
-
-def tool_search_datasets(query: str, max_results: int = 10) -> list:
-    pass
-
-def tool_search_conferences(query: str, max_results: int = 10) -> list:
-    pass
 
 tools = {
     "search_web": {
@@ -687,7 +672,7 @@ tools = {
             "max_results": {"type": "integer", "description": "Maximum number of articles to return."}
         },
     },
-    "search_scholar": {
+    "    "search_scholar": {
         "function": tool_search_scholar,
         "description": "Searches Google Scholar for academic publications.",
         "parameters": {
@@ -702,7 +687,7 @@ tools = {
             "url": {"type": "string", "description": "The URL of the article to extract"}
         },
     },
-    "summarize_paper": {
+        "summarize_paper": {
         "function": tool_summarize_paper,
         "description": "Summarizes the content of an academic paper.",
         "parameters": {
@@ -753,7 +738,7 @@ tools = {
         "description": "Identifies contradictions across multiple insights.",
         "parameters": {
             "insights": {"type": "array", "description": "Collection of insights to analyze for contradictions."},
-        },
+                },
     },
     "identify_focus_areas": {
         "function": tool_identify_focus_areas,
@@ -787,38 +772,6 @@ tools = {
             "entities": {"type": "array", "description": "Key entities to investigate."},
             "focus_areas": {"type": "array", "description": "Additional areas to focus on."}
         }
-    },
-    "search_patents": {
-        "function": tool_search_patents,
-        "description": "Searches patent databases globally",
-        "parameters": {
-            "query": {"type": "string", "description": "Patent search query"},
-            "max_results": {"type": "integer", "description": "Maximum number of patents to return"}
-        }
-    },
-    "search_clinical_trials": {
-        "function": tool_search_clinical_trials,
-        "description": "Search ClinicalTrials.gov and WHO ICTRP",
-        "parameters": {
-            "query": {"type": "string", "description": "Search query for ClinicalTrials.gov and WHO ICTRP"},
-            "max_results": {"type": "integer", "description": "Maximum number of results to return"}
-        }
-    },
-    "search_datasets": {
-        "function": tool_search_datasets,
-        "description": "Search academic datasets from repositories like Kaggle, UCI, etc.",
-        "parameters": {
-            "query": {"type": "string", "description": "Search query for academic datasets"},
-            "max_results": {"type": "integer", "description": "Maximum number of results to return"}
-        }
-    },
-    "search_conferences": {
-        "function": tool_search_conferences,
-        "description": "Search major conference proceedings",
-        "parameters": {
-            "query": {"type": "string", "description": "Search query for conference proceedings"},
-            "max_results": {"type": "integer", "description": "Maximum number of results to return"}
-        }
     }
 }
 
@@ -845,18 +798,12 @@ Instructions:
 Select the BEST tool and parameters for the current research stage. Output valid JSON. If no tool is appropriate, respond with {}.
 Only use provided tools. Be strategic about which tool to use next based on the research progress so far.
 You MUST be methodical.  Think step-by-step:
-
-Plan: If it's the very beginning, extract key entities, identify focus areas, and then draft a research plan.
-
-Search:  Use a variety of search tools.  Start with broad searches, then narrow down.  Use specific search tools (arXiv, PubMed, Scholar) for relevant topics.
-
-Analyze:  Reason deeply about search results, and critique your reasoning.  Identify contradictions. Filter and use FAISS index for relevant information.
-
-Refine: If results are poor, generate better search queries. Adjust focus areas.
-
-Iterate: Repeat steps 2-4, focusing on different entities and aspects.
-
-Synthesize:  Finally, summarize the findings, addressing contradictions.
+1.  **Plan:** If it's the very beginning, extract key entities, identify focus areas, and then draft a research plan.
+2.  **Search:**  Use a variety of search tools.  Start with broad searches, then narrow down.  Use specific search tools (arXiv, PubMed, Scholar) for relevant topics.
+3.  **Analyze:**  Reason deeply about search results, and critique your reasoning.  Identify contradictions. Filter and use FAISS index for relevant information.
+4.  **Refine:** If results are poor, generate *better* search queries. Adjust focus areas.
+5.  **Iterate:** Repeat steps 2-4, focusing on different entities and aspects.
+6.  **Synthesize:**  Finally, summarize the findings, addressing contradictions.
 Example:
 {"tool": "search_web", "parameters": {"query": "Eiffel Tower location"}}
 Output:
@@ -866,7 +813,7 @@ Output:
 def deep_research(prompt):
     task_description = "You are an advanced research assistant, designed to be as comprehensive as possible.  Use available tools iteratively, focus on different aspects, explore promising leads thoroughly, critically evaluate your findings, and build up a comprehensive understanding of the research topic. Utilize the FAISS index to avoid redundant searches and to build a persistent knowledge base."
     research_data = load_research_data()
-    paper_summaries = load_paper_summaries()
+    paper_summaries = load_paper_summaries()  # Load paper summaries
 
     context = research_data.get('context', [])
     all_insights = research_data.get('all_insights', [])
@@ -881,59 +828,62 @@ def deep_research(prompt):
     seen_snippets = set(research_data.get('seen_snippets', []))
     contradictions = research_data.get('contradictions', [])
     research_session_id = research_data.get('research_session_id', str(uuid4()))
-
+    
     global index
     if research_data:
-        logger.info("Restoring FAISS Index from loaded data.")
+         logger.info("Restoring FAISS Index from loaded data.")
     else:
         index.reset()
         logger.info("Initialized a fresh FAISS Index")
 
     key_entities_with_descriptions = tool_extract_key_entities(prompt=prompt)
-    key_entities = [e.split(":")[0].strip() for e in key_entities_with_descriptions]
+    key_entities = [e.split(":")[0].strip() for e in key_entities_with_descriptions] # Extract just entity names
     if key_entities:
         context.append(f"Identified key entities: {key_entities}")
         intermediate_output += f"Identified key entities for focused research: {key_entities_with_descriptions}\n"
         yield "Identifying key entities... (Completed)"
 
+    # Initialize progress tracking for each entity.
     entity_progress = {entity: {'queries': [], 'insights': []} for entity in key_entities}
-    entity_progress['general'] = {'queries': [], 'insights': []}
+    entity_progress['general'] = {'queries': [], 'insights': []}  # For general, non-entity-specific searches
     for entity in key_entities + ['general']:
-        if entity in research_data:
+        if entity in research_data:  # Load existing progress
             entity_progress[entity]['queries'] = research_data[entity]['queries']
             entity_progress[entity]['insights'] = research_data[entity]['insights']
 
-    if not focus_areas:
+    if not focus_areas:  # Corrected placement: outside the loop
         initial_focus_areas = tool_identify_focus_areas(prompt=prompt)
         yield "Identifying initial focus areas...(Completed)"
         research_plan = tool_draft_research_plan(prompt=prompt, entities=key_entities, focus_areas=initial_focus_areas)
         yield "Drafting initial research plan...(Completed)"
-        context.append(f"Initial Research Plan: {research_plan[:200]}...")
+        context.append(f"Initial Research Plan: {research_plan[:200]}...") # Add plan to context
         intermediate_output += f"Initial Research Plan:\n{research_plan}\n\n"
         focus_areas = initial_focus_areas
 
 
     for i in range(MAX_ITERATIONS):
-        if key_entities and i > 0:
-            entities_to_process = key_entities + ['general']
+        # Entity-focused iteration strategy
+        if key_entities and i > 0:  # Cycle through entities *after* initial setup
+            entities_to_process = key_entities + ['general']  # Include 'general' for broad searches
             current_entity = entities_to_process[i % len(entities_to_process)]
         else:
-            current_entity = 'general'
+            current_entity = 'general'  # Start with general research.
 
         context.append(f"Current focus: {current_entity}")
         
-        if i > 0:
+        # FAISS Retrieval
+        if i > 0:  # Use FAISS *after* the first iteration (once we have data)
             faiss_results_indices = search_faiss_index(prompt if current_entity == 'general' else f"{prompt} {current_entity}")
             faiss_context = []
             for idx in faiss_results_indices:
-                if idx < len(all_insights):
+                if idx < len(all_insights):  # Check index bounds
                     faiss_context.append(f"Previously found insight: {all_insights[idx]}")
             if faiss_context:
-                context.extend(faiss_context)
+                context.extend(faiss_context) # Add FAISS context
                 intermediate_output += f"Iteration {i+1} - Retrieved {len(faiss_context)} relevant items from FAISS index.\n"
-            
+                
 
-        if i == 0:
+        if i == 0: #Initial broad search
             initial_query = tool_generate_search_query(prompt=prompt)
             yield f"Generating initial search query... (Iteration {i+1})"
             if initial_query:
@@ -996,68 +946,69 @@ def deep_research(prompt):
                     for future in as_completed(futures):
                         search_results.extend(future.result())
 
-            yield f"Searching for information on entity: {current_entity}... (Iteration {i+1})"
-            filtered_search_results = filter_results(search_results,
+                yield f"Searching for information on entity: {current_entity}... (Iteration {i+1})"
+                filtered_search_results = filter_results(search_results,
                                                         f"{prompt} {current_entity}",
-                                                        previous_snippets=seen_snippets)
-
-            if filtered_search_results:
-                context.append(f"Entity Search for {current_entity}: {len(filtered_search_results)} results")
-
-                entity_reasoning = tool_reason(
-                    prompt=f"{prompt} focusing on {current_entity}",
-                    search_results=filtered_search_results,
-                    reasoning_context=entity_progress[current_entity]['insights'],
-                    focus_areas=focus_areas
-                )
-                yield f"Reasoning about entity: {current_entity}... (Iteration {i+1})"
+                                                        previous_snippets=seen_snippets) # Pass existing snippets
 
-                if entity_reasoning:
-                    all_insights.append(entity_reasoning)
-                    entity_progress[current_entity]['insights'].append(entity_reasoning)
-
-                    if current_entity not in entity_specific_insights:
-                        entity_specific_insights[current_entity] = []
-                    entity_specific_insights[current_entity].append(entity_reasoning)
-
-                    context.append(f"Reasoning about {current_entity}: {entity_reasoning[:200]}...")
-                    add_to_faiss_index(entity_reasoning)
-            else:
-                failed_queries.append(entity_query)
-                context.append(f"Entity query for {current_entity} yielded no relevant results")
+                if filtered_search_results:
+                    context.append(f"Entity Search for {current_entity}: {len(filtered_search_results)} results")
+
+                    entity_reasoning = tool_reason(
+                        prompt=f"{prompt} focusing on {current_entity}",
+                        search_results=filtered_search_results,
+                        reasoning_context=entity_progress[current_entity]['insights'], # Use entity-specific context
+                        focus_areas=focus_areas
+                    )
+                    yield f"Reasoning about entity: {current_entity}... (Iteration {i+1})"
+
+                    if entity_reasoning:
+                        all_insights.append(entity_reasoning)
+                        entity_progress[current_entity]['insights'].append(entity_reasoning)
+
+                        if current_entity not in entity_specific_insights:
+                            entity_specific_insights[current_entity] = []
+                        entity_specific_insights[current_entity].append(entity_reasoning)
+
+                        context.append(f"Reasoning about {current_entity}: {entity_reasoning[:200]}...")
+                        add_to_faiss_index(entity_reasoning)
+                else:
+                    failed_queries.append(entity_query)
+                    context.append(f"Entity query for {current_entity} yielded no relevant results")
 
         llm_prompt = create_prompt(task_description, prompt, tools, context)
-        llm_response = hf_inference(MAIN_LLM_MODEL, llm_prompt, stream=True)
+        llm_response = hf_inference(MAIN_LLM_MODEL, llm_prompt, stream=True)  # Use streaming
 
         if isinstance(llm_response, dict) and "error" in llm_response:
             intermediate_output += f"LLM Error: {llm_response['error']}\n"
-            yield f"LLM Error (Iteration {i+1}): {llm_response['error']}"
+            yield f"LLM Error (Iteration {i+1}): {llm_response['error']}"  # Display error in output
             continue
 
+        # Process streaming response
         response_text = ""
         try:
             for chunk in llm_response:
                 if chunk.choices and chunk.choices[0].delta and chunk.choices[0].delta.content:
                     response_text += chunk.choices[0].delta.content
-                    yield f"Iteration {i+1} - Thinking... {response_text}"
+                    yield f"Iteration {i+1} - Thinking... {response_text}" # Real time output
 
         except Exception as e:
             intermediate_output += f"Streaming Error: {str(e)}\n"
-            yield f"Streaming Error (Iteration {i+1}): {str(e)}"
+            yield f"Streaming Error (Iteration {i+1}): {str(e)}" #Error
             continue
 
         try:
-            response_json = json.loads(response_text)
+            response_json = json.loads(response_text)  # Parse the JSON response.
             intermediate_output += f"Iteration {i+1} - Focus: {current_entity} - Action: {response_text}\n"
         except json.JSONDecodeError:
             intermediate_output += f"Iteration {i+1} - LLM Response (Invalid JSON): {response_text[:100]}...\n"
-            context.append(f"Invalid JSON: {response_text[:100]}...")
+            context.append(f"Invalid JSON: {response_text[:100]}...") # Add invalid JSON to context
             continue
 
         tool_name = response_json.get("tool")
         parameters = response_json.get("parameters", {})
 
-        if not tool_name:
+        if not tool_name: #LLM didn't return a tool.  End the process if we are past halfway.
             if all_insights:
                 if i > MAX_ITERATIONS // 2:
                     break
@@ -1080,7 +1031,7 @@ def deep_research(prompt):
                 yield f"Iteration {i+1} - Generated search query: {result}"
 
                 if current_entity != 'general':
-                    entity_progress[current_entity]['queries'].append(result)
+                    entity_progress[current_entity]['queries'].append(result) # Add entity-specific
 
                 previous_queries.append(result)
 
@@ -1092,12 +1043,13 @@ def deep_research(prompt):
 
                 filtered_result = filter_results(result, search_prompt, previous_snippets=seen_snippets)
 
-                result = filtered_result
+                result = filtered_result  # Work with filtered results
 
-                if not result and 'query' in parameters:
+                if not result and 'query' in parameters: # Add query to failures if nothing returned.
                     failed_queries.append(parameters['query'])
 
             elif tool_name == "reason":
+                # Ensure correct reasoning context is passed.
                 if current_entity != 'general' and 'reasoning_context' not in parameters:
                     parameters['reasoning_context'] = entity_progress[current_entity]['insights']
                 elif 'reasoning_context' not in parameters:
@@ -1110,9 +1062,9 @@ def deep_research(prompt):
                         parameters['prompt'] = prompt
 
                 if 'search_results' not in parameters:
-                    parameters['search_results'] = []
+                    parameters['search_results'] = [] #Avoid errors if no search results.
 
-                if 'focus_areas' not in parameters and focus_areas:
+                if 'focus_areas' not in parameters and focus_areas: # Avoid overwriting focus_areas if already set
                     parameters['focus_areas'] = focus_areas
 
                 result = tool["function"](**parameters)
@@ -1121,20 +1073,20 @@ def deep_research(prompt):
                 if current_entity != 'general':
                     entity_progress[current_entity]['insights'].append(result)
                     if current_entity not in entity_specific_insights:
-                        entity_specific_insights[current_entity] = []
+                         entity_specific_insights[current_entity] = []
                     entity_specific_insights[current_entity].append(result)
                 else:
-                    reasoning_context.append(result)
+                    reasoning_context.append(result) #Add to general context.
                 add_to_faiss_index(result)
                 all_insights.append(result)
 
             elif tool_name == "critique_reasoning":
-                if 'previous_critiques' not in parameters:
+                if 'previous_critiques' not in parameters: #Pass in the previous critiques.
                     parameters['previous_critiques'] = previous_critiques
 
                 if all_insights:
                     if 'reasoning_output' not in parameters:
-                        parameters['reasoning_output'] = all_insights[-1]
+                        parameters['reasoning_output'] = all_insights[-1]  #Critique the most recent insight.
                     if 'prompt' not in parameters:
                         parameters['prompt'] = prompt
 
@@ -1149,7 +1101,7 @@ def deep_research(prompt):
                 result = tool["function"](**parameters)
                 yield f"Iteration {i+1} - Identifying contradictions..."
                 if result:
-                    contradictions = result
+                    contradictions = result  # Keep track of contradictions.
                     context.append(f"Identified contradictions: {result}")
 
             elif tool_name == "identify_focus_areas":
@@ -1159,8 +1111,8 @@ def deep_research(prompt):
                 yield f"Iteration {i+1} - Identifying focus areas..."
                 if result:
                     old_focus = set(focus_areas)
-                    focus_areas = result
-                    failed_areas.extend([area for area in old_focus if area not in result])
+                    focus_areas = result  # Update focus areas
+                    failed_areas.extend([area for area in old_focus if area not in result])  #Track failed areas
                     context.append(f"New focus areas: {result}")
 
             elif tool_name == "extract_article":
@@ -1168,20 +1120,21 @@ def deep_research(prompt):
                 yield f"Iteration {i+1} - Extracting article content..."
                 if result:
                     context.append(f"Extracted article content from {parameters['url']}: {result[:200]}...")
+                    # Reason specifically about the extracted article.
                     reasoning_about_article = tool_reason(prompt=prompt, search_results=[{"title": "Extracted Article", "snippet": result, "url": parameters['url']}])
                     if reasoning_about_article:
                         all_insights.append(reasoning_about_article)
                         add_to_faiss_index(reasoning_about_article)
-        
+            
             elif tool_name == "summarize_paper":
                 result = tool["function"](**parameters)
                 yield f"Iteration {i+1} - Summarizing paper..."
                 if result:
-                    paper_summaries[parameters['paper_text'][:100]] = result
+                    paper_summaries[parameters['paper_text'][:100]] = result  # Store by a snippet of the text
                     save_paper_summaries(paper_summaries)
                     context.append(f"Summarized paper: {result[:200]}...")
-                    add_to_faiss_index(result)
-                    all_insights.append(result)
+                    add_to_faiss_index(result) # Add the summary itself to FAISS.
+                    all_insights.append(result) #Add summary to insights for later summarization.
 
             elif tool_name == "meta_analyze":
                 if 'entity_insights' not in parameters:
@@ -1191,13 +1144,13 @@ def deep_research(prompt):
                 result = tool["function"](**parameters)
                 yield f"Iteration {i+1} - Performing meta-analysis..."
                 if result:
-                    all_insights.append(result)
+                    all_insights.append(result)  # Add meta-analysis to overall insights.
                     context.append(f"Meta-analysis across entities: {result[:200]}...")
                     add_to_faiss_index(result)
 
 
             elif tool_name == "draft_research_plan":
-                result = "Research plan already generated."
+                result = "Research plan already generated."  # Avoid re-generating.
 
             else:
                 result = tool["function"](**parameters)
@@ -1208,6 +1161,7 @@ def deep_research(prompt):
 
             intermediate_output += f"Iteration {i+1} - Result: {result_str}\n"
 
+            # Add tool use to context, limit context length
             result_context = result_str
             if len(result_str) > 300:
                 result_context = result_str[:300] + "..."
@@ -1219,6 +1173,7 @@ def deep_research(prompt):
             intermediate_output += f"Iteration {i+1} - Error: {str(e)}\n"
             continue
 
+        #Save data
         research_data = {
             'context': context,
             'all_insights': all_insights,
@@ -1234,28 +1189,31 @@ def deep_research(prompt):
             'research_session_id': research_session_id
         }
         for entity in entity_progress:
-            research_data[entity] = entity_progress[entity]
+             research_data[entity] = entity_progress[entity] #save the individual entity
         save_research_data(research_data, index)
 
+
+    # Perform meta-analysis *before* final summarization, if we have enough entity-specific insights.
     if len(entity_specific_insights) > 1 and len(all_insights) > 2:
         meta_analysis = tool_meta_analyze(entity_insights=entity_specific_insights, prompt=prompt)
         if meta_analysis:
             all_insights.append(meta_analysis)
             intermediate_output += f"Final Meta-Analysis: {meta_analysis[:500]}...\n"
-            add_to_faiss_index(meta_analysis)
+            add_to_faiss_index(meta_analysis)  # Add to FAISS
 
     if all_insights:
-        final_result = tool_summarize(all_insights, prompt, contradictions)
+        final_result = tool_summarize(all_insights, prompt, contradictions) # Summarize all insights.
     else:
         final_result = "Could not find meaningful information despite multiple attempts."
 
+
     full_output = f"**Research Prompt:** {prompt}\n\n"
 
     if key_entities_with_descriptions:
-        full_output += f"**Key Entities Identified:**\n"
-        for entity in key_entities_with_descriptions:
-            full_output += f"- {entity}\n"
-        full_output += "\n"
+         full_output += f"**Key Entities Identified:**\n"
+         for entity in key_entities_with_descriptions:
+             full_output += f"- {entity}\n"
+         full_output += "\n"
 
     full_output += "**Research Process:**\n" + intermediate_output + "\n"
 
@@ -1272,124 +1230,54 @@ def deep_research(prompt):
     full_output += f"Total iterations: {i+1}\n"
     full_output += f"Total insights generated: {len(all_insights)}\n"
 
-    yield full_output
+    yield full_output  # Final output
+    
 
 custom_css = """
-/* Modern Research Interface */
-.research-container {
-    display: grid;
-    grid-template-columns: 2fr 1fr;
-    gap: 20px;
-    padding: 20px;
-    background: #1a1a1a;
+.gradio-container {
+    background-color: #f7f9fc;
 }
-
-.main-content {
-    background: #2d2d2d;
-    border-radius: 10px;
-    padding: 20px;
+.output-box {
+    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
+    line-height: 1.5;
+    font-size: 14px;
+    white-space: pre-wrap; /* Preserve newlines and spacing */
 }
-
-.sidebar {
-    background: #2d2d2d;
-    border-radius: 10px;
-    padding: 20px;
+h3 {
+    color: #2c3e50;
+    font-weight: 600;
 }
-
-/* Progress Tracking */
-.progress-container {
-    background: #333;
-    border-radius: 8px;
-    padding: 15px;
-    margin-bottom: 15px;
-}
-
-.tool-indicator {
-    display: flex;
-    align-items: center;
-    gap: 10px;
-    padding: 8px;
-    background: #444;
-    border-radius: 5px;
-    margin-bottom: 8px;
-}
-
-.tool-icon {
-    width: 24px;
-    height: 24px;
-}
-
-.tool-name {
-    color: #4CAF50;
-    font-weight: bold;
-}
-
-/* Enhanced Output Formatting */
-.research-output {
-    font-family: 'Inter', sans-serif;
-    line-height: 1.6;
-}
-
-.research-output h2 {
-    color: #4CAF50;
-    border-bottom: 2px solid #4CAF50;
-    padding-bottom: 5px;
-}
-
-.research-output code {
-    background: #333;
-    padding: 2px 6px;
-    border-radius: 4px;
-}
-
-/* Statistics Panel */
-.stats-panel {
-    background: #333;
-    border-radius: 8px;
-    padding: 15px;
-    margin-top: 15px;
-}
-
-.stat-item {
-    display: flex;
-    justify-content: space-between;
-    padding: 8px 0;
-    border-bottom: 1px solid #444;
+.footer {
+    text-align: center;
+    margin-top: 20px;
+    color: #7f8c8d;
+    font-size: 0.9em;
 }
 """
 
 iface = gr.Interface(
     fn=deep_research,
     inputs=[
-        gr.Textbox(lines=5, placeholder="Enter your research question...", label="Research Question"),
-        gr.Dropdown(
-            choices=["Quick", "Standard", "Comprehensive", "Exhaustive"],
-            label="Research Depth",
-            value="Standard"
-        ),
-        gr.CheckboxGroup(
-            choices=["Academic Papers", "Patents", "News", "Clinical Trials", "Datasets"],
-            label="Source Types",
-            value=["Academic Papers", "News"]
-        ),
-        gr.Slider(
-            minimum=1,
-            maximum=24,
-            value=2,
-            label="Time Limit (hours)"
-        )
+        gr.Textbox(lines=5, placeholder="Enter your research question...", label="Research Question")
     ],
-    outputs=[
-        gr.Markdown(label="Research Results", elem_classes=["research-output"]),
-        gr.JSON(label="Progress Statistics", elem_classes=["stats-panel"]),
-        gr.HTML(label="Active Tools", elem_classes=["tool-indicator"])
+    outputs=gr.Markdown(label="Research Results",  elem_classes=["output-box"]), #Changed to markdown
+    title="Advanced Multi-Stage Research Assistant",
+    description="""This tool performs deep, multi-faceted research, leveraging multiple search engines, 
+                   specialized academic databases, and advanced AI models. It incorporates a persistent knowledge 
+                   base using FAISS indexing to avoid redundant searches and build upon previous findings.  Progress is shown in real-time.""",
+    examples=[
+        ["What are the key factors affecting urban tree survival and how do they vary between developing and developed countries?"],
+        ["Compare and contrast the economic policies of China and the United States over the past two decades, analyzing their impacts on global trade."],
+        ["What are the most promising approaches to quantum computing and what are their respective advantages and limitations?"],
+        ["Analyze the environmental and social impacts of lithium mining for electric vehicle batteries."],
+        ["How has artificial intelligence influenced medical diagnostics in the past five years, and what are the ethical considerations?"]
     ],
-    title="Advanced Research Institution Assistant",
-    description="""Enterprise-grade research system with real-time progress tracking, 
-    comprehensive source coverage, and advanced analysis capabilities.""",
-    theme=gr.themes.Base(primary_hue="green"),
-    css=custom_css
+    theme="default",
+    cache_examples=False,
+    css=custom_css,
+    allow_flagging="never",
+    live=True #for real time streaming
 )
 
 if __name__ == "__main__":
-    iface.launch(share=False)
+    iface.launch(share=False)