run-llama--llama_index/docs/examples/vector_stores/RocksetIndexDemo.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/vector_stores/RocksetIndexDemo.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Rockset Vector Store\n",
    "\n",
    "As a real-time search and analytics database, Rockset uses indexing to deliver scalable and performant personalization, product search, semantic search, chatbot applications, and more.\n",
    "Since Rockset is purpose-built for real-time, you can build these responsive applications on constantly updating, streaming data. \n",
    "By integrating Rockset with LlamaIndex, you can easily use LLMs on your own real-time data for production-ready vector search applications.\n",
    "\n",
    "We'll walk through a demonstration of how to use Rockset as a vector store in LlamaIndex. \n",
    "\n",
    "## Tutorial\n",
    "In this example, we'll use OpenAI's `text-embedding-ada-002` model to generate embeddings and Rockset as vector store to store embeddings.\n",
    "We'll ingest text from a file and ask questions about the content.\n",
    "\n",
    "### Setting Up Your Environment\n",
    "1. Create a [collection](https://rockset.com/docs/collections) from the Rockset console with the [Write API](https://rockset.com/docs/write-api/) as your source.\n",
    "Name your collection `llamaindex_demo`. Configure the following [ingest transformation](https://rockset.com/docs/ingest-transformation) \n",
    "with [`VECTOR_ENFORCE`](https://rockset.com/docs/vector-functions) to define your embeddings field and take advantage of performance and storage optimizations:\n",
    "```sql\n",
    "SELECT \n",
    "    _input.* EXCEPT(_meta), \n",
    "    VECTOR_ENFORCE(\n",
    "        _input.embedding,\n",
    "        1536,\n",
    "        'float'\n",
    "    ) as embedding\n",
    "FROM _input\n",
    "```\n",
    "\n",
    "2. Create an [API key](https://rockset.com/docs/iam) from the Rockset console and set the `ROCKSET_API_KEY` environment variable.\n",
    "Find your API server [here](http://rockset.com/docs/rest-api#introduction) and set the `ROCKSET_API_SERVER` environment variable. \n",
    "Set the `OPENAI_API_KEY` environment variable.\n",
    "\n",
    "3. Install the dependencies.\n",
    "```shell\n",
    "pip3 install llama_index rockset \n",
    "```\n",
    "\n",
    "4. LlamaIndex allows you to ingest data from a variety of sources. \n",
    "For this example, we'll read from a text file named `constitution.txt`, which is a transcript of the American Constitution, found [here](https://www.archives.gov/founding-docs/constitution-transcript). \n",
    "\n",
    "### Data ingestion \n",
    "Use LlamaIndex's `SimpleDirectoryReader` class to convert the text file to a list of `Document` objects."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-llms-openai\n",
    "%pip install llama-index-vector-stores-rocksetdb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import SimpleDirectoryReader\n",
    "\n",
    "docs = SimpleDirectoryReader(\n",
    "    input_files=[\"{path to}/consitution.txt\"]\n",
    ").load_data()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Instantiate the LLM and service context."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import Settings\n",
    "from llama_index.llms.openai import OpenAI\n",
    "\n",
    "Settings.llm = OpenAI(temperature=0.8, model=\"gpt-3.5-turbo\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Instantiate the vector store and storage context."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import StorageContext\n",
    "from llama_index.vector_stores.rocksetdb import RocksetVectorStore\n",
    "\n",
    "vector_store = RocksetVectorStore(collection=\"llamaindex_demo\")\n",
    "storage_context = StorageContext.from_defaults(vector_store=vector_store)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Add documents to the `llamaindex_demo` collection and create an index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import VectorStoreIndex\n",
    "\n",
    "index = VectorStoreIndex.from_documents(\n",
    "    docs,\n",
    "    storage_context=storage_context,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Querying\n",
    "Ask a question about your document and generate a response."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = index.as_query_engine().query(\"What is the duty of the president?\")\n",
    "\n",
    "print(str(response))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "Run the program.\n",
    "```text\n",
    "$ python3 main.py\n",
    "The duty of the president is to faithfully execute the Office of President of the United States, preserve, protect and defend the Constitution of the United States, serve as the Commander in Chief of the Army and Navy, grant reprieves and pardons for offenses against the United States (except in cases of impeachment), make treaties and appoint ambassadors and other public ministers, take care that the laws be faithfully executed, and commission all the officers of the United States.\n",
    "```\n",
    "\n",
    "## Metadata Filtering\n",
    "Metadata filtering allows you to retrieve relevant documents that match specific filters.\n",
    "\n",
    "1. Add nodes to your vector store and create an index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.vector_stores.rocksetdb import RocksetVectorStore\n",
    "from llama_index.core import VectorStoreIndex, StorageContext\n",
    "from llama_index.core.vector_stores.types import NodeWithEmbedding\n",
    "from llama_index.core.schema import TextNode\n",
    "\n",
    "nodes = [\n",
    "    NodeWithEmbedding(\n",
    "        node=TextNode(\n",
    "            text=\"Apples are blue\",\n",
    "            metadata={\"type\": \"fruit\"},\n",
    "        ),\n",
    "        embedding=[],\n",
    "    )\n",
    "]\n",
    "index = VectorStoreIndex(\n",
    "    nodes,\n",
    "    storage_context=StorageContext.from_defaults(\n",
    "        vector_store=RocksetVectorStore(collection=\"llamaindex_demo\")\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2. Define metadata filters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.vector_stores import ExactMatchFilter, MetadataFilters\n",
    "\n",
    "filters = MetadataFilters(\n",
    "    filters=[ExactMatchFilter(key=\"type\", value=\"fruit\")]\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "3. Retrieve relevant documents that satisfy the filters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "retriever = index.as_retriever(filters=filters)\n",
    "retriever.retrieve(\"What colors are apples?\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating an Index from an Existing Collection\n",
    "You can create indices with data from existing collections."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import VectorStoreIndex\n",
    "from llama_index.vector_stores.rocksetdb import RocksetVectorStore\n",
    "\n",
    "vector_store = RocksetVectorStore(collection=\"llamaindex_demo\")\n",
    "\n",
    "index = VectorStoreIndex.from_vector_store(vector_store)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating an Index from a New Collection\n",
    "You can also create a new Rockset collection to use as a vector store."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.vector_stores.rocksetdb import RocksetVectorStore\n",
    "\n",
    "vector_store = RocksetVectorStore.with_new_collection(\n",
    "    collection=\"llamaindex_demo\",  # name of new collection\n",
    "    dimensions=1536,  # specifies length of vectors in ingest tranformation (optional)\n",
    "    # other RocksetVectorStore args\n",
    ")\n",
    "\n",
    "index = VectorStoreIndex(\n",
    "    nodes,\n",
    "    storage_context=StorageContext.from_defaults(vector_store=vector_store),\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configuration\n",
    "* **collection**: Name of the collection to query (required).\n",
    "\n",
    "```python\n",
    "RocksetVectorStore(collection=\"my_collection\")\n",
    "```\n",
    "\n",
    "* **workspace**: Name of the workspace containing the collection. Defaults to `\"commons\"`.\n",
    "```python\n",
    "RocksetVectorStore(worksapce=\"my_workspace\")\n",
    "```\n",
    "\n",
    "* **api_key**: The API key to use to authenticate Rockset requests. Ignored if `client` is passed in. Defaults to the `ROCKSET_API_KEY` environment variable.\n",
    "```python\n",
    "RocksetVectorStore(api_key=\"<my key>\")\n",
    "```\n",
    "\n",
    "* **api_server**: The API server to use for Rockset requests. Ignored if `client` is passed in. Defaults to the `ROCKSET_API_KEY` environment variable or `\"https://api.use1a1.rockset.com\"` if the `ROCKSET_API_SERVER` is not set.\n",
    "```python\n",
    "from rockset import Regions\n",
    "RocksetVectorStore(api_server=Regions.euc1a1)\n",
    "```\n",
    "\n",
    "* **client**: Rockset client object to use to execute Rockset requests. If not specified, a client object is internally constructed with the `api_key` parameter (or `ROCKSET_API_SERVER` environment variable) and the `api_server` parameter (or `ROCKSET_API_SERVER` environment variable).\n",
    "```python\n",
    "from rockset import RocksetClient\n",
    "RocksetVectorStore(client=RocksetClient(api_key=\"<my key>\"))\n",
    "```\n",
    "\n",
    "* **embedding_col**: The name of the database field containing embeddings. Defaults to `\"embedding\"`.\n",
    "```python\n",
    "RocksetVectorStore(embedding_col=\"my_embedding\")\n",
    "```\n",
    "\n",
    "* **metadata_col**: The name of the database field containing node data. Defaults to `\"metadata\"`.\n",
    "```python\n",
    "RocksetVectorStore(metadata_col=\"node\")\n",
    "```\n",
    "\n",
    "* **distance_func**: The metric to measure vector relationship. Defaults to cosine similarity.\n",
    "```python\n",
    "RocksetVectorStore(distance_func=RocksetVectorStore.DistanceFunc.DOT_PRODUCT)\n",
    "```"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}