third_party/Wikipedia/wikipedia-search-cookbook.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Iteratively Searching Wikipedia with Claude"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[DISCLAIMER: This notebook was created using Claude 2 models and is considered legacy.]\n",
    "\n",
    "Some questions can't be answered by Claude off the top of Claude's head. Maybe they're about current events. Maybe you have an intensely detailed question that Claude hasn't memorized the answer to. No worries! With some prompting and scaffolding, Claude can search the web to find answers. In this notebook, we will create a virtual research assistant who has the ability to search Wikipedia to find answers to your question. The same approach can be used to allow Claude to search the broader web, or a set of documents you provide.\n",
    "\n",
    "What is the approach? Broadly it falls under the category of \"tool use\". We create a search tool, tell Claude about it, and let it go to work. In pseudocode:\n",
    "\n",
    "1. Prompt Claude with a description of the search tool, how it's best used, and how to \"call\" it (by issuing a special string).\n",
    "2. Tell Claude your question.\n",
    "3. Claude produces tokens like normal. If it produces the special string, terminate the token production stream, and issue a query to a search API.\n",
    "4. Construct a new prompt which consists of the prompt from step 1, plus everything Claude generated up to the search call string, plus the results of the API call.\n",
    "5. Repeat until Claude decides it's done.\n",
    "\n",
    "Let's zoom in on the prompts for tool use and retrieval."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prompts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "You will be asked a question by a human user. You have access to the following tool to help answer the question. <tool_description> Search Engine Tool * The search engine will exclusively search over Wikipedia for pages similar to your query. It returns for each page its title and full page content. Use this tool if you want to get up-to-date and comprehensive information on a topic to help answer queries. Queries should be as atomic as possible -- they only need to address one part of the user's question. For example, if the user's query is \"what is the color of a basketball?\", your search query should be \"basketball\". Here's another example: if the user's question is \"Who created the first neural network?\", your first query should be \"neural network\". As you can see, these queries are quite short. Think keywords, not phrases. * At any time, you can make a call to the search engine using the following syntax: <search_query>query_word</search_query>. * You'll then get results back in <search_result> tags.</tool_description>\n"
     ]
    }
   ],
   "source": [
    "# Tool Description Prompt\n",
    "wikipedia_prompt = \"\"\"You will be asked a question by a human user. You have access to the following tool to help answer the question. <tool_description> Search Engine Tool * The search engine will exclusively search over Wikipedia for pages similar to your query. It returns for each page its title and full page content. Use this tool if you want to get up-to-date and comprehensive information on a topic to help answer queries. Queries should be as atomic as possible -- they only need to address one part of the user's question. For example, if the user's query is \"what is the color of a basketball?\", your search query should be \"basketball\". Here's another example: if the user's question is \"Who created the first neural network?\", your first query should be \"neural network\". As you can see, these queries are quite short. Think keywords, not phrases. * At any time, you can make a call to the search engine using the following syntax: <search_query>query_word</search_query>. * You'll then get results back in <search_result> tags.</tool_description>\"\"\"\n",
    "print(wikipedia_prompt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that there is a lot of advice in this prompt about how to search Wikipedia properly. We're all used to just typing random nonsense into Google and getting decent results because the query parsing logic is so good. Wikipedia search is not like that. As an example: consider the query \"What's the best way to purchase potatoes in the United Arab Emirates\". The [top hits for this on Wikipedia](https://en.wikipedia.org/w/index.php?search=What%27s+the+best+way+to+purchase+potatoes+in+the+United+Arab+Emirates&title=Special:Search&profile=advanced&fulltext=1&ns0=1) are for Slavery in the United States, 1973 Oil Crisis, Wendy's, and Tim Horton's (??). Meanwhile Google correctly takes you straight to Carrefour UAE."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Another difference is that Wikipedia search returns entire pages. With vector search, you might be getting narrower chunks, so you might want to ask for more results, use a more specific query, or both. The big-picture takeaway is that your results can vary a lot on your choices here so pay attention!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Before beginning to research the user's question, first think for a moment inside <scratchpad> tags about what information is necessary for a well-informed answer. If the user's question is complex, you may need to decompose the query into multiple subqueries and execute them individually. Sometimes the search engine will return empty search results, or the search results may not contain the information you need. In such cases, feel free to try again with a different query. \n",
      "\n",
      "After each call to the Search Engine Tool, reflect briefly inside <search_quality></search_quality> tags about whether you now have enough information to answer, or whether more information is needed. If you have all the relevant information, write it in <information></information> tags, WITHOUT actually answering the question. Otherwise, issue a new search.\n",
      "\n",
      "Here is the user's question: <question>{query}</question> Remind yourself to make short queries in your scratchpad as you plan out your strategy.\n"
     ]
    }
   ],
   "source": [
    "retrieval_prompt = \"\"\"Before beginning to research the user's question, first think for a moment inside <scratchpad> tags about what information is necessary for a well-informed answer. If the user's question is complex, you may need to decompose the query into multiple subqueries and execute them individually. Sometimes the search engine will return empty search results, or the search results may not contain the information you need. In such cases, feel free to try again with a different query. \n",
    "\n",
    "After each call to the Search Engine Tool, reflect briefly inside <search_quality></search_quality> tags about whether you now have enough information to answer, or whether more information is needed. If you have all the relevant information, write it in <information></information> tags, WITHOUT actually answering the question. Otherwise, issue a new search.\n",
    "\n",
    "Here is the user's question: <question>{query}</question> Remind yourself to make short queries in your scratchpad as you plan out your strategy.\"\"\"\n",
    "print(retrieval_prompt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "We use a scratchpad here for the normal chain-of-thought reasons -- it makes Claude come up with a coherent plan to answer the question. The search quality reflection is used to induce Claude to be persistent and not jump the gun by answering the question before gathering all the relevant information. But why are we telling Claude to synthesize the information and not answer right away?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Here is a user query: <query>{query}</query>. Here is some relevant information: <information>{information}</information>. Please answer the question using the relevant information.\n"
     ]
    }
   ],
   "source": [
    "answer_prompt = \"Here is a user query: <query>{query}</query>. Here is some relevant information: <information>{information}</information>. Please answer the question using the relevant information.\"\n",
    "print(answer_prompt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By extracting the information and presenting it to Claude in a new query, we allow Claude to focus all its attention on synthesizing the information into the right answer. Without this step, we found that Claude would sometimes precommit to an answer and then \"justify\" it with the search results, rather than allowing the results to guide it."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now follows a bunch of code that implements the pseudocode for searching + retrieving + reprompting.\n",
    "\n",
    "### Search Implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dataclasses import dataclass\n",
    "from abc import ABC, abstractmethod\n",
    "import wikipedia, re\n",
    "from anthropic import Anthropic, HUMAN_PROMPT, AI_PROMPT\n",
    "from typing import Tuple, Optional\n",
    "\n",
    "@dataclass\n",
    "class SearchResult:\n",
    "    \"\"\"\n",
    "    A single search result.\n",
    "    \"\"\"\n",
    "    content: str\n",
    "\n",
    "class SearchTool:\n",
    "    \"\"\"\n",
    "    A search tool that can run a query and return a formatted string of search results.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__():\n",
    "        pass\n",
    "\n",
    "    @abstractmethod\n",
    "    def raw_search(self, query: str, n_search_results_to_use: int) -> list[SearchResult]:\n",
    "        \"\"\"\n",
    "        Runs a query using the searcher, then returns the raw search results without formatting.\n",
    "\n",
    "        :param query: The query to run.\n",
    "        :param n_search_results_to_use: The number of results to return.\n",
    "        \"\"\"\n",
    "        raise NotImplementedError()\n",
    "    \n",
    "    @abstractmethod\n",
    "    def process_raw_search_results(\n",
    "        self, results: list[SearchResult],\n",
    "    ) -> list[str]:\n",
    "        \"\"\"\n",
    "        Extracts the raw search content from the search results and returns a list of strings that can be passed to Claude.\n",
    "\n",
    "        :param results: The search results to extract.\n",
    "        \"\"\"\n",
    "        raise NotImplementedError()\n",
    "    \n",
    "    def search_results_to_string(self, extracted: list[str]) -> str:\n",
    "        \"\"\"\n",
    "        Joins and formats the extracted search results as a string.\n",
    "\n",
    "        :param extracted: The extracted search results to format.\n",
    "        \"\"\"\n",
    "        result = \"\\n\".join(\n",
    "            [\n",
    "                f'<item index=\"{i+1}\">\\n<page_content>\\n{r}\\n</page_content>\\n</item>'\n",
    "                for i, r in enumerate(extracted)\n",
    "            ]\n",
    "        )\n",
    "        return result\n",
    "\n",
    "    def wrap_search_results(self, extracted: list[str]) -> str:\n",
    "        \"\"\"\n",
    "        Formats the extracted search results as a string, including the <search_results> tags.\n",
    "\n",
    "        :param extracted: The extracted search results to format.\n",
    "        \"\"\"\n",
    "        return f\"\\n<search_results>\\n{self.search_results_to_string(extracted)}\\n</search_results>\"\n",
    "    \n",
    "    def search(self, query: str, n_search_results_to_use: int) -> str:\n",
    "        raw_search_results = self.raw_search(query, n_search_results_to_use)\n",
    "        processed_search_results = self.process_raw_search_results(raw_search_results)\n",
    "        displayable_search_results = self.wrap_search_results(processed_search_results)\n",
    "        return displayable_search_results "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class WikipediaSearchResult(SearchResult):\n",
    "    title: str\n",
    "    \n",
    "class WikipediaSearchTool(SearchTool):\n",
    "\n",
    "    def __init__(self,\n",
    "                 truncate_to_n_tokens: Optional[int] = 5000):\n",
    "        self.truncate_to_n_tokens = truncate_to_n_tokens\n",
    "        if truncate_to_n_tokens is not None:\n",
    "            self.tokenizer = Anthropic().get_tokenizer()\n",
    "\n",
    "    def raw_search(self, query: str, n_search_results_to_use: int) -> list[WikipediaSearchResult]:\n",
    "        search_results = self._search(query, n_search_results_to_use)\n",
    "        return search_results\n",
    "    \n",
    "    def process_raw_search_results(self, results: list[WikipediaSearchResult]) -> list[str]:\n",
    "        processed_search_results = [f'Page Title: {result.title.strip()}\\nPage Content:\\n{self.truncate_page_content(result.content)}' for result in results]\n",
    "        return processed_search_results\n",
    "\n",
    "    def truncate_page_content(self, page_content: str) -> str:\n",
    "        if self.truncate_to_n_tokens is None:\n",
    "            return page_content.strip()\n",
    "        else:\n",
    "            return self.tokenizer.decode(self.tokenizer.encode(page_content).ids[:self.truncate_to_n_tokens]).strip()\n",
    "        \n",
    "    def _search(self, query: str, n_search_results_to_use: int) -> list[WikipediaSearchResult]:\n",
    "        results: list[str] = wikipedia.search(query)\n",
    "        search_results: list[WikipediaSearchResult] = []\n",
    "        for result in results:\n",
    "            if len(search_results) >= n_search_results_to_use:\n",
    "                break\n",
    "            try:\n",
    "                page = wikipedia.page(result)\n",
    "                print(page.url)\n",
    "            except:\n",
    "                # The Wikipedia API is a little flaky, so we just skip over pages that fail to load\n",
    "                continue\n",
    "            content = page.content\n",
    "            title = page.title\n",
    "            search_results.append(WikipediaSearchResult(content=content, title=title))\n",
    "        return search_results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_between_tags(tag: str, string: str, strip: bool = True) -> list[str]:\n",
    "    ext_list = re.findall(f\"<{tag}\\s?>(.+?)</{tag}\\s?>\", string, re.DOTALL)\n",
    "    if strip:\n",
    "        ext_list = [e.strip() for e in ext_list]\n",
    "    return ext_list\n",
    "\n",
    "class ClientWithRetrieval(Anthropic):\n",
    "\n",
    "    def __init__(self, search_tool: SearchTool, verbose: bool = True, *args, **kwargs):\n",
    "        super().__init__(*args, **kwargs)\n",
    "        self.search_tool = search_tool\n",
    "        self.verbose = verbose\n",
    "\n",
    "    # Helper methods\n",
    "    def _search_query_stop(self, partial_completion: str, n_search_results_to_use: int) -> Tuple[list[SearchResult], str]:\n",
    "        search_query = extract_between_tags('search_query', partial_completion + '</search_query>') \n",
    "        if search_query is None:\n",
    "            raise Exception(f'Completion with retrieval failed as partial completion returned mismatched <search_query> tags.')\n",
    "        print(f'Running search query against SearchTool: {search_query}')\n",
    "        search_results = self.search_tool.raw_search(search_query, n_search_results_to_use)\n",
    "        extracted_search_results = self.search_tool.process_raw_search_results(search_results)\n",
    "        formatted_search_results = self.search_tool.wrap_search_results(extracted_search_results)\n",
    "        return search_results, formatted_search_results\n",
    "    \n",
    "    def retrieve(self,\n",
    "                       query: str,\n",
    "                       model: str,\n",
    "                       n_search_results_to_use: int = 3,\n",
    "                       stop_sequences: list[str] = [HUMAN_PROMPT],\n",
    "                       max_tokens_to_sample: int = 1000,\n",
    "                       max_searches_to_try: int = 5,\n",
    "                       temperature: float = 1.0) -> tuple[list[SearchResult], str]:\n",
    "        \n",
    "        prompt = f\"{HUMAN_PROMPT} {wikipedia_prompt} {retrieval_prompt.format(query=query)}{AI_PROMPT}\"\n",
    "        starting_prompt = prompt\n",
    "        print(\"Starting prompt:\", starting_prompt)\n",
    "        token_budget = max_tokens_to_sample\n",
    "        all_raw_search_results: list[SearchResult] = []\n",
    "        for tries in range(max_searches_to_try):\n",
    "            partial_completion = self.completions.create(prompt = prompt,\n",
    "                                                     stop_sequences=stop_sequences + ['</search_query>'],\n",
    "                                                     model=model,\n",
    "                                                     max_tokens_to_sample = token_budget,\n",
    "                                                     temperature = temperature)\n",
    "            partial_completion, stop_reason, stop_seq = partial_completion.completion, partial_completion.stop_reason, partial_completion.stop\n",
    "            print(partial_completion)\n",
    "            token_budget -= self.count_tokens(partial_completion)\n",
    "            prompt += partial_completion\n",
    "            if stop_reason == 'stop_sequence' and stop_seq == '</search_query>':\n",
    "                print(f'Attempting search number {tries}.')\n",
    "                raw_search_results, formatted_search_results = self._search_query_stop(partial_completion, n_search_results_to_use)\n",
    "                prompt += '</search_query>' + formatted_search_results\n",
    "                all_raw_search_results += raw_search_results\n",
    "            else:\n",
    "                break\n",
    "        final_model_response = prompt[len(starting_prompt):]\n",
    "        return all_raw_search_results, final_model_response\n",
    "    \n",
    "    # Main methods\n",
    "    def completion_with_retrieval(self,\n",
    "                                        query: str,\n",
    "                                        model: str,\n",
    "                                        n_search_results_to_use: int = 3,\n",
    "                                        stop_sequences: list[str] = [HUMAN_PROMPT],\n",
    "                                        max_tokens_to_sample: int = 1000,\n",
    "                                        max_searches_to_try: int = 5,\n",
    "                                        temperature: float = 1.0) -> str:\n",
    "        \n",
    "        _, retrieval_response = self.retrieve(query, model=model,\n",
    "                                                 n_search_results_to_use=n_search_results_to_use, stop_sequences=stop_sequences,\n",
    "                                                 max_tokens_to_sample=max_tokens_to_sample,\n",
    "                                                 max_searches_to_try=max_searches_to_try,\n",
    "                                                 temperature=temperature)\n",
    "        information = extract_between_tags('information', retrieval_response)[-1]\n",
    "        prompt = f\"{HUMAN_PROMPT} {answer_prompt.format(query=query, information=information)}{AI_PROMPT}\"\n",
    "        print(\"Summarizing:\\n\", prompt)\n",
    "        answer = self.completions.create(\n",
    "            prompt = prompt, model=model, temperature=temperature, max_tokens_to_sample=1000\n",
    "        ).completion\n",
    "        return answer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Running a Query"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We're ready to execute a query! Let's pick something:\n",
    "- recent, so it's less likely to be in Claude's training data, and\n",
    "- compound/complex so it requires multiple searches."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting prompt: \n",
      "\n",
      "Human: You will be asked a question by a human user. You have access to the following tool to help answer the question. <tool_description> Search Engine Tool * The search engine will exclusively search over Wikipedia for pages similar to your query. It returns for each page its title and full page content. Use this tool if you want to get up-to-date and comprehensive information on a topic to help answer queries. Queries should be as atomic as possible -- they only need to address one part of the user's question. For example, if the user's query is \"what is the color of a basketball?\", your search query should be \"basketball\". Here's another example: if the user's question is \"Who created the first neural network?\", your first query should be \"neural network\". As you can see, these queries are quite short. Think keywords, not phrases. * At any time, you can make a call to the search engine using the following syntax: <search_query>query_word</search_query>. * You'll then get results back in <search_result> tags.</tool_description> Before beginning to research the user's question, first think for a moment inside <scratchpad> tags about what information is necessary for a well-informed answer. If the user's question is complex, you may need to decompose the query into multiple subqueries and execute them individually. Sometimes the search engine will return empty search results, or the search results may not contain the information you need. In such cases, feel free to try again with a different query. \n",
      "\n",
      "After each call to the Search Engine Tool, reflect briefly inside <search_quality></search_quality> tags about whether you now have enough information to answer, or whether more information is needed. If you have all the relevant information, write it in <information></information> tags, WITHOUT actually answering the question. Otherwise, issue a new search.\n",
      "\n",
      "Here is the user's question: <question>Which movie came out first: Oppenheimer, or Are You There God It's Me Margaret?</question> Remind yourself to make short queries in your scratchpad as you plan out your strategy.\n",
      "\n",
      "Assistant:\n",
      " <scratchpad>\n",
      "To answer this question, I need to find the release dates for the two movies:\n",
      "- Oppenheimer release date\n",
      "- Are You There God It's Me Margaret release date\n",
      "I can search for each movie title individually to get the release date.\n",
      "</scratchpad>\n",
      "\n",
      "<search_query>Oppenheimer movie\n",
      "Attempting search number 0.\n",
      "Running search query against SearchTool: ['Oppenheimer movie']\n",
      "https://en.wikipedia.org/wiki/Oppenheimer_(film)\n",
      "\n",
      "\n",
      "The search results indicate that Oppenheimer is scheduled for theatrical release on July 21, 2023. This provides the release date for Oppenheimer.\n",
      "\n",
      "<search_quality>The search results directly provided the release date for Oppenheimer, so I now have enough information to answer this part of the question.</search_quality>\n",
      "\n",
      "<search_query>Are You There God It's Me Margaret movie\n",
      "Attempting search number 1.\n",
      "Running search query against SearchTool: [\"Are You There God It's Me Margaret movie\"]\n",
      "https://en.wikipedia.org/wiki/Are_You_There_God%3F_It%27s_Me,_Margaret.\n",
      "\n",
      "\n",
      "The search results indicate that the film adaptation of Are You There God? It's Me, Margaret was released on April 28, 2023. This provides the release date for Are You There God? It's Me, Margaret. \n",
      "\n",
      "<search_quality>The search results directly stated the release date for the Are You There God? It's Me, Margaret movie adaptation, so I now have enough information to fully answer the question.</search_quality>\n",
      "\n",
      "<information>\n",
      "- Oppenheimer was released on July 21, 2023\n",
      "- Are You There God? It's Me, Margaret was released on April 28, 2023\n",
      "</information>\n",
      "\n",
      "Based on the release dates found through my searches, Oppenheimer came out first, being released on July 21, 2023, while Are You There God? It's Me, Margaret was released later on April 28, 2023.\n",
      "Summarizing:\n",
      " \n",
      "\n",
      "Human:Here is a user query: <query>Which movie came out first: Oppenheimer, or Are You There God It's Me Margaret?</query>. Here is some relevant information: <information>- Oppenheimer was released on July 21, 2023\n",
      "- Are You There God? It's Me, Margaret was released on April 28, 2023</information>. Please answer the question using the relevant information.\n",
      "\n",
      "Assistant:\n",
      " Based on the information provided, Are You There God? It's Me, Margaret was released first, on April 28, 2023. Oppenheimer was released later, on July 21, 2023.\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "# Create a searcher\n",
    "wikipedia_search_tool = WikipediaSearchTool()\n",
    "ANTHROPIC_SEARCH_MODEL = \"claude-2\"\n",
    "\n",
    "client = ClientWithRetrieval(api_key=os.environ['ANTHROPIC_API_KEY'], verbose=True, search_tool = wikipedia_search_tool)\n",
    "\n",
    "query = \"Which movie came out first: Oppenheimer, or Are You There God It's Me Margaret?\"\n",
    "\n",
    "augmented_response = client.completion_with_retrieval(\n",
    "    query=query,\n",
    "    model=ANTHROPIC_SEARCH_MODEL,\n",
    "    n_search_results_to_use=1,\n",
    "    max_searches_to_try=5,\n",
    "    max_tokens_to_sample=1000,\n",
    "    temperature=0)\n",
    "print(augmented_response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Cool, Claude was able to make a plan, execute the queries, and synthesize the information into an accurate answer. Note: without the extra information extraction step, Claude would sometimes determine the release dates of the movies correctly but then get the ordering wrong in its final answer. Let's do another."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting prompt: \n",
      "\n",
      "Human: You will be asked a question by a human user. You have access to the following tool to help answer the question. <tool_description> Search Engine Tool * The search engine will exclusively search over Wikipedia for pages similar to your query. It returns for each page its title and full page content. Use this tool if you want to get up-to-date and comprehensive information on a topic to help answer queries. Queries should be as atomic as possible -- they only need to address one part of the user's question. For example, if the user's query is \"what is the color of a basketball?\", your search query should be \"basketball\". Here's another example: if the user's question is \"Who created the first neural network?\", your first query should be \"neural network\". As you can see, these queries are quite short. Think keywords, not phrases. * At any time, you can make a call to the search engine using the following syntax: <search_query>query_word</search_query>. * You'll then get results back in <search_result> tags.</tool_description> Before beginning to research the user's question, first think for a moment inside <scratchpad> tags about what information is necessary for a well-informed answer. If the user's question is complex, you may need to decompose the query into multiple subqueries and execute them individually. Sometimes the search engine will return empty search results, or the search results may not contain the information you need. In such cases, feel free to try again with a different query. \n",
      "\n",
      "After each call to the Search Engine Tool, reflect briefly inside <search_quality></search_quality> tags about whether you now have enough information to answer, or whether more information is needed. If you have all the relevant information, write it in <information></information> tags, WITHOUT actually answering the question. Otherwise, issue a new search.\n",
      "\n",
      "Here is the user's question: <question>Who won the 2023 NBA championship? Who was that team's best player in the year 2009?</question> Remind yourself to make short queries in your scratchpad as you plan out your strategy.\n",
      "\n",
      "Assistant:\n",
      " <scratchpad>\n",
      "To answer this question, I need to find:\n",
      "1. The team that won the 2023 NBA championship\n",
      "2. The best player on that team in 2009\n",
      "I can search for these things separately.\n",
      "</scratchpad>\n",
      "\n",
      "<search_query>2023 nba championship winner\n",
      "Attempting search number 0.\n",
      "Running search query against SearchTool: ['2023 nba championship winner']\n",
      "https://en.wikipedia.org/wiki/List_of_NBA_champions\n",
      "\n",
      "\n",
      "<search_quality>The search results contain the team that won the 2023 NBA championship, so I have the information I need to answer the first part of the question.</search_quality>\n",
      "\n",
      "<information>\n",
      "The Denver Nuggets won the 2023 NBA championship.\n",
      "</information>\n",
      "\n",
      "<search_query>denver nuggets best player 2009\n",
      "Attempting search number 1.\n",
      "Running search query against SearchTool: ['denver nuggets best player 2009']\n",
      "https://en.wikipedia.org/wiki/2009%E2%80%9310_Denver_Nuggets_season\n",
      "\n",
      "\n",
      "<search_quality>The search results indicate that Carmelo Anthony was the Nuggets' best player in 2009, so I now have all the information needed to fully answer the question.</search_quality>\n",
      "\n",
      "<information>\n",
      "- The Denver Nuggets won the 2023 NBA championship.\n",
      "- Carmelo Anthony was the Nuggets' best player in 2009.\n",
      "</information>\n",
      "Summarizing:\n",
      " \n",
      "\n",
      "Human:Here is a user query: <query>Who won the 2023 NBA championship? Who was that team's best player in the year 2009?</query>. Here is some relevant information: <information>- The Denver Nuggets won the 2023 NBA championship.\n",
      "- Carmelo Anthony was the Nuggets' best player in 2009.</information>. Please answer the question using the relevant information.\n",
      "\n",
      "Assistant:\n",
      " <response>\n",
      "Based on the provided information:\n",
      "- The Denver Nuggets won the 2023 NBA championship. \n",
      "- Carmelo Anthony was the Nuggets' best player in 2009.\n",
      "</response>\n"
     ]
    }
   ],
   "source": [
    "augmented_response = client.completion_with_retrieval(\n",
    "    query=\"Who won the 2023 NBA championship? Who was that team's best player in the year 2009?\",\n",
    "    model=ANTHROPIC_SEARCH_MODEL,\n",
    "    n_search_results_to_use=1,\n",
    "    max_searches_to_try=5,\n",
    "    max_tokens_to_sample=1000,\n",
    "    temperature=0)\n",
    "print(augmented_response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And there you have it! You may notice that the search tool code is nice and abstract and can be adapted to use a search API of your choice with minor modifications. Just remember to explain to Claude any tips it needs to use the tool well. You can even give Claude some few-shot examples of ideal query plans and query structure to improve performance further."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.11"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}