{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "Link to the original blog post: https://jina.ai/news/implementing-a-chat-history-rag-with-jina-ai-and-milvus-lite"
      ],
      "metadata": {
        "id": "4fsDDMiPlL37"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Implementing a Chat History RAG with Jina AI and Milvus Lite"
      ],
      "metadata": {
        "id": "iEO9F3VFmlKq"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Install the prerequisites"
      ],
      "metadata": {
        "id": "NWEh5VGXlOBA"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install -U pymilvus\n",
        "!pip install -U \"pymilvus[model]\"\n",
        "!pip install langchain\n",
        "!pip install langchain-community"
      ],
      "metadata": {
        "id": "IW5fp9auPFzU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Restart the kernel after installing the prerequisites"
      ],
      "metadata": {
        "id": "qEdd3v9fN3c5"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Download the dataset"
      ],
      "metadata": {
        "id": "XIWdHmLiO9qN"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "\n",
        "if not os.path.exists(\"chat_history.json\"):\n",
        "    !wget https://raw.githubusercontent.com/jina-ai/workshops/main/notebooks/embeddings/milvus/chat_history.json"
      ],
      "metadata": {
        "id": "V_cIRkm5O7rw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Add the [Jina AI API key](https://jina.ai/embeddings) and [Hugging Face token](https://huggingface.co/docs/hub/en/security-tokens) as environment variables\n",
        "\n",
        "Make sure that your Hugging Face token is set to READ to access the Hugging Face Hub on which [Mixtral 7B Instruct](https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1) is hosted."
      ],
      "metadata": {
        "id": "rjkLyi2ulRro"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import getpass\n",
        "\n",
        "os.environ[\"JINAAI_API_KEY\"] = getpass.getpass(prompt=\"Jina AI API Key: \")\n",
        "os.environ[\"HUGGINGFACEHUB_API_TOKEN\"] = getpass.getpass(prompt=\"Hugging Face Token: \")"
      ],
      "metadata": {
        "id": "F61RMqu4wsbs"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Create the Milvus Lite collection"
      ],
      "metadata": {
        "id": "P782IdpUlbT3"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from pymilvus import MilvusClient, DataType\n",
        "\n",
        "# Specify a local file name as uri parameter of MilvusClient to use Milvus Lite\n",
        "client = MilvusClient(\"milvus_jina.db\")\n",
        "\n",
        "schema = MilvusClient.create_schema(\n",
        "    auto_id=True,\n",
        "    enable_dynamic_field=True,\n",
        ")\n",
        "\n",
        "schema.add_field(field_name=\"id\", datatype=DataType.INT64, description=\"The Primary Key\", is_primary=True)\n",
        "schema.add_field(field_name=\"embedding\", datatype=DataType.FLOAT_VECTOR, description=\"The Embedding Vector\", dim=768)\n",
        "\n",
        "index_params = client.prepare_index_params()\n",
        "index_params.add_index(field_name=\"embedding\", metric_type=\"COSINE\", index_type=\"AUTOINDEX\")\n",
        "\n",
        "client.create_collection(collection_name=\"milvus_jina\", schema=schema, index_params=index_params)"
      ],
      "metadata": {
        "id": "CQz1yiyljUUM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Parse the chat history data"
      ],
      "metadata": {
        "id": "zW2aUAFLlec4"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import json\n",
        "\n",
        "with open(\"chat_history.json\", \"r\", encoding=\"utf-8\") as file:\n",
        "    chat_data = json.load(file)\n",
        "\n",
        "messages = []\n",
        "metadatas = []\n",
        "\n",
        "for channel in chat_data:\n",
        "  chat_history = channel[\"chat_history\"]\n",
        "  chat_topic = channel[\"topic\"]\n",
        "  chat_expert = channel[\"expert_user\"]\n",
        "  for message in chat_history:\n",
        "    text = f\"\"\"{message[\"user\"]}: {message[\"message\"]}\"\"\"\n",
        "    messages.append(text)\n",
        "    meta = {\n",
        "        \"time_stamp\": message[\"time_stamp\"],\n",
        "        \"file_name\": message[\"file_name\"],\n",
        "        \"parent_message_nr\": message[\"parent_message_nr\"],\n",
        "        \"channel\": chat_topic,\n",
        "        \"expert\": True if message[\"user\"] == chat_expert else False\n",
        "    }\n",
        "    metadatas.append(meta)"
      ],
      "metadata": {
        "id": "UlmpyORbG91o"
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Embed the data using Jina Embeddings v2"
      ],
      "metadata": {
        "id": "JKYQQoLslhKZ"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "Q79J9Dh5M_5w"
      },
      "outputs": [],
      "source": [
        "from pymilvus.model.dense import JinaEmbeddingFunction\n",
        "\n",
        "jina_ef = JinaEmbeddingFunction(\"jina-embeddings-v2-base-en\")\n",
        "\n",
        "embeddings = jina_ef.encode_documents(messages)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Index the data in the Milvus Lite collection"
      ],
      "metadata": {
        "id": "ba4BRQBXlilh"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "collection_data = [{\n",
        "    \"message\": message,\n",
        "    \"embedding\": embedding,\n",
        "    \"metadata\": metadata\n",
        "} for message, embedding, metadata in zip(messages, embeddings, metadatas)]\n",
        "\n",
        "data = client.insert(\n",
        "    collection_name=\"milvus_jina\",\n",
        "    data=collection_data\n",
        ")"
      ],
      "metadata": {
        "id": "dU2uUmqCD3Qm"
      },
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Query the chat history"
      ],
      "metadata": {
        "id": "2WmlZXihlkhR"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "query = \"Who knows the most about encryption protocols in my team?\""
      ],
      "metadata": {
        "id": "xYEjI13Lk5Zm"
      },
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Retrieve relevant messages and rerank them using Jina Reranker v1"
      ],
      "metadata": {
        "id": "-b4PRrj_oSyt"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from pymilvus.model.reranker import JinaRerankFunction\n",
        "\n",
        "query_vectors = jina_ef.encode_queries([query])\n",
        "\n",
        "results = client.search(\n",
        "    collection_name=\"milvus_jina\",\n",
        "    data=query_vectors,\n",
        "    limit=5,\n",
        ")\n",
        "\n",
        "results = results[0]\n",
        "\n",
        "ids = [results[i][\"id\"] for i in range(len(results))]\n",
        "\n",
        "results = client.get(\n",
        "    collection_name=\"milvus_jina\",\n",
        "    ids=ids,\n",
        "    output_fields=[\"id\", \"message\", \"metadata\"]\n",
        ")\n",
        "\n",
        "jina_rf = JinaRerankFunction(\"jina-reranker-v1-base-en\")\n",
        "\n",
        "documents = [results[i][\"message\"] for i in range(len(results))]\n",
        "reranked_documents = jina_rf(query, documents)\n",
        "\n",
        "reranked_messages = []\n",
        "for reranked_document in reranked_documents:\n",
        "  idx = reranked_document.index\n",
        "  reranked_messages.append(results[idx])"
      ],
      "metadata": {
        "id": "r9JNxW5jMWjx"
      },
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Generate the response using Mixtral 7B Instruct"
      ],
      "metadata": {
        "id": "OF3OiUSVohdL"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from langchain.prompts import PromptTemplate\n",
        "from langchain_community.llms import HuggingFaceEndpoint\n",
        "\n",
        "llm = HuggingFaceEndpoint(repo_id=\"mistralai/Mixtral-8x7B-Instruct-v0.1\")\n",
        "\n",
        "prompt = \"\"\"<s>[INST] Context information is below.\\\\n\n",
        "        It includes the five most relevant messages to the query, sorted based on their relevance to the query.\\\\n\n",
        "        ---------------------\\\\n\n",
        "        {context_str}\\\\\\\\n\n",
        "        ---------------------\\\\n\n",
        "        Given the context information and not prior knowledge,\n",
        "        answer the query. Please be brief, concise, and complete.\\\\n\n",
        "        If the context information does not contain an answer to the query,\n",
        "        respond with \\\\\"No information\\\\\".\\\\n\n",
        "        Query: {query_str}[/INST] </s>\"\"\"\n",
        "\n",
        "prompt = PromptTemplate(template=prompt, input_variables=[\"query_str\", \"context_str\"])\n",
        "\n",
        "llm_chain = prompt | llm\n",
        "\n",
        "answer = llm_chain.invoke({\"query_str\":query, \"context_str\":reranked_messages})\n",
        "\n",
        "print(f\"\\n\\nANSWER:\\n\\n{answer}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Zp_akv_3ZoNy",
        "outputId": "6dc3921b-e6bc-495f-f9a5-cd350eba1738"
      },
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/langchain_core/_api/deprecation.py:119: LangChainDeprecationWarning: The class `HuggingFaceEndpoint` was deprecated in LangChain 0.0.37 and will be removed in 0.3. An updated version of the class exists in the from langchain-huggingface package and should be used instead. To use it run `pip install -U from langchain-huggingface` and import as `from from langchain_huggingface import llms import HuggingFaceEndpoint`.\n",
            "  warn_deprecated(\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "The token has not been saved to the git credentials helper. Pass `add_to_git_credential=True` in this function directly or `--add-to-git-credential` if using via `huggingface-cli` if you want to set the git credential as well.\n",
            "Token is valid (permission: read).\n",
            "Your token has been saved to /root/.cache/huggingface/token\n",
            "Login successful\n",
            "\n",
            "\n",
            "ANSWER",
            "\n",
            "\n",
            "Based on the context information, User5 seems to be the most knowledgeable about encryption protocols in your team. They have mentioned that the new protocols enhance data security significantly, especially for cloud deployments.\n"
          ]
        }
      ]
    }
  ]
}