Update app.py

app.py

@@ -12,42 +12,90 @@ import logging
 import numpy as np
 from collections import deque
 from huggingface_hub import InferenceClient
+import requests
+import arxiv
+import scholarly
+import pymed
+import wikipedia
+from newspaper import Article
+import pickle
+import faiss
+import threading
+from concurrent.futures import ThreadPoolExecutor, as_completed
+import tiktoken
 
-# Set up logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-# Configuration
 HF_API_KEY = os.environ.get("HF_API_KEY")
 if not HF_API_KEY:
     raise ValueError("Please set the HF_API_KEY environment variable.")
 
-# Initialize Hugging Face Inference Client
 client = InferenceClient(provider="hf-inference", api_key=HF_API_KEY)
 
-# Model endpoints
-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"
+MAIN_LLM_MODEL = "meta-llama/Llama-3-70b-instruct"
+REASONING_LLM_MODEL = "anthropic/claude-3-opus-20240229"
+CRITIC_LLM_MODEL = "google/gemini-1.5-pro"
+ENSEMBLE_MODELS = [MAIN_LLM_MODEL, REASONING_LLM_MODEL, CRITIC_LLM_MODEL]
 
-MAX_ITERATIONS = 12
-TIMEOUT = 60
+MAX_ITERATIONS = 20
+TIMEOUT = 120
 RETRY_DELAY = 5
-NUM_RESULTS = 10
+NUM_RESULTS = 15
 SIMILARITY_THRESHOLD = 0.15
-MAX_CONTEXT_ITEMS = 20
-MAX_HISTORY_ITEMS = 5
+MAX_CONTEXT_ITEMS = 30
+MAX_HISTORY_ITEMS = 8
+MAX_FULL_TEXT_LENGTH = 10000
+FAISS_INDEX_PATH = "research_index.faiss"
+RESEARCH_DATA_PATH = "research_data.pkl"
 
-# Load multiple embedding models for different purposes
 try:
-    main_similarity_model = SentenceTransformer('all-mpnet-base-v2')
+    main_similarity_model = SentenceTransformer('sentence-transformers/all-mpnet-base-v2')
     concept_similarity_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
+    document_similarity_model = SentenceTransformer('sentence-transformers/multi-qa-mpnet-base-dot-v1')
+
+    embedding_dim = document_similarity_model.get_sentence_embedding_dimension()
+    if os.path.exists(FAISS_INDEX_PATH):
+        index = faiss.read_index(FAISS_INDEX_PATH)
+        logger.info(f"Loaded FAISS index from {FAISS_INDEX_PATH}")
+    else:
+        index = faiss.IndexFlatIP(embedding_dim)
+        logger.info("Created a new FAISS index.")
 except Exception as e:
-    logger.error(f"Failed to load SentenceTransformer models: {e}")
-    main_similarity_model = None
-    concept_similarity_model = None
+    logger.error(f"Failed to load models or initialize FAISS: {e}")
+    raise
+
+def get_token_count(text):
+    try:
+        encoding = tiktoken.get_encoding("cl100k_base")
+        return len(encoding.encode(text))
+    except:
+        return len(text.split()) * 1.3
 
-def hf_inference(model_name, prompt, max_tokens=500, retries=5):
+def save_research_data(data, index):
+    try:
+        with open(RESEARCH_DATA_PATH, "wb") as f:
+            pickle.dump(data, f)
+        faiss.write_index(index, FAISS_INDEX_PATH)
+        logger.info(f"Research data and index saved to {RESEARCH_DATA_PATH} and {FAISS_INDEX_PATH}")
+    except Exception as e:
+        logger.error(f"Error saving research data: {e}")
+
+def load_research_data():
+    if os.path.exists(RESEARCH_DATA_PATH):
+        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
+        except Exception as e:
+            logger.error(f"Error loading research data: {e}")
+            return {}
+    else:
+        logger.info("No existing research data found.")
+        return {}
+
+def hf_inference(model_name, prompt, max_tokens=2000, retries=5):
     for attempt in range(retries):
         try:
             messages = [{"role": "user", "content": prompt}]
@@ -64,22 +112,185 @@ def hf_inference(model_name, prompt, max_tokens=500, retries=5):
             time.sleep(RETRY_DELAY * (1 + attempt))
     return {"error": "Request failed after multiple retries."}
 
+def ensemble_inference(prompt, models=ENSEMBLE_MODELS, max_tokens=1500):
+    results = []
+    with ThreadPoolExecutor(max_workers=len(models)) as executor:
+        future_to_model = {executor.submit(hf_inference, model, prompt, max_tokens): model for model in models}
+        for future in as_completed(future_to_model):
+            model = future_to_model[future]
+            try:
+                result = future.result()
+                if "generated_text" in result:
+                    results.append({"model": model, "text": result["generated_text"]})
+            except Exception as e:
+                logger.error(f"Error with model {model}: {e}")
+
+    if not results:
+        return {"error": "All models failed to generate responses"}
+
+    if len(results) == 1:
+        return {"generated_text": results[0]["text"]}
+
+    synthesis_prompt = "Synthesize these expert responses into a single coherent answer:\n\n"
+    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)
+    if "generated_text" in synthesis:
+        return synthesis
+    else:
+        return {"generated_text": max(results, key=lambda x: len(x["text"]))["text"]}
+
 def tool_search_web(query: str, num_results: int = NUM_RESULTS, safesearch: str = "moderate",
-                    time_filter: str = "", 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:
-            results = [r for r in ddgs.text(query, max_results=num_results, safesearch=safesearch,
-                                           time=time_filter, region=region, hreflang=language)]
+            kwargs = {
+                "keywords": query,
+                "max_results": num_results,
+                "safesearch": safesearch,
+                "region": region,
+                "hreflang": language,
+            }
+            if time_filter:
+                if time_filter in ['d', 'w', 'm', 'y']:
+                    kwargs["time"] = time_filter
+
+            results = [r for r in ddgs.text(**kwargs)]
             if results:
-                return [{"title": r["title"], "snippet": r["body"], "url": r["href"]} for r in results]
+                return [{"title": r["title"], "snippet": r["body"], "url": r["href"]}] for r in results]
             else:
+                if time_filter and "time" in kwargs:
+                    del kwargs["time"]
+                    results = [r for r in ddgs.text(**kwargs)]
+                    if results:
+                        return [{"title": r["title"], "snippet": r["body"], "url": r["href"]}] for r in results]
                 return []
     except Exception as e:
         logger.error(f"DuckDuckGo search error: {e}")
+        try:
+            with DDGS() as ddgs:
+                results = [r for r in ddgs.text(
+                    keywords=query,
+                    max_results=num_results,
+                    safesearch=safesearch,
+                    region=region,
+                    hreflang=language
+                )]
+                if results:
+                    return [{"title": r["title"], "snippet": r["body"], "url": r["href"]} for r in results]
+        except Exception as e2:
+            logger.error(f"Fallback DuckDuckGo search also failed: {e2}")
         return []
 
+def tool_search_arxiv(query: str, max_results: int = 5) -> list:
+    try:
+        client = arxiv.Client()
+        search = arxiv.Search(
+            query=query,
+            max_results=max_results,
+            sort_by=arxiv.SortCriterion.Relevance
+        )
+        results = []
+        for paper in client.results(search):
+            results.append({
+                "title": paper.title,
+                "snippet": paper.summary[:500] + "..." if len(paper.summary) > 500 else paper.summary,
+                "url": paper.pdf_url,
+                "authors": ", ".join(author.name for author in paper.authors),
+                "published": paper.published.strftime("%Y-%m-%d") if paper.published else "Unknown",
+                "source": "arXiv"
+            })
+        return results
+    except Exception as e:
+        logger.error(f"arXiv search error: {e}")
+        return []
+
+def tool_search_pubmed(query: str, max_results: int = 5) -> list:
+    try:
+        pubmed = pymed.PubMed(tool="ResearchAssistant", email="[email protected]")
+        results = list(pubmed.query(query, max_results=max_results))
+
+        output = []
+        for article in results:
+            try:
+                data = article.toDict()
+                output.append({
+                    "title": data.get("title", "No title"),
+                    "snippet": data.get("abstract", "No abstract")[:500] + "..." if data.get("abstract", "") and len(data.get("abstract", "")) > 500 else data.get("abstract", "No abstract"),
+                    "url": f"https://pubmed.ncbi.nlm.nih.gov/{data.get('pubmed_id')}/",
+                    "authors": ", ".join(author.get("name", "") for author in data.get("authors", [])),
+                    "published": data.get("publication_date", "Unknown"),
+                    "source": "PubMed"
+                })
+            except:
+                continue
+        return output
+    except Exception as e:
+        logger.error(f"PubMed search error: {e}")
+        return []
+
+def tool_search_wikipedia(query: str, max_results: int = 3) -> list:
+    try:
+        search_results = wikipedia.search(query, results=max_results)
+        results = []
+
+        for title in search_results:
+            try:
+                page = wikipedia.page(title)
+                summary = page.summary
+                snippet = summary[:500] + "..." if len(summary) > 500 else summary
+                results.append({
+                    "title": page.title,
+                    "snippet": snippet,
+                    "url": page.url,
+                    "source": "Wikipedia"
+                })
+            except (wikipedia.exceptions.DisambiguationError, wikipedia.exceptions.PageError):
+                continue
+
+        return results
+    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)
+        results = []
+        for _ in range(max_results):
+            try:
+                result = next(search_query)
+                results.append({
+                    "title": result.get("bib", {}).get("title", "No title"),
+                    "snippet": result.get("bib", {}).get("abstract", "No abstract")[:500] + "..." if result.get("bib", {}).get("abstract") else result.get("bib", {}).get("abstract", "No abstract"),
+                    "url": result.get("pub_url", "#"),
+                    "authors": ", ".join(result.get("bib", {}).get("author", [])),
+                    "published": result.get("bib", {}).get("pub_year", "Unknown"),
+                    "source": "Google Scholar"
+                })
+            except StopIteration:
+                break
+            except Exception as e:
+                logger.warning(f"Error processing Scholar result: {e}")
+                continue
+        return results
+    except Exception as e:
+        logger.error(f"Google Scholar search error: {e}")
+        return []
+
+def extract_article_content(url: str) -> str:
+    try:
+        article = Article(url)
+        article.download()
+        article.parse()
+        return article.text
+    except Exception as e:
+        logger.error(f"Failed to extract article content from {url}: {e}")
+        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."
 
@@ -89,8 +300,22 @@ def tool_reason(prompt: str, search_results: list, reasoning_context: list = [],
     if focus_areas:
         reasoning_input += f"Focus particularly on these aspects: {', '.join(focus_areas)}\n\n"
 
+    results_by_source = {}
     for i, result in enumerate(search_results):
-        reasoning_input += f"- Result {i + 1}: Title: {result['title']}, Snippet: {result['snippet']}\n"
+        source = result.get('source', 'Web Search')
+        if source not in results_by_source:
+            results_by_source[source] = []
+        results_by_source[source].append((i, result))
+
+    for source, results in results_by_source.items():
+        reasoning_input += f"\n--- {source} Results ---\n"
+        for i, result in results:
+            reasoning_input += f"- Result {i + 1}: Title: {result['title']}\n  Snippet: {result['snippet']}\n"
+            if 'authors' in result:
+                reasoning_input += f"  Authors: {result['authors']}\n"
+            if 'published' in result:
+                reasoning_input += f"  Published: {result['published']}\n"
+            reasoning_input += "\n"
 
     if reasoning_context:
         recent_context = reasoning_context[-MAX_HISTORY_ITEMS:]
@@ -99,9 +324,9 @@ def tool_reason(prompt: str, search_results: list, reasoning_context: list = [],
     if critique:
         reasoning_input += f"\n\nRecent critique to address: {critique}\n"
 
-    reasoning_input += "\nProvide a thorough, nuanced analysis that builds upon previous reasoning if applicable. Consider multiple perspectives and potential contradictions in the search results."
+    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."
 
-    reasoning_output = hf_inference(REASONING_LLM_MODEL, reasoning_input)
+    reasoning_output = ensemble_inference(reasoning_input)
 
     if isinstance(reasoning_output, dict) and "generated_text" in reasoning_output:
         return reasoning_output["generated_text"].strip()
@@ -114,14 +339,27 @@ def tool_summarize(insights: list, prompt: str, contradictions: list = []) -> st
         return "No insights to summarize."
 
     summarization_input = f"Synthesize the following insights into a cohesive and comprehensive summary regarding: '{prompt}'\n\n"
-    summarization_input += "\n\n".join(insights[-MAX_HISTORY_ITEMS:])
+
+    max_tokens = 12000
+    selected_insights = []
+    token_count = get_token_count(summarization_input) + get_token_count("\n\n".join(contradictions))
+
+    for insight in reversed(insights):
+        insight_tokens = get_token_count(insight)
+        if token_count + insight_tokens < max_tokens:
+            selected_insights.insert(0, insight)
+            token_count += insight_tokens
+        else:
+            break
+
+    summarization_input += "\n\n".join(selected_insights)
 
     if contradictions:
         summarization_input += "\n\nAddress these specific contradictions:\n" + "\n".join(contradictions)
 
-    summarization_input += "\n\nProvide a well-structured summary that:\n1. Presents the main findings\n2. Acknowledges limitations and uncertainties\n3. Highlights areas of consensus and disagreement\n4. Suggests potential directions for further inquiry"
+    summarization_input += "\n\nProvide a well-structured summary that:\n1. Presents the main findings\n2. Acknowledges limitations and uncertainties\n3. Highlights areas of consensus and disagreement\n4. Suggests potential directions for further inquiry\n5. Evaluates the strength of evidence for key claims"
 
-    summarization_output = hf_inference(MAIN_LLM_MODEL, summarization_input)
+    summarization_output = ensemble_inference(summarization_input)
 
     if isinstance(summarization_output, dict) and "generated_text" in summarization_output:
         return summarization_output["generated_text"].strip()
@@ -130,7 +368,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:
@@ -143,7 +381,7 @@ def tool_generate_search_query(prompt: str, previous_queries: list = [],
     if focus_areas:
         query_gen_input += f"Focus particularly on these aspects: {', '.join(focus_areas)}\n"
 
-    query_gen_input += "Refine the search query based on previous queries, aiming for more precise results.\n"
+    query_gen_input += "Refine the search query based on previous queries, aiming for more precise results. Consider using advanced search operators like site:, filetype:, intitle:, etc. when appropriate. Make sure the query is well-formed for academic and scientific search engines.\n"
     query_gen_input += "Search Query:"
 
     query_gen_output = hf_inference(MAIN_LLM_MODEL, query_gen_input)
@@ -155,13 +393,13 @@ 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:
         critique_input += "Previous critiques that should be addressed:\n" + "\n".join(previous_critiques[-MAX_HISTORY_ITEMS:]) + "\n\n"
 
-    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."
+    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)
 
@@ -175,8 +413,20 @@ def tool_identify_contradictions(insights: list) -> list:
     if len(insights) < 2:
         return []
 
-    contradiction_input = "Identify specific contradictions in these insights:\n\n" + "\n\n".join(insights[-MAX_HISTORY_ITEMS:])
-    contradiction_input += "\n\nList each contradiction as a separate numbered point. If no contradictions exist, respond with 'No contradictions found.'"
+    max_tokens = 12000
+    selected_insights = []
+    token_count = 0
+
+    for insight in reversed(insights):
+        insight_tokens = get_token_count(insight)
+        if token_count + insight_tokens < max_tokens:
+            selected_insights.insert(0, insight)
+            token_count += insight_tokens
+        else:
+            break
+
+    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)
 
@@ -192,27 +442,46 @@ 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:
-        focus_input += "And these existing insights:\n" + "\n".join(insights[-3:]) + "\n\n"
+        recent_insights = insights[-5:] if len(insights) > 5 else insights
+        focus_input += "And these existing insights:\n" + "\n".join(recent_insights) + "\n\n"
 
     if failed_areas:
         focus_input += f"These focus areas didn't yield useful results: {', '.join(failed_areas)}\n\n"
 
-    focus_input += "Identify 2-3 specific aspects that should be investigated further to get a complete understanding. Be precise and prioritize underexplored areas."
+    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)
 
     if isinstance(focus_output, dict) and "generated_text" in focus_output:
         result = focus_output["generated_text"].strip()
         areas = re.findall(r'(?:^|\n)(?:\d+\.|\*|\-)\s*(.*?)(?=(?:\n(?:\d+\.|\*|\-|$))|$)', result)
-        return [area.strip() for area in areas if area.strip()][:3]
+        return [area.strip() for area in areas if area.strip()][:5]
 
     logger.error(f"Failed to identify focus areas: {focus_output}")
     return []
 
+def add_to_faiss_index(text: str):
+    """Adds the embedding of the given text to the FAISS index."""
+    embedding = document_similarity_model.encode(text, convert_to_tensor=True)
+    embedding_np = embedding.cpu().numpy().reshape(1, -1)  # Ensure 2D array
+    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)
+    index.add(embedding_np)
+
+def search_faiss_index(query: str, top_k: int = 5) -> List[str]:
+    """Searches the FAISS index for the most similar texts to the query."""
+    query_embedding = document_similarity_model.encode(query, convert_to_tensor=True)
+    query_embedding_np = query_embedding.cpu().numpy().reshape(1, -1)
+    faiss.normalize_L2(query_embedding_np)
+    distances, indices = index.search(query_embedding_np, top_k)
+    return indices[0].tolist() # Return indices
+
 def filter_results(search_results, prompt, previous_snippets=None):
     if not main_similarity_model or not search_results:
         return search_results
@@ -233,11 +502,14 @@ def filter_results(search_results, prompt, previous_snippets=None):
 
             result_embedding = main_similarity_model.encode(combined_text, convert_to_tensor=True)
             cosine_score = util.pytorch_cos_sim(prompt_embedding, result_embedding)[0][0].item()
-
+            
             if cosine_score >= SIMILARITY_THRESHOLD:
                 result['relevance_score'] = cosine_score
                 filtered_results.append(result)
                 seen_snippets.add(result['snippet'])
+                # Add snippet to FAISS index
+                add_to_faiss_index(result['snippet'])
+
 
         filtered_results.sort(key=lambda x: x.get('relevance_score', 0), reverse=True)
         return filtered_results
@@ -247,14 +519,14 @@ def filter_results(search_results, prompt, previous_snippets=None):
         return search_results
 
 def tool_extract_key_entities(prompt: str) -> list:
-    entity_input = f"Extract the key entities (people, organizations, concepts, technologies, etc.) from this research prompt that should be investigated individually:\n\n{prompt}\n\nList only the most important 3-5 entities, one per line."
+    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."
 
     entity_output = hf_inference(MAIN_LLM_MODEL, entity_input)
 
     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[:5]
+        return entities[:7]
 
     logger.error(f"Failed to extract key entities: {entity_output}")
     return []
@@ -269,9 +541,9 @@ def tool_meta_analyze(entity_insights: Dict[str, list], prompt: str) -> str:
         if insights:
             meta_input += f"\n--- {entity} ---\n" + insights[-1] + "\n"
 
-    meta_input += "\nProvide a high-level synthesis that identifies:\n1. Common themes across entities\n2. Important differences\n3. How these entities interact or influence each other\n4. The broader implications for the original research question"
+    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 = hf_inference(MAIN_LLM_MODEL, meta_input)
+    meta_output = ensemble_inference(meta_input)
 
     if isinstance(meta_output, dict) and "generated_text" in meta_output:
         return meta_output["generated_text"].strip()
@@ -279,6 +551,41 @@ def tool_meta_analyze(entity_insights: Dict[str, list], prompt: str) -> str:
     logger.error(f"Failed to perform meta-analysis: {meta_output}")
     return "Could not generate a meta-analysis due to an error."
 
+def tool_draft_research_plan(prompt: str, entities: list, focus_areas: list = []) -> str:
+    plan_input = f"Create a detailed research plan for investigating this question: '{prompt}'\n\n"
+
+    if entities:
+        plan_input += "Key entities to investigate:\n" + "\n".join(entities) + "\n\n"
+
+    if focus_areas:
+        plan_input += "Additional focus areas:\n" + "\n".join(focus_areas) + "\n\n"
+
+    plan_input += "The research plan should include:\n"
+    plan_input += "1. Main research questions and sub-questions\n"
+    plan_input += "2. Methodology for investigating each aspect\n"
+    plan_input += "3. Potential sources and databases to consult\n"
+    plan_input += "4. Suggested sequence of investigation\n"
+    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)
+
+    if isinstance(plan_output, dict) and "generated_text" in plan_output:
+        return plan_output["generated_text"].strip()
+
+    logger.error(f"Failed to generate research plan: {plan_output}")
+    return "Could not generate a research plan due to an error."
+
+def tool_extract_article(url: str) -> str:
+    content = extract_article_content(url)
+    if not content:
+        return f"Could not extract content from {url}"
+
+    if len(content) > MAX_FULL_TEXT_LENGTH:
+        content = content[:MAX_FULL_TEXT_LENGTH] + "... [content truncated]"
+
+    return content
+
 tools = {
     "search_web": {
         "function": tool_search_web,
@@ -291,6 +598,45 @@ tools = {
             "language": {"type": "string", "description": "Optional language code."}
         },
     },
+    "search_arxiv": {
+        "function": tool_search_arxiv,
+        "description": "Searches arXiv for scientific papers.",
+        "parameters": {
+            "query": {"type": "string", "description": "The search query for scientific papers."},
+            "max_results": {"type": "integer", "description": "Maximum number of papers to return."}
+        },
+    },
+    "search_pubmed": {
+        "function": tool_search_pubmed,
+        "description": "Searches PubMed for medical and scientific literature.",
+        "parameters": {
+            "query": {"type": "string", "description": "The search query for medical literature."},
+            "max_results": {"type": "integer", "description": "Maximum number of articles to return."}
+        },
+    },
+    "search_wikipedia": {
+        "function": tool_search_wikipedia,
+        "description": "Searches Wikipedia for information.",
+        "parameters": {
+            "query": {"type": "string", "description": "The search query for Wikipedia."},
+            "max_results": {"type": "integer", "description": "Maximum number of articles to return."}
+        },
+    },
+    "search_scholar": {
+        "function": tool_search_scholar,
+        "description": "Searches Google Scholar for academic publications.",
+        "parameters": {
+            "query": {"type": "string", "description": "The search query for Google Scholar."},
+            "max_results": {"type": "integer", "description": "Maximum number of articles to return."}
+        }
+    },
+    "extract_article": {
+        "function": tool_extract_article,
+        "description": "Extracts the main content from a web article URL",
+        "parameters": {
+            "url": {"type": "string", "description": "The URL of the article to extract"}
+        },
+    },
     "reason": {
         "function": tool_reason,
         "description": "Analyzes and reasons about information.",
@@ -360,6 +706,15 @@ tools = {
             "entity_insights": {"type": "object", "description": "Dictionary mapping entities to their insights."},
             "prompt": {"type": "string", "description": "The original research prompt."}
         },
+    },
+    "draft_research_plan": {
+        "function": tool_draft_research_plan,
+        "description": "Creates a detailed research plan.",
+        "parameters": {
+            "prompt": {"type": "string", "description": "The research question/prompt."},
+            "entities": {"type": "array", "description": "Key entities to investigate."},
+            "focus_areas": {"type": "array", "description": "Additional areas to focus on."}
+        }
     }
 }
 
@@ -385,6 +740,16 @@ Available Tools:
 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:
+
+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:
@@ -392,28 +757,52 @@ Output:
     return prompt
 
 def deep_research(prompt):
-    task_description = "You are an advanced research assistant that can perform deep, multi-stage analysis. Use available tools iteratively, focus on different aspects, follow promising leads, and critically evaluate your findings."
-    context = []
-    all_insights = []
-    entity_specific_insights = {}
-    intermediate_output = ""
-    previous_queries = []
-    failed_queries = []
-    reasoning_context = []
-    previous_critiques = []
-    focus_areas = []
-    failed_areas = []
-    seen_snippets = set()
-    contradictions = []
-    research_session_id = str(uuid4())
-
-    key_entities = tool_extract_key_entities(prompt=prompt)
+    task_description = "You are an advanced research assistant. Use available tools iteratively, focus on different aspects, follow promising leads, critically evaluate your findings, and build up a comprehensive understanding. Utilize the FAISS index to avoid redundant searches and build a persistent knowledge base."
+    research_data = load_research_data()
+    
+    context = research_data.get('context', [])
+    all_insights = research_data.get('all_insights', [])
+    entity_specific_insights = research_data.get('entity_specific_insights', {})
+    intermediate_output = ""  # For Gradio display
+    previous_queries = research_data.get('previous_queries', [])
+    failed_queries = research_data.get('failed_queries', [])
+    reasoning_context = research_data.get('reasoning_context', [])
+    previous_critiques = research_data.get('previous_critiques', [])
+    focus_areas = research_data.get('focus_areas', [])
+    failed_areas = research_data.get('failed_areas', [])
+    seen_snippets = set(research_data.get('seen_snippets', []))
+    contradictions = research_data.get('contradictions', [])
+    research_session_id = research_data.get('research_session_id', str(uuid4()))
+    
+    # Restore or initialize FAISS index
+    global index
+    if research_data:
+         logger.info("Restoring FAISS Index from loaded data.")
+    else:
+        index.reset() #Start Fresh
+        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]
     if key_entities:
         context.append(f"Identified key entities: {key_entities}")
-        intermediate_output += f"Identified key entities for focused research: {key_entities}\n"
+        intermediate_output += f"Identified key entities for focused research: {key_entities_with_descriptions}\n"
 
     entity_progress = {entity: {'queries': [], 'insights': []} for entity in key_entities}
     entity_progress['general'] = {'queries': [], 'insights': []}
+    for entity in key_entities + ['general']:
+        if entity in research_data:
+            entity_progress[entity]['queries'] = research_data[entity]['queries']
+            entity_progress[entity]['insights'] = research_data[entity]['insights']
+
+    if i == 0:
+        initial_focus_areas = tool_identify_focus_areas(prompt=prompt)
+        research_plan = tool_draft_research_plan(prompt=prompt, entities=key_entities, focus_areas=initial_focus_areas)
+        context.append(f"Initial Research Plan: {research_plan[:200]}...")
+        intermediate_output += f"Initial Research Plan:\n{research_plan}\n\n"
+        focus_areas = initial_focus_areas
+    elif not focus_areas:
+        focus_areas = tool_identify_focus_areas(prompt=prompt, insights=all_insights, failed_areas=failed_areas)
 
     for i in range(MAX_ITERATIONS):
         if key_entities and i > 0:
@@ -423,17 +812,39 @@ def deep_research(prompt):
             current_entity = 'general'
 
         context.append(f"Current focus: {current_entity}")
+        
+        # FAISS similarity search before web/arxiv/pubmed searches
+        if i > 0: # Don't do it on first iteration
+            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):
+                    faiss_context.append(f"Previously found insight: {all_insights[idx]}")
+            if faiss_context:
+                context.extend(faiss_context)
+                intermediate_output += f"Iteration {i+1} - Retrieved {len(faiss_context)} relevant items from FAISS index.\n"
+                
 
         if i == 0:
             initial_query = tool_generate_search_query(prompt=prompt)
             if initial_query:
                 previous_queries.append(initial_query)
                 entity_progress['general']['queries'].append(initial_query)
-                search_results = tool_search_web(query=initial_query)
-                filtered_search_results = filter_results(search_results, prompt)
 
-                for result in filtered_search_results:
-                    seen_snippets.add(result['snippet'])
+                with ThreadPoolExecutor(max_workers=5) as executor:
+                    futures = [
+                        executor.submit(tool_search_web, query=initial_query, num_results=NUM_RESULTS),
+                        executor.submit(tool_search_arxiv, query=initial_query, max_results=5),
+                        executor.submit(tool_search_pubmed, query=initial_query, max_results=5),
+                        executor.submit(tool_search_wikipedia, query=initial_query, max_results=3),
+                        executor.submit(tool_search_scholar, query=initial_query, max_results=5)
+                    ]
+
+                    search_results = []
+                    for future in as_completed(futures):
+                        search_results.extend(future.result())
+
+                filtered_search_results = filter_results(search_results, prompt)
 
                 if filtered_search_results:
                     context.append(f"Initial Search Results: {len(filtered_search_results)} items found")
@@ -443,6 +854,7 @@ def deep_research(prompt):
                         entity_progress['general']['insights'].append(reasoning_output)
                         reasoning_context.append(reasoning_output)
                         context.append(f"Initial Reasoning: {reasoning_output[:200]}...")
+                        add_to_faiss_index(reasoning_output) # Add reasoning to FAISS
                 else:
                     failed_queries.append(initial_query)
                     context.append(f"Initial query yielded no relevant results: {initial_query}")
@@ -458,14 +870,23 @@ def deep_research(prompt):
                 previous_queries.append(entity_query)
                 entity_progress[current_entity]['queries'].append(entity_query)
 
-                search_results = tool_search_web(query=entity_query)
+                with ThreadPoolExecutor(max_workers=5) as executor:
+                    futures = [
+                        executor.submit(tool_search_web, query=entity_query, num_results=NUM_RESULTS//2),
+                        executor.submit(tool_search_arxiv, query=entity_query, max_results=3),
+                        executor.submit(tool_search_pubmed, query=entity_query, max_results=3),
+                        executor.submit(tool_search_wikipedia, query=entity_query, max_results=2),
+                        executor.submit(tool_search_scholar, query=entity_query, max_results=3)
+                    ]
+
+                    search_results = []
+                    for future in as_completed(futures):
+                        search_results.extend(future.result())
+
                 filtered_search_results = filter_results(search_results,
                                                         f"{prompt} {current_entity}",
                                                         previous_snippets=seen_snippets)
 
-                for result in filtered_search_results:
-                    seen_snippets.add(result['snippet'])
-
                 if filtered_search_results:
                     context.append(f"Entity Search for {current_entity}: {len(filtered_search_results)} results")
 
@@ -485,6 +906,7 @@ def deep_research(prompt):
                         entity_specific_insights[current_entity].append(entity_reasoning)
 
                         context.append(f"Reasoning about {current_entity}: {entity_reasoning[:200]}...")
+                        add_to_faiss_index(entity_reasoning)  # Add to FAISS
                 else:
                     failed_queries.append(entity_query)
                     context.append(f"Entity query for {current_entity} yielded no relevant results")
@@ -538,6 +960,19 @@ def deep_research(prompt):
 
                 previous_queries.append(result)
 
+            elif tool_name in ["search_web", "search_arxiv", "search_pubmed", "search_wikipedia", "search_scholar"]:
+                result = tool["function"](**parameters)
+                search_prompt = prompt
+                if current_entity != 'general':
+                    search_prompt = f"{prompt} focusing on {current_entity}"
+
+                filtered_result = filter_results(result, search_prompt, previous_snippets=seen_snippets)
+
+                result = filtered_result
+
+                if not result and 'query' in parameters:
+                    failed_queries.append(parameters['query'])
+
             elif tool_name == "reason":
                 if current_entity != 'general' and 'reasoning_context' not in parameters:
                     parameters['reasoning_context'] = entity_progress[current_entity]['insights']
@@ -565,25 +1000,9 @@ def deep_research(prompt):
                     entity_specific_insights[current_entity].append(result)
                 else:
                     reasoning_context.append(result)
-
+                add_to_faiss_index(result) # Add reasoning to FAISS
                 all_insights.append(result)
 
-            elif tool_name == "search_web":
-                result = tool_search_web(**parameters)
-                filtered_result = filter_results(result,
-                                               prompt if current_entity == 'general' else f"{prompt} {current_entity}",
-                                               previous_snippets=seen_snippets)
-
-                for r in filtered_result:
-                    seen_snippets.add(r['snippet'])
-
-                result = filtered_result
-
-                if not result:
-                    query = parameters.get('query', '')
-                    if query:
-                        failed_queries.append(query)
-
             elif tool_name == "critique_reasoning":
                 if 'previous_critiques' not in parameters:
                     parameters['previous_critiques'] = previous_critiques
@@ -616,6 +1035,16 @@ def deep_research(prompt):
                     failed_areas.extend([area for area in old_focus if area not in result])
                     context.append(f"New focus areas: {result}")
 
+            elif tool_name == "extract_article":
+                result = tool["function"](**parameters)
+                if result:
+                    context.append(f"Extracted article content from {parameters['url']}: {result[:200]}...")
+                    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) # Add to FAISS
+
+
             elif tool_name == "meta_analyze":
                 if 'entity_insights' not in parameters:
                     parameters['entity_insights'] = entity_specific_insights
@@ -625,6 +1054,10 @@ def deep_research(prompt):
                 if result:
                     all_insights.append(result)
                     context.append(f"Meta-analysis across entities: {result[:200]}...")
+                    add_to_faiss_index(result)  # Add to FAISS
+
+            elif tool_name == "draft_research_plan":
+                result = "Research plan already generated."
 
             else:
                 result = tool["function"](**parameters)
@@ -646,11 +1079,32 @@ def deep_research(prompt):
             intermediate_output += f"Iteration {i+1} - Error: {str(e)}\n"
             continue
 
+        # Save research data after each iteration
+        research_data = {
+            'context': context,
+            'all_insights': all_insights,
+            'entity_specific_insights': entity_specific_insights,
+            'previous_queries': previous_queries,
+            'failed_queries': failed_queries,
+            'reasoning_context': reasoning_context,
+            'previous_critiques': previous_critiques,
+            'focus_areas': focus_areas,
+            'failed_areas': failed_areas,
+            'seen_snippets': list(seen_snippets),  # Convert set to list for pickling
+            'contradictions': contradictions,
+            'research_session_id': research_session_id
+        }
+        for entity in entity_progress:
+             research_data[entity] = entity_progress[entity]
+        save_research_data(research_data, index)
+
+
     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)
 
     if all_insights:
         final_result = tool_summarize(all_insights, prompt, contradictions)
@@ -659,8 +1113,11 @@ def deep_research(prompt):
 
     full_output = f"**Research Prompt:** {prompt}\n\n"
 
-    if key_entities:
-        full_output += f"**Key Entities Identified:** {', '.join(key_entities)}\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 += "**Research Process:**\n" + intermediate_output + "\n"
 
@@ -679,7 +1136,8 @@ def deep_research(prompt):
 
     return full_output
 
-# Create CSS for a more professional look
+# Gradio Interface
+
 custom_css = """
 .gradio-container {
     background-color: #f7f9fc;
@@ -687,6 +1145,7 @@ custom_css = """
 .output-box {
     font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
     line-height: 1.5;
+    font-size: 14px; /* Increased font size */
 }
 h3 {
     color: #2c3e50;
@@ -700,7 +1159,6 @@ h3 {
 }
 """
 
-# Create the Gradio interface with enhanced UI
 iface = gr.Interface(
     fn=deep_research,
     inputs=[
@@ -708,11 +1166,9 @@ iface = gr.Interface(
     ],
     outputs=gr.Textbox(lines=30, placeholder="Research results will appear here...", label="Research Results", elem_classes=["output-box"]),
     title="Advanced Multi-Stage Research Assistant",
-    description="""This tool performs deep, multi-faceted research by:
-1. Breaking down complex topics into key entities and aspects
-2. Iteratively searching, reasoning, and critiquing findings
-3. Exploring different perspectives and addressing contradictions
-4. Synthesizing insights across multiple information sources""",
+    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.""",
     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."],
@@ -723,18 +1179,8 @@ iface = gr.Interface(
     theme="default",
     cache_examples=False,
     css=custom_css,
-    flagging_mode='never',
-    analytics_enabled=False,
+    allow_flagging="never",  # Disable flagging
 )
 
-# Add footer with additional information (Optional, good for context)
-footer_html = """
-<div class="footer">
-    <p>This research assistant performs advanced multi-stage analysis using natural language processing and web search.</p>
-    <p>Results should be verified with additional sources. Not suitable for medical, legal, or emergency use.</p>
-</div>
-"""
-
-
-# Launch the interface
-iface.launch(share=False)
+if __name__ == "__main__":
+    iface.launch(share=False)