run-llama--llama_index/docs/examples/managed/PostgresMLDemo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "adf7d63d",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/managed/vectaraDemo.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db0855d0",
   "metadata": {},
   "source": [
    "# PostgresML Managed Index\n",
    "In this notebook we are going to show how to use [PostgresML](https://postgresml.org) with LlamaIndex."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cfe2497c",
   "metadata": {},
   "source": [
    "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74158792",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index-indices-managed-postgresml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6019e01a",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "359797ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.indices.managed.postgresml import PostgresMLIndex\n",
    "\n",
    "from llama_index.core import SimpleDirectoryReader\n",
    "\n",
    "# Need this as asyncio can get pretty wild with notebooks and this prevents event loop errors\n",
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7010b1d-d1bb-4f08-9309-a328bb4ea396",
   "metadata": {},
   "source": [
    "### Loading documents\n",
    "Load the `paul_graham_essay.txt` document."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "wsPWym0jC-1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir data\n",
    "!curl -o data/paul_graham_essay.txt https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c154dd4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "documents = SimpleDirectoryReader(\"data\").load_data()\n",
    "print(f\"documents loaded into {len(documents)} document objects\")\n",
    "print(f\"Document ID of first doc is {documents[0].doc_id}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0232fd1",
   "metadata": {},
   "source": [
    "### Upsert the documents into your PostgresML database"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "EDQLEUHiKvRk",
   "metadata": {},
   "source": [
    "First let's set the url to our PostgresML database. If you don't have a url yet, you can make one for free here: https://postgresml.org/signup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "V1lCc5VlEhan",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's set some secrets we need\n",
    "from google.colab import userdata\n",
    "\n",
    "PGML_DATABASE_URL = userdata.get(\"PGML_DATABASE_URL\")\n",
    "\n",
    "# If you don't have those secrets set, uncomment the lines below and run them instead\n",
    "# Make sure to replace {REPLACE_ME} with your keys\n",
    "# PGML_DATABASE_URL = \"{REPLACE_ME}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8dfdb0a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "index = PostgresMLIndex.from_documents(\n",
    "    documents,\n",
    "    collection_name=\"llama-index-example-demo\",\n",
    "    pgml_database_url=PGML_DATABASE_URL,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ee4473a-094f-4d0a-a825-e1213db07240",
   "metadata": {},
   "source": [
    "### Query the Postgresml Index\n",
    "We can now ask questions using the PostgresMLIndex retriever."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb174ec3",
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"What did the author write about?\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52878fd2",
   "metadata": {},
   "source": [
    "We can use a retriever to list search our documents:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a2bcc07",
   "metadata": {},
   "outputs": [],
   "source": [
    "retriever = index.as_retriever()\n",
    "response = retriever.retrieve(query)\n",
    "texts = [t.node.text for t in response]\n",
    "\n",
    "print(\"The Nodes:\")\n",
    "print(response)\n",
    "print(\"\\nThe Texts\")\n",
    "print(texts)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "DtmyAoURXLBO",
   "metadata": {},
   "source": [
    "PostgresML allows for easy re-reranking in the same query as doing retrieval:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "xn1wj-9OXK1g",
   "metadata": {},
   "outputs": [],
   "source": [
    "retriever = index.as_retriever(\n",
    "    limit=2,  # Limit to returning the 2 most related Nodes\n",
    "    rerank={\n",
    "        \"model\": \"mixedbread-ai/mxbai-rerank-base-v1\",  # Use the mxbai-rerank-base model for reranking\n",
    "        \"num_documents_to_rerank\": 100,  # Rerank up to 100 results returned from the vector search\n",
    "    },\n",
    ")\n",
    "response = retriever.retrieve(query)\n",
    "texts = [t.node.text for t in response]\n",
    "\n",
    "print(\"The Nodes:\")\n",
    "print(response)\n",
    "print(\"\\nThe Texts\")\n",
    "print(texts)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9fbf0c93",
   "metadata": {},
   "source": [
    "with the as_query_engine(), we can ask questions and get the response in one query:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "890f7133",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_engine = index.as_query_engine()\n",
    "response = query_engine.query(query)\n",
    "\n",
    "print(\"The Response:\")\n",
    "print(response)\n",
    "print(\"\\nThe Source Nodes:\")\n",
    "print(response.get_formatted_sources())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc0874ad",
   "metadata": {},
   "source": [
    "Note that the \"response\" object above includes both the summary text but also the source documents used to provide this response (citations). Notice the source nodes are all from the same document. That is because we only uploaded one document which PostgresML automatically split before embedding for us. All parameters can be controlled. See the documentation for more information."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "-JAlUGSNMA20",
   "metadata": {},
   "source": [
    "We can enable streaming by passing `streaming=True` when we create our query_engine.\n",
    "\n",
    "**NOTE: Streaming is painfully slow on google collab due to their internet connectivity.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "Ui9wM_w_L1Z9",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_engine = index.as_query_engine(streaming=True)\n",
    "results = query_engine.query(query)\n",
    "for text in results.response_gen:\n",
    "    print(text, end=\"\", flush=True)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}