jupyter-legacy/local_open_deep_researcher_via_playwright.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install nest_asyncio asyncio aiohttp ollama playwright docling"
   ]
  },
  {
   "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",
    "import time\n",
    "from pathlib import Path\n",
    "from collections import defaultdict\n",
    "from urllib.parse import urlparse\n",
    "import mimetypes\n",
    "from playwright.async_api import async_playwright\n",
    "from docling.document_converter import DocumentConverter\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",
    "DEFAULT_MODEL = \"mistral-small\"\n",
    "REASON_MODEL = \"deepseek-r1:14b\"\n",
    "\n",
    "BASE_SEARXNG_URL = \"http://localhost:4000/search\"\n",
    "\n",
    "\n",
    "# ---------------------\n",
    "# Concurrency control\n",
    "# ---------------------\n",
    "concurrent_limit = 3  # Max concurrent tasks for website visiting\n",
    "cool_down = 10.0  # Cooldown time per domain in seconds to avoid rate limits\n",
    "CHROME_PORT = 9222  # Port for connecting to a running Chrome via `chromium --remote-debugging-port=9222 --user-data-dir=/path/to/profile`\n",
    "global_semaphore = asyncio.Semaphore(concurrent_limit)\n",
    "domain_locks = defaultdict(asyncio.Lock)  # domain -> asyncio.Lock\n",
    "domain_next_allowed_time = defaultdict(lambda: 0.0)  # domain -> float (epoch time)\n",
    "\n",
    "# ---------------------\n",
    "# Parsing settings\n",
    "# ---------------------\n",
    "TEMP_PDF_DIR = Path(\"./temp_pdf\")  # Directory for storing downloaded PDFs\n",
    "BROWSE_LITE = 0 # whether to parse webpage with reader-lm and parse pdf with docling or not\n",
    "PDF_MAX_PAGES=30\n",
    "PDF_MAX_FILESIZE=20971520\n",
    "TIMEOUT_PDF=75\n",
    "MAX_HTML_LENGTH = 5120\n",
    "MAX_EVAL_TIME = 15  # Maximum seconds for JavaScript execution to clean HTML"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# --------------------------\n",
    "# Local AI and Browser use\n",
    "# --------------------------\n",
    "\n",
    "async def call_ollama_async(session, messages, model=DEFAULT_MODEL, max_tokens=20000):\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, options=dict(num_predict=max_tokens)):\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",
    "def is_pdf_url(url):\n",
    "    parsed_url = urlparse(url)\n",
    "    if parsed_url.path.lower().endswith(\".pdf\"):\n",
    "        return True\n",
    "    mime_type, _ = mimetypes.guess_type(url)\n",
    "    return mime_type == \"application/pdf\"\n",
    "\n",
    "def get_domain(url):\n",
    "    parsed = urlparse(url)\n",
    "    return parsed.netloc.lower()\n",
    "\n",
    "# Global lock to ensure only one PDF processing task runs at a time\n",
    "pdf_processing_lock = asyncio.Lock()\n",
    "\n",
    "async def process_pdf(pdf_path):\n",
    "    \"\"\"\n",
    "    Converts a local PDF file to text using Docling.\n",
    "    Ensures only one PDF processing task runs at a time to prevent GPU OutOfMemoryError.\n",
    "    \"\"\"\n",
    "    converter = DocumentConverter()\n",
    "\n",
    "    async def docling_task():\n",
    "        return converter.convert(str(pdf_path), max_num_pages=PDF_MAX_PAGES, max_file_size=PDF_MAX_FILESIZE)\n",
    "\n",
    "    # Ensure no other async task runs while processing the PDF\n",
    "    async with pdf_processing_lock:\n",
    "        try:\n",
    "            return await asyncio.wait_for(docling_task(), TIMEOUT_PDF)  # Enforce timeout\n",
    "        except asyncio.TimeoutError:\n",
    "            return \"Parser unable to parse the resource within defined time.\"\n",
    "\n",
    "async def download_pdf(page, url):\n",
    "    \"\"\"\n",
    "    Downloads a PDF from a webpage using Playwright and saves it locally.\n",
    "    \"\"\"\n",
    "    pdf_filename = TEMP_PDF_DIR / f\"{hash(url)}.pdf\"\n",
    "\n",
    "    async def intercept_request(request):\n",
    "        \"\"\"Intercepts request to log PDF download attempts.\"\"\"\n",
    "        if request.resource_type == \"document\":\n",
    "            print(f\"Downloading PDF: {request.url}\")\n",
    "\n",
    "    # Attach the request listener\n",
    "    page.on(\"request\", intercept_request)\n",
    "\n",
    "    try:\n",
    "        await page.goto(url, timeout=30000)  # 30 seconds timeout\n",
    "        await page.wait_for_load_state(\"networkidle\")\n",
    "        \n",
    "        # Attempt to save PDF (works for direct links)\n",
    "        await page.pdf(path=str(pdf_filename))\n",
    "        return pdf_filename\n",
    "\n",
    "    except Exception as e:\n",
    "        print(f\"Error downloading PDF {url}: {e}\")\n",
    "        return None\n",
    "\n",
    "async def get_cleaned_html(page):\n",
    "    \"\"\"\n",
    "    Extracts cleaned HTML from a page while enforcing a timeout.\n",
    "    \"\"\"\n",
    "    try:\n",
    "        cleaned_html = await asyncio.wait_for(\n",
    "            page.evaluate(\"\"\"\n",
    "                () => {\n",
    "                    let clone = document.cloneNode(true);\n",
    "                    \n",
    "                    // Remove script, style, and noscript elements\n",
    "                    clone.querySelectorAll('script, style, noscript').forEach(el => el.remove());\n",
    "                    \n",
    "                    // Optionally remove navigation and footer\n",
    "                    clone.querySelectorAll('nav, footer, aside').forEach(el => el.remove());\n",
    "                    \n",
    "                    return clone.body.innerHTML;\n",
    "                }\n",
    "            \"\"\"),\n",
    "            timeout=MAX_EVAL_TIME\n",
    "        )\n",
    "        return cleaned_html\n",
    "    except asyncio.TimeoutError:\n",
    "        return \"Parser unable to extract HTML within defined time.\"\n",
    "\n",
    "\n",
    "async def fetch_webpage_text_async(session, url):\n",
    "    \"\"\"\n",
    "    Fetches webpage HTML using Playwright, processes it using Ollama reader-lm:1.5b,\n",
    "    or downloads and processes a PDF using Docling.\n",
    "    Respects concurrency limits and per-domain cooldown.\n",
    "    \"\"\"\n",
    "    domain = get_domain(url)\n",
    "\n",
    "    async with global_semaphore:  # Global concurrency limit\n",
    "        async with domain_locks[domain]:  # Ensure only one request per domain at a time\n",
    "            now = time.time()\n",
    "            if now < domain_next_allowed_time[domain]:\n",
    "                await asyncio.sleep(domain_next_allowed_time[domain] - now)  # Respect per-domain cooldown\n",
    "\n",
    "            async with async_playwright() as p:\n",
    "                # Attempt to connect to an already running Chrome instance\n",
    "                try:\n",
    "                    browser = await p.chromium.connect_over_cdp(f\"http://localhost:{CHROME_PORT}\")\n",
    "                except Exception as e:\n",
    "                    print(f\"Error connecting to Chrome on port {CHROME_PORT}: {e}\")\n",
    "                    return f\"Failed to connect to Chrome on port {CHROME_PORT}\"\n",
    "\n",
    "                context = await browser.new_context()\n",
    "                page = await context.new_page()\n",
    "\n",
    "                # PDFs\n",
    "                if is_pdf_url(url):\n",
    "                    if BROWSE_LITE:\n",
    "                        result = \"PDF parsing is disabled in lite browsing mode.\"\n",
    "                    else:\n",
    "                        pdf_path = await download_pdf(page, url)\n",
    "                        if pdf_path:\n",
    "                            text = await process_pdf(pdf_path)\n",
    "                            result = f\"# PDF Content\\n{text}\"\n",
    "                        else:\n",
    "                            result = \"Failed to download or process PDF\"\n",
    "                else:\n",
    "                    try:\n",
    "                        await page.goto(url, timeout=30000) # 30 seconds timeout to wait for loading\n",
    "                        title = await page.title() or \"Untitled Page\"\n",
    "                        if BROWSE_LITE:\n",
    "                            # Extract main content using JavaScript inside Playwright\n",
    "                            main_content = await page.evaluate(\"\"\"\n",
    "                                () => {\n",
    "                                    let mainEl = document.querySelector('main') || document.body;\n",
    "                                    return mainEl.innerText.trim();\n",
    "                                }\n",
    "                            \"\"\")\n",
    "                            result = f\"# {title}\\n{main_content}\"\n",
    "                        else:\n",
    "                            # Clean HTML before sending to reader-lm\n",
    "                            cleaned_html = await get_cleaned_html(page)\n",
    "                            cleaned_html = cleaned_html[:MAX_HTML_LENGTH] # Enforce a Maximum length for a webpage\n",
    "                            messages = [{\"role\": \"user\", \"content\": cleaned_html}]\n",
    "                            markdown_text = await call_ollama_async(session, messages, model=\"reader-lm:0.5b\", max_tokens=int(1.25*MAX_HTML_LENGTH)) # Don't get stuck when exceed reader-lm ctx\n",
    "                            result = f\"# {title}\\n{markdown_text}\"\n",
    "\n",
    "                    except Exception as e:\n",
    "                        print(f\"Error fetching webpage: {e}\")\n",
    "                        result = f\"Failed to fetch {url}\"\n",
    "\n",
    "                await browser.close()\n",
    "\n",
    "            # Update next allowed time for this domain (cool down time per domain)\n",
    "            domain_next_allowed_time[domain] = time.time() + cool_down\n",
    "\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "GJTo96a7DGUz",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# ============================\n",
    "# Asynchronous Helper Functions\n",
    "# ============================\n",
    "\n",
    "async def generate_search_queries_async(session, user_query):\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\"User Query: {user_query}\\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 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, 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}\\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",
    "        \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",
    "\n",
    "\n",
    "async def generate_final_report_async(session, user_query, 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 and a bibliography list with their title and url 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\\n{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",
    "        new_search_queries = await generate_search_queries_async(session, user_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 -----\n",
    "            if iteration + 1 < iteration_limit:\n",
    "                new_search_queries = await get_new_search_queries_async(session, user_query, 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 = await generate_final_report_async(session, user_query, 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
}