notebooks/langchain/langchain-vector-store-using-elser.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Similarity search with Langchain and SparseVectorRetrievalStrategy(ELSER Model)\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/elastic/elasticsearch-labs/blob/main/notebooks/langchain/langchain-vector-store-using-elser.ipynb)\n",
    "\n",
    "\n",
    "This workbook demonstrates similiarity search using [SparseVectorRetrievalStrategy](https://api.python.langchain.com/en/latest/vectorstores/langchain.vectorstores.elasticsearch.SparseRetrievalStrategy.html#langchain.vectorstores.elasticsearch.SparseRetrievalStrategy) (ELSER). First,  we  split the documents into chunks using `langchain` and then index into elasticsearch through [`ElasticsearchStore.from_documents`](https://api.python.langchain.com/en/latest/vectorstores/langchain.vectorstores.elasticsearch.ElasticsearchStore.html#langchain.vectorstores.elasticsearch.ElasticsearchStore.from_documents). \n",
    "\n",
    "The [SparseVectorRetrievalStrategy](https://api.python.langchain.com/en/latest/vectorstores/langchain.vectorstores.elasticsearch.SparseRetrievalStrategy.html#langchain.vectorstores.elasticsearch.SparseRetrievalStrategy) converts each document into tokens and would be stored in `vector` field with datatype `rank_features`. Hence, the inference is handled within elasticsearch.\n",
    "\n",
    "The `similarity_search` performs semantic search using `text_expansion` and sets the query.\n",
    "\n",
    "We will then see how to filter the metadata within query. \n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install packages and import modules\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# install packages\n",
    "!python3 -m pip install -qU langchain openai elasticsearch tiktoken\n",
    "\n",
    "# import modules\n",
    "from getpass import getpass\n",
    "from langchain.vectorstores import ElasticsearchStore\n",
    "from urllib.request import urlopen\n",
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "import json"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Connect to Elasticsearch\n",
    "\n",
    "ℹ️ We're using an Elastic Cloud deployment of Elasticsearch for this notebook. If you don't have an Elastic Cloud deployment, sign up [here](https://cloud.elastic.co/registration?utm_source=github&utm_content=elasticsearch-labs-notebook) for a free trial. \n",
    "\n",
    "We'll use the **Cloud ID** to identify our deployment, because we are using Elastic Cloud deployment. To find the Cloud ID for your deployment, go to https://cloud.elastic.co/deployments and select your deployment.\n",
    "\n",
    "\n",
    "We will use [ElasticsearchStore](https://api.python.langchain.com/en/latest/vectorstores/langchain.vectorstores.elasticsearch.ElasticsearchStore.html) to connect to our elastic cloud deployment, This would help create and index data easily. We will be using text embedding from ELSER model `.elser_model_2`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# https://www.elastic.co/search-labs/tutorials/install-elasticsearch/elastic-cloud#finding-your-cloud-id\n",
    "ELASTIC_CLOUD_ID = getpass(\"Elastic Cloud ID: \")\n",
    "\n",
    "# https://www.elastic.co/search-labs/tutorials/install-elasticsearch/elastic-cloud#creating-an-api-key\n",
    "ELASTIC_API_KEY = getpass(\"Elastic Api Key: \")\n",
    "\n",
    "# https://platform.openai.com/api-keys\n",
    "OPENAI_API_KEY = getpass(\"OpenAI API key: \")\n",
    "\n",
    "vector_store = ElasticsearchStore(\n",
    "  es_cloud_id=ELASTIC_CLOUD_ID, \n",
    "  es_api_key=ELASTIC_API_KEY,\n",
    "  index_name= \"workplace_index\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download the dataset \n",
    "\n",
    "Let's download the sample dataset and deserialize the document."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "url = \"https://raw.githubusercontent.com/elastic/elasticsearch-labs/main/example-apps/chatbot-rag-app/data/data.json\"\n",
    "\n",
    "response = urlopen(url)\n",
    "\n",
    "workplace_docs = json.loads(response.read())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Split Documents into Passages\n",
    "\n",
    "\n",
    "We will chunk these documents into 500 token passages with an overlap of 0 tokens using a simple splitter. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metadata = []\n",
    "content = []\n",
    "\n",
    "for doc in workplace_docs:\n",
    "  content.append(doc[\"content\"])\n",
    "  metadata.append({\n",
    "      \"name\": doc[\"name\"],\n",
    "      \"summary\": doc[\"summary\"],\n",
    "      \"rolePermissions\":doc[\"rolePermissions\"]\n",
    "  })\n",
    "\n",
    "text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(chunk_size=512, chunk_overlap=256)\n",
    "docs = text_splitter.create_documents(content, metadatas=metadata)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Index data into elasticsearch\n",
    "\n",
    "Next, we will index data to elasticsearch using [ElasticsearchStore.from_documents](https://api.python.langchain.com/en/latest/vectorstores/langchain.vectorstores.elasticsearch.ElasticsearchStore.html#langchain.vectorstores.elasticsearch.ElasticsearchStore.from_documents). We will use Cloud ID,  Password and Index name values set in the `Create cloud deployment` step.\n",
    "\n",
    "In the instance, we will set `strategy` to [ElasticsearchStore.SparseVectorRetrievalStrategy()](https://api.python.langchain.com/en/latest/vectorstores/langchain.vectorstores.elasticsearch.SparseRetrievalStrategy.html#langchain.vectorstores.elasticsearch.SparseRetrievalStrategy)\n",
    "\n",
    "Note: Before we begin indexing, ensure you have [downloaded and deployed ELSER model](https://www.elastic.co/guide/en/machine-learning/current/ml-nlp-elser.html#download-deploy-elser) in your deployment and is running in ml node. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "documents = vector_store.from_documents(\n",
    "    docs, \n",
    "    es_cloud_id=ELASTIC_CLOUD_ID,\n",
    "    es_api_key=ELASTIC_API_KEY,\n",
    "    index_name=\"workplace_index\",\n",
    "    strategy=ElasticsearchStore.SparseVectorRetrievalStrategy(model_id = \".elser_model_2\")\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Results functions\n",
    "Next, we will create a small function to show the results of our query in human-readable outputs. This function would be used in our examples to display the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def showResults(output):\n",
    "  print(\"Total results: \", len(output))\n",
    "  for index in range(len(output)):\n",
    "    print(output[index])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Querying the dataset with similarity_search\n",
    "\n",
    "Now that we have indexed our sample data to elasticsearch, we will perform a similarity search on query - `How does the compensation work?`. By default returns top `4` documents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total results:  4\n",
      "page_content='Performance-Based Compensation:\\nIn addition to the defined compensation bands, we emphasize a performance-based compensation model. Performance evaluations will be conducted regularly, and employees exceeding performance expectations will be eligible for bonuses, incentives, and salary increases. This approach rewards high achievers and motivates employees to excel in their roles.' metadata={'summary': 'This document outlines a compensation framework for IT teams. It includes job levels, compensation bands, and performance-based incentives to ensure fair and competitive wages. Regular market benchmarking will be conducted to adjust the bands according to industry trends.', 'rolePermissions': ['manager'], 'name': 'Compensation Framework For It Teams'}\n",
      "page_content='Performance-Based Compensation:\\nIn addition to the defined compensation bands, we emphasize a performance-based compensation model. Performance evaluations will be conducted regularly, and employees exceeding performance expectations will be eligible for bonuses, incentives, and salary increases. This approach rewards high achievers and motivates employees to excel in their roles.' metadata={'summary': 'This document outlines a compensation framework for IT teams. It includes job levels, compensation bands, and performance-based incentives to ensure fair and competitive wages. Regular market benchmarking will be conducted to adjust the bands according to industry trends.', 'rolePermissions': ['manager'], 'name': 'Compensation Framework For It Teams'}\n",
      "page_content='Purpose:\\nThe purpose of this compensation bands strategy is to:\\na. Define clear guidelines for salary ranges based on job levels and market benchmarks.\\nb. Support equitable compensation practices across different teams.\\nc. Encourage employee growth and performance.\\nd. Enable effective budgeting and resource allocation.' metadata={'summary': 'This document outlines a compensation framework for IT teams. It includes job levels, compensation bands, and performance-based incentives to ensure fair and competitive wages. Regular market benchmarking will be conducted to adjust the bands according to industry trends.', 'rolePermissions': ['manager'], 'name': 'Compensation Framework For It Teams'}\n",
      "page_content='Purpose:\\nThe purpose of this compensation bands strategy is to:\\na. Define clear guidelines for salary ranges based on job levels and market benchmarks.\\nb. Support equitable compensation practices across different teams.\\nc. Encourage employee growth and performance.\\nd. Enable effective budgeting and resource allocation.' metadata={'summary': 'This document outlines a compensation framework for IT teams. It includes job levels, compensation bands, and performance-based incentives to ensure fair and competitive wages. Regular market benchmarking will be conducted to adjust the bands according to industry trends.', 'rolePermissions': ['manager'], 'name': 'Compensation Framework For It Teams'}\n"
     ]
    }
   ],
   "source": [
    "results = documents.similarity_search(\"How does the compensation work?\")\n",
    "showResults(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Querying the dataset show top 10 documents\n",
    "\n",
    "Now we will set `k=10` and try same query to see top 10 documents.  \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total results:  10\n",
      "page_content='Performance-Based Compensation:\\nIn addition to the defined compensation bands, we emphasize a performance-based compensation model. Performance evaluations will be conducted regularly, and employees exceeding performance expectations will be eligible for bonuses, incentives, and salary increases. This approach rewards high achievers and motivates employees to excel in their roles.' metadata={'summary': 'This document outlines a compensation framework for IT teams. It includes job levels, compensation bands, and performance-based incentives to ensure fair and competitive wages. Regular market benchmarking will be conducted to adjust the bands according to industry trends.', 'rolePermissions': ['manager'], 'name': 'Compensation Framework For It Teams'}\n",
      "page_content='Performance-Based Compensation:\\nIn addition to the defined compensation bands, we emphasize a performance-based compensation model. Performance evaluations will be conducted regularly, and employees exceeding performance expectations will be eligible for bonuses, incentives, and salary increases. This approach rewards high achievers and motivates employees to excel in their roles.' metadata={'summary': 'This document outlines a compensation framework for IT teams. It includes job levels, compensation bands, and performance-based incentives to ensure fair and competitive wages. Regular market benchmarking will be conducted to adjust the bands according to industry trends.', 'rolePermissions': ['manager'], 'name': 'Compensation Framework For It Teams'}\n",
      "page_content='Purpose:\\nThe purpose of this compensation bands strategy is to:\\na. Define clear guidelines for salary ranges based on job levels and market benchmarks.\\nb. Support equitable compensation practices across different teams.\\nc. Encourage employee growth and performance.\\nd. Enable effective budgeting and resource allocation.' metadata={'summary': 'This document outlines a compensation framework for IT teams. It includes job levels, compensation bands, and performance-based incentives to ensure fair and competitive wages. Regular market benchmarking will be conducted to adjust the bands according to industry trends.', 'rolePermissions': ['manager'], 'name': 'Compensation Framework For It Teams'}\n",
      "page_content='Purpose:\\nThe purpose of this compensation bands strategy is to:\\na. Define clear guidelines for salary ranges based on job levels and market benchmarks.\\nb. Support equitable compensation practices across different teams.\\nc. Encourage employee growth and performance.\\nd. Enable effective budgeting and resource allocation.' metadata={'summary': 'This document outlines a compensation framework for IT teams. It includes job levels, compensation bands, and performance-based incentives to ensure fair and competitive wages. Regular market benchmarking will be conducted to adjust the bands according to industry trends.', 'rolePermissions': ['manager'], 'name': 'Compensation Framework For It Teams'}\n",
      "page_content='Market Benchmarking:\\nTo ensure our compensation remains competitive, regular market benchmarking will be conducted. This involves analyzing industry salary trends, regional compensation data, and market demand for specific roles. The findings will inform periodic adjustments to our compensation bands to maintain alignment with the market.' metadata={'summary': 'This document outlines a compensation framework for IT teams. It includes job levels, compensation bands, and performance-based incentives to ensure fair and competitive wages. Regular market benchmarking will be conducted to adjust the bands according to industry trends.', 'rolePermissions': ['manager'], 'name': 'Compensation Framework For It Teams'}\n",
      "page_content='Market Benchmarking:\\nTo ensure our compensation remains competitive, regular market benchmarking will be conducted. This involves analyzing industry salary trends, regional compensation data, and market demand for specific roles. The findings will inform periodic adjustments to our compensation bands to maintain alignment with the market.' metadata={'summary': 'This document outlines a compensation framework for IT teams. It includes job levels, compensation bands, and performance-based incentives to ensure fair and competitive wages. Regular market benchmarking will be conducted to adjust the bands according to industry trends.', 'rolePermissions': ['manager'], 'name': 'Compensation Framework For It Teams'}\n",
      "page_content=\"Introduction:\\nThis document outlines the compensation bands strategy for the various teams within our IT company. The goal is to establish a fair and competitive compensation structure that aligns with industry standards, rewards performance, and attracts top talent. By implementing this strategy, we aim to foster employee satisfaction and retention while ensuring the company's overall success.\" metadata={'summary': 'This document outlines a compensation framework for IT teams. It includes job levels, compensation bands, and performance-based incentives to ensure fair and competitive wages. Regular market benchmarking will be conducted to adjust the bands according to industry trends.', 'rolePermissions': ['manager'], 'name': 'Compensation Framework For It Teams'}\n",
      "page_content=\"Introduction:\\nThis document outlines the compensation bands strategy for the various teams within our IT company. The goal is to establish a fair and competitive compensation structure that aligns with industry standards, rewards performance, and attracts top talent. By implementing this strategy, we aim to foster employee satisfaction and retention while ensuring the company's overall success.\" metadata={'summary': 'This document outlines a compensation framework for IT teams. It includes job levels, compensation bands, and performance-based incentives to ensure fair and competitive wages. Regular market benchmarking will be conducted to adjust the bands according to industry trends.', 'rolePermissions': ['manager'], 'name': 'Compensation Framework For It Teams'}\n",
      "page_content='Compensation Bands:\\nBased on the job levels, the following compensation bands have been established:\\na. Entry-Level Band: This band encompasses salary ranges for employees in entry-level positions. It aims to provide competitive compensation for individuals starting their careers within the company.' metadata={'summary': 'This document outlines a compensation framework for IT teams. It includes job levels, compensation bands, and performance-based incentives to ensure fair and competitive wages. Regular market benchmarking will be conducted to adjust the bands according to industry trends.', 'rolePermissions': ['manager'], 'name': 'Compensation Framework For It Teams'}\n",
      "page_content='Compensation Bands:\\nBased on the job levels, the following compensation bands have been established:\\na. Entry-Level Band: This band encompasses salary ranges for employees in entry-level positions. It aims to provide competitive compensation for individuals starting their careers within the company.' metadata={'summary': 'This document outlines a compensation framework for IT teams. It includes job levels, compensation bands, and performance-based incentives to ensure fair and competitive wages. Regular market benchmarking will be conducted to adjust the bands according to industry trends.', 'rolePermissions': ['manager'], 'name': 'Compensation Framework For It Teams'}\n"
     ]
    }
   ],
   "source": [
    "results = documents.similarity_search(\"How does the compensation work?\", k=10)\n",
    "showResults(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Querying the dataset with filtering Metadata\n",
    "We will now add metadata filtering by Keyword at query time, to match `rolePermissions` as `manager`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total results:  4\n",
      "page_content='Performance-Based Compensation:\\nIn addition to the defined compensation bands, we emphasize a performance-based compensation model. Performance evaluations will be conducted regularly, and employees exceeding performance expectations will be eligible for bonuses, incentives, and salary increases. This approach rewards high achievers and motivates employees to excel in their roles.' metadata={'summary': 'This document outlines a compensation framework for IT teams. It includes job levels, compensation bands, and performance-based incentives to ensure fair and competitive wages. Regular market benchmarking will be conducted to adjust the bands according to industry trends.', 'rolePermissions': ['manager'], 'name': 'Compensation Framework For It Teams'}\n",
      "page_content='Performance-Based Compensation:\\nIn addition to the defined compensation bands, we emphasize a performance-based compensation model. Performance evaluations will be conducted regularly, and employees exceeding performance expectations will be eligible for bonuses, incentives, and salary increases. This approach rewards high achievers and motivates employees to excel in their roles.' metadata={'summary': 'This document outlines a compensation framework for IT teams. It includes job levels, compensation bands, and performance-based incentives to ensure fair and competitive wages. Regular market benchmarking will be conducted to adjust the bands according to industry trends.', 'rolePermissions': ['manager'], 'name': 'Compensation Framework For It Teams'}\n",
      "page_content='Compensation Bands:\\nBased on the job levels, the following compensation bands have been established:\\na. Entry-Level Band: This band encompasses salary ranges for employees in entry-level positions. It aims to provide competitive compensation for individuals starting their careers within the company.' metadata={'summary': 'This document outlines a compensation framework for IT teams. It includes job levels, compensation bands, and performance-based incentives to ensure fair and competitive wages. Regular market benchmarking will be conducted to adjust the bands according to industry trends.', 'rolePermissions': ['manager'], 'name': 'Compensation Framework For It Teams'}\n",
      "page_content='Compensation Bands:\\nBased on the job levels, the following compensation bands have been established:\\na. Entry-Level Band: This band encompasses salary ranges for employees in entry-level positions. It aims to provide competitive compensation for individuals starting their careers within the company.' metadata={'summary': 'This document outlines a compensation framework for IT teams. It includes job levels, compensation bands, and performance-based incentives to ensure fair and competitive wages. Regular market benchmarking will be conducted to adjust the bands according to industry trends.', 'rolePermissions': ['manager'], 'name': 'Compensation Framework For It Teams'}\n"
     ]
    }
   ],
   "source": [
    "results = documents.similarity_search(\"How does the compensation work\", filter=[{ 'match': { \"metadata.rolePermissions\": \"manager\" }}])\n",
    "showResults(results)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.11.4 64-bit",
   "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.11.4"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "b0fa6594d8f4cbf19f97940f81e996739fb7646882a419484c72d19e05852a7e"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}