langchain/docs/extras/integrations/document_loaders/rockset.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Rockset\n",
    "\n",
    "> Rockset is a real-time analytics database which enables queries on massive, semi-structured data without operational burden. With Rockset, ingested data is queryable within one second and analytical queries against that data typically execute in milliseconds. Rockset is compute optimized, making it suitable for serving high concurrency applications in the sub-100TB range (or larger than 100s of TBs with rollups).\n",
    "\n",
    "This notebook demonstrates how to use Rockset as a document loader in langchain. To get started, make sure you have a Rockset account and an API key available.\n",
    "\n",
    "\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up the environment\n",
    "\n",
    "1. Go to the [Rockset console](https://console.rockset.com/apikeys) and get an API key. Find your API region from the [API reference](https://rockset.com/docs/rest-api/#introduction). For the purpose of this notebook, we will assume you're using Rockset from `Oregon(us-west-2)`.\n",
    "2. Set your the environment variable `ROCKSET_API_KEY`.\n",
    "3. Install the Rockset python client, which will be used by langchain to interact with the Rockset database."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "shellscript"
    }
   },
   "outputs": [],
   "source": [
    "$ pip3 install rockset"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Loading Documents\n",
    "The Rockset integration with LangChain allows you to load documents from Rockset collections with SQL queries. In order to do this you must construct a `RocksetLoader` object. Here is an example snippet that initializes a `RocksetLoader`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.document_loaders import RocksetLoader\n",
    "from rockset import RocksetClient, Regions, models\n",
    "\n",
    "loader = RocksetLoader(\n",
    "    RocksetClient(Regions.usw2a1, \"<api key>\"),\n",
    "    models.QueryRequestSql(query=\"SELECT * FROM langchain_demo LIMIT 3\"),  # SQL query\n",
    "    [\"text\"],  # content columns\n",
    "    metadata_keys=[\"id\", \"date\"],  # metadata columns\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, you can see that the following query is run:\n",
    "\n",
    "```sql\n",
    "SELECT * FROM langchain_demo LIMIT 3\n",
    "```\n",
    "\n",
    "The `text` column in the collection is used as the page content, and the record's `id` and `date` columns are used as metadata (if you do not pass anything into `metadata_keys`, the whole Rockset document will be used as metadata). \n",
    "\n",
    "To execute the query and access an iterator over the resulting `Document`s, run:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loader.lazy_load()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To execute the query and access all resulting `Document`s at once, run:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loader.load()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is an example response of `loader.load()`:\n",
    "```python\n",
    "[\n",
    "    Document(\n",
    "        page_content=\"Lorem ipsum dolor sit amet, consectetur adipiscing elit. Maecenas a libero porta, dictum ipsum eget, hendrerit neque. Morbi blandit, ex ut suscipit viverra, enim velit tincidunt tellus, a tempor velit nunc et ex. Proin hendrerit odio nec convallis lobortis. Aenean in purus dolor. Vestibulum orci orci, laoreet eget magna in, commodo euismod justo.\", \n",
    "        metadata={\"id\": 83209, \"date\": \"2022-11-13T18:26:45.000000Z\"}\n",
    "    ),\n",
    "    Document(\n",
    "        page_content=\"Integer at finibus odio. Nam sit amet enim cursus lacus gravida feugiat vestibulum sed libero. Aenean eleifend est quis elementum tincidunt. Curabitur sit amet ornare erat. Nulla id dolor ut magna volutpat sodales fringilla vel ipsum. Donec ultricies, lacus sed fermentum dignissim, lorem elit aliquam ligula, sed suscipit sapien purus nec ligula.\", \n",
    "        metadata={\"id\": 89313, \"date\": \"2022-11-13T18:28:53.000000Z\"}\n",
    "    ),\n",
    "    Document(\n",
    "        page_content=\"Morbi tortor enim, commodo id efficitur vitae, fringilla nec mi. Nullam molestie faucibus aliquet. Praesent a est facilisis, condimentum justo sit amet, viverra erat. Fusce volutpat nisi vel purus blandit, et facilisis felis accumsan. Phasellus luctus ligula ultrices tellus tempor hendrerit. Donec at ultricies leo.\", \n",
    "        metadata={\"id\": 87732, \"date\": \"2022-11-13T18:49:04.000000Z\"}\n",
    "    )\n",
    "]\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using multiple columns as content\n",
    "\n",
    "You can choose to use multiple columns as content:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.document_loaders import RocksetLoader\n",
    "from rockset import RocksetClient, Regions, models\n",
    "\n",
    "loader = RocksetLoader(\n",
    "    RocksetClient(Regions.usw2a1, \"<api key>\"),\n",
    "    models.QueryRequestSql(query=\"SELECT * FROM langchain_demo LIMIT 1 WHERE id=38\"),\n",
    "    [\"sentence1\", \"sentence2\"],  # TWO content columns\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Assuming the \"sentence1\" field is `\"This is the first sentence.\"` and the \"sentence2\" field is `\"This is the second sentence.\"`, the `page_content` of the resulting `Document` would be:\n",
    "\n",
    "```\n",
    "This is the first sentence.\n",
    "This is the second sentence.\n",
    "```\n",
    "\n",
    "You can define you own function to join content columns by setting the `content_columns_joiner` argument in the `RocksetLoader` constructor. `content_columns_joiner` is a method that takes in a `List[Tuple[str, Any]]]` as an argument, representing a list of tuples of (column name, column value). By default, this is a method that joins each column value with a new line.\n",
    "\n",
    "For example, if you wanted to join sentence1 and sentence2 with a space instead of a new line, you could set `content_columns_joiner` like so:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "RocksetLoader(\n",
    "    RocksetClient(Regions.usw2a1, \"<api key>\"),\n",
    "    models.QueryRequestSql(query=\"SELECT * FROM langchain_demo LIMIT 1 WHERE id=38\"),\n",
    "    [\"sentence1\", \"sentence2\"],\n",
    "    content_columns_joiner=lambda docs: \" \".join(\n",
    "        [doc[1] for doc in docs]\n",
    "    ),  # join with space instead of /n\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `page_content` of the resulting `Document` would be:\n",
    "\n",
    "```\n",
    "This is the first sentence. This is the second sentence.\n",
    "```\n",
    "\n",
    "Oftentimes you want to include the column name in the `page_content`. You can do that like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "RocksetLoader(\n",
    "    RocksetClient(Regions.usw2a1, \"<api key>\"),\n",
    "    models.QueryRequestSql(query=\"SELECT * FROM langchain_demo LIMIT 1 WHERE id=38\"),\n",
    "    [\"sentence1\", \"sentence2\"],\n",
    "    content_columns_joiner=lambda docs: \"\\n\".join(\n",
    "        [f\"{doc[0]}: {doc[1]}\" for doc in docs]\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This would result in the following `page_content`:\n",
    "\n",
    "```\n",
    "sentence1: This is the first sentence.\n",
    "sentence2: This is the second sentence.\n",
    "```"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11.4"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}