Claude_Researcher.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyORAqWyBDR5YL8GM+6Xtyy1",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/mshumer/ai-researcher/blob/main/Claude_Researcher.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install serpapi\n",
        "\n",
        "import requests\n",
        "import json\n",
        "import time\n",
        "import ast\n",
        "import re\n",
        "\n",
        "ANTHROPIC_API_KEY = \"YOUR KEY HERE\"  # Replace with your Anthropic API key\n",
        "SERP_API_KEY = \"YOUR KEY HERE\"  # Replace with your SERP API key"
      ],
      "metadata": {
        "id": "aCGIopPW_RMi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def remove_first_line(test_string):\n",
        "    if test_string.startswith(\"Here\") and test_string.split(\"\\n\")[0].strip().endswith(\":\"):\n",
        "        return re.sub(r'^.*\\n', '', test_string, count=1)\n",
        "    return test_string\n",
        "\n",
        "def generate_text(prompt, model=\"claude-3-haiku-20240307\", max_tokens=2000, temperature=0.7):\n",
        "    headers = {\n",
        "        \"x-api-key\": ANTHROPIC_API_KEY,\n",
        "        \"anthropic-version\": \"2023-06-01\",\n",
        "        \"content-type\": \"application/json\"\n",
        "    }\n",
        "    data = {\n",
        "        \"model\": model,\n",
        "        \"max_tokens\": max_tokens,\n",
        "        \"temperature\": temperature,\n",
        "        \"system\": \"You are a world-class researcher. Analyze the given information and generate a well-structured report.\",\n",
        "        \"messages\": [{\"role\": \"user\", \"content\": prompt}],\n",
        "    }\n",
        "    response = requests.post(\"https://api.anthropic.com/v1/messages\", headers=headers, json=data)\n",
        "    print(response.json())\n",
        "    response_text = response.json()['content'][0]['text']\n",
        "    print(remove_first_line(response_text.strip()))\n",
        "    return remove_first_line(response_text.strip())\n",
        "\n",
        "def search_web(search_term):\n",
        "    url = f\"https://serpapi.com/search.json?q={search_term}&api_key={SERP_API_KEY}\"\n",
        "    response = requests.get(url)\n",
        "    data = response.json()\n",
        "    print(data)\n",
        "    return data\n",
        "\n",
        "def generate_subtopic_report(subtopic):\n",
        "    search_data = []\n",
        "    all_queries = []\n",
        "\n",
        "    print(f\"Generating initial search queries for subtopic: {subtopic}...\")\n",
        "    initial_queries_prompt = f\"Generate 3 search queries to gather information on the subtopic '{subtopic}'. Return your queries in a Python-parseable list. Return nothing but the list. Do so in one line. Start your response with [\\\"\"\n",
        "    initial_queries = ast.literal_eval('[' + generate_text(initial_queries_prompt).split('[')[1])\n",
        "    print(initial_queries)\n",
        "    all_queries.extend(initial_queries)\n",
        "\n",
        "    for i in range(3):\n",
        "        print(f\"Performing search round {i+1} for subtopic: {subtopic}...\")\n",
        "        for query in initial_queries:\n",
        "            search_results = search_web(query)\n",
        "            search_data.append(search_results)\n",
        "\n",
        "        print(f\"Generating additional search queries for subtopic: {subtopic}...\")\n",
        "        additional_queries_prompt = f\"Here are the search results so far for the subtopic '{subtopic}':\\n\\n{str(search_data)}\\n\\n---\\n\\nHere are all the search queries you have used so far for this subtopic:\\n\\n{str(all_queries)}\\n\\n---\\n\\nBased on the search results and previous queries, generate 3 new and unique search queries to expand the knowledge on the subtopic '{subtopic}'. Return your queries in a Python-parseable list. Return nothing but the list. Do so in one line. Start your response with [\\\"\"\n",
        "        additional_queries = ast.literal_eval('[' + generate_text(additional_queries_prompt).split('[')[1])\n",
        "        initial_queries = additional_queries\n",
        "        all_queries.extend(additional_queries)\n",
        "\n",
        "    print(f\"Generating initial report for subtopic: {subtopic}...\")\n",
        "    report_prompt = f\"When writing your report, make it incredibly detailed, thorough, specific, and well-structured. Use Markdown for formatting. Analyze the following search data and generate a comprehensive report on the subtopic '{subtopic}':\\n\\n{str(search_data)}\"\n",
        "    report = generate_text(report_prompt, max_tokens=4000)\n",
        "\n",
        "    for i in range(3):\n",
        "        print(f\"Analyzing report and generating additional searches (Round {i+1}) for subtopic: {subtopic}...\")\n",
        "        analysis_prompt = f\"Analyze the following report on the subtopic '{subtopic}' and identify areas that need more detail or further information:\\n\\n{report}\\n\\n---\\n\\nHere are all the search queries you have used so far for this subtopic:\\n\\n{str(all_queries)}\\n\\n---\\n\\nGenerate 3 new and unique search queries to fill in the gaps and provide more detail on the identified areas. Return your queries in a Python-parseable list. Return nothing but the list. Do so in one line. Start your response with [\\\"\"\n",
        "        additional_queries = ast.literal_eval('[' + generate_text(analysis_prompt).split('[')[1])\n",
        "        all_queries.extend(additional_queries)\n",
        "\n",
        "        round_search_data = []\n",
        "        for query in additional_queries:\n",
        "            search_results = search_web(query)\n",
        "            round_search_data.append(search_results)\n",
        "\n",
        "        print(f\"Updating report with additional information (Round {i+1}) for subtopic: {subtopic}...\")\n",
        "        update_prompt = f\"Update the following report on the subtopic '{subtopic}' by incorporating the new information from the additional searches. Keep all existing information... only add new information:\\n\\n{report}\\n\\n---\\n\\nAdditional search data for this round:\\n\\n{str(round_search_data)}\\n\\n---\\n\\nGenerate an updated report that includes the new information and provides more detail in the identified areas. Use Markdown for formatting.\"\n",
        "        report = generate_text(update_prompt, max_tokens=4000)\n",
        "\n",
        "    print(f\"Generating boss feedback for subtopic: {subtopic}...\")\n",
        "    feedback_prompt = f\"Imagine you are the boss reviewing the following report on the subtopic '{subtopic}':\\n\\n{report}\\n\\n---\\n\\nProvide constructive feedback on what information is missing or needs further elaboration in the report. Be specific and detailed in your feedback.\"\n",
        "    feedback = generate_text(feedback_prompt, max_tokens=1000)\n",
        "\n",
        "    print(f\"Generating final round of searches based on feedback for subtopic: {subtopic}...\")\n",
        "    final_queries_prompt = f\"Based on the following feedback from the boss regarding the subtopic '{subtopic}':\\n\\n{feedback}\\n\\n---\\n\\nGenerate 3 search queries to find the missing information and address the areas that need further elaboration. Return your queries in a Python-parseable list. Return nothing but the list. Do so in one line. Start your response with [\\\"\"\n",
        "    final_queries = ast.literal_eval('[' + generate_text(final_queries_prompt).split('[')[1])\n",
        "    all_queries.extend(final_queries)\n",
        "\n",
        "    final_search_data = []\n",
        "    for query in final_queries:\n",
        "        search_results = search_web(query)\n",
        "        final_search_data.append(search_results)\n",
        "\n",
        "    print(f\"Updating report with final information for subtopic: {subtopic}...\")\n",
        "    final_update_prompt = f\"Update the following report on the subtopic '{subtopic}' by incorporating the new information from the final round of searches based on the boss's feedback:\\n\\n{report}\\n\\n---\\n\\nFinal search data:\\n\\n{str(final_search_data)}\\n\\n---\\n\\nGenerate the final report that addresses the boss's feedback and includes the missing information. Use Markdown for formatting.\"\n",
        "    final_report = generate_text(final_update_prompt, max_tokens=4000)\n",
        "\n",
        "    print(f\"Final report generated for subtopic: {subtopic}!\")\n",
        "    return final_report\n",
        "\n",
        "def generate_comprehensive_report(topic, subtopic_reports):\n",
        "    print(\"Generating comprehensive report...\")\n",
        "    comprehensive_report_prompt = f\"Generate a comprehensive report on the topic '{topic}' by combining the following reports on various subtopics:\\n\\n{subtopic_reports}\\n\\n---\\n\\nEnsure that the final report is well-structured, coherent, and covers all the important aspects of the topic. Make sure that it includes EVERYTHING in each of the reports, in a better structured, more info-heavy manner. Nothing -- absolutely nothing, should be left out. If you forget to include ANYTHING from any of the previous reports, you will face the consequences. Include a table of contents. Leave nothing out. Use Markdown for formatting.\"\n",
        "    comprehensive_report = generate_text(comprehensive_report_prompt, model=\"claude-3-opus-20240229\", max_tokens=4000)\n",
        "\n",
        "    print(\"Comprehensive report generated!\")\n",
        "    return comprehensive_report\n",
        "\n",
        "# User input\n",
        "research_topic = input(\"Enter the research topic: \")\n",
        "\n",
        "# Generate subtopic checklist\n",
        "subtopic_checklist_prompt = f\"Generate a detailed checklist of subtopics to research for the topic '{research_topic}'. Return your checklist in a Python-parseable list. Return nothing but the list. Do so in one line. Maximum five sub-topics. Start your response with [\\\"\"\n",
        "subtopic_checklist = ast.literal_eval('[' + generate_text(subtopic_checklist_prompt).split('[')[1])\n",
        "print(f\"Subtopic Checklist: {subtopic_checklist}\")\n",
        "\n",
        "# Generate reports for each subtopic\n",
        "subtopic_reports = []\n",
        "for subtopic in subtopic_checklist:\n",
        "    subtopic_report = generate_subtopic_report(subtopic)\n",
        "    subtopic_reports.append(subtopic_report)\n",
        "\n",
        "# Combine subtopic reports into a comprehensive report\n",
        "comprehensive_report = generate_comprehensive_report(research_topic, \"\\n\\n\".join(subtopic_reports))\n",
        "\n",
        "# Save the comprehensive report to a file\n",
        "with open(\"comprehensive_report.txt\", \"w\") as file:\n",
        "    file.write(comprehensive_report)\n",
        "\n",
        "print(\"Comprehensive report saved as 'comprehensive_report.txt'.\")"
      ],
      "metadata": {
        "id": "KHJcqwWe4rgN"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}