{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LangChain and pgvector: Up and Running\n",
    "[LangChain](https://langchain.com) is the most popular framework for building LLM applications and agents. This notebook is an introduction to building LLM applications with the LangChain framework, using PostgreSQL and pgvector as a vector database for embeddings data.\n",
    "\n",
    "We'll use the example of creating a chatbot to answer questions about the blog posts from the Timescale blog to illustrate the following concepts:\n",
    "- How to prepare your documents for insertion into PostgreSQL and pgvector using LangChain document transformer TextSplitter\n",
    "- How to create embeddings from your data using the OpenAI embeddings model and insert them into PostgreSQL and pgvector.\n",
    "- How to use embeddings retrieved from a vector database to augment LLM generation. \n",
    "\n",
    "This is a great first step for more advanced LangChain projects in Python -- for example, creating a chatbot for your company documentation or an application to answer questions from uploaded PDFs.\n",
    "\n",
    "Let's get started!"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup and Configuration\n",
    "- Signup for an OpenAI Developer Account and create an API Key. See [OpenAI's developer platform](https://platform.openai.com/overview).\n",
    "- Install Python\n",
    "- Install and configure a python virtual environment. We recommend [Pyenv](https://github.com/pyenv/pyenv)\n",
    "- Install the requirements for this notebook using the following command:\n",
    "\n",
    "```\n",
    "pip install -r requirements.txt\n",
    "```\n",
    "\n",
    "Or if you already have langchain installed, run: \n",
    "```pip install --upgrade langchain```\n",
    "\n",
    "Protip: Add your OpenAI API Key to the provided shell script (named ```setenv.sh```) and run it to set your environment variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "# Run export OPENAI_API_KEY=sk-YOUR_OPENAI_API_KEY...\n",
    "# Get openAI api key by reading local .env file\n",
    "from dotenv import load_dotenv, find_dotenv\n",
    "_ = load_dotenv(find_dotenv())\n",
    "OPENAI_API_KEY  = os.environ['OPENAI_API_KEY']"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we need a way for LangChain to interact with PostgreSQL and pgvector. This is acheived by importing the PGVector class from the langchain.vectorstores package as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.vectorstores.pgvector import PGVector"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we'll construct our connection string for LangChain to connect to our PostgreSQL database.\n",
    "\n",
    " Because LangChain uses SQLAlchemy to connect to SQL databases like PostgreSQL, we need to create our connection string programmatically, reading each of the components of the string (host, database name, password, port etc) from our environment variables.\n",
    "\n",
    "In this example, we'll use a PostgreSQL database with pgvector installed that's hosted on Timescale. You can create your own cloud PostgreSQL database in minutes [at this link](https://console.cloud.timescale.com/signup) to follow along. If you're using a Timescale database, you can find all this information in the \"Cheat Sheet\" file you download when you first create your new database service.\n",
    "\n",
    "Alternatively, you can also use a local PostgreSQL database if you prefer.\n",
    "\n",
    "Protip: Add your database credentials to the provided shell script (named ```setenv.sh```)  and run it to set your environment variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build the PGVector Connection String from params\n",
    "\n",
    "# Found in the credential cheat-sheet or \"Connection Info\" in the Timescale console\n",
    "# In terminal, run: export VAR_NAME=value for each of the values below\n",
    "host= os.environ['TIMESCALE_HOST']\n",
    "port= os.environ['TIMESCALE_PORT']\n",
    "user= os.environ['TIMESCALE_USER']\n",
    "password= os.environ['TIMESCALE_PASSWORD']\n",
    "dbname= os.environ['TIMESCALE_DBNAME']\n",
    "\n",
    "# We use postgresql rather than postgres in the conn string since LangChain uses sqlalchemy under the hood\n",
    "# You can remove the ?sslmode=require if you have a local PostgreSQL instance running without SSL\n",
    "CONNECTION_STRING = f\"postgresql+psycopg2://{user}:{password}@{host}:{port}/{dbname}?sslmode=require\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 1: Split a CSV file into smaller chunks while preserving associated metadata\n",
    "In this section, we will parse our CSV file into smaller chunks for similarity search and retrieval, with help from LangChains TokenTextSplitter.\n",
    "\n",
    "First let's take a look at the CSV file we'll be working with"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "df = pd.read_csv('blog_posts_data.csv')\n",
    "df.head()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As shown above, this is a CSV file of blog posts about Timescale use cases. \n",
    "\n",
    "Ordinarily, we would use the langchain [CSVLoader](https://python.langchain.com/docs/modules/data_connection/document_loaders/how_to/csv) to load the contents of a CSV file, but in this case, we need to pre-process the content column of our CSV to be able to create embeddings for each blog post within the token limits of the OpenAI embeddings API.\n",
    "\n",
    "We also need a way to split the text of content column of the CSV while retaining the associated metadata with that text (i.e the blog title and URL)."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "LangChain has a number of built-in document transformers that make it easy to split, combine, filter, and otherwise manipulate documents.\n",
    "\n",
    "We'll use LangChain's [Token Text Splitter](https://python.langchain.com/docs/modules/data_connection/document_transformers/text_splitters/split_by_token) to help us split up the content column of our CSV into chunks of a specified token amount. You an alternatively use the [Recursive Character Text Splitter](https://python.langchain.com/docs/modules/data_connection/document_transformers/text_splitters/character_text_splitter), if you'd rather split text by number of characters rather than tokens.\n",
    "\n",
    "We will split the text into chunks of around 512 tokens, with a 20% or 103 token overlap."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tiktoken\n",
    "from langchain.text_splitter import TokenTextSplitter\n",
    "\n",
    "# We need to split the text into chunks of 512 tokens, with 20% token overlap\n",
    "text_splitter = TokenTextSplitter(chunk_size=512,chunk_overlap=103)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Helper func: calculate number of tokens\n",
    "def num_tokens_from_string(string: str, encoding_name = \"cl100k_base\") -> int:\n",
    "    if not string:\n",
    "        return 0\n",
    "    # Returns the number of tokens in a text string\n",
    "    encoding = tiktoken.get_encoding(encoding_name)\n",
    "    num_tokens = len(encoding.encode(string))\n",
    "    return num_tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#list for smaller chunked text and metadata\n",
    "new_list = []\n",
    "# Create a new list by splitting up text into token sizes of around 512 tokens\n",
    "for i in range(len(df.index)):\n",
    "    text = df['content'][i]\n",
    "    token_len = num_tokens_from_string(text)\n",
    "    if token_len <= 512:\n",
    "        new_list.append([df['title'][i], df['content'][i], df['url'][i]])\n",
    "    else:\n",
    "        #split text into 512 token chunks using text splitter\n",
    "        split_text = text_splitter.split_text(text)\n",
    "        for j in range(len(split_text)):\n",
    "            new_list.append([df['title'][i], split_text[j], df['url'][i]])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's take a look at how the content looks after being split:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_new = pd.DataFrame(new_list, columns=['title', 'content', 'url'])\n",
    "df_new.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Quick check on how many items in our new list\n",
    "print(len(new_list))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Optional: save to new csv for easy reloading\n",
    "df_new.to_csv('blog_posts_data_chunked.csv', index=False)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 2: Insert embeddings into PostgreSQL and pgvector"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have our original CSV split up into smaller chunks and the associated metadata preserved, we will use the LangChain [Pandas Data Frame Loader](https://python.langchain.com/docs/modules/data_connection/document_loaders/integrations/pandas_dataframe) to load data from our new pandas data frame and insert it into our PostgreSQL database with pgvector installed.\n",
    "\n",
    "Note that we must specify which column in the Data Frame contains the text that we'll create embeddings for."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#load documents from Pandas dataframe for insertion into database\n",
    "from langchain.document_loaders import DataFrameLoader\n",
    "\n",
    "# page_content_column is the column name in the dataframe that contains the we'll create embeddings for\n",
    "loader = DataFrameLoader(df_new, page_content_column = 'content')\n",
    "docs = loader.load()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll use the OpenAI embeddings model for our documents, so let's import the OpenAIEmbeddings module from the langchain.embeddings package and create an instance of it. \n",
    "\n",
    "This instance can be used to generate embeddings for text data using the OpenAI API.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.embeddings import OpenAIEmbeddings\n",
    "embeddings = OpenAIEmbeddings()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before we create embeddings for all the data in our DataFrame, let's briefly overview how creating an embedding works. \n",
    "\n",
    "Here's how we create an embedding for a string:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create OpenAI embedding using LangChain's OpenAIEmbeddings class\n",
    "query_string = \"PostgreSQL is my favorite database\"\n",
    "embed = embeddings.embed_query(query_string)\n",
    "print(len(embed)) # Should be 1536, the dimensinality of the OpenAI model's embeddings\n",
    "print(embed[:5]) # Should be a list of floats"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now for the main event, we'll connect to our PostgreSQL database and store the documents we loaded along with their embeddings.\n",
    "\n",
    "Thanks to LangChain, creating the embeddings and storing the data in our PostgreSQL database is a one command operation!\n",
    "\n",
    "We pass in the following arguments:\n",
    "- ```documents```: The documents we loaded from the Pandas Data Frame.\n",
    "- ```embedding```: Our instance of the OpenAI embeddings class, which is the model we'll use the create the embeddings. \n",
    "-  ```collection_name```: The name of the table we want our embeddings and metadata to live in\n",
    "- ```distance_strategy```: The distance strategy we wan to use to calculate the distance between vectors, in our case we'll use cosine distance\n",
    "- ```connection_string```: The connection string to our PostgreSQL database which we constructed in the setup section"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a PGVector instance to house the documents and embeddings\n",
    "from langchain.vectorstores.pgvector import DistanceStrategy\n",
    "db = PGVector.from_documents(\n",
    "    documents= docs,\n",
    "    embedding = embeddings,\n",
    "    collection_name= \"blog_posts\",\n",
    "    distance_strategy = DistanceStrategy.COSINE,\n",
    "    connection_string=CONNECTION_STRING)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that our data is in the database, let's perform a similarity search to fetch the documents most similar to a query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.schema import Document\n",
    "\n",
    "# Query for which we want to find semantically similar documents\n",
    "query = \"Tell me about how Edeva uses Timescale?\"\n",
    "\n",
    "#Fetch the k=3 most similar documents\n",
    "docs =  db.similarity_search(query, k=3)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The query on our database returns a list of LangChain Documents, let's learn how to interact with those documents below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "type(docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "type(docs[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Interact with a document returned from the similarity search on pgvector\n",
    "doc = docs[0]\n",
    "\n",
    "# Access the document's content\n",
    "doc_content = doc.page_content\n",
    "# Access the document's metadata object\n",
    "doc_metadata = doc.metadata\n",
    "\n",
    "print(\"Content snippet:\" + doc_content[:500])\n",
    "print(\"Document title: \" + doc_metadata['title'])\n",
    "print(\"Document url: \" + doc_metadata['url'])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 3: Question Answering with Retrieval Augmented Generation\n",
    "Next let's tie everything we've learned together and build a simple example of using LangChain for questions answering using an LLM from OpenAI and the most relevant documents the question from our database. \n",
    "\n",
    "This technique is called Retrieval Augmented Generation and works as follows:\n",
    "- Create an embedding vector for the user question.\n",
    "- Use pgvector to perform a vector similarity search and retrieve the k nearest neighbors to the question embedding from our database of embedding vectors representing the blog content. In our example, we’ll use k=3, finding the three most similar embedding vectors and associated content.\n",
    "- Supply the content retrieved from the database as additional context to the model and ask it to perform a completion task to answer the user question."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To more easily retrieve documents from our PostgreSQL vector database, we'll use a LangChain [retriever](https://python.langchain.com/docs/modules/data_connection/retrievers/). \n",
    "\n",
    "In LangChain, a retriever is an interface that returns documents given an unstructured query. A retriever's main purpose is only to return (or retrieve) documents. \n",
    "\n",
    "We will use a [vector store-backed retriever](https://python.langchain.com/docs/modules/data_connection/retrievers/how_to/vectorstore) which is a retriever that uses a vector store to retrieve documents. It is a lightweight wrapper around the Vector Store class to make it conform to the Retriever interface. It uses the search methods implemented by a vector store, like similarity search and MMR, to query the texts in the vector store."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create retriever from database\n",
    "# We specify the number of results we want to retrieve (k=3)\n",
    "retriever = db.as_retriever(\n",
    "    search_kwargs={\"k\": 3}\n",
    "    )"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we'll import the LLM we want to use to generate a response to our question. In this case we'll use [OpenAI's GPT-3.5 model](https://platform.openai.com/docs/models/) with a 16k token context window, so that we won't have any trouble fitting in retrieved documents as context in addition to the user question."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chat_models import ChatOpenAI\n",
    "llm = ChatOpenAI(temperature = 0.0, model = 'gpt-3.5-turbo-16k')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next up, we'll use one of the most useful chain's in LangChain, the [Retrieval Q+A chain](https://python.langchain.com/docs/modules/chains/popular/vector_db_qa), which is used for question answering over an a vector database (vector store or index as its also known.)\n",
    "\n",
    "\n",
    "We'll combine it with a [stuff chain](https://python.langchain.com/docs/modules/chains/document/stuff) which takes a list of documents, inserts them all into a prompt (_stuffs_ them in) and passes that prompt to an LLM.\n",
    "\n",
    "And for the final ingredient, let's formulate a question we want to the model to answer with the help from the documents in our database and pass it to our chain to process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains import RetrievalQA\n",
    "qa_stuff = RetrievalQA.from_chain_type(\n",
    "    llm=llm, \n",
    "    chain_type=\"stuff\", \n",
    "    retriever=retriever,\n",
    "    verbose=True,\n",
    ")\n",
    "\n",
    "query =  \"How does Edeva use continuous aggregates?\"\n",
    "response = qa_stuff.run(query)\n",
    "\n",
    "from IPython.display import Markdown, display\n",
    "display(Markdown(response))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bonus: Cite your sources with LangChain and pgvector for RAG"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For even more advanced functionality, you might want your answer to include the sources used to give users peace of mind. Here's how you can do that with the RetrievalQA chain using the ```return_source_documents``` argument:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# New chain to return context and sources\n",
    "qa_stuff_with_sources = RetrievalQA.from_chain_type(\n",
    "    llm=llm, \n",
    "    chain_type=\"stuff\", \n",
    "    retriever=retriever,\n",
    "    return_source_documents=True,\n",
    "    verbose=True,\n",
    ")\n",
    "\n",
    "query =  \"How does Edeva use continuous aggregates?\"\n",
    "\n",
    "# To run the query, we use a different syntax since we're returning more than just the response text\n",
    "responses = qa_stuff_with_sources({\"query\": query})"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's inspect the source documents that got returned:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "responses[\"source_documents\"]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And finally, let's print the out the result with the source document cited:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "source_documents = responses[\"source_documents\"]\n",
    "source_content = [doc.page_content for doc in source_documents]\n",
    "source_metadata = [doc.metadata for doc in source_documents]\n",
    "\n",
    "# Construct a single string with the LLM output and the source titles and urls\n",
    "def construct_result_with_sources():\n",
    "    result = responses['result']\n",
    "    result += \"\\n\\n\"\n",
    "    result += \"Sources used:\"\n",
    "    for i in range(len(source_content)):\n",
    "        result += \"\\n\\n\"\n",
    "        result += source_metadata[i]['title']\n",
    "        result += \"\\n\\n\"\n",
    "        result += source_metadata[i]['url']\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(Markdown(construct_result_with_sources()))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The cite your sources functionality is useful because it can help explain unexpected responses from the model due to irrelevant but highly similar documents being retrieved from the database."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Next steps:\n",
    "- Check out [Conversational Retrieval QA Chain](https://python.langchain.com/docs/modules/chains/popular/chat_vector_db) for how to added memory and use what you learned above in a chatbot conversation setting\n",
    "- Use [Chainlit](https://github.com/Chainlit/chainlit) to build your own LLM chatbot in python (Timescale tutorial coming soon!)\n",
    "- Learn about how pgvector finds approximate nearest neighbors in this blog post: [What Are ivfflat Indexes in pgvector and How Do They Work](https://www.timescale.com/blog/nearest-neighbor-indexes-what-are-ivfflat-indexes-in-pgvector-and-how-do-they-work/)"
   ]
  }
 ],
 "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.8.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
