{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial notebook\n",
    "\n",
    "This notebook shares base snippets to use Abacus.AI API with chat and AI agent functionality\n",
    "API comes in handy to ease up some recurring tasks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Loading Documents\n",
    "When documents are uploaded into the platform, they are uploaded as a special Class type \"BlobInput\"."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Testing BlobInput Locally"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-22T08:57:33.707610Z",
     "iopub.status.busy": "2024-07-22T08:57:33.707215Z",
     "iopub.status.idle": "2024-07-22T08:57:33.823496Z",
     "shell.execute_reply": "2024-07-22T08:57:33.822805Z",
     "shell.execute_reply.started": "2024-07-22T08:57:33.707580Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from abacusai import ApiClient\n",
    "client = ApiClient()\n",
    "\n",
    "# Your project id consists of numbers and letters id, \n",
    "# Can be found as a part of the browser URL or the project's main page. \n",
    "# Needed for some API calls. \n",
    "# For this example, this should be the AI Agent Project\n",
    "project_id = 'your_project_id' \n",
    "try:\n",
    "    client.describe_project(project_id)\n",
    "except:\n",
    "    raise Exception('Provide your current project ID')\n",
    "\n"
   ]
  },
  {
   "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",
    "document = BlobInput.from_local_file(\"test.docx\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#document.contents is a bytes string of the document\n",
    "extracted_doc_data = client.extract_document_data(document.contents)\n",
    "\n",
    "print(extracted_doc_data.pages[0]) # Text from page 0\n",
    "print(extracted_doc_data.embedded_text) # All text from the document"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2024-07-19T11:57:40.640115Z",
     "iopub.status.idle": "2024-07-19T11:57:40.640754Z",
     "shell.execute_reply": "2024-07-19T11:57:40.640509Z",
     "shell.execute_reply.started": "2024-07-19T11:57:40.640484Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "len(extracted_doc_data.pages)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2024-07-19T11:57:40.642525Z",
     "iopub.status.idle": "2024-07-19T11:57:40.642957Z",
     "shell.execute_reply": "2024-07-19T11:57:40.642750Z",
     "shell.execute_reply.started": "2024-07-19T11:57:40.642728Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Here, we will create the upload for the dataset from the notebook as an Abacus dataset that we will be able to use later.\n",
    "# Docstore is special table format for document storage\n",
    "\n",
    "upload = client.create_dataset_from_upload(\n",
    "    table_name='my_documents_'+client.describe_user().email.split('@')[0], #name should be unique inside the organisation\n",
    "    file_format='DOCX',\n",
    "    is_documentset=True\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-18T16:23:49.025773Z",
     "iopub.status.busy": "2024-07-18T16:23:49.025348Z",
     "iopub.status.idle": "2024-07-18T16:23:49.776343Z",
     "shell.execute_reply": "2024-07-18T16:23:49.775579Z",
     "shell.execute_reply.started": "2024-07-18T16:23:49.025744Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset(dataset_id='93c789714',\n",
       "  source_type='UPLOAD',\n",
       "  data_source='re://datasets/93c789714',\n",
       "  created_at='2024-07-18T16:23:47+00:00',\n",
       "  ephemeral=False,\n",
       "  feature_group_table_name='my_documents_bogdan',\n",
       "  incremental=False,\n",
       "  is_documentset=True,\n",
       "  extract_bounding_boxes=False,\n",
       "  merge_file_schemas=True,\n",
       "  reference_only_documentset=False,\n",
       "  latest_dataset_version=DatasetVersion(dataset_version='18b974e6a',\n",
       "  status='CONVERTING',\n",
       "  dataset_id='93c789714',\n",
       "  size=136504,\n",
       "  created_at='2024-07-18T16:23:47+00:00',\n",
       "  merge_file_schemas=True),\n",
       "  document_processing_config=DocumentProcessingConfig(extract_bounding_boxes=False, ocr_mode='DEFAULT', use_full_ocr=None, remove_header_footer=False, remove_watermarks=True, convert_to_markdown=False))"
      ]
     },
     "execution_count": 68,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with open(\"test.docx\", \"rb\") as file:\n",
    "\n",
    "    file_uploaded = upload.upload_file(file)\n",
    "    file_uploaded.wait_for_import()\n",
    "\n",
    "file_uploaded"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-19T12:01:07.897492Z",
     "iopub.status.busy": "2024-07-19T12:01:07.897085Z",
     "iopub.status.idle": "2024-07-19T12:01:08.014131Z",
     "shell.execute_reply": "2024-07-19T12:01:08.013392Z",
     "shell.execute_reply.started": "2024-07-19T12:01:07.897463Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " feature group found\n"
     ]
    }
   ],
   "source": [
    "# Here we verify our upload and check the structure of file we created in docstore\n",
    "\n",
    "try:\n",
    "    feature_group = client.describe_feature_group_by_table_name('my_documents_'+client.describe_user().email.split('@')[0])\n",
    "    print(' feature group found')\n",
    "except:\n",
    "    feature_group = file_uploaded.describe_feature_group()\n",
    "    if not feature_group.list_versions():\n",
    "        print(\"creating first version\")\n",
    "        feature_group.create_version()\n",
    "    feature_group.wait_for_materialization()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-19T12:01:09.179593Z",
     "iopub.status.busy": "2024-07-19T12:01:09.179205Z",
     "iopub.status.idle": "2024-07-19T12:01:10.773613Z",
     "shell.execute_reply": "2024-07-19T12:01:10.772818Z",
     "shell.execute_reply.started": "2024-07-19T12:01:09.179565Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>doc_id</th>\n",
       "      <th>page_infos</th>\n",
       "      <th>file_path</th>\n",
       "      <th>file_size_bytes</th>\n",
       "      <th>file_checksum</th>\n",
       "      <th>file_description</th>\n",
       "      <th>mime_type</th>\n",
       "      <th>page_count</th>\n",
       "      <th>token_count</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>18b974e6a-000000000-5c536d98323d4a361309f5516c...</td>\n",
       "      <td>{'first_page': 0, 'last_page': 57}</td>\n",
       "      <td>uploaded_data.docx</td>\n",
       "      <td>136504</td>\n",
       "      <td>SHA512_256:5f4100c619746dc53c62b706540e759ac86...</td>\n",
       "      <td>Microsoft Word 2007+</td>\n",
       "      <td>application/vnd.openxmlformats-officedocument....</td>\n",
       "      <td>58</td>\n",
       "      <td>17598</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                              doc_id  \\\n",
       "0  18b974e6a-000000000-5c536d98323d4a361309f5516c...   \n",
       "\n",
       "                           page_infos           file_path  file_size_bytes  \\\n",
       "0  {'first_page': 0, 'last_page': 57}  uploaded_data.docx           136504   \n",
       "\n",
       "                                       file_checksum      file_description  \\\n",
       "0  SHA512_256:5f4100c619746dc53c62b706540e759ac86...  Microsoft Word 2007+   \n",
       "\n",
       "                                           mime_type  page_count  token_count  \n",
       "0  application/vnd.openxmlformats-officedocument....          58        17598  "
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = feature_group.load_as_pandas()\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-19T12:01:10.775611Z",
     "iopub.status.busy": "2024-07-19T12:01:10.774972Z",
     "iopub.status.idle": "2024-07-19T12:01:10.780482Z",
     "shell.execute_reply": "2024-07-19T12:01:10.779846Z",
     "shell.execute_reply.started": "2024-07-19T12:01:10.775579Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'18b974e6a-000000000-5c536d98323d4a361309f5516c6282e003a385e11822616975ed720f8d473ba4'"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df['doc_id'][0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Extracting the documents text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-18T16:27:06.403338Z",
     "iopub.status.busy": "2024-07-18T16:27:06.403055Z",
     "iopub.status.idle": "2024-07-18T16:27:12.595804Z",
     "shell.execute_reply": "2024-07-18T16:27:12.595177Z",
     "shell.execute_reply.started": "2024-07-18T16:27:06.403312Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "\n",
    "data_from_docstore = client.get_docstore_document_data(df['doc_id'][0]) # Get data from a document stored in the docstore"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Getting document text from uploaded data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-18T16:41:18.951733Z",
     "iopub.status.busy": "2024-07-18T16:41:18.951323Z",
     "iopub.status.idle": "2024-07-18T16:41:19.890519Z",
     "shell.execute_reply": "2024-07-18T16:41:19.889614Z",
     "shell.execute_reply.started": "2024-07-18T16:41:18.951705Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['metadata', 'tokens', 'pages', 'doc_id', 'embedded_text', 'extracted_text'])"
      ]
     },
     "execution_count": 74,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# To access docstore later, or when it was created outside of this notebook, we may use the name or id of it by functions describe_feature_group_by_table_name or describe_feature_group, respectively\n",
    "\n",
    "df = client.describe_feature_group_by_table_name(feature_group.name).load_as_pandas_documents(doc_id_column = 'doc_id',document_column = 'page_infos')\n",
    "df['page_infos'][0].keys()\n",
    "# dict_keys(['pages', 'tokens', 'metadata', 'extracted_text'])\n",
    "\n",
    "#pages: This is the embedded text from the document on a per page level\n",
    "#extracted_text: This is the OCR extracted text from the document"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating RAG systems on the fly\n",
    "\n",
    "How to create RAG system \"on the fly\" with an uploaded document"
   ]
  },
  {
   "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 = [df['doc_id'][0]],\n",
    "        # blobs={\"document\": document.contents},\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": 79,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-18T16:41:41.891390Z",
     "iopub.status.busy": "2024-07-18T16:41:41.890693Z",
     "iopub.status.idle": "2024-07-18T16:41:41.954327Z",
     "shell.execute_reply": "2024-07-18T16:41:41.953701Z",
     "shell.execute_reply.started": "2024-07-18T16:41:41.891362Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# First we connect our docstore to our project\n",
    "\n",
    "client.add_feature_group_to_project(\n",
    "    feature_group_id=feature_group.id,\n",
    "    project_id=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": 80,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-18T16:41:41.955884Z",
     "iopub.status.busy": "2024-07-18T16:41:41.955576Z",
     "iopub.status.idle": "2024-07-18T16:41:41.960320Z",
     "shell.execute_reply": "2024-07-18T16:41:41.959690Z",
     "shell.execute_reply.started": "2024-07-18T16:41:41.955858Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'1246377b2c'"
      ]
     },
     "execution_count": 80,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "feature_group.id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-18T16:41:41.961462Z",
     "iopub.status.busy": "2024-07-18T16:41:41.961181Z",
     "iopub.status.idle": "2024-07-18T16:41:47.887608Z",
     "shell.execute_reply": "2024-07-18T16:41:47.886850Z",
     "shell.execute_reply.started": "2024-07-18T16:41:41.961437Z"
    },
    "tags": []
   },
   "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')])"
      ]
     },
     "execution_count": 81,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "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": 82,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-18T16:41:47.888874Z",
     "iopub.status.busy": "2024-07-18T16:41:47.888586Z",
     "iopub.status.idle": "2024-07-18T16:41:47.892105Z",
     "shell.execute_reply": "2024-07-18T16:41:47.891417Z",
     "shell.execute_reply.started": "2024-07-18T16:41:47.888847Z"
    },
    "tags": []
   },
   "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='15ed76a6a8', feature_group_id='98a8d9cce', feature_group_type='DOCUMENTS')\n",
    "# client.set_feature_mapping(project_id,feature_group.id,feature_name='doc_id',feature_mapping='DOCUMENT_ID')\n",
    "# client.set_feature_mapping(project_id,feature_group.id,feature_name='page_infos',feature_mapping='DOCUMENT')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-19T12:01:48.962128Z",
     "iopub.status.busy": "2024-07-19T12:01:48.961697Z",
     "iopub.status.idle": "2024-07-19T12:01:49.195026Z",
     "shell.execute_reply": "2024-07-19T12:01:49.194395Z",
     "shell.execute_reply.started": "2024-07-19T12:01:48.962086Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Creating a document retriever\n",
    "\n",
    "document_retriever = client.create_document_retriever(\n",
    "    project_id=project_id,\n",
    "    name='demo_document_retriever__'+client.describe_user().email.split('@')[0],\n",
    "    feature_group_id=feature_group.id\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Accessing document retriever that is already crreated\n",
    "\n",
    "# dr = client.describe_document_retriever_by_name('demo_document_retriever_'+client.describe_user().email.split('@')[0])\n",
    "# dr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-19T13:02:32.750744Z",
     "iopub.status.busy": "2024-07-19T13:02:32.750341Z",
     "iopub.status.idle": "2024-07-19T13:02:32.877158Z",
     "shell.execute_reply": "2024-07-19T13:02:32.876301Z",
     "shell.execute_reply.started": "2024-07-19T13:02:32.750715Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[]"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "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",
    "r.get_matching_documents(query = \"WHATEVER_YOU_WANT_TO_ASK\", filters = {\"document_identification\":['AXIP-4440']})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-19T14:58:12.356794Z",
     "iopub.status.busy": "2024-07-19T14:58:12.356373Z",
     "iopub.status.idle": "2024-07-19T14:58:12.896285Z",
     "shell.execute_reply": "2024-07-19T14:58:12.895452Z",
     "shell.execute_reply.started": "2024-07-19T14:58:12.356765Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "10"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "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": 32,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-19T15:05:15.506942Z",
     "iopub.status.busy": "2024-07-19T15:05:15.506540Z",
     "iopub.status.idle": "2024-07-19T15:05:16.143817Z",
     "shell.execute_reply": "2024-07-19T15:05:16.143043Z",
     "shell.execute_reply.started": "2024-07-19T15:05:15.506910Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[]"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Example of getting no results\n",
    "\n",
    "res2 = document_retriever.get_matching_documents(\"planting potatoes on a mars\", required_phrases=['mars'])\n",
    "res2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Calling a Large Language Model\n",
    "You can use the `evalute_prompt` method to call the LLM of your choice:\n",
    "- prompt: This is the actual message that the model receives from the user\n",
    "- system_message: These are the instructions that the model will follow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-19T13:33:27.110785Z",
     "iopub.status.busy": "2024-07-19T13:33:27.110381Z",
     "iopub.status.idle": "2024-07-19T13:33:31.256828Z",
     "shell.execute_reply": "2024-07-19T13:33:31.255945Z",
     "shell.execute_reply.started": "2024-07-19T13:33:27.110755Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Athens\n"
     ]
    }
   ],
   "source": [
    "r = client.evaluate_prompt(prompt = \"What is the capital of Greece?\", system_message = \"You should answer all questions with a single word.\", llm_name = \"OPENAI_GPT4O\")\n",
    "\n",
    "# Response:\n",
    "print(r.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calling a Large Language Model and specifying some output schema\n",
    "You can also use the `json_response_schema` to specify the output of the model in a pre-defined manner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-19T15:08:40.269078Z",
     "iopub.status.busy": "2024-07-19T15:08:40.268645Z",
     "iopub.status.idle": "2024-07-19T15:08:44.355169Z",
     "shell.execute_reply": "2024-07-19T15:08:44.354149Z",
     "shell.execute_reply.started": "2024-07-19T15:08:40.269046Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'learning_objectives': ['Understand the components and functions of car batteries',\n",
       "  'Learn how to maintain and troubleshoot car batteries',\n",
       "  'Gain knowledge about the different types of car doors and their mechanisms',\n",
       "  'Learn how to repair and replace car doors',\n",
       "  'Understand the principles and components of car suspension systems',\n",
       "  'Learn how to diagnose and fix common issues in car suspension systems']}"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import json\n",
    "\n",
    "r = client.evaluate_prompt(prompt = \"In this course, you will learn about car batteries, car doors, and car suspension system\",\n",
    "                           # system_message = \"OPTIONAL, but good to have\", \n",
    "                           llm_name = 'OPENAI_GPT4O',\n",
    "                           json_response_schema = {\"learning_objectives\": {\"type\": \"list\", \"description\": \"A list of learning objectives\", \"is_required\": True}}\n",
    ")\n",
    "learning_objectives = json.loads(r.content)\n",
    "learning_objectives"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating a simple AI Agent with workflows"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-19T15:17:10.645602Z",
     "iopub.status.busy": "2024-07-19T15:17:10.645200Z",
     "iopub.status.idle": "2024-07-19T15:17:10.649548Z",
     "shell.execute_reply": "2024-07-19T15:17:10.648754Z",
     "shell.execute_reply.started": "2024-07-19T15:17:10.645573Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from abacusai import (\n",
    "    AgentInterface,\n",
    "    WorkflowGraph,\n",
    "    WorkflowGraphEdge,\n",
    "    WorkflowGraphNode,\n",
    "    WorkflowNodeInputMapping,\n",
    "    WorkflowNodeInputSchema,\n",
    "    WorkflowNodeInputType,\n",
    "    WorkflowNodeOutputMapping,\n",
    "    WorkflowNodeOutputSchema,\n",
    "    WorkflowNodeOutputType,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### For this agent you can select one of preselected charecters to answer questions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-19T16:33:28.364370Z",
     "iopub.status.busy": "2024-07-19T16:33:28.363964Z",
     "iopub.status.idle": "2024-07-19T16:33:28.368666Z",
     "shell.execute_reply": "2024-07-19T16:33:28.367948Z",
     "shell.execute_reply.started": "2024-07-19T16:33:28.364340Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def agent_function(nlp_query, character):\n",
    "    \"\"\"\n",
    "        Args:\n",
    "            nlp_query (Any): Data row to predict on/with or to pass to the agent for execution\n",
    "        Returns:\n",
    "            The result which can be any json serializable python type\n",
    "    \"\"\"\n",
    "    from abacusai import ApiClient\n",
    "\n",
    "    # Let agent respond like your favorite character.\n",
    "    char = character or 'Sherlock Holmes'\n",
    "    response = ApiClient().evaluate_prompt(prompt=nlp_query, system_message=f'Respond like {char}. Prepend your name.')\n",
    "    return str(response.content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-19T16:33:28.718690Z",
     "iopub.status.busy": "2024-07-19T16:33:28.718280Z",
     "iopub.status.idle": "2024-07-19T16:33:30.911606Z",
     "shell.execute_reply": "2024-07-19T16:33:30.910908Z",
     "shell.execute_reply.started": "2024-07-19T16:33:28.718661Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Homer Simpson: Mmm, donuts...'"
      ]
     },
     "execution_count": 96,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent_function('what is your favorite food','Homer Simpson')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-19T16:33:38.887253Z",
     "iopub.status.busy": "2024-07-19T16:33:38.886850Z",
     "iopub.status.idle": "2024-07-19T16:33:38.891021Z",
     "shell.execute_reply": "2024-07-19T16:33:38.890220Z",
     "shell.execute_reply.started": "2024-07-19T16:33:38.887219Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "package_requirements = []  # e.g. ['numpy==1.2.3', 'pandas>=1.4.0']\n",
    "description = None\n",
    "memory = 16\n",
    "enable_binary_input = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "WorkflowGraphNode is one block of creation of AI Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-19T16:34:33.287084Z",
     "iopub.status.busy": "2024-07-19T16:34:33.286482Z",
     "iopub.status.idle": "2024-07-19T16:34:33.294142Z",
     "shell.execute_reply": "2024-07-19T16:34:33.293351Z",
     "shell.execute_reply.started": "2024-07-19T16:34:33.287054Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "workflow_graph_node = WorkflowGraphNode(\n",
    "    name=\"input_text\",\n",
    "    function=agent_function,\n",
    "    input_mappings=[\n",
    "        WorkflowNodeInputMapping(\n",
    "            name=\"nlp_query\",\n",
    "            variable_type=WorkflowNodeInputType.USER_INPUT,\n",
    "            # variable_source=\"obi van\"\n",
    "        ),\n",
    "        WorkflowNodeInputMapping(\n",
    "            name=\"character\",\n",
    "            variable_type=WorkflowNodeInputType.USER_INPUT,\n",
    "            # variable_source=\"obi van\"\n",
    "        ),\n",
    "    ],\n",
    "    input_schema = WorkflowNodeInputSchema(\n",
    "        json_schema={\n",
    "            \"type\": \"object\",\n",
    "            \"title\": \"Get character responce\",\n",
    "            \"required\": [\"nlp_query\", \"character\"],\n",
    "            \"properties\": {\n",
    "                \"nlp_query\": {\"type\": \"string\", \"title\": \"Your question\"},\n",
    "                \"character\": {\n",
    "                            \"type\": \"string\",\n",
    "                            \"title\": \"Characters\",\n",
    "                            \"enum\": [\"Sherlock Holmes\", \"Elon Musk\", \"Homer Simpson\"],\n",
    "                            \"default\": \"Homer Simpson\"\n",
    "                            # \"enumNames\": [\"Sherlock Holmes\", \"Elon Musk\", \"Homer Simpson\"]\n",
    "                            }\n",
    "                # \"table_name\": {\"type\": \"string\", \"title\": \"Table Name\"},\n",
    "                # \"document_column_name\": {\"type\": \"string\", \"title\": \"Document Column Name\"},\n",
    "                # \"chunk_size\": {\"type\": \"integer\", \"title\": \"Chunk Size\"},\n",
    "                # \"text_encoder\": {\"type\": \"string\", \"title\": \"Text Encoder\", \"enum\": [e.value for e in VectorStoreTextEncoder]},\n",
    "            },\n",
    "        },\n",
    "        # ui_schema={\n",
    "        #     \"text_encoder\": {\"ui:widget\": \"select\"},\n",
    "        # }\n",
    "    ),\n",
    "    output_mappings=[\n",
    "        WorkflowNodeOutputMapping(\n",
    "            name=\"str_out\",\n",
    "            variable_type=WorkflowNodeOutputType.STRING\n",
    "        ),\n",
    "    ],\n",
    "    output_schema=WorkflowNodeOutputSchema({\n",
    "        \"type\": \"object\",\n",
    "        \"properties\": {\n",
    "            \"str_out\": {\"type\": \"string\", \"title\": \"Response\"},\n",
    "        },\n",
    "    })\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "WorkflowGraph is final graph of all nodes and edges that create an AI Agent Logic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-19T16:34:35.190444Z",
     "iopub.status.busy": "2024-07-19T16:34:35.190040Z",
     "iopub.status.idle": "2024-07-19T16:34:35.194012Z",
     "shell.execute_reply": "2024-07-19T16:34:35.193341Z",
     "shell.execute_reply.started": "2024-07-19T16:34:35.190412Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "workflow_graph = WorkflowGraph(\n",
    "    nodes=[\n",
    "        workflow_graph_node,\n",
    "    ],\n",
    "    edges=[],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-19T16:34:35.886585Z",
     "iopub.status.busy": "2024-07-19T16:34:35.886176Z",
     "iopub.status.idle": "2024-07-19T16:34:35.949178Z",
     "shell.execute_reply": "2024-07-19T16:34:35.948341Z",
     "shell.execute_reply.started": "2024-07-19T16:34:35.886553Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Agent(name='example_agent',\n",
       "   agent_id='3c981ea4e',\n",
       "   created_at='2024-06-24T10:47:22+00:00',\n",
       "   project_id={'projectId': '45d76db9c', 'problemType': 'ai_agent'},\n",
       "   agent_config={'ENABLE_BINARY_INPUT': True},\n",
       "   agent_execution_config={'character': ['Elon Musk', 'Joe Biden']},\n",
       "   latest_agent_version=AgentVersion(agent_version='325916946',\n",
       "   status='COMPLETE',\n",
       "   publishing_completed_at='2024-06-24T10:48:09+00:00')),\n",
       " Agent(name='example_agent',\n",
       "   agent_id='d741a5db6',\n",
       "   created_at='2024-06-25T16:24:16+00:00',\n",
       "   project_id={'projectId': '45d76db9c', 'problemType': 'ai_agent'},\n",
       "   agent_config={},\n",
       "   latest_agent_version=AgentVersion(agent_version='75cdd1928',\n",
       "   status='COMPLETE',\n",
       "   publishing_completed_at='2024-07-01T11:20:44+00:00'),\n",
       "   workflow_graph=WorkflowGraph(nodes=[WorkflowGraphNode()], edges=[], primary_start_node='input_text')),\n",
       " Agent(name='Example_Character_Agent',\n",
       "   agent_id='1000067952',\n",
       "   created_at='2024-07-19T15:48:24+00:00',\n",
       "   project_id={'projectId': '45d76db9c', 'problemType': 'ai_agent'},\n",
       "   agent_config={},\n",
       "   latest_agent_version=AgentVersion(agent_version='f9cc62cc0',\n",
       "   status='COMPLETE',\n",
       "   publishing_completed_at='2024-07-19T16:31:30+00:00'),\n",
       "   workflow_graph=WorkflowGraph(nodes=[WorkflowGraphNode()], edges=[], primary_start_node='input_text'))]"
      ]
     },
     "execution_count": 101,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "client.list_agents(project_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "There are 2 main types of AI Agents AgentInterface.DEFAULT and AgentInterface.CHAT\n",
    "\n",
    "- AgentInterface.DEFAULT is an AI agent that uses forms to fill in and work like an app\n",
    "- AgentInterface.CHAT reproduces experiense of chat with LLM with logic that you may create for it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-19T16:34:41.086643Z",
     "iopub.status.busy": "2024-07-19T16:34:41.086228Z",
     "iopub.status.idle": "2024-07-19T16:35:42.676266Z",
     "shell.execute_reply": "2024-07-19T16:35:42.675477Z",
     "shell.execute_reply.started": "2024-07-19T16:34:41.086610Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Agent(name='Example_Character_Agent',\n",
       "  agent_id='1000067952',\n",
       "  created_at='2024-07-19T15:48:24+00:00',\n",
       "  project_id={'problemType': 'ai_agent', 'allProjectModels': None, 'projectId': '45d76db9c', 'ingressName': None, 'starred': 0, 'tags': None, 'useCase': 'ai_agent', 'name': 'AI_agent_bogdan', 'ingressType': None, 'createdAt': '2024-06-19T12:30:10+00:00', 'deployments': None, 'systemCreated': False, 'info': None, 'updatedAt': '2024-06-19T12:30:10+00:00'},\n",
       "  notebook_id='5441509c0',\n",
       "  agent_config={},\n",
       "  code_source=CodeSource(source_type='TEXT',\n",
       "  source_code='def agent_function(nlp_query, character):\\n    \"\"\"\\n        Args:\\n            nlp_query (Any): Data row to predict on/with or to pass to the agent for execution\\n        Returns:\\n            The result which can be any json serializable python type\\n    \"\"\"\\n    from abacusai import ApiClient\\n\\n    # Let agent respond like your favorite character.\\n    char = character or \\'Sherlock Holmes\\'\\n    response = ApiClient().evaluate_prompt(prompt=nlp_query, system_message=f\\'Respond like {char}. Prepend your name.\\')\\n    return str(response.content)\\n\\n',\n",
       "  package_requirements=[],\n",
       "  status='PENDING',\n",
       "  module_dependencies=[]),\n",
       "  latest_agent_version=AgentVersion(agent_version='117e736ee8',\n",
       "  status='COMPLETE',\n",
       "  agent_id='1000067952',\n",
       "  agent_config={},\n",
       "  publishing_started_at='2024-07-19T16:35:00+00:00',\n",
       "  publishing_completed_at='2024-07-19T16:35:13+00:00',\n",
       "  pending_deployment_ids=['69d0022c6'],\n",
       "  failed_deployment_ids=[],\n",
       "  code_source=CodeSource(source_type='TEXT',\n",
       "  source_code='def agent_function(nlp_query, character):\\n    \"\"\"\\n        Args:\\n            nlp_query (Any): Data row to predict on/with or to pass to the agent for execution\\n        Returns:\\n            The result which can be any json serializable python type\\n    \"\"\"\\n    from abacusai import ApiClient\\n\\n    # Let agent respond like your favorite character.\\n    char = character or \\'Sherlock Holmes\\'\\n    response = ApiClient().evaluate_prompt(prompt=nlp_query, system_message=f\\'Respond like {char}. Prepend your name.\\')\\n    return str(response.content)\\n\\n',\n",
       "  package_requirements=[],\n",
       "  status='COMPLETE',\n",
       "  publishing_msg={'warningInfo': ''},\n",
       "  module_dependencies=[]),\n",
       "  workflow_graph=WorkflowGraph(nodes=[WorkflowGraphNode()], edges=[], primary_start_node='input_text')),\n",
       "  workflow_graph=WorkflowGraph(nodes=[WorkflowGraphNode()], edges=[], primary_start_node='input_text'))"
      ]
     },
     "execution_count": 103,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from abacusai import ApiClient\n",
    "client = ApiClient()\n",
    "agent_interface: AgentInterface = AgentInterface.DEFAULT\n",
    "if 'agent_function' not in vars():\n",
    "    raise Exception('Please define agent function with name - agent_function')\n",
    "\n",
    "if not [x for x in client.list_agents(project_id) if 'Example_Character_Agent'==x.name]:\n",
    "    agent = client.create_agent(project_id=project_id,\n",
    "                                # function_source_code=agent_function, agent_function_name='agent_function',\n",
    "                                name='Example_Character_Agent', \n",
    "                                package_requirements=package_requirements,\n",
    "                                description=description,\n",
    "                                # enable_binary_input=enable_binary_input, memory=memory,\n",
    "                                workflow_graph=workflow_graph, agent_interface=agent_interface)\n",
    "    agent.wait_for_publish()\n",
    "    deployment = client.create_deployment(model_id=agent.agent_id)\n",
    "    deployment.wait_for_deployment()\n",
    "\n",
    "else:\n",
    "    agent = client.update_agent(model_id=agent.id,\n",
    "                            # function_source_code=agent_function, agent_function_name='agent_function',\n",
    "                            # name='example_agent',\n",
    "                            # package_requirements=package_requirements,\n",
    "                            # description=description,\n",
    "                            # enable_binary_input=enable_binary_input, memory=memory,\n",
    "                            workflow_graph=workflow_graph, agent_interface=agent_interface)\n",
    "    agent.wait_for_publish()\n",
    "\n",
    "agent"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Deployment of AI Agent\n",
    "\n",
    "Additionally to creation of model it should be deployed\n",
    "Ai agent later may be reached through Deployments > Predictions Dash inside this project\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client.list_deployment_tokens(project_id)\n",
    "# use client.create_deployment_token() if you have no tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-19T16:01:11.274557Z",
     "iopub.status.busy": "2024-07-19T16:01:11.273771Z",
     "iopub.status.idle": "2024-07-19T16:01:56.721141Z",
     "shell.execute_reply": "2024-07-19T16:01:56.720353Z",
     "shell.execute_reply.started": "2024-07-19T16:01:11.274520Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Deployment(deployment_id='69d0022c6',\n",
       "  name='Example_Character_Agent Deployment',\n",
       "  status='ACTIVE',\n",
       "  description='',\n",
       "  deployed_at='2024-07-19T16:01:54+00:00',\n",
       "  created_at='2024-07-19T16:01:11+00:00',\n",
       "  project_id='45d76db9c',\n",
       "  model_id='1000067952',\n",
       "  model_version='3da4b4f46',\n",
       "  calls_per_second=5,\n",
       "  auto_deploy=True,\n",
       "  skip_metrics_check=False,\n",
       "  algo_name='AI Agent',\n",
       "  regions=[{'name': 'Us East 1', 'value': 'us-east-1'}],\n",
       "  batch_streaming_updates=False,\n",
       "  algorithm='2efe1d48f',\n",
       "  model_deployment_config={'otherModelsForDataClusterTypes': {}, 'streamingFeatureGroupDetails': [], 'modelTrainingType': None})"
      ]
     },
     "execution_count": 69,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "deployment = client.create_deployment(model_id=agent.agent_id)\n",
    "deployment.wait_for_deployment()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# You can use below command, to get the response from a deployed ChatLLM model / This ChatLLM model might be using RAG under the hood.\n",
    "r = client.get_chat_response(deployment_token=client.list_deployment_tokens(project_id)[0], deployment_id='fddsfff', messages=[{\"is_user\":True,\"text\":\"What is the meaning of life?\"}])\n",
    "print(r.keys())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Agent JSON Schema\n",
    "Below instructions are only relevant for when you are creating an \"AI Agent\". The \"json_schema\" variable allows you to create a custom UX that the user can use to interact with the agent. You can find a playground for the json schema here: https://rjsf-team.github.io/react-jsonschema-form/.\n",
    "\n",
    "Below is an example of a json_schema that allows user to:\n",
    "1. Upload a document\n",
    "2. Select an option"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Additional parameters for the JSON Schema allows you to create the UX\n",
    "json_schema ={\n",
    "\"type\": \"object\",\n",
    "\"title\": \"Upload Document and select task\",\n",
    "\"required\": [\"document\", \"options\"],\n",
    "\"properties\": {\n",
    "    \"document\": {\n",
    "        \"type\": \"string\",\n",
    "        \"title\": \"Upload Document\",\n",
    "        \"format\": \"data-url\"},\n",
    "    \"options\": {\n",
    "        \"type\": \"string\",\n",
    "        \"title\": \"Options\",\n",
    "        \"enum\": [\"extract_rfp_questions\", \"complete_rfp_questions\"],\n",
    "        \"enumNames\": [\"Extract RFP Questions\", \"Complete RFP Questions\"]\n",
    "        }\n",
    "}\n",
    "}\n",
    "\n",
    "# The WorkflowNodeOutputSchema allows you to setup the output to be of data-url so that the user can download.\n",
    "output_schema=WorkflowNodeOutputSchema(\n",
    "    json_schema={\n",
    "        \"type\": \"object\",\n",
    "        \"properties\": {\n",
    "            \"processed_doc\": {\n",
    "                \"type\": \"string\",\n",
    "                \"title\": \"Responses\",\n",
    "                \"format\": \"data-url\",\n",
    "            }\n",
    "        },\n",
    "    }\n",
    ")\n",
    "\n",
    "# Here is how the Agent Response should look like:\n",
    "return AgentResponse(processed_document_doc=Blob(doc_bytes,\"application/vnd.openxmlformats-officedocument.wordprocessingml.document\",filename=f\"result.docx\",))"
   ]
  }
 ],
 "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
