{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1ad7250ddd99fba9",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "# Tencent Cloud VectorDB\n",
    "\n",
    "> [Tencent Cloud VectorDB](https://cloud.tencent.com/document/product/1709) is a fully managed, self-developed, enterprise-level distributed database    service designed for storing, retrieving, and analyzing multi-dimensional vector data.\n",
    "\n",
    "In the walkthrough, we'll demo the `SelfQueryRetriever` with a Tencent Cloud VectorDB."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "209652d4ab38ba7f",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## create a TencentVectorDB instance\n",
    "First we'll want to create a TencentVectorDB and seed it with some data. We've created a small demo set of documents that contain summaries of movies.\n",
    "\n",
    "**Note:** The self-query retriever requires you to have `lark` installed (`pip install lark`) along with integration-specific requirements."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b68da3303b0625f2",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-29T02:39:28.887634Z",
     "start_time": "2024-03-29T02:39:27.277978Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\r\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.2.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.0\u001b[0m\r\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\r\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip install --upgrade --quiet tcvectordb langchain-openai tiktoken lark"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1113af6008f3f3d",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "We want to use `OpenAIEmbeddings` so we have to get the OpenAI API Key."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c243e15bcf72d539",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-29T02:40:59.788206Z",
     "start_time": "2024-03-29T02:40:59.783798Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = getpass.getpass(\"OpenAI API Key:\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5277a4dba027bb8",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "create a TencentVectorDB instance and seed it with some data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "fd0c70c0be7d7130",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-29T02:42:28.467682Z",
     "start_time": "2024-03-29T02:42:21.255335Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from langchain_community.vectorstores.tencentvectordb import (\n",
    "    ConnectionParams,\n",
    "    MetaField,\n",
    "    TencentVectorDB,\n",
    ")\n",
    "from langchain_core.documents import Document\n",
    "from tcvectordb.model.enum import FieldType\n",
    "\n",
    "meta_fields = [\n",
    "    MetaField(name=\"year\", data_type=\"uint64\", index=True),\n",
    "    MetaField(name=\"rating\", data_type=\"string\", index=False),\n",
    "    MetaField(name=\"genre\", data_type=FieldType.String, index=True),\n",
    "    MetaField(name=\"director\", data_type=FieldType.String, index=True),\n",
    "]\n",
    "\n",
    "docs = [\n",
    "    Document(\n",
    "        page_content=\"The Shawshank Redemption is a 1994 American drama film written and directed by Frank Darabont.\",\n",
    "        metadata={\n",
    "            \"year\": 1994,\n",
    "            \"rating\": \"9.3\",\n",
    "            \"genre\": \"drama\",\n",
    "            \"director\": \"Frank Darabont\",\n",
    "        },\n",
    "    ),\n",
    "    Document(\n",
    "        page_content=\"The Godfather is a 1972 American crime film directed by Francis Ford Coppola.\",\n",
    "        metadata={\n",
    "            \"year\": 1972,\n",
    "            \"rating\": \"9.2\",\n",
    "            \"genre\": \"crime\",\n",
    "            \"director\": \"Francis Ford Coppola\",\n",
    "        },\n",
    "    ),\n",
    "    Document(\n",
    "        page_content=\"The Dark Knight is a 2008 superhero film directed by Christopher Nolan.\",\n",
    "        metadata={\n",
    "            \"year\": 2008,\n",
    "            \"rating\": \"9.0\",\n",
    "            \"genre\": \"science fiction\",\n",
    "            \"director\": \"Christopher Nolan\",\n",
    "        },\n",
    "    ),\n",
    "    Document(\n",
    "        page_content=\"Inception is a 2010 science fiction action film written and directed by Christopher Nolan.\",\n",
    "        metadata={\n",
    "            \"year\": 2010,\n",
    "            \"rating\": \"8.8\",\n",
    "            \"genre\": \"science fiction\",\n",
    "            \"director\": \"Christopher Nolan\",\n",
    "        },\n",
    "    ),\n",
    "    Document(\n",
    "        page_content=\"The Avengers is a 2012 American superhero film based on the Marvel Comics superhero team of the same name.\",\n",
    "        metadata={\n",
    "            \"year\": 2012,\n",
    "            \"rating\": \"8.0\",\n",
    "            \"genre\": \"science fiction\",\n",
    "            \"director\": \"Joss Whedon\",\n",
    "        },\n",
    "    ),\n",
    "    Document(\n",
    "        page_content=\"Black Panther is a 2018 American superhero film based on the Marvel Comics character of the same name.\",\n",
    "        metadata={\n",
    "            \"year\": 2018,\n",
    "            \"rating\": \"7.3\",\n",
    "            \"genre\": \"science fiction\",\n",
    "            \"director\": \"Ryan Coogler\",\n",
    "        },\n",
    "    ),\n",
    "]\n",
    "\n",
    "vector_db = TencentVectorDB.from_documents(\n",
    "    docs,\n",
    "    None,\n",
    "    connection_params=ConnectionParams(\n",
    "        url=\"http://10.0.X.X\",\n",
    "        key=\"eC4bLRy2va******************************\",\n",
    "        username=\"root\",\n",
    "        timeout=20,\n",
    "    ),\n",
    "    collection_name=\"self_query_movies\",\n",
    "    meta_fields=meta_fields,\n",
    "    drop_old=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3810b731a981a957",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Creating our self-querying retriever\n",
    "Now we can instantiate our retriever. To do this we'll need to provide some information upfront about the metadata fields that our documents support and a short description of the document contents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7095b68ea997468c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-29T02:42:37.901230Z",
     "start_time": "2024-03-29T02:42:36.836827Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from langchain.chains.query_constructor.base import AttributeInfo\n",
    "from langchain.retrievers.self_query.base import SelfQueryRetriever\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "metadata_field_info = [\n",
    "    AttributeInfo(\n",
    "        name=\"genre\",\n",
    "        description=\"The genre of the movie\",\n",
    "        type=\"string\",\n",
    "    ),\n",
    "    AttributeInfo(\n",
    "        name=\"year\",\n",
    "        description=\"The year the movie was released\",\n",
    "        type=\"integer\",\n",
    "    ),\n",
    "    AttributeInfo(\n",
    "        name=\"director\",\n",
    "        description=\"The name of the movie director\",\n",
    "        type=\"string\",\n",
    "    ),\n",
    "    AttributeInfo(\n",
    "        name=\"rating\", description=\"A 1-10 rating for the movie\", type=\"string\"\n",
    "    ),\n",
    "]\n",
    "document_content_description = \"Brief summary of a movie\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "cbbf7e54054bb3aa",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-29T02:42:45.187071Z",
     "start_time": "2024-03-29T02:42:45.138462Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "llm = ChatOpenAI(temperature=0, model=\"gpt-4\", max_tokens=4069)\n",
    "retriever = SelfQueryRetriever.from_llm(\n",
    "    llm, vector_db, document_content_description, metadata_field_info, verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65ff2054be9d5236",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Test it out\n",
    "And now we can try actually using our retriever!\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "267e2a68f26505b1",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-29T02:42:51.526470Z",
     "start_time": "2024-03-29T02:42:48.328191Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": "[Document(page_content='The Dark Knight is a 2008 superhero film directed by Christopher Nolan.', metadata={'year': 2008, 'rating': '9.0', 'genre': 'science fiction', 'director': 'Christopher Nolan'}),\n Document(page_content='The Avengers is a 2012 American superhero film based on the Marvel Comics superhero team of the same name.', metadata={'year': 2012, 'rating': '8.0', 'genre': 'science fiction', 'director': 'Joss Whedon'}),\n Document(page_content='Black Panther is a 2018 American superhero film based on the Marvel Comics character of the same name.', metadata={'year': 2018, 'rating': '7.3', 'genre': 'science fiction', 'director': 'Ryan Coogler'}),\n Document(page_content='The Godfather is a 1972 American crime film directed by Francis Ford Coppola.', metadata={'year': 1972, 'rating': '9.2', 'genre': 'crime', 'director': 'Francis Ford Coppola'})]"
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# This example only specifies a relevant query\n",
    "retriever.invoke(\"movies about a superhero\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "3afd98ca20782dda",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-29T02:42:55.179002Z",
     "start_time": "2024-03-29T02:42:53.057022Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": "[Document(page_content='The Avengers is a 2012 American superhero film based on the Marvel Comics superhero team of the same name.', metadata={'year': 2012, 'rating': '8.0', 'genre': 'science fiction', 'director': 'Joss Whedon'}),\n Document(page_content='Black Panther is a 2018 American superhero film based on the Marvel Comics character of the same name.', metadata={'year': 2018, 'rating': '7.3', 'genre': 'science fiction', 'director': 'Ryan Coogler'})]"
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# This example only specifies a filter\n",
    "retriever.invoke(\"movies that were released after 2010\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "9974f641e11abfe8",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-29T02:42:58.472620Z",
     "start_time": "2024-03-29T02:42:56.131594Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": "[Document(page_content='The Avengers is a 2012 American superhero film based on the Marvel Comics superhero team of the same name.', metadata={'year': 2012, 'rating': '8.0', 'genre': 'science fiction', 'director': 'Joss Whedon'}),\n Document(page_content='Black Panther is a 2018 American superhero film based on the Marvel Comics character of the same name.', metadata={'year': 2018, 'rating': '7.3', 'genre': 'science fiction', 'director': 'Ryan Coogler'})]"
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# This example specifies both a relevant query and a filter\n",
    "retriever.invoke(\"movies about a superhero which were released after 2010\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be593d3a6c508517",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Filter k\n",
    "\n",
    "We can also use the self query retriever to specify `k`: the number of documents to fetch.\n",
    "\n",
    "We can do this by passing `enable_limit=True` to the constructor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "e255b69c937fa424",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-29T02:43:02.779337Z",
     "start_time": "2024-03-29T02:43:02.759900Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "retriever = SelfQueryRetriever.from_llm(\n",
    "    llm,\n",
    "    vector_db,\n",
    "    document_content_description,\n",
    "    metadata_field_info,\n",
    "    verbose=True,\n",
    "    enable_limit=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "45674137c7f8a9d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-29T02:43:07.357830Z",
     "start_time": "2024-03-29T02:43:04.854323Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": "[Document(page_content='The Dark Knight is a 2008 superhero film directed by Christopher Nolan.', metadata={'year': 2008, 'rating': '9.0', 'genre': 'science fiction', 'director': 'Christopher Nolan'}),\n Document(page_content='The Avengers is a 2012 American superhero film based on the Marvel Comics superhero team of the same name.', metadata={'year': 2012, 'rating': '8.0', 'genre': 'science fiction', 'director': 'Joss Whedon'})]"
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "retriever.invoke(\"what are two movies about a superhero\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}