apitable_langchain_function.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents import initialize_agent, Tool, AgentType\n",
    "from langchain.chat_models import ChatOpenAI\n",
    "import os\n",
    "import json\n",
    "from apitable import Apitable\n",
    "import tiktoken\n",
    "from typing import Optional\n",
    "from pydantic import BaseModel, Field\n",
    "from typing import Dict, List\n",
    "from langchain.callbacks import get_openai_callback\n",
    "from langchain.tools import BaseTool\n",
    "from typing import Type"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = ChatOpenAI(temperature=0, model=\"gpt-3.5-turbo-0613\")\n",
    "apitable_api_token = os.getenv(\"APITABLE_API_TOKEN\")\n",
    "apitable = Apitable(token=apitable_api_token)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "enc = tiktoken.encoding_for_model(\"gpt-3.5-turbo\")\n",
    "\n",
    "def count_tokens(s):\n",
    "    return len(enc.encode(s))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def trans_key(field_key_map, key: str):\n",
    "    \"\"\"\n",
    "    When there is a field mapping, convert the mapped key to the actual key\n",
    "    \"\"\"\n",
    "    if key in [\"_id\", \"recordId\"]:\n",
    "        return key\n",
    "    if field_key_map:\n",
    "        _key = field_key_map.get(key, key)\n",
    "        return _key\n",
    "    return key\n",
    "\n",
    "def query_parse(field_key_map, **kwargs) -> str:\n",
    "    query_list = []\n",
    "    for k, v in kwargs.items():\n",
    "        # Handling null\n",
    "        if v is None:\n",
    "            v = \"BLANK()\"\n",
    "        # Handling string\n",
    "        elif isinstance(v, str):\n",
    "            v = f'\"{v}\"'\n",
    "        elif isinstance(v, bool):\n",
    "            v = \"TRUE()\" if v else \"FALSE()\"\n",
    "        # Handling array type values, multiple select, members?\n",
    "        elif isinstance(v, list):\n",
    "            v = f'\"{\", \".join(v)}\"'\n",
    "        query_list.append(f\"{{{trans_key(field_key_map, k)}}}={v}\")\n",
    "    if len(query_list) == 1:\n",
    "        return query_list[0]\n",
    "    else:\n",
    "        qs = \",\".join(query_list)\n",
    "        return f\"AND({qs})\"\n",
    "\n",
    "def get_spaces(question: str):\n",
    "    spaces = apitable.spaces.all()\n",
    "    return [json.loads(space.json()) for space in spaces]\n",
    "\n",
    "def get_nodes(space_id: str):\n",
    "    nodes = apitable.space(space_id=space_id).nodes.all()\n",
    "    return [json.loads(node.json()) for node in nodes]\n",
    "\n",
    "def get_fields(datasheet_id: str):\n",
    "    fields = apitable.datasheet(datasheet_id).fields.all()\n",
    "    return [json.loads(field.json()) for field in fields]\n",
    "\n",
    "def create_fields(space_id: str, datasheet_id: str, field_data: dict[str, str]):\n",
    "    field = (\n",
    "        apitable.space(space_id)\n",
    "        .datasheet(datasheet_id)\n",
    "        .fields.create(field_data)\n",
    "    )\n",
    "    return field.json()\n",
    "\n",
    "def get_records(datasheet_id:str, filter_condition: Optional[dict]= None, sort_condition: Optional[list] = None, maxRecords_condition: Optional[int] = None):\n",
    "    dst = apitable.datasheet(datasheet_id)\n",
    "    query_kwargs = {}\n",
    "    if filter_condition:\n",
    "        query_formula = query_parse(filter_condition)\n",
    "        query_kwargs[\"filterByFormula\"] = query_formula\n",
    "    if sort_condition:\n",
    "        query_kwargs[\"sort\"] = sort_condition\n",
    "    if maxRecords_condition:\n",
    "        query_kwargs[\"maxRecords\"] = maxRecords_condition\n",
    "    records = dst.records.all(**query_kwargs)\n",
    "    parsed_records = [record.json() for record in records]\n",
    "    if count_tokens(str(parsed_records)) > 1000:\n",
    "        parsed_records_str = (\n",
    "            \"Found \"\n",
    "            + str(len(parsed_records))\n",
    "            + \" records, too many to show. Try to limit\"\n",
    "        )\n",
    "        return parsed_records_str\n",
    "    else:\n",
    "        return parsed_records"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SortCondition(BaseModel):\n",
    "    field: str = Field(description=\"Field name\")\n",
    "    order: str = Field(description=\"Sort order\", enum=[\"desc\", \"asc\"])\n",
    "\n",
    "class GetRecordsInput(BaseModel):\n",
    "    \"\"\"Inputs for get records\"\"\"\n",
    "    datasheet_id: str = Field(description=\"The ID of the datasheet to retrieve records from.\")\n",
    "    filter_condition: Optional[Dict[str, str]] = Field(\n",
    "        description=\"\"\"\n",
    "            Find records that meet specific conditions.\n",
    "            This object should contain a key-value pair where the key is the field name and the value is the lookup value. For instance: {\"title\": \"test\"}.\n",
    "            \"\"\"\n",
    "    )\n",
    "    sort_condition: Optional[List[SortCondition]] = Field(min_items=1,\n",
    "        description=\"Sort returned records by specific field\"\n",
    "    )\n",
    "    maxRecords_condition: Optional[int] = Field(\n",
    "        description=\"Limit the number of returned values.\"\n",
    "    )\n",
    "\n",
    "class GetRecordsTool(BaseTool):\n",
    "    name = \"get_records\"\n",
    "    description = \"\"\"\n",
    "        Useful for retrieving data in a datasheet\n",
    "        \"\"\"\n",
    "    args_schema: Type[BaseModel] = GetRecordsInput\n",
    "\n",
    "    def _run(self, datasheet_id:str, filter_condition: Optional[dict]= None, sort_condition: Optional[list] = None, maxRecords_condition: Optional[int] = None):\n",
    "        try:\n",
    "            records = get_records(datasheet_id, filter_condition, sort_condition, maxRecords_condition)\n",
    "            return records\n",
    "        except Exception as e:\n",
    "            if str(e) == \"The sorted field does not exist\":\n",
    "                return f\"Function excute failedL: '{e}', please try another function to get right field name.\"\n",
    "            elif str(e) == \"api_param_formula_error\":\n",
    "                return  f\"Function excute failed: '{e}', please try to make right filter_condition.\"\n",
    "            else:\n",
    "                return f\"Function excute failed: '{e}', please try another function.\"\n",
    "\n",
    "    def _arun(self):\n",
    "        raise NotImplementedError(\"get_records does not support async\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "tools = [\n",
    "    Tool(\n",
    "        name = \"get_spaces\",\n",
    "        func=get_spaces,\n",
    "        description=\"useful for accessing all spaces the user has access to. Input should be in the form of a question containing full context\",\n",
    "    ),\n",
    "    Tool(\n",
    "        name=\"get_nodes\",\n",
    "        func=get_nodes,\n",
    "        description=\"useful for retrieving all file nodes in a space\"\n",
    "    ),\n",
    "    Tool(\n",
    "        name=\"get_fields\",\n",
    "        func=get_fields,\n",
    "        description=\"useful for retrieving fields in a datasheet, if you need get\"\n",
    "    ),\n",
    "    GetRecordsTool()\n",
    "]\n",
    "\n",
    "mrkl = initialize_agent(tools, llm, agent=AgentType.OPENAI_FUNCTIONS, verbose=True)\n",
    "\n",
    "def main(prompt: str):\n",
    "    with get_openai_callback() as cb:\n",
    "        mrkl.run(prompt)\n",
    "        print(f\"Total Tokens: {cb.total_tokens}\")\n",
    "        print(f\"Prompt Tokens: {cb.prompt_tokens}\")\n",
    "        print(f\"Completion Tokens: {cb.completion_tokens}\")\n",
    "        print(f\"Successful Requests: {cb.successful_requests}\")\n",
    "        print(f\"Total Cost (USD): ${cb.total_cost}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new  chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3m\n",
      "Invoking: `get_spaces` with `What spaces do I have?`\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3m[{'id': 'spcS0eZxZ8mSA', 'name': 'APITable Ltd.', 'isAdmin': None}, {'id': 'spctqtTZpssYw', 'name': \"xukecheng's Space\", 'isAdmin': True}, {'id': 'spcGSVizcwRYF', 'name': 'Gmail Space', 'isAdmin': None}, {'id': 'spcV5VkzU71Lf', 'name': 'test03-bug验证升级', 'isAdmin': None}]\u001b[0m\u001b[32;1m\u001b[1;3mYou have the following spaces:\n",
      "\n",
      "1. APITable Ltd.\n",
      "2. xukecheng's Space\n",
      "3. Gmail Space\n",
      "4. test03-bug验证升级\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n",
      "Total Tokens: 642\n",
      "Prompt Tokens: 585\n",
      "Completion Tokens: 57\n",
      "Successful Requests: 2\n",
      "Total Cost (USD): $0.0009915\n"
     ]
    }
   ],
   "source": [
    "main(\"What spaces do I have?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new  chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3m\n",
      "Invoking: `get_spaces` with `xukecheng`\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3m[{'id': 'spcS0eZxZ8mSA', 'name': 'APITable Ltd.', 'isAdmin': None}, {'id': 'spctqtTZpssYw', 'name': \"xukecheng's Space\", 'isAdmin': True}, {'id': 'spcGSVizcwRYF', 'name': 'Gmail Space', 'isAdmin': None}, {'id': 'spcV5VkzU71Lf', 'name': 'test03-bug验证升级', 'isAdmin': None}]\u001b[0m\u001b[32;1m\u001b[1;3m\n",
      "Invoking: `get_nodes` with `spctqtTZpssYw`\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[33;1m\u001b[1;3m[{'id': 'dstZTEueDaWedFWgbh', 'name': 'Project management', 'type': 'Datasheet', 'icon': '', 'isFav': False}, {'id': 'fodHl5i2LeMBH', 'name': 'Make Test', 'type': 'Folder', 'icon': '', 'isFav': False}, {'id': 'fomh7eepvbqkBrwCvW', 'name': 'Form', 'type': 'Form', 'icon': '', 'isFav': False}, {'id': 'dstFQKumRsAp4p5RBE', 'name': 'AITEST', 'type': 'Datasheet', 'icon': '', 'isFav': False}, {'id': 'dstWuVdBJNDzpjn8ck', 'name': 'email_oss', 'type': 'Datasheet', 'icon': '', 'isFav': False}, {'id': 'dsti6VpNpuKQpHVSnh', 'name': 'APITable MAUs', 'type': 'Datasheet', 'icon': '', 'isFav': False}]\u001b[0m\u001b[32;1m\u001b[1;3m\n",
      "Invoking: `get_records` with `{'datasheet_id': 'dsti6VpNpuKQpHVSnh', 'sort_condition': [{'field': 'Timestamp', 'order': 'desc'}], 'maxRecords_condition': 1}`\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3mFunction excute failed: 'The sorted field \"Timestamp\" does not exist', please try another tool.\u001b[0m\u001b[32;1m\u001b[1;3mI apologize, but it seems that there is an issue retrieving the latest value of APITable MAU in xukecheng's space. The error message suggests that the \"Timestamp\" field does not exist in the datasheet. Please check the field name or try using another tool to retrieve the latest value.\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n",
      "Total Tokens: 2180\n",
      "Prompt Tokens: 2012\n",
      "Completion Tokens: 168\n",
      "Successful Requests: 4\n",
      "Total Cost (USD): $0.003354\n"
     ]
    }
   ],
   "source": [
    "main(\"Tell me the latest value of APITable MAU in xukecheng's space\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}