notebooks/integrations/gemini/vector-search-gemini-elastic.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c52e30d1-cb29-4e70-af4a-9c953fcb0f2e",
   "metadata": {},
   "source": [
    "# Quickstart: Vector search using Gemini Embeddings and Elasticsearch\n",
    "\n",
    "This tutorial demonstrates how to use the [Gemini API](https://ai.google.dev/docs) to create [embeddings](https://ai.google.dev/docs/embeddings_guide) and store them in Elasticsearch. Elasticsearch will enable us to perform vector search (Knn) to find similar documents."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88303061-f357-43d8-8b63-c4f79e9a1746",
   "metadata": {},
   "source": [
    "## setup\n",
    "\n",
    "* Elastic Credentials - Create [Cloud deployment](https://www.elastic.co/search-labs/tutorials/install-elasticsearch/elastic-cloud) to get all Elastic credentials (`ELASTIC_CLOUD_ID`, `ELASTIC_API_KEY`).\n",
    "\n",
    "* `GOOGLE_API_KEY` - To use the Gemini API, you need an API key. [Follow](https://ai.google.dev/tutorials/setup) to create a key with one click in Google AI Studio."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76ca723c-6148-4682-a5ae-486e73cb2b94",
   "metadata": {},
   "source": [
    "## Install packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef1f1e52-f892-489f-8947-3e4698f5f5c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "pip install -q -U google-generativeai elasticsearch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d86d3fa-4ca0-41b6-a4bc-81bacf26bf02",
   "metadata": {},
   "source": [
    "## Import packages and credentials"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb62d8fb-6c34-44fd-bc94-18b644422ee8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import google.generativeai as genai\n",
    "import google.ai.generativelanguage as glm\n",
    "from elasticsearch import Elasticsearch, helpers\n",
    "from getpass import getpass\n",
    "\n",
    "GOOGLE_API_KEY = getpass(\"Google API Key :\")\n",
    "ELASTIC_API_KEY = getpass(\"Elastic API Key :\")\n",
    "ELASTIC_CLOUD_ID = getpass(\"Elastic Cloud ID :\")\n",
    "elastic_index_name = \"gemini-demo\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b22dc16-c0a0-48f0-979d-5d21c17bd264",
   "metadata": {},
   "source": [
    "## Embedding generation\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca56532d-7c82-4e2b-aecf-2173520d3696",
   "metadata": {},
   "outputs": [],
   "source": [
    "genai.configure(api_key=GOOGLE_API_KEY)\n",
    "\n",
    "title = \"Climate in India\"\n",
    "sample_text = \"India generally experiences a hot summer from March to June, with temperatures often exceeding 40°C in central and northern regions. Monsoon season, from June to September, brings heavy rainfall, especially in the western coast and northeastern areas. Post-monsoon months, October and November, mark a transition with decreasing rainfall. Winter, from December to February, varies in temperature across the country, with colder conditions in the north and milder weather in the south. India's diverse climate is influenced by its geographical features, resulting in regional \"\n",
    "\n",
    "model = \"models/embedding-001\"\n",
    "embedding = genai.embed_content(\n",
    "    model=model, content=sample_text, task_type=\"retrieval_document\", title=title\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6239eda7-3bed-43dd-a6a8-a8369b907d5c",
   "metadata": {},
   "source": [
    "## Connecting Elasticsearch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7cbade18-3049-46f1-8d3e-5b22d4aade5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "es = Elasticsearch(cloud_id=ELASTIC_CLOUD_ID, api_key=ELASTIC_API_KEY)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20d070c8-9e19-48a3-bc3b-5f22067eb63f",
   "metadata": {},
   "source": [
    "## Index document with Elasticsearch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e02ca81e-7caa-4505-95c6-3c6be7843c8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "doc = {\"text\": sample_text, \"text_embedding\": embedding[\"embedding\"]}\n",
    "\n",
    "resp = es.index(index=elastic_index_name, document=doc)\n",
    "\n",
    "print(resp)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afa0d371-afbf-4f98-9cd1-ee457839f323",
   "metadata": {},
   "source": [
    "## Searching for document with Elasticsearch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d71eeacc-d0c8-4035-b052-a1c03300aec0",
   "metadata": {},
   "outputs": [],
   "source": [
    "q = \"How's weather in India?\"\n",
    "\n",
    "embedding = genai.embed_content(model=model, content=q, task_type=\"retrieval_query\")\n",
    "\n",
    "resp = es.search(\n",
    "    index=elastic_index_name,\n",
    "    knn={\n",
    "        \"field\": \"text_embedding\",\n",
    "        \"query_vector\": embedding[\"embedding\"],\n",
    "        \"k\": 10,\n",
    "        \"num_candidates\": 100,\n",
    "    },\n",
    ")\n",
    "\n",
    "\n",
    "for result in resp[\"hits\"][\"hits\"]:\n",
    "    pretty_output = f\"\\n\\nID: {result['_id']}\\n\\nText: {result['_source']['text']}\\n\\nEmbedding: {result['_source']['text_embedding']}\"\n",
    "    print(pretty_output)"
   ]
  }
 ],
 "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.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}