{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import uuid\n",
    "import math\n",
    "import pinecone\n",
    "import pandas as pd\n",
    "from environs import Env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = Env()\n",
    "env.read_env()\n",
    "\n",
    "PINECONE_API_KEY = env.str('PINECONE_API_KEY')\n",
    "PINECONE_ENVIRONMENT = env.str('PINECONE_ENVIRONMENT')\n",
    "PINECONE_INDEX_NAME = env.str('PINECONE_INDEX_NAME', default='embedding-768')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pinecone.init(api_key=PINECONE_API_KEY, environment=PINECONE_ENVIRONMENT)\n",
    "\n",
    "def get_or_create_pinecone_index(index_name):\n",
    "    if index_name not in pinecone.list_indexes():\n",
    "        pinecone.create_index(\n",
    "            index_name,\n",
    "            dimension=768,\n",
    "            metric='cosine',\n",
    "            metadata_config={'indexed': ['version']}\n",
    "        )\n",
    "    return pinecone.Index(index_name)\n",
    "\n",
    "pinecone_index = get_or_create_pinecone_index(PINECONE_INDEX_NAME)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_pinecone_vectors(df: pd.DataFrame):\n",
    "    vecs = []\n",
    "    for i, r in df.iterrows():\n",
    "        metadata = r.to_dict()\n",
    "        del metadata['vector']\n",
    "        vecs.append({\n",
    "            \"id\": str(uuid.uuid4()),\n",
    "            \"values\": r.vector.tolist(),\n",
    "            \"metadata\": metadata\n",
    "        })\n",
    "    return vecs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def bulk_upsert_to_pinecone(df, ns):\n",
    "    step = 100\n",
    "    for i in range(math.ceil(len(df) / step)):\n",
    "        sliced_df = df[i * step: (i + 1) * step]\n",
    "        pinecone_index.upsert(vectors=get_pinecone_vectors(sliced_df), namespace=ns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_pickle('/data/documents.pkl').reset_index(drop=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bulk_upsert_to_pinecone(df[df.source == 'tidb'], 'tidb-v7.1-2023-04-24')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bulk_upsert_to_pinecone(df[df.source == 'tidbcloud'], 'tidbcloud-2023-04-24')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bulk_upsert_to_pinecone(df[df.source == 'community'], 'community-2023-04-24')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bulk_upsert_to_pinecone(df[df.source == 'knowledge-base-cn'], 'knowledge-base-cn-2023-04-24')"
   ]
  }
 ],
 "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
}