examples/qa/qa_quick_start.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "99f717ef-4cba-4300-b258-0b1c248cb873",
   "metadata": {},
   "source": [
    "# QA Quick Start\n",
    "\n",
    "In this QA Quick Start notebook, we take a first look at semantic ingestion, RAG, and retrieval, presenting basic usage patterns.\n",
    "\n",
    "For more advanced QA examples, check out [QA Deep Dive](./qa_deep_dive.ipynb).\n",
    "\n",
    "### Access required\n",
    "\n",
    "The content of this notebook requires access to Deep Search capabilities which are not\n",
    "available on the public access system.\n",
    "\n",
    "[Contact us](https://ds4sd.github.io) if you are interested in exploring\n",
    "these Deep Search capabilities.\n",
    "\n",
    "\n",
    "### GenAI Integration required\n",
    "\n",
    "When interacting with the virtual assistant, Deep Search requires a connection to a Generative AI API. Currently, we support connections to [watsonx.ai](https://www.ibm.com/products/watsonx-ai) or the IBM-internal GenAI platform BAM.\n",
    "\n",
    "Deep Search allows custom GenAI configurations for each project.\n",
    "In the following example you will require to work in a project which has such GenAI capabilities activated."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "256aef50-71a1-4278-9b22-17cb99a6566e",
   "metadata": {},
   "source": [
    "### Set notebook parameters\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5b244bdd-1b52-41ff-b63e-9a203570d210",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from dotenv import load_dotenv\n",
    "from pydantic import TypeAdapter\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "PROFILE_NAME = os.environ.get(\"DS_NB_PROFILE\")  # profile to use; defaults to active one\n",
    "PROJ_KEY = os.environ[\"DS_NB_PROJ_KEY\"]  # project to use\n",
    "INDEX_KEY = os.environ[\"DS_NB_QA_IDX_KEY\"]\n",
    "DOC_HASH = os.environ.get(\"DS_NB_QA_DOC_HASH\")  # set only when targeting a specific doc\n",
    "QUESTION = os.environ[\"DS_NB_QUESTION\"]\n",
    "\n",
    "# whether to ingest incrementally:\n",
    "SKIP_INGESTED_DOCS = TypeAdapter(bool).validate_python(\n",
    "    os.environ.get(\"DS_NB_SKIP_INGESTED_DOCS\", True)\n",
    ")\n",
    "RETR_K = os.environ.get(\"DS_NB_RETR_K\", 3)  # number of search results to retrieve\n",
    "GEN_TIMEOUT = os.environ.get(\"DS_NB_GEN_TIMEOUT\", 10)  # generation timeout in seconds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5269060-bb5f-4fe3-9b64-547202db6714",
   "metadata": {},
   "source": [
    "### Import example dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5d236ea0-db1c-4171-8e11-cdd0bad69d66",
   "metadata": {},
   "outputs": [],
   "source": [
    "import rich\n",
    "from deepsearch.cps.client.api import CpsApi\n",
    "from deepsearch.cps.client.components.documents import create_private_data_source\n",
    "from deepsearch.cps.queries import RAGQuery, SemanticQuery\n",
    "from deepsearch.cps.queries.results import RAGResult, SearchResult"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "293c249b-6018-46f2-b4d8-795f994d4729",
   "metadata": {},
   "source": [
    "### Connect to Deep Search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9c108432-a285-4c7b-a996-008ac3ff3d7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "api = CpsApi.from_env(profile_name=PROFILE_NAME)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a655e8ce",
   "metadata": {},
   "source": [
    "### Prepare data source\n",
    "\n",
    "The cell below shows how to configure a private data source, i.e. either a whole private collection (in which case `document_hash` should be `None` or omitted) or a given doc within one.\n",
    "\n",
    "For more details on data sources check out [QA Deep Dive](./qa_deep_dive.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "adc0f7a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_source = create_private_data_source(\n",
    "    proj_key=PROJ_KEY,\n",
    "    index_key=INDEX_KEY,\n",
    "    document_hash=DOC_HASH,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94119e3d",
   "metadata": {},
   "source": [
    "### Ingestion\n",
    "\n",
    "If your data source has not yet been semantically indexed, you can ingest it into the vector DB as shown below. Otherwise you can skip this step.\n",
    "\n",
    "Ingestion of already indexed docs is controlled via param `skip_ingested_docs`. \n",
    "\n",
    "Particularly when indexing whole collections, note that the larger the data source, the longer the ingestion duration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "947e87c5",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/pva/work/github.com/DS4SD/deepsearch-examples/.venv/lib/python3.10/site-packages/pydantic/main.py:347: UserWarning: Pydantic serializer warnings:\n",
      "  Expected `list[str]` but got `_LiteralGenericAlias` - serialized value may not be as expected\n",
      "  return self.__pydantic_serializer__.to_python(\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'ing_out': {}}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "task = api.documents.semantic_ingest(\n",
    "    project=PROJ_KEY,\n",
    "    data_source=data_source,\n",
    "    skip_ingested_docs=SKIP_INGESTED_DOCS,\n",
    ")\n",
    "\n",
    "# wait for the ingestion task to finish\n",
    "api.tasks.wait_for(PROJ_KEY, task.task_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb0e8542",
   "metadata": {},
   "source": [
    "### RAG\n",
    "\n",
    "The cell below demonstrates basic RAG usage.\n",
    "\n",
    "For more advanced usage and parametrization, check out [QA Deep Dive](./qa_deep_dive.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b960b309",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Where was the first European IBM lab located?\n",
       "</pre>\n"
      ],
      "text/plain": [
       "Where was the first European IBM lab located?\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">RAGResult</span><span style=\"font-weight: bold\">(</span>\n",
       "    <span style=\"color: #808000; text-decoration-color: #808000\">answers</span>=<span style=\"font-weight: bold\">[</span>\n",
       "        <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">RAGAnswerItem</span><span style=\"font-weight: bold\">(</span>\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">answer</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'The first European IBM lab was located in Adliswil, Switzerland, near Zurich.'</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">grounding</span>=<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">RAGGroundingInfo</span><span style=\"font-weight: bold\">(</span>\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">retr_items</span>=<span style=\"font-weight: bold\">[</span>\n",
       "                    <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">SearchResultItem</span><span style=\"font-weight: bold\">(</span>\n",
       "                        <span style=\"color: #808000; text-decoration-color: #808000\">doc_hash</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'b30bc667a324ae111d025526563b674a8d3fd869bc07c8fd204aa95b05d41f0c'</span>,\n",
       "                        <span style=\"color: #808000; text-decoration-color: #808000\">chunk</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'Switzerland\\nIn 1956, IBM opened their first European research laboratory in </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">Adliswil, Switzerland, near Zurich. The lab moved to its own campus in neighboring Rüschlikon in 1962. The Zurich </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">lab is staffed by a multicultural and interdisciplinary team of a few hundred permanent research staff members, </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">graduate students and post-doctoral fellows, representing about 45 nationalities. Collocated with the lab is a </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">Client Center (formerly the Industry Solutions Lab), an executive briefing facility demonstrating technology </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">prototypes and solutions.'</span>,\n",
       "                        <span style=\"color: #808000; text-decoration-color: #808000\">main_path</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'main-text.71'</span>,\n",
       "                        <span style=\"color: #808000; text-decoration-color: #808000\">path_group</span>=<span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'main-text.69'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'main-text.71'</span><span style=\"font-weight: bold\">]</span>,\n",
       "                        <span style=\"color: #808000; text-decoration-color: #808000\">source_is_text</span>=<span style=\"color: #00ff00; text-decoration-color: #00ff00; font-style: italic\">True</span>\n",
       "                    <span style=\"font-weight: bold\">)</span>\n",
       "                <span style=\"font-weight: bold\">]</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">gen_ctx_paths</span>=<span style=\"font-weight: bold\">[</span>\n",
       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'main-text.60'</span>,\n",
       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'main-text.61'</span>,\n",
       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'main-text.62'</span>,\n",
       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'main-text.63'</span>,\n",
       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'main-text.64'</span>,\n",
       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'main-text.65'</span>,\n",
       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'main-text.66'</span>,\n",
       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'main-text.67'</span>,\n",
       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'main-text.68'</span>,\n",
       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'main-text.69'</span>,\n",
       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'main-text.70'</span>,\n",
       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'main-text.71'</span>,\n",
       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'main-text.72'</span>,\n",
       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'main-text.73'</span>,\n",
       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'main-text.74'</span>,\n",
       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'main-text.75'</span>,\n",
       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'main-text.76'</span>,\n",
       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'main-text.77'</span>,\n",
       "                    <span style=\"color: #008000; text-decoration-color: #008000\">'main-text.78'</span>\n",
       "                <span style=\"font-weight: bold\">]</span>\n",
       "            <span style=\"font-weight: bold\">)</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">prompt</span>=<span style=\"color: #800080; text-decoration-color: #800080; font-style: italic\">None</span>\n",
       "        <span style=\"font-weight: bold\">)</span>\n",
       "    <span style=\"font-weight: bold\">]</span>,\n",
       "    <span style=\"color: #808000; text-decoration-color: #808000\">search_result_items</span>=<span style=\"font-weight: bold\">[</span>\n",
       "        <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">SearchResultItem</span><span style=\"font-weight: bold\">(</span>\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">doc_hash</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'b30bc667a324ae111d025526563b674a8d3fd869bc07c8fd204aa95b05d41f0c'</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">chunk</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'Switzerland\\nIn 1956, IBM opened their first European research laboratory in Adliswil, </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">Switzerland, near Zurich. The lab moved to its own campus in neighboring Rüschlikon in 1962. The Zurich lab is </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">staffed by a multicultural and interdisciplinary team of a few hundred permanent research staff members, graduate </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">students and post-doctoral fellows, representing about 45 nationalities. Collocated with the lab is a Client Center</span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">(formerly the Industry Solutions Lab), an executive briefing facility demonstrating technology prototypes and </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">solutions.'</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">main_path</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'main-text.71'</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">path_group</span>=<span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'main-text.69'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'main-text.71'</span><span style=\"font-weight: bold\">]</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">source_is_text</span>=<span style=\"color: #00ff00; text-decoration-color: #00ff00; font-style: italic\">True</span>\n",
       "        <span style=\"font-weight: bold\">)</span>,\n",
       "        <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">SearchResultItem</span><span style=\"font-weight: bold\">(</span>\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">doc_hash</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'b30bc667a324ae111d025526563b674a8d3fd869bc07c8fd204aa95b05d41f0c'</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">chunk</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'Switzerland\\nIBM Research-Zurich (previously called IBM Zurich Research Laboratory, ZRL) is the </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">European branch of IBM Research. It was opened in 1956 and is located in Rüschlikon, near Zurich, Switzerland.'</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">main_path</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'main-text.70'</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">path_group</span>=<span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'main-text.69'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'main-text.70'</span><span style=\"font-weight: bold\">]</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">source_is_text</span>=<span style=\"color: #00ff00; text-decoration-color: #00ff00; font-style: italic\">True</span>\n",
       "        <span style=\"font-weight: bold\">)</span>,\n",
       "        <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">SearchResultItem</span><span style=\"font-weight: bold\">(</span>\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">doc_hash</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'b30bc667a324ae111d025526563b674a8d3fd869bc07c8fd204aa95b05d41f0c'</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">chunk</span>=<span style=\"color: #008000; text-decoration-color: #008000\">\"History\\nThe roots of today's IBM Research began with the 1945 opening of the Watson Scientific </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">Computing Laboratory at Columbia University. $^{[4]}$ This was the first IBM laboratory devoted to pure science and</span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">later expanded into additional IBM Research locations in Westchester County, New York, starting in the 1950s, </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">$^{[5][6]}$ including the Thomas J. Watson Research Center in 1961. [5][6]\"</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">main_path</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'main-text.8'</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">path_group</span>=<span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'main-text.7'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'main-text.8'</span><span style=\"font-weight: bold\">]</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">source_is_text</span>=<span style=\"color: #00ff00; text-decoration-color: #00ff00; font-style: italic\">True</span>\n",
       "        <span style=\"font-weight: bold\">)</span>\n",
       "    <span style=\"font-weight: bold\">]</span>\n",
       "<span style=\"font-weight: bold\">)</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mRAGResult\u001b[0m\u001b[1m(\u001b[0m\n",
       "    \u001b[33manswers\u001b[0m=\u001b[1m[\u001b[0m\n",
       "        \u001b[1;35mRAGAnswerItem\u001b[0m\u001b[1m(\u001b[0m\n",
       "            \u001b[33manswer\u001b[0m=\u001b[32m'The first European IBM lab was located in Adliswil, Switzerland, near Zurich.'\u001b[0m,\n",
       "            \u001b[33mgrounding\u001b[0m=\u001b[1;35mRAGGroundingInfo\u001b[0m\u001b[1m(\u001b[0m\n",
       "                \u001b[33mretr_items\u001b[0m=\u001b[1m[\u001b[0m\n",
       "                    \u001b[1;35mSearchResultItem\u001b[0m\u001b[1m(\u001b[0m\n",
       "                        \u001b[33mdoc_hash\u001b[0m=\u001b[32m'b30bc667a324ae111d025526563b674a8d3fd869bc07c8fd204aa95b05d41f0c'\u001b[0m,\n",
       "                        \u001b[33mchunk\u001b[0m=\u001b[32m'Switzerland\\nIn 1956, IBM opened their first European research laboratory in \u001b[0m\n",
       "\u001b[32mAdliswil, Switzerland, near Zurich. The lab moved to its own campus in neighboring Rüschlikon in 1962. The Zurich \u001b[0m\n",
       "\u001b[32mlab is staffed by a multicultural and interdisciplinary team of a few hundred permanent research staff members, \u001b[0m\n",
       "\u001b[32mgraduate students and post-doctoral fellows, representing about 45 nationalities. Collocated with the lab is a \u001b[0m\n",
       "\u001b[32mClient Center \u001b[0m\u001b[32m(\u001b[0m\u001b[32mformerly the Industry Solutions Lab\u001b[0m\u001b[32m)\u001b[0m\u001b[32m, an executive briefing facility demonstrating technology \u001b[0m\n",
       "\u001b[32mprototypes and solutions.'\u001b[0m,\n",
       "                        \u001b[33mmain_path\u001b[0m=\u001b[32m'main-text.71'\u001b[0m,\n",
       "                        \u001b[33mpath_group\u001b[0m=\u001b[1m[\u001b[0m\u001b[32m'main-text.69'\u001b[0m, \u001b[32m'main-text.71'\u001b[0m\u001b[1m]\u001b[0m,\n",
       "                        \u001b[33msource_is_text\u001b[0m=\u001b[3;92mTrue\u001b[0m\n",
       "                    \u001b[1m)\u001b[0m\n",
       "                \u001b[1m]\u001b[0m,\n",
       "                \u001b[33mgen_ctx_paths\u001b[0m=\u001b[1m[\u001b[0m\n",
       "                    \u001b[32m'main-text.60'\u001b[0m,\n",
       "                    \u001b[32m'main-text.61'\u001b[0m,\n",
       "                    \u001b[32m'main-text.62'\u001b[0m,\n",
       "                    \u001b[32m'main-text.63'\u001b[0m,\n",
       "                    \u001b[32m'main-text.64'\u001b[0m,\n",
       "                    \u001b[32m'main-text.65'\u001b[0m,\n",
       "                    \u001b[32m'main-text.66'\u001b[0m,\n",
       "                    \u001b[32m'main-text.67'\u001b[0m,\n",
       "                    \u001b[32m'main-text.68'\u001b[0m,\n",
       "                    \u001b[32m'main-text.69'\u001b[0m,\n",
       "                    \u001b[32m'main-text.70'\u001b[0m,\n",
       "                    \u001b[32m'main-text.71'\u001b[0m,\n",
       "                    \u001b[32m'main-text.72'\u001b[0m,\n",
       "                    \u001b[32m'main-text.73'\u001b[0m,\n",
       "                    \u001b[32m'main-text.74'\u001b[0m,\n",
       "                    \u001b[32m'main-text.75'\u001b[0m,\n",
       "                    \u001b[32m'main-text.76'\u001b[0m,\n",
       "                    \u001b[32m'main-text.77'\u001b[0m,\n",
       "                    \u001b[32m'main-text.78'\u001b[0m\n",
       "                \u001b[1m]\u001b[0m\n",
       "            \u001b[1m)\u001b[0m,\n",
       "            \u001b[33mprompt\u001b[0m=\u001b[3;35mNone\u001b[0m\n",
       "        \u001b[1m)\u001b[0m\n",
       "    \u001b[1m]\u001b[0m,\n",
       "    \u001b[33msearch_result_items\u001b[0m=\u001b[1m[\u001b[0m\n",
       "        \u001b[1;35mSearchResultItem\u001b[0m\u001b[1m(\u001b[0m\n",
       "            \u001b[33mdoc_hash\u001b[0m=\u001b[32m'b30bc667a324ae111d025526563b674a8d3fd869bc07c8fd204aa95b05d41f0c'\u001b[0m,\n",
       "            \u001b[33mchunk\u001b[0m=\u001b[32m'Switzerland\\nIn 1956, IBM opened their first European research laboratory in Adliswil, \u001b[0m\n",
       "\u001b[32mSwitzerland, near Zurich. The lab moved to its own campus in neighboring Rüschlikon in 1962. The Zurich lab is \u001b[0m\n",
       "\u001b[32mstaffed by a multicultural and interdisciplinary team of a few hundred permanent research staff members, graduate \u001b[0m\n",
       "\u001b[32mstudents and post-doctoral fellows, representing about 45 nationalities. Collocated with the lab is a Client Center\u001b[0m\n",
       "\u001b[32m(\u001b[0m\u001b[32mformerly the Industry Solutions Lab\u001b[0m\u001b[32m)\u001b[0m\u001b[32m, an executive briefing facility demonstrating technology prototypes and \u001b[0m\n",
       "\u001b[32msolutions.'\u001b[0m,\n",
       "            \u001b[33mmain_path\u001b[0m=\u001b[32m'main-text.71'\u001b[0m,\n",
       "            \u001b[33mpath_group\u001b[0m=\u001b[1m[\u001b[0m\u001b[32m'main-text.69'\u001b[0m, \u001b[32m'main-text.71'\u001b[0m\u001b[1m]\u001b[0m,\n",
       "            \u001b[33msource_is_text\u001b[0m=\u001b[3;92mTrue\u001b[0m\n",
       "        \u001b[1m)\u001b[0m,\n",
       "        \u001b[1;35mSearchResultItem\u001b[0m\u001b[1m(\u001b[0m\n",
       "            \u001b[33mdoc_hash\u001b[0m=\u001b[32m'b30bc667a324ae111d025526563b674a8d3fd869bc07c8fd204aa95b05d41f0c'\u001b[0m,\n",
       "            \u001b[33mchunk\u001b[0m=\u001b[32m'Switzerland\\nIBM Research-Zurich \u001b[0m\u001b[32m(\u001b[0m\u001b[32mpreviously called IBM Zurich Research Laboratory, ZRL\u001b[0m\u001b[32m)\u001b[0m\u001b[32m is the \u001b[0m\n",
       "\u001b[32mEuropean branch of IBM Research. It was opened in 1956 and is located in Rüschlikon, near Zurich, Switzerland.'\u001b[0m,\n",
       "            \u001b[33mmain_path\u001b[0m=\u001b[32m'main-text.70'\u001b[0m,\n",
       "            \u001b[33mpath_group\u001b[0m=\u001b[1m[\u001b[0m\u001b[32m'main-text.69'\u001b[0m, \u001b[32m'main-text.70'\u001b[0m\u001b[1m]\u001b[0m,\n",
       "            \u001b[33msource_is_text\u001b[0m=\u001b[3;92mTrue\u001b[0m\n",
       "        \u001b[1m)\u001b[0m,\n",
       "        \u001b[1;35mSearchResultItem\u001b[0m\u001b[1m(\u001b[0m\n",
       "            \u001b[33mdoc_hash\u001b[0m=\u001b[32m'b30bc667a324ae111d025526563b674a8d3fd869bc07c8fd204aa95b05d41f0c'\u001b[0m,\n",
       "            \u001b[33mchunk\u001b[0m=\u001b[32m\"History\u001b[0m\u001b[32m\\nThe roots of today's IBM Research began with the 1945 opening of the Watson Scientific \u001b[0m\n",
       "\u001b[32mComputing Laboratory at Columbia University. $^\u001b[0m\u001b[32m{\u001b[0m\u001b[32m[\u001b[0m\u001b[32m4\u001b[0m\u001b[32m]\u001b[0m\u001b[32m}\u001b[0m\u001b[32m$ This was the first IBM laboratory devoted to pure science and\u001b[0m\n",
       "\u001b[32mlater expanded into additional IBM Research locations in Westchester County, New York, starting in the 1950s, \u001b[0m\n",
       "\u001b[32m$^\u001b[0m\u001b[32m{\u001b[0m\u001b[32m[\u001b[0m\u001b[32m5\u001b[0m\u001b[32m]\u001b[0m\u001b[32m[\u001b[0m\u001b[32m6\u001b[0m\u001b[32m]\u001b[0m\u001b[32m}\u001b[0m\u001b[32m$ including the Thomas J. Watson Research Center in 1961. \u001b[0m\u001b[32m[\u001b[0m\u001b[32m5\u001b[0m\u001b[32m]\u001b[0m\u001b[32m[\u001b[0m\u001b[32m6\u001b[0m\u001b[32m]\u001b[0m\u001b[32m\"\u001b[0m,\n",
       "            \u001b[33mmain_path\u001b[0m=\u001b[32m'main-text.8'\u001b[0m,\n",
       "            \u001b[33mpath_group\u001b[0m=\u001b[1m[\u001b[0m\u001b[32m'main-text.7'\u001b[0m, \u001b[32m'main-text.8'\u001b[0m\u001b[1m]\u001b[0m,\n",
       "            \u001b[33msource_is_text\u001b[0m=\u001b[3;92mTrue\u001b[0m\n",
       "        \u001b[1m)\u001b[0m\n",
       "    \u001b[1m]\u001b[0m\n",
       "\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "query = RAGQuery(\n",
    "    question=QUESTION,\n",
    "    project=PROJ_KEY,\n",
    "    data_source=data_source,\n",
    "    retr_k=RETR_K,  # optional\n",
    "    gen_timeout=GEN_TIMEOUT,  # optional\n",
    ")\n",
    "\n",
    "api_output = api.queries.run(query)\n",
    "result = RAGResult.from_api_output(api_output)\n",
    "\n",
    "rich.print(QUESTION)\n",
    "rich.print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "feef40f4",
   "metadata": {},
   "source": [
    "### Semantic retrieval\n",
    "\n",
    "In certain cases, a user may only be interested in the semantic retrieval part, instead of the whole RAG pipeline.\n",
    "\n",
    "Basic semantic retrieval usage is shown below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "492fd066",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Where was the first European IBM lab located?\n",
       "</pre>\n"
      ],
      "text/plain": [
       "Where was the first European IBM lab located?\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">SearchResult</span><span style=\"font-weight: bold\">(</span>\n",
       "    <span style=\"color: #808000; text-decoration-color: #808000\">search_result_items</span>=<span style=\"font-weight: bold\">[</span>\n",
       "        <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">SearchResultItem</span><span style=\"font-weight: bold\">(</span>\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">doc_hash</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'b30bc667a324ae111d025526563b674a8d3fd869bc07c8fd204aa95b05d41f0c'</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">chunk</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'Switzerland\\nIn 1956, IBM opened their first European research laboratory in Adliswil, </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">Switzerland, near Zurich. The lab moved to its own campus in neighboring Rüschlikon in 1962. The Zurich lab is </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">staffed by a multicultural and interdisciplinary team of a few hundred permanent research staff members, graduate </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">students and post-doctoral fellows, representing about 45 nationalities. Collocated with the lab is a Client Center</span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">(formerly the Industry Solutions Lab), an executive briefing facility demonstrating technology prototypes and </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">solutions.'</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">main_path</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'main-text.71'</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">path_group</span>=<span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'main-text.69'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'main-text.71'</span><span style=\"font-weight: bold\">]</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">source_is_text</span>=<span style=\"color: #00ff00; text-decoration-color: #00ff00; font-style: italic\">True</span>\n",
       "        <span style=\"font-weight: bold\">)</span>,\n",
       "        <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">SearchResultItem</span><span style=\"font-weight: bold\">(</span>\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">doc_hash</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'b30bc667a324ae111d025526563b674a8d3fd869bc07c8fd204aa95b05d41f0c'</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">chunk</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'Switzerland\\nIBM Research-Zurich (previously called IBM Zurich Research Laboratory, ZRL) is the </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">European branch of IBM Research. It was opened in 1956 and is located in Rüschlikon, near Zurich, Switzerland.'</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">main_path</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'main-text.70'</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">path_group</span>=<span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'main-text.69'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'main-text.70'</span><span style=\"font-weight: bold\">]</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">source_is_text</span>=<span style=\"color: #00ff00; text-decoration-color: #00ff00; font-style: italic\">True</span>\n",
       "        <span style=\"font-weight: bold\">)</span>,\n",
       "        <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">SearchResultItem</span><span style=\"font-weight: bold\">(</span>\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">doc_hash</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'b30bc667a324ae111d025526563b674a8d3fd869bc07c8fd204aa95b05d41f0c'</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">chunk</span>=<span style=\"color: #008000; text-decoration-color: #008000\">\"History\\nThe roots of today's IBM Research began with the 1945 opening of the Watson Scientific </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">Computing Laboratory at Columbia University. $^{[4]}$ This was the first IBM laboratory devoted to pure science and</span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">later expanded into additional IBM Research locations in Westchester County, New York, starting in the 1950s, </span>\n",
       "<span style=\"color: #008000; text-decoration-color: #008000\">$^{[5][6]}$ including the Thomas J. Watson Research Center in 1961. [5][6]\"</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">main_path</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'main-text.8'</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">path_group</span>=<span style=\"font-weight: bold\">[</span><span style=\"color: #008000; text-decoration-color: #008000\">'main-text.7'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'main-text.8'</span><span style=\"font-weight: bold\">]</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">source_is_text</span>=<span style=\"color: #00ff00; text-decoration-color: #00ff00; font-style: italic\">True</span>\n",
       "        <span style=\"font-weight: bold\">)</span>\n",
       "    <span style=\"font-weight: bold\">]</span>\n",
       "<span style=\"font-weight: bold\">)</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mSearchResult\u001b[0m\u001b[1m(\u001b[0m\n",
       "    \u001b[33msearch_result_items\u001b[0m=\u001b[1m[\u001b[0m\n",
       "        \u001b[1;35mSearchResultItem\u001b[0m\u001b[1m(\u001b[0m\n",
       "            \u001b[33mdoc_hash\u001b[0m=\u001b[32m'b30bc667a324ae111d025526563b674a8d3fd869bc07c8fd204aa95b05d41f0c'\u001b[0m,\n",
       "            \u001b[33mchunk\u001b[0m=\u001b[32m'Switzerland\\nIn 1956, IBM opened their first European research laboratory in Adliswil, \u001b[0m\n",
       "\u001b[32mSwitzerland, near Zurich. The lab moved to its own campus in neighboring Rüschlikon in 1962. The Zurich lab is \u001b[0m\n",
       "\u001b[32mstaffed by a multicultural and interdisciplinary team of a few hundred permanent research staff members, graduate \u001b[0m\n",
       "\u001b[32mstudents and post-doctoral fellows, representing about 45 nationalities. Collocated with the lab is a Client Center\u001b[0m\n",
       "\u001b[32m(\u001b[0m\u001b[32mformerly the Industry Solutions Lab\u001b[0m\u001b[32m)\u001b[0m\u001b[32m, an executive briefing facility demonstrating technology prototypes and \u001b[0m\n",
       "\u001b[32msolutions.'\u001b[0m,\n",
       "            \u001b[33mmain_path\u001b[0m=\u001b[32m'main-text.71'\u001b[0m,\n",
       "            \u001b[33mpath_group\u001b[0m=\u001b[1m[\u001b[0m\u001b[32m'main-text.69'\u001b[0m, \u001b[32m'main-text.71'\u001b[0m\u001b[1m]\u001b[0m,\n",
       "            \u001b[33msource_is_text\u001b[0m=\u001b[3;92mTrue\u001b[0m\n",
       "        \u001b[1m)\u001b[0m,\n",
       "        \u001b[1;35mSearchResultItem\u001b[0m\u001b[1m(\u001b[0m\n",
       "            \u001b[33mdoc_hash\u001b[0m=\u001b[32m'b30bc667a324ae111d025526563b674a8d3fd869bc07c8fd204aa95b05d41f0c'\u001b[0m,\n",
       "            \u001b[33mchunk\u001b[0m=\u001b[32m'Switzerland\\nIBM Research-Zurich \u001b[0m\u001b[32m(\u001b[0m\u001b[32mpreviously called IBM Zurich Research Laboratory, ZRL\u001b[0m\u001b[32m)\u001b[0m\u001b[32m is the \u001b[0m\n",
       "\u001b[32mEuropean branch of IBM Research. It was opened in 1956 and is located in Rüschlikon, near Zurich, Switzerland.'\u001b[0m,\n",
       "            \u001b[33mmain_path\u001b[0m=\u001b[32m'main-text.70'\u001b[0m,\n",
       "            \u001b[33mpath_group\u001b[0m=\u001b[1m[\u001b[0m\u001b[32m'main-text.69'\u001b[0m, \u001b[32m'main-text.70'\u001b[0m\u001b[1m]\u001b[0m,\n",
       "            \u001b[33msource_is_text\u001b[0m=\u001b[3;92mTrue\u001b[0m\n",
       "        \u001b[1m)\u001b[0m,\n",
       "        \u001b[1;35mSearchResultItem\u001b[0m\u001b[1m(\u001b[0m\n",
       "            \u001b[33mdoc_hash\u001b[0m=\u001b[32m'b30bc667a324ae111d025526563b674a8d3fd869bc07c8fd204aa95b05d41f0c'\u001b[0m,\n",
       "            \u001b[33mchunk\u001b[0m=\u001b[32m\"History\u001b[0m\u001b[32m\\nThe roots of today's IBM Research began with the 1945 opening of the Watson Scientific \u001b[0m\n",
       "\u001b[32mComputing Laboratory at Columbia University. $^\u001b[0m\u001b[32m{\u001b[0m\u001b[32m[\u001b[0m\u001b[32m4\u001b[0m\u001b[32m]\u001b[0m\u001b[32m}\u001b[0m\u001b[32m$ This was the first IBM laboratory devoted to pure science and\u001b[0m\n",
       "\u001b[32mlater expanded into additional IBM Research locations in Westchester County, New York, starting in the 1950s, \u001b[0m\n",
       "\u001b[32m$^\u001b[0m\u001b[32m{\u001b[0m\u001b[32m[\u001b[0m\u001b[32m5\u001b[0m\u001b[32m]\u001b[0m\u001b[32m[\u001b[0m\u001b[32m6\u001b[0m\u001b[32m]\u001b[0m\u001b[32m}\u001b[0m\u001b[32m$ including the Thomas J. Watson Research Center in 1961. \u001b[0m\u001b[32m[\u001b[0m\u001b[32m5\u001b[0m\u001b[32m]\u001b[0m\u001b[32m[\u001b[0m\u001b[32m6\u001b[0m\u001b[32m]\u001b[0m\u001b[32m\"\u001b[0m,\n",
       "            \u001b[33mmain_path\u001b[0m=\u001b[32m'main-text.8'\u001b[0m,\n",
       "            \u001b[33mpath_group\u001b[0m=\u001b[1m[\u001b[0m\u001b[32m'main-text.7'\u001b[0m, \u001b[32m'main-text.8'\u001b[0m\u001b[1m]\u001b[0m,\n",
       "            \u001b[33msource_is_text\u001b[0m=\u001b[3;92mTrue\u001b[0m\n",
       "        \u001b[1m)\u001b[0m\n",
       "    \u001b[1m]\u001b[0m\n",
       "\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "query = SemanticQuery(\n",
    "    question=QUESTION,\n",
    "    project=PROJ_KEY,\n",
    "    data_source=data_source,\n",
    "    retr_k=RETR_K,  # optional\n",
    ")\n",
    "api_output = api.queries.run(query)\n",
    "result = SearchResult.from_api_output(api_output)\n",
    "\n",
    "rich.print(QUESTION)\n",
    "rich.print(result)"
   ]
  }
 ],
 "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}