{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "####  Document Retriever\n",
    "\n",
    "A document retriever is Abacus own vector database. It can be used to:\n",
    "1. Create embeddings of documents\n",
    "2. Retrieve document chunks that are semantically close to a phrase that the user passes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### How to create RAG on the fly with a local document"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here, we upload training file from the current location of the notebook\n",
    "# You can add files to Jupyter Notebook by drag and drop\n",
    "from abacusai.client import BlobInput\n",
    "import abacusai\n",
    "client = abacusai.ApiClient('YOUR_API_KEY')\n",
    "document = BlobInput.from_local_file(\"YOUR_DOCUMENT.pdf/word\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Returns chunks of documents that are relevant to the query and can be used to feed an LLM\n",
    "# Example for blob in memory of notebook\n",
    "\n",
    "relevant_snippets = client.get_relevant_snippets(\n",
    "        blobs={\"document\": document.contents},\n",
    "        query=\"What are the key terms\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Returns chunks of documents that are relevant to the query and can be used to feed an LLM\n",
    "# Example for document in the docstore\n",
    "\n",
    "relevant_snippets = client.get_relevant_snippets(\n",
    "        doc_ids = ['YOUR_DOC_ID_1','YOUR_DOC_ID_2'],\n",
    "        query=\"What are the key terms\")\n",
    "\n",
    "relevant_snippets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Using A document Retriever as a standalone deployment\n",
    "You can also use a documen retriever, even if a ChatLLM model is not trained!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# First we connect our docstore to our project\n",
    "\n",
    "client.add_feature_group_to_project(\n",
    "    feature_group_id='YOUR_FEATURE_GROUP_ID_WITH_DOCUMENTS'\n",
    "    project_id='YOUR_PROJECT_ID',\n",
    "    feature_group_type='DOCUMENTS'  # Optional, defaults to 'CUSTOM_TABLE'. But important to set 'DOCUMENTS' as it will enable Document retriver to work properly with it\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "InferredFeatureMappings(error='',\n",
       "  feature_mappings=[FeatureMapping(feature_mapping='DOCUMENT_ID',\n",
       "  feature_name='doc_id'), FeatureMapping(feature_mapping='DOCUMENT',\n",
       "  feature_name='file_description')])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "ifm = client.infer_feature_mappings(project_id='PROJECT_ID',feature_group_id='FEATURE_GROUP_ID')\n",
    "\n",
    "# ifm = client.infer_feature_mappings(project_id='15ed76a6a8',feature_group_id='98a8d9cce')\n",
    "ifm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This blocs of code might be useful to fix featuregroup for docstore usage by document retrievers\n",
    "\n",
    "# client.set_feature_group_type(project_id='YOUR_PROJECT_ID', feature_group_id='98a8d9cce', feature_group_type='DOCUMENTS')\n",
    "# client.set_feature_mapping(project_id='YOUR_PROJECT_ID',feature_group_id = 'YOUR_FEATURE_GROUP_ID',feature_name='doc_id',feature_mapping='DOCUMENT_ID')\n",
    "# client.set_feature_mapping(project_id='YOUR_PROJECT_ID',feature_group_id = 'YOUR_FEATURE_GROUP_ID',feature_name='page_infos',feature_mapping='DOCUMENT')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Creating a document retriever\n",
    "\n",
    "document_retriever = client.create_document_retriever(\n",
    "    project_id=project_id,\n",
    "    name='NAME_OF_YOUR_DOCUMENT_RETRIEVER',\n",
    "    feature_group_id='YOUR_FEATURE_GROUP_ID'\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Accessing document retriever that is already crreated\n",
    "\n",
    "# dr = client.describe_document_retriever_by_name('DOCUMENT_RETRIEVER_NAME')\n",
    "# dr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "r = client.describe_document_retriever(document_retriever.id)\n",
    "# Filters allow you to filter the documents that the doc retriever can use on the fly, using some columns of the training feature group that was used as input to the doc retriever.\n",
    "# Filters are also available when using .get_chat_reponse\n",
    "\n",
    "client.get_matching_documents(document_retriever_id='DOCUMENT_RETRIEVER_ID', \n",
    "                              query='WHATEVER_YOU_NEED_TO_ASK',limit= 10,\n",
    "                              filters={\"state\": [\"MICHIGAN\",\"NATIONAL\"]})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "10"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Examples of document retriever usage\n",
    "\n",
    "res = document_retriever.get_matching_documents(\"Agreement of the Parties\")\n",
    "len(res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Example of getting no results\n",
    "\n",
    "res2 = document_retriever.get_matching_documents(\"planting potatoes on a mars\", required_phrases=['mars'])\n",
    "res2"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
