jupyter-legacy/local_open_deep_researcher_with_planning.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install nest_asyncio asyncio aiohttp ollama"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "y7cTpP9rDZW-",
    "outputId": "5a443ad2-7a8d-4fef-f315-12108c28f1a2"
   },
   "outputs": [],
   "source": [
    "import nest_asyncio\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import asyncio\n",
    "import aiohttp\n",
    "import json\n",
    "import re\n",
    "import ast\n",
    "\n",
    "# ---------------------------\n",
    "# Local AI\n",
    "# ---------------------------\n",
    "from ollama import AsyncClient\n",
    "OLLAMA_BASE_URL='http://localhost:11434'\n",
    "\n",
    "# ---------------------------\n",
    "# Configuration Constants\n",
    "# ---------------------------\n",
    "JINA_API_KEY = \"jina_xxxx\" # Replace with your Jina API key\n",
    "\n",
    "JINA_BASE_URL = \"https://r.jina.ai/\"\n",
    "\n",
    "DEFAULT_MODEL = \"mistral-small\"\n",
    "REASON_MODEL = \"deepseek-r1:14b\"\n",
    "\n",
    "BASE_SEARXNG_URL = \"http://localhost:4000/search\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "GJTo96a7DGUz",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# ============================\n",
    "# Asynchronous Helper Functions\n",
    "# ============================\n",
    "\n",
    "async def call_ollama_async(session, messages, model=DEFAULT_MODEL):\n",
    "    \"\"\"\n",
    "    Asynchronously call the Ollama chat completion API with the provided messages.\n",
    "    Returns the content of the assistant’s reply.\n",
    "    \"\"\"\n",
    "    try:\n",
    "        client = AsyncClient(host=OLLAMA_BASE_URL)\n",
    "        response_content = \"\"\n",
    "\n",
    "        async for part in await client.chat(model=model, messages=messages, stream=True):\n",
    "            if 'message' in part and 'content' in part['message']:\n",
    "                response_content += part['message']['content']\n",
    "\n",
    "        return response_content if response_content else None\n",
    "    except Exception as e:\n",
    "        print(f\"Error calling Ollama API: {e}\")\n",
    "        return None\n",
    "\n",
    "\n",
    "async def make_initial_searching_plan_async(session, user_query):\n",
    "    \"\"\"\n",
    "    Ask the reasoning LLMs to produce a research plan based on the user’s query.\n",
    "    \"\"\"\n",
    "    prompt = (\n",
    "        \"You are an advanced reasoning LLM that specializes in structuring and refining research plans. Based on the given user query,\"\n",
    "        \"you will generate a comprehensive research plan that expands on the topic, identifies key areas of investigation, and breaks down\"\n",
    "        \" the research process into actionable steps for a search agent to execute.\\n\"\n",
    "        \"Process:\\n\"\n",
    "        \"Expand the Query: 1. Clarify and enrich the user’s query by considering related aspects, possible interpretations,\"\n",
    "        \"and necessary contextual details. 2.Identify any ambiguities and resolve them by assuming the most logical and useful framing of the problem.\\n\"\n",
    "        \"Identify Key Research Areas: 1. Break down the expanded query into core themes, subtopics, or dimensions of investigation.\"\n",
    "        \"2.Determine what information is necessary to provide a comprehensive answer.\\n\"\n",
    "        \"Define Research Steps: 1. Outline a structured plan with clear steps that guide the search agent on how to gather information.\"\n",
    "        \"2. Specify which sources or types of data are most relevant (e.g., academic papers, government reports, news sources, expert opinions).\"\n",
    "        \"3. Prioritize steps based on importance and logical sequence.\\n\"\n",
    "        \"Suggest Search Strategies: 1.Recommend search terms, keywords, and boolean operators to optimize search efficiency.\"\n",
    "        \"2. Identify useful databases, journals, and sources where high-quality information can be found.\"\n",
    "        \"3. Suggest methodologies for verifying credibility and synthesizing findings.\\n\"\n",
    "        \"NO EXPLANATIONS, write plans ONLY!\\n\"\n",
    "    )\n",
    "    messages = [\n",
    "        {\"role\": \"system\", \"content\": \"You are an advanced reasoning LLM that guides a following search agent to search for relevant information for a research.\"},\n",
    "        {\"role\": \"user\", \"content\": f\"User Query: {user_query}\\n\\n{prompt}\"}\n",
    "    ]\n",
    "    response = await call_ollama_async(session, messages, model=REASON_MODEL)\n",
    "    if response:\n",
    "        try:\n",
    "            # Remove <think>...</think> tags and their content if they exist\n",
    "            cleaned_response = re.sub(r\"<think>.*?</think>\", \"\", response, flags=re.DOTALL).strip()\n",
    "            print(f\"Plan:{cleaned_response}\")\n",
    "            return cleaned_response\n",
    "        except Exception as e:\n",
    "            print(f\"Error processing response: {e}\")\n",
    "    return []\n",
    "\n",
    "async def judge_search_result_and_future_plan_aync(session, user_query, original_plan, context_combined):\n",
    "    \"\"\"\n",
    "    Ask the reasoning LLMs to judge the result of the search attempt and produce a plan for next interation.\n",
    "    \"\"\"\n",
    "    prompt = (\n",
    "    \"You are an advanced reasoning LLM that specializes in evaluating research results and refining search strategies. \"\n",
    "    \"Your task is to analyze the search agent’s findings, assess their relevance and completeness, \"\n",
    "    \"and generate a structured plan for the next search iteration. Your goal is to ensure a thorough and efficient research process \"\n",
    "    \"that ultimately provides a comprehensive answer to the user’s query. But still, if you think everything is enough, you can tell search agent to stop\\n\"\n",
    "    \"Process:\\n\"\n",
    "    \"1. **Evaluate Search Results:**\\n\"\n",
    "    \"   - Analyze the retrieved search results to determine their relevance, credibility, and completeness.\\n\"\n",
    "    \"   - Identify missing information, knowledge gaps, or weak sources.\\n\"\n",
    "    \"   - Assess whether the search results sufficiently address the key research areas from the original plan.\\n\"\n",
    "    \"   - If everything is enough, tell search agent to stop with your reason\\n\"\n",
    "    \"2. **Determine Next Steps:**\\n\"\n",
    "    \"   - Based on gaps identified, refine or expand the research focus.\\n\"\n",
    "    \"   - Suggest additional search directions or alternative sources to explore.\\n\"\n",
    "    \"   - If necessary, propose adjustments to search strategies, including keyword modifications, new sources, or filtering techniques.\\n\"\n",
    "    \"3. **Generate an Updated Research Plan:**\\n\"\n",
    "    \"   - Provide a structured step-by-step plan for the next search iteration.\\n\"\n",
    "    \"   - Clearly outline what aspects need further investigation and where the search agent should focus next.\\n\"\n",
    "    \"NO EXPLANATIONS, write plans ONLY!\\n\"\n",
    "    \"Now, based on the above information and instruction, evaluate the search results and generate a refined research plan for the next iteration.\"\n",
    "    )\n",
    "\n",
    "    messages = [\n",
    "        {\"role\": \"system\", \"content\": \"You are an advanced reasoning LLM that guides a following search agent to search for relevant information for a research.\"},\n",
    "        {\"role\": \"user\", \"content\": f\"User Query: {user_query}\\n\\n Original Research Plan: {original_plan} \\n\\nExtracted Relevant Contexts by the search agent:\\n{context_combined} \\n\\n{prompt}\"}\n",
    "    ]\n",
    "    response = await call_ollama_async(session, messages, model=REASON_MODEL)\n",
    "    if response:\n",
    "        try:\n",
    "            # Remove <think>...</think> tags and their content if they exist\n",
    "            cleaned_response = re.sub(r\"<think>.*?</think>\", \"\", response, flags=re.DOTALL).strip()\n",
    "            print(f\"Next Plan:{cleaned_response}\")\n",
    "            return cleaned_response\n",
    "        except Exception as e:\n",
    "            print(f\"Error processing response: {e}\")\n",
    "    return []\n",
    "\n",
    "async def generate_writing_plan_aync(session, user_query, aggregated_contexts):\n",
    "    \"\"\"\n",
    "    Ask the reasoning LLM to generate a structured writing plan for the final report based on the aggregated research findings.\n",
    "    \"\"\"\n",
    "    prompt = (\n",
    "        \"You are an advanced reasoning LLM that specializes in structuring comprehensive research reports. \"\n",
    "        \"Your task is to analyze the aggregated research findings from previous searches and generate a structured plan \"\n",
    "        \"for writing the final report. The goal is to ensure that the report is well-organized, complete, and effectively presents the findings.\\n\\n\"\n",
    "        \"Process:\\n\"\n",
    "        \"1. **Analyze Aggregated Findings:**\\n\"\n",
    "        \"   - Review the provided research findings for relevance, accuracy, and coherence.\\n\"\n",
    "        \"   - Identify key insights, supporting evidence, and significant conclusions.\\n\"\n",
    "        \"   - Highlight any inconsistencies or areas that may need further clarification.\\n\\n\"\n",
    "        \"2. **Determine Report Structure:**\\n\"\n",
    "        \"   - Define a clear outline with logical sections that effectively communicate the research results.\\n\"\n",
    "        \"   - Ensure the structure follows a coherent flow, such as Introduction, Key Findings, Analysis, Conclusion, and References.\\n\"\n",
    "        \"   - Specify where different pieces of evidence should be integrated within the report.\\n\\n\"\n",
    "        \"3. **Provide Writing Guidelines:**\\n\"\n",
    "        \"   - Suggest the tone and style appropriate for the report (e.g., formal, analytical, concise).\\n\"\n",
    "        \"   - Recommend how to synthesize multiple sources into a coherent narrative.\\n\"\n",
    "        \"   - If any section lacks sufficient information, indicate where additional elaboration may be needed.\\n\\n\"\n",
    "        \"NO EXPLANATIONS, write plans ONLY!\\n\"\n",
    "        \"Now, based on the above instructions, generate a structured plan for writing the final research report.\\n\"\n",
    "    )\n",
    "\n",
    "    messages = [\n",
    "        {\"role\": \"system\", \"content\": \"You are an advanced reasoning LLM that structures research findings into a well-organized final report.\"},\n",
    "        {\"role\": \"user\", \"content\": f\"User Query: {user_query}\\n\\nAggregated Research Findings:\\n{aggregated_contexts}\\n\\n{prompt}\"}\n",
    "    ]\n",
    "    response = await call_ollama_async(session, messages, model=REASON_MODEL)\n",
    "    if response:\n",
    "        try:\n",
    "            # Remove <think>...</think> tags and their content if they exist\n",
    "            cleaned_response = re.sub(r\"<think>.*?</think>\", \"\", response, flags=re.DOTALL).strip()\n",
    "            print(f\"Writing Plan: {cleaned_response}\")\n",
    "            return cleaned_response\n",
    "        except Exception as e:\n",
    "            print(f\"Error processing response: {e}\")\n",
    "    return []\n",
    "\n",
    "    \n",
    "async def generate_search_queries_async(session, query_plan):\n",
    "    \"\"\"\n",
    "    Ask the LLM to produce up to four precise search queries (in Python list format)\n",
    "    based on the user’s query.\n",
    "    \"\"\"\n",
    "    prompt = (\n",
    "        \"You are an expert research assistant. Given the user's query, generate up to four distinct, \"\n",
    "        \"precise search queries that would help gather comprehensive information on the topic. \"\n",
    "        \"Return ONLY a Python list of strings, for example: ['query1', 'query2', 'query3'].\"\n",
    "    )\n",
    "    messages = [\n",
    "        {\"role\": \"system\", \"content\": \"You are a helpful and precise research assistant.\"},\n",
    "        {\"role\": \"user\", \"content\": f\"{query_plan}\\n\\n{prompt}\"}\n",
    "    ]\n",
    "    response = await call_ollama_async(session, messages)\n",
    "    if response:\n",
    "        cleaned = response.strip()\n",
    "        \n",
    "        # Remove triple backticks and language specifier if present\n",
    "        cleaned = re.sub(r\"```(?:\\w+)?\\n(.*?)\\n```\", r\"\\1\", cleaned, flags=re.DOTALL).strip()\n",
    "\n",
    "        # First, try to directly evaluate the cleaned string\n",
    "        try:\n",
    "            new_queries = ast.literal_eval(cleaned)\n",
    "            if isinstance(new_queries, list):\n",
    "                return new_queries\n",
    "        except Exception as e:\n",
    "            # Direct evaluation failed; try to extract the list part from the string.\n",
    "            match = re.search(r'(\\[.*\\])', cleaned, re.DOTALL)\n",
    "            if match:\n",
    "                list_str = match.group(1)\n",
    "                try:\n",
    "                    new_queries = ast.literal_eval(list_str)\n",
    "                    if isinstance(new_queries, list):\n",
    "                        return new_queries\n",
    "                except Exception as e_inner:\n",
    "                    print(\"Error parsing extracted list:\", e_inner, \"\\nExtracted text:\", list_str)\n",
    "                    return []\n",
    "            print(\"Error parsing new search queries:\", e, \"\\nResponse:\", response)\n",
    "            return []\n",
    "    return []\n",
    "\n",
    "\n",
    "async def perform_search_async(session, query):\n",
    "    params = {\n",
    "        \"q\": query,\n",
    "        \"format\": \"json\"\n",
    "    }\n",
    "    try:\n",
    "        async with session.get(BASE_SEARXNG_URL, params=params) as resp:\n",
    "            if resp.status == 200:\n",
    "                results = await resp.json()\n",
    "                if \"results\" in results:\n",
    "                    # Extract the URLs from the returned results list.\n",
    "                    links = [item.get(\"url\") for item in results[\"results\"] if \"url\" in item]\n",
    "                    return links\n",
    "                else:\n",
    "                    print(\"No results in SearXNG response.\")\n",
    "                    return []\n",
    "            else:\n",
    "                text = await resp.text()\n",
    "                print(f\"SearXNG error: {resp.status} - {text}\")\n",
    "                return []\n",
    "    except Exception as e:\n",
    "        print(\"Error performing SearXNG search:\", e)\n",
    "        return []\n",
    "\n",
    "\n",
    "async def fetch_webpage_text_async(session, url):\n",
    "    \"\"\"\n",
    "    Asynchronously retrieve the text content of a webpage using Jina.\n",
    "    The URL is appended to the Jina endpoint.\n",
    "    \"\"\"\n",
    "    full_url = f\"{JINA_BASE_URL}{url}\"\n",
    "    headers = {\n",
    "        \"Authorization\": f\"Bearer {JINA_API_KEY}\"\n",
    "    }\n",
    "    try:\n",
    "        async with session.get(full_url, headers=headers) as resp:\n",
    "            if resp.status == 200:\n",
    "                return await resp.text()\n",
    "            else:\n",
    "                text = await resp.text()\n",
    "                print(f\"Jina fetch error for {url}: {resp.status} - {text}\")\n",
    "                return \"\"\n",
    "    except Exception as e:\n",
    "        print(\"Error fetching webpage text with Jina:\", e)\n",
    "        return \"\"\n",
    "\n",
    "\n",
    "async def is_page_useful_async(session, user_query, page_text, page_url):\n",
    "    \"\"\"\n",
    "    Ask the LLM if the provided webpage content is useful for answering the user's query.\n",
    "    The LLM must reply with exactly \"Yes\" or \"No\".\n",
    "    \"\"\"\n",
    "    prompt = (\n",
    "        \"You are a critical research evaluator. Given the user's query and the content of a webpage, \"\n",
    "        \"determine if the webpage contains information relevant and useful for addressing the query. \"\n",
    "        \"Respond with exactly one word: 'Yes' if the page is useful, or 'No' if it is not. Do not include any extra text.\"\n",
    "    )\n",
    "    messages = [\n",
    "        {\"role\": \"system\", \"content\": \"You are a strict and concise evaluator of research relevance.\"},\n",
    "        {\"role\": \"user\", \"content\": f\"User Query: {user_query}\\n\\nWeb URL: {page_url}\\n\\nWebpage Content (first 20000 characters):\\n{page_text[:20000]}\\n\\n{prompt}\"}\n",
    "    ]\n",
    "    response = await call_ollama_async(session, messages)\n",
    "    if response:\n",
    "        answer = response.strip()\n",
    "        if answer in [\"Yes\", \"No\"]:\n",
    "            return answer\n",
    "        else:\n",
    "            # Fallback: try to extract Yes/No from the response.\n",
    "            if \"Yes\" in answer:\n",
    "                return \"Yes\"\n",
    "            elif \"No\" in answer:\n",
    "                return \"No\"\n",
    "    return \"No\"\n",
    "\n",
    "\n",
    "async def extract_relevant_context_async(session, user_query, search_query, page_text, page_url):\n",
    "    \"\"\"\n",
    "    Given the original query, the search query used, and the page content,\n",
    "    have the LLM extract all information relevant for answering the query.\n",
    "    \"\"\"\n",
    "    prompt = (\n",
    "        \"You are an expert information extractor. Given the user's query, the search query that led to this page, \"\n",
    "        \"and the webpage content, extract all pieces of information that are relevant to answering the user's query. \"\n",
    "        \"Return only the relevant context as plain text without commentary.\"\n",
    "    )\n",
    "    messages = [\n",
    "        {\"role\": \"system\", \"content\": \"You are an expert in extracting and summarizing relevant information.\"},\n",
    "        {\"role\": \"user\", \"content\": f\"User Query: {user_query}\\nSearch Query: {search_query}\\n\\nWeb URL: {page_url}\\n\\nWebpage Content (first 20000 characters):\\n{page_text[:20000]}\\n\\n{prompt}\"}\n",
    "    ]\n",
    "    response = await call_ollama_async(session, messages)\n",
    "    if response:\n",
    "        return response.strip()\n",
    "    return \"\"\n",
    "\n",
    "\n",
    "async def get_new_search_queries_async(session, user_query, new_research_plan, previous_search_queries, all_contexts):\n",
    "    \"\"\"\n",
    "    Based on the original query, the previously used search queries, and all the extracted contexts,\n",
    "    ask the LLM whether additional search queries are needed. If yes, return a Python list of up to four queries;\n",
    "    if the LLM thinks research is complete, it should return \"<done>\".\n",
    "    \"\"\"\n",
    "    context_combined = \"\\n\".join(all_contexts)\n",
    "    prompt = (\n",
    "        \"You are an analytical research assistant. Based on the original query, the search queries performed so far, \"\n",
    "        \"the next step plan by a planning agent and the extracted contexts from webpages, determine if further research is needed. \"\n",
    "        \"If further research is needed, ONLY provide up to four new search queries as a Python list IN ONE LINE (for example, \"\n",
    "        \"['new query1', 'new query2']) in PLAIN text, NEVER wrap in code env. If you believe no further research is needed, respond with exactly <done>.\"\n",
    "        \"\\nREMEMBER: Output ONLY a Python list or the token <done> WITHOUT any additional text or explanations.\"\n",
    "    )\n",
    "    messages = [\n",
    "        {\"role\": \"system\", \"content\": \"You are a systematic research planner.\"},\n",
    "        {\"role\": \"user\", \"content\": f\"User Query: {user_query}\\nPrevious Search Queries: {previous_search_queries}\\n\\nExtracted Relevant Contexts:\\n{context_combined}Next Research Plan by planning agent:{new_research_plan}\\n\\n{prompt}\"}\n",
    "    ]\n",
    "    response = await call_ollama_async(session, messages)\n",
    "    if response:\n",
    "        cleaned = response.strip()\n",
    "        if cleaned == \"<done>\":\n",
    "            return \"<done>\"\n",
    "        # Remove triple backticks and language specifier if present\n",
    "        cleaned = re.sub(r\"```(?:\\w+)?\\n(.*?)\\n```\", r\"\\1\", cleaned, flags=re.DOTALL).strip()\n",
    "        # First, try to directly evaluate the cleaned string\n",
    "        try:\n",
    "            new_queries = ast.literal_eval(cleaned)\n",
    "            if isinstance(new_queries, list):\n",
    "                return new_queries\n",
    "        except Exception as e:\n",
    "            # Direct evaluation failed; try to extract the list part from the string.\n",
    "            match = re.search(r'(\\[.*\\])', cleaned, re.DOTALL)\n",
    "            if match:\n",
    "                list_str = match.group(1)\n",
    "                try:\n",
    "                    new_queries = ast.literal_eval(list_str)\n",
    "                    if isinstance(new_queries, list):\n",
    "                        return new_queries\n",
    "                except Exception as e_inner:\n",
    "                    print(\"Error parsing extracted list:\", e_inner, \"\\nExtracted text:\", list_str)\n",
    "                    return []\n",
    "            print(\"Error parsing new search queries:\", e, \"\\nResponse:\", response)\n",
    "            return []\n",
    "    return []\n",
    "\n",
    "\n",
    "\n",
    "async def generate_final_report_async(session, user_query, report_planning, all_contexts):\n",
    "    \"\"\"\n",
    "    Generate the final comprehensive report using all gathered contexts.\n",
    "    \"\"\"\n",
    "    context_combined = \"\\n\".join(all_contexts)\n",
    "    prompt = (\n",
    "        \"You are an expert researcher and report writer. Based on the gathered contexts below and the original query, \"\n",
    "        \"write a comprehensive, well-structured, and detailed report that addresses the query thoroughly. \"\n",
    "        \"Include all relevant insights and conclusions without extraneous commentary.\"\n",
    "        \"Math equations should use proper LaTeX syntax in markdown format, with $\\\\LaTeX{}$ for inline, $$\\\\LaTeX{}$$ for block.\"\n",
    "        \"Properly cite all the sources inline from 'Gathered Relevant Contexts' with [cite_number],\"\n",
    "        \"and a corresponding bibliography list with their urls in markdown format in the end.\"\n",
    "    )\n",
    "    messages = [\n",
    "        {\"role\": \"system\", \"content\": \"You are a skilled report writer.\"},\n",
    "        {\"role\": \"user\", \"content\": f\"User Query: {user_query}\\n\\nGathered Relevant Contexts:\\n{context_combined}\\n\\nWriting plan from a planning agent:\\n{report_planning}\\n\\nWriting Instructions:{prompt}\"}\n",
    "    ]\n",
    "    report = await call_ollama_async(session, messages)\n",
    "    return report\n",
    "\n",
    "\n",
    "async def process_link(session, link, user_query, search_query):\n",
    "    \"\"\"\n",
    "    Process a single link: fetch its content, judge its usefulness, and if useful, extract the relevant context.\n",
    "    \"\"\"\n",
    "    print(f\"Fetching content from: {link}\")\n",
    "    page_text = await fetch_webpage_text_async(session, link)\n",
    "    if not page_text:\n",
    "        return None\n",
    "    usefulness = await is_page_useful_async(session, user_query, page_text, link)\n",
    "    print(f\"Page usefulness for {link}: {usefulness}\")\n",
    "    if usefulness == \"Yes\":\n",
    "        context = await extract_relevant_context_async(session, user_query, search_query, page_text, link)\n",
    "        if context:\n",
    "            print(f\"Extracted context from {link} (first 200 chars): {context[:200]}\")\n",
    "            context=\"url:\"+link+\"\\ncontext:\"+context\n",
    "            return context\n",
    "    return None\n",
    "\n",
    "\n",
    "# =========================\n",
    "# Main Asynchronous Routine\n",
    "# =========================\n",
    "\n",
    "async def async_main():\n",
    "    user_query = input(\"Enter your research query/topic: \").strip()\n",
    "    iter_limit_input = input(\"Enter maximum number of iterations (default 10): \").strip()\n",
    "    max_search_items = input(\"Enter maximum number of URLs to go through per iteration per query (default 4): \").strip()\n",
    "    iteration_limit = int(iter_limit_input) if iter_limit_input.isdigit() else 10\n",
    "    max_search_items = int(max_search_items) if max_search_items.isdigit() else 4\n",
    "\n",
    "    aggregated_contexts = []    # All useful contexts from every iteration\n",
    "    all_search_queries = []     # Every search query used across iterations\n",
    "    iteration = 0\n",
    "\n",
    "    async with aiohttp.ClientSession() as session:\n",
    "        # ----- INITIAL SEARCH QUERIES -----\n",
    "        reseach_plan = await make_initial_searching_plan_async(session, user_query)\n",
    "        initial_query = \"User Query:\" + user_query + \"\\n\\nResaerch Plan:\" + reseach_plan\n",
    "        new_search_queries = await generate_search_queries_async(session, initial_query)\n",
    "        if not new_search_queries:\n",
    "            print(\"No search queries were generated by the LLM. Exiting.\")\n",
    "            return\n",
    "        all_search_queries.extend(new_search_queries)\n",
    "\n",
    "        # ----- ITERATIVE RESEARCH LOOP -----\n",
    "        while iteration < iteration_limit:\n",
    "            print(f\"\\n=== Iteration {iteration + 1} ===\")\n",
    "            iteration_contexts = []\n",
    "\n",
    "            # For each search query, perform searxng searches concurrently.\n",
    "            search_tasks = [asyncio.create_task(perform_search_async(session, query)) for query in new_search_queries]\n",
    "            # Wait for all searches to complete\n",
    "            full_search_results = await asyncio.gather(*search_tasks)\n",
    "            # Ensure only max_search_items per query are saved\n",
    "            search_results = [result[:max_search_items] for result in full_search_results]# limit for local models\n",
    "\n",
    "            # Aggregate all unique links from all search queries of this iteration.\n",
    "            # Map each unique link to the search query that produced it.\n",
    "            unique_links = {}\n",
    "            for idx, links in enumerate(search_results):\n",
    "                query = new_search_queries[idx]\n",
    "                for link in links:\n",
    "                    if link not in unique_links:\n",
    "                        unique_links[link] = query\n",
    "\n",
    "            print(f\"Aggregated {len(unique_links)} unique links from this iteration.\")\n",
    "\n",
    "            # Process each link concurrently: fetch, judge, and extract context.\n",
    "            link_tasks = [\n",
    "                process_link(session, link, user_query, unique_links[link])\n",
    "                for link in unique_links\n",
    "            ]\n",
    "            link_results = await asyncio.gather(*link_tasks)\n",
    "\n",
    "            # Collect non-None contexts.\n",
    "            for res in link_results:\n",
    "                if res:\n",
    "                    iteration_contexts.append(res)\n",
    "\n",
    "            if iteration_contexts:\n",
    "                aggregated_contexts.extend(iteration_contexts)\n",
    "            else:\n",
    "                print(\"No useful contexts were found in this iteration.\")\n",
    "\n",
    "            # ----- ASK THE LLM IF MORE SEARCHES ARE NEEDED AND NOT THE FINIAL ROUND -----\n",
    "            if iteration + 1 < iteration_limit:\n",
    "                new_research_plan = await judge_search_result_and_future_plan_aync(session, user_query, reseach_plan, aggregated_contexts)\n",
    "                new_search_queries = await get_new_search_queries_async(session, user_query, new_research_plan, all_search_queries, aggregated_contexts)\n",
    "                if new_search_queries == \"<done>\":\n",
    "                    print(\"LLM indicated that no further research is needed.\")\n",
    "                    break\n",
    "                elif new_search_queries:\n",
    "                    print(\"LLM provided new search queries:\", new_search_queries)\n",
    "                    all_search_queries.extend(new_search_queries)\n",
    "                else:\n",
    "                    print(\"LLM did not provide any new search queries. Ending the loop.\")\n",
    "                    break\n",
    "\n",
    "            iteration += 1\n",
    "\n",
    "        # ----- FINAL REPORT -----\n",
    "        print(\"\\nGenerating final report...\")\n",
    "        final_report_planning = await generate_writing_plan_aync(session, user_query, aggregated_contexts)\n",
    "        final_report = await generate_final_report_async(session, user_query, final_report_planning, aggregated_contexts)\n",
    "        print(\"\\n==== FINAL REPORT ====\\n\")\n",
    "        print(final_report)\n",
    "\n",
    "\n",
    "def main():\n",
    "    asyncio.run(async_main())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "46Q5XpapDJZT",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    main()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}