mirror of
https://github.com/hwchase17/langchain.git
synced 2025-09-24 12:01:54 +00:00
docs: add YugabyteDB Distributed SQL database (#32571)
- **Description:** The `langchain-yugabytedb` package implementations of core LangChain abstractions using `YugabyteDB` Distributed SQL Database. YugabyteDB is a cloud-native distributed PostgreSQL-compatible database that combines strong consistency with ultra-resilience, seamless scalability, geo-distribution, and highly flexible data locality to deliver business-critical, transactional applications. [YugabyteDB](https://www.yugabyte.com/ai/) combines the power of the `pgvector` PostgreSQL extension with an inherently distributed architecture. This future-proofed foundation helps you build GenAI applications using RAG retrieval that demands high-performance vector search. - [ ] **tests and docs**: 1. `langchain-yugabytedb` [github](https://github.com/yugabyte/langchain-yugabytedb) repo. 2. YugabyteDB VectorStore example notebook showing its use. It lives in `langchain/docs/docs/integrations/vectorstores/yugabytedb.ipynb` directory. 3. Running `langchain-yugabytedb` unit tests - Setting up a Development Environment This document details how to set up a local development environment that will allow you to contribute changes to the project. Acquire sources and create virtualenv. ```shell git clone https://github.com/yugabyte/langchain-yugabytedb cd langchain-yugabytedb uv venv --python=3.13 source .venv/bin/activate ``` Install package in editable mode. ```shell uv pip install pipx pipx install poetry poetry install uv pip install pytest pytest_asyncio pytest-timeout langchain-core langchain_tests sqlalchemy psycopg psycopg-binary numpy pgvector ``` Start YugabyteDB RF-1 Universe. ```shell docker run -d --name yugabyte_node01 --hostname yugabyte01 \ -p 7000:7000 -p 9000:9000 -p 15433:15433 -p 5433:5433 -p 9042:9042 \ yugabytedb/yugabyte:2.25.2.0-b359 bin/yugabyted start --background=false \ --master_flags="allowed_preview_flags_csv=ysql_yb_enable_advisory_locks,ysql_yb_enable_advisory_locks=true" \ --tserver_flags="allowed_preview_flags_csv=ysql_yb_enable_advisory_locks,ysql_yb_enable_advisory_locks=true" docker exec -it yugabyte_node01 bin/ysqlsh -h yugabyte01 -c "CREATE extension vector;" ``` Invoke test cases. ```shell pytest -vvv tests/unit_tests/yugabytedb_tests ```
This commit is contained in:
committed by
GitHub
parent
03f0ebd93e
commit
e6b5ff213a
664
docs/docs/integrations/vectorstores/yugabytedb.ipynb
Normal file
664
docs/docs/integrations/vectorstores/yugabytedb.ipynb
Normal file
@@ -0,0 +1,664 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "raw",
|
||||
"id": "1957f5cb",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"---\n",
|
||||
"sidebar_label: YugabyteDB\n",
|
||||
"---"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "ef1f0986",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# YugabyteDBVectorStore\n",
|
||||
"\n",
|
||||
"This notebook covers how to get started with the YugabyteDB vector store in langchain, using the `langchain-yugabytedb` package.\n",
|
||||
"\n",
|
||||
"YugabyteDB is a cloud-native distributed PostgreSQL-compatible database that combines strong consistency with ultra-resilience, seamless scalability, geo-distribution, and highly flexible data locality to deliver business-critical, transactional applications.\n",
|
||||
"\n",
|
||||
"[YugabyteDB](https://www.yugabyte.com/ai/) combines the power of the `pgvector` PostgreSQL extension with an inherently distributed architecture. This future-proofed foundation helps you build GenAI applications using RAG retrieval that demands high-performance vector search.\n",
|
||||
"\n",
|
||||
"YugabyteDB’s unique approach to vector indexing addresses the limitations of single-node PostgreSQL systems when dealing with large-scale vector datasets.\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"## Setup\n",
|
||||
"\n",
|
||||
"### Minimum Version\n",
|
||||
"`langchain-yugabytedb` module requires YugabyteDB `v2025.1.0.0` or higher.\n",
|
||||
"\n",
|
||||
"### Connecting to YugabyteDB database\n",
|
||||
"\n",
|
||||
"In order to get started with `YugabyteDBVectorStore`, lets start a local YugabyteDB node for development purposes - \n",
|
||||
"\n",
|
||||
"### Start YugabyteDB RF-1 Universe."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "5a8147d9",
|
||||
"metadata": {
|
||||
"vscode": {
|
||||
"languageId": "plaintext"
|
||||
}
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"\n",
|
||||
"docker run -d --name yugabyte_node01 --hostname yugabyte01 \\\n",
|
||||
" -p 7000:7000 -p 9000:9000 -p 15433:15433 -p 5433:5433 -p 9042:9042 \\\n",
|
||||
" yugabytedb/yugabyte:2.25.2.0-b359 bin/yugabyted start --background=false \\\n",
|
||||
" --master_flags=\"allowed_preview_flags_csv=ysql_yb_enable_advisory_locks,ysql_yb_enable_advisory_locks=true\" \\\n",
|
||||
" --tserver_flags=\"allowed_preview_flags_csv=ysql_yb_enable_advisory_locks,ysql_yb_enable_advisory_locks=true\"\n",
|
||||
"\n",
|
||||
"docker exec -it yugabyte_node01 bin/ysqlsh -h yugabyte01 -c \"CREATE extension vector;\""
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "541e4507",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"For production deployment, performance benchmarking, or deploying a true multi-node on multi-host setup, see Deploy [YugabyteDB](https://docs.yugabyte.com/stable/deploy/)."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "36fdc060",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Installation"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "432f461c",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%pip install --upgrade --quiet langchain\n",
|
||||
"%pip install --upgrade --quiet langchain-openai langchain-community tiktoken\n",
|
||||
"%pip install --upgrade --quiet psycopg-binary"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "64e28aa6",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%pip install -qU \"langchain-yugabytedb\""
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "3f9951f4",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Set your YugabyteDB Values\n",
|
||||
"\n",
|
||||
"YugabyteDB clients connect to the cluster using a PostgreSQL compliant connection string. YugabyteDB connection parameters are provided below."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "b4715d61",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"YUGABYTEDB_USER = \"yugabyte\" # @param {type: \"string\"}\n",
|
||||
"YUGABYTEDB_PASSWORD = \"\" # @param {type: \"string\"}\n",
|
||||
"YUGABYTEDB_HOST = \"localhost\" # @param {type: \"string\"}\n",
|
||||
"YUGABYTEDB_PORT = \"5433\" # @param {type: \"string\"}\n",
|
||||
"YUGABYTEDB_DB = \"yugabyte\" # @param {type: \"string\"}"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "93df377e",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Initialization\n",
|
||||
"\n",
|
||||
"### Environment Setup\n",
|
||||
"\n",
|
||||
"This notebook uses the OpenAI API through `OpenAIEmbeddings`. We suggest obtaining an OpenAI API key and export it as an environment variable with the name `OPENAI_API_KEY`.\n",
|
||||
"\n",
|
||||
"### Connecting to YugabyteDB Universe"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "dc37144c-208d-4ab3-9f3a-0407a69fe052",
|
||||
"metadata": {
|
||||
"tags": []
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from langchain_yugabytedb import YBEngine, YugabyteDBVectorStore\n",
|
||||
"from langchain_openai import OpenAIEmbeddings\n",
|
||||
"\n",
|
||||
"TABLE_NAME = \"my_doc_collection\"\n",
|
||||
"VECTOR_SIZE = 1536\n",
|
||||
"\n",
|
||||
"CONNECTION_STRING = (\n",
|
||||
" f\"postgresql+asyncpg://{YUGABYTEDB_USER}:{YUGABYTEDB_PASSWORD}@{YUGABYTEDB_HOST}\"\n",
|
||||
" f\":{YUGABYTEDB_PORT}/{YUGABYTEDB_DB}\"\n",
|
||||
")\n",
|
||||
"engine = YBEngine.from_connection_string(url=CONNECTION_STRING)\n",
|
||||
"\n",
|
||||
"embeddings = OpenAIEmbeddings()\n",
|
||||
"engine.init_vectorstore_table(\n",
|
||||
" table_name=TABLE_NAME,\n",
|
||||
" vector_size=VECTOR_SIZE,\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"yugabyteDBVectorStore = YugabyteDBVectorStore.create_sync(\n",
|
||||
" engine=engine,\n",
|
||||
" table_name=TABLE_NAME,\n",
|
||||
" embedding_service=embeddings,\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "ac6071d4",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Manage vector store\n",
|
||||
"\n",
|
||||
"### Add items to vector store"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "17f5efc0",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from langchain_core.documents import Document\n",
|
||||
"\n",
|
||||
"docs = [\n",
|
||||
" Document(page_content=\"Apples and oranges\"),\n",
|
||||
" Document(page_content=\"Cars and airplanes\"),\n",
|
||||
" Document(page_content=\"Train\"),\n",
|
||||
"]\n",
|
||||
"\n",
|
||||
"yugabyteDBVectorStore.add_documents(docs)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "7b92b5f6",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"['b40e7f47-3a4e-4b88-b6e2-cb3465dde6bd', '275823d2-1a47-440d-904b-c07b132fd72b', 'f0c5a9bc-1456-40fe-906b-4e808d601470']"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "dcf1b905",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Delete items from vector store\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "ef61e188",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"yugabyteDBVectorStore.delete(ids=[\"275823d2-1a47-440d-904b-c07b132fd72b\"])"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "f8f751e1",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Update items from vector store\n",
|
||||
"\n",
|
||||
"Note: Update operation is not supported by YugabyteDBVectorStore."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "c3620501",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Query vector store\n",
|
||||
"\n",
|
||||
"Once your vector store has been created and the relevant documents have been added you will most likely wish to query it during the running of your chain or agent. \n",
|
||||
"\n",
|
||||
"### Query directly\n",
|
||||
"\n",
|
||||
"Performing a simple similarity search can be done as follows:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "aa0a16fa",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"query = \"I'd like a fruit.\"\n",
|
||||
"docs = yugabyteDBVectorStore.similarity_search(query)\n",
|
||||
"print(docs)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "3ed9d733",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"If you want to execute a similarity search and receive the corresponding scores you can run:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "5efd2eaa",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"query = \"I'd like a fruit.\"\n",
|
||||
"docs = yugabyteDBVectorStore.similarity_search(query, k=1)\n",
|
||||
"print(docs)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "0c235cdc",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Query by turning into retriever\n",
|
||||
"\n",
|
||||
"You can also transform the vector store into a retriever for easier usage in your chains. \n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "f3460093",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"retriever = yugabyteDBVectorStore.as_retriever(search_kwargs={\"k\": 1})\n",
|
||||
"retriever.invoke(\"I'd like a fruit.\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "5e657ae6",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## ChatMessageHistory\n",
|
||||
"\n",
|
||||
"The chat message history abstraction helps to persist chat message history in a YugabyteDB table.\n",
|
||||
"\n",
|
||||
"`YugabyteDBChatMessageHistory` is parameterized using a table_name and a session_id.\n",
|
||||
"\n",
|
||||
"The table_name is the name of the table in the database where the chat messages will be stored.\n",
|
||||
"\n",
|
||||
"The session_id is a unique identifier for the chat session. It can be assigned by the caller using uuid.uuid4()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "0677c927",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import uuid\n",
|
||||
"\n",
|
||||
"from langchain_core.messages import SystemMessage, AIMessage, HumanMessage\n",
|
||||
"from langchain_yugabytedb import YugabyteDBChatMessageHistory\n",
|
||||
"import psycopg\n",
|
||||
"\n",
|
||||
"# Establish a synchronous connection to the database\n",
|
||||
"# (or use psycopg.AsyncConnection for async)\n",
|
||||
"conn_info = \"dbname=yugabyte user=yugabyte host=localhost port=5433\"\n",
|
||||
"sync_connection = psycopg.connect(conn_info)\n",
|
||||
"\n",
|
||||
"# Create the table schema (only needs to be done once)\n",
|
||||
"table_name = \"chat_history\"\n",
|
||||
"YugabyteDBChatMessageHistory.create_tables(sync_connection, table_name)\n",
|
||||
"\n",
|
||||
"session_id = str(uuid.uuid4())\n",
|
||||
"\n",
|
||||
"# Initialize the chat history manager\n",
|
||||
"chat_history = YugabyteDBChatMessageHistory(\n",
|
||||
" table_name, session_id, sync_connection=sync_connection\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"# Add messages to the chat history\n",
|
||||
"chat_history.add_messages(\n",
|
||||
" [\n",
|
||||
" SystemMessage(content=\"Meow\"),\n",
|
||||
" AIMessage(content=\"woof\"),\n",
|
||||
" HumanMessage(content=\"bark\"),\n",
|
||||
" ]\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"print(chat_history.messages)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "901c75dc",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Usage for retrieval-augmented generation\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"One of the primary advantages of the vector stores is to provide contextual data to the LLMs. LLMs often are trained with stale data and might not have the relevant domain specific knowledge which results in halucinations in LLMs responses. Take the following example - \n",
|
||||
"\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "d0f51eb3",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import getpass\n",
|
||||
"from langchain_openai import ChatOpenAI\n",
|
||||
"from langchain_core.messages import HumanMessage, SystemMessage, AIMessage\n",
|
||||
"\n",
|
||||
"my_api_key = getpass.getpass(\"Enter your API Key: \")\n",
|
||||
"\n",
|
||||
"llm = ChatOpenAI(model=\"gpt-3.5-turbo\", temperature=0.7, api_key=my_api_key)\n",
|
||||
"# Start with a system message to set the persona/behavior of the AI\n",
|
||||
"messages = [\n",
|
||||
" SystemMessage(\n",
|
||||
" content=\"You are a helpful and friendly assistant named 'YugaAI'. You love to answer questions about YugabyteDB and distributed sql.\"\n",
|
||||
" ),\n",
|
||||
" # First human turn\n",
|
||||
" HumanMessage(content=\"Hi YugaAI! Where's the headquarters of YugabyteDB?\"),\n",
|
||||
"]\n",
|
||||
"\n",
|
||||
"print(\"--- First Interaction ---\")\n",
|
||||
"print(f\"Human: {messages[1].content}\") # Print the human message\n",
|
||||
"response1 = llm.invoke(messages)\n",
|
||||
"print(f\"YugaAI: {response1.content}\")\n",
|
||||
"\n",
|
||||
"print(\"\\n--- Second Interaction ---\")\n",
|
||||
"print(f\"Human: {messages[2].content}\") # Print the new human message\n",
|
||||
"response2 = llm.invoke(messages) # Send the *entire* message history\n",
|
||||
"print(f\"YugaAI: {response2.content}\")\n",
|
||||
"\n",
|
||||
"# Add the second AI response to the history\n",
|
||||
"messages.append(AIMessage(content=response2.content))\n",
|
||||
"\n",
|
||||
"# --- 5. Another Turn with a different topic ---\n",
|
||||
"messages.append(\n",
|
||||
" HumanMessage(\n",
|
||||
" content=\"Can you tell me the current preview release version of YugabyteDB?\"\n",
|
||||
" )\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"print(\"\\n--- Third Interaction ---\")\n",
|
||||
"print(f\"Human: {messages[4].content}\") # Print the new human message\n",
|
||||
"response3 = llm.invoke(messages) # Send the *entire* message history\n",
|
||||
"print(f\"YugaAI: {response3.content}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "8500e27b",
|
||||
"metadata": {
|
||||
"vscode": {
|
||||
"languageId": "Log"
|
||||
}
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"--- First Interaction ---\n",
|
||||
"Human: Hi YugaAI! Where's the headquarters of YugabyteDB?\n",
|
||||
"YugaAI: Hello! YugabyteDB's headquarters is located in Sunnyvale, California, USA.\n",
|
||||
"\n",
|
||||
"--- Second Interaction ---\n",
|
||||
"Human: And what are YugabyteDB's supported APIs?\n",
|
||||
"YugaAI: YugabyteDB's headquarters is located in Sunnyvale, California, USA.\n",
|
||||
"\n",
|
||||
"YugabyteDB supports several APIs, including:\n",
|
||||
"1. YSQL (PostgreSQL-compatible SQL)\n",
|
||||
"2. YCQL (Cassandra-compatible query language)\n",
|
||||
"3. YEDIS (Redis-compatible key-value store)\n",
|
||||
"\n",
|
||||
"These APIs allow developers to interact with YugabyteDB using familiar interfaces and tools.\n",
|
||||
"\n",
|
||||
"--- Third Interaction ---\n",
|
||||
"Human: Can you tell me the current preview release version of YugabyteDB?\n",
|
||||
"YugaAI: The current preview release version of YugabyteDB is 2.11.0. This version includes new features, improvements, and bug fixes that are being tested by the community before the official stable release."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "7fe4301c",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"The current preview release of YugabyteDB is `v2.25.2.0`, however LLMs is providing stale information which is 2-3 years old. This is where the vector stores complement the LLMs by providing a way to store and retrive relevant information."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "d4e19220",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Construct a RAG for providing contextual information\n",
|
||||
"\n",
|
||||
"We will provide the relevant information to the LLMs by reading the YugabyteDB documentation. Let's first read the YugabyteDB docs and add data into YugabyteDB Vectorstore by loading, splitting and chuncking data from a html source. We will then store the vector embeddings generated by OpenAI embeddings into YugabyteDB Vectorstore.\n",
|
||||
"\n",
|
||||
"#### Generate Embeddings "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "05ec4ebc",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import getpass\n",
|
||||
"from langchain_community.document_loaders import WebBaseLoader\n",
|
||||
"from langchain_text_splitters import CharacterTextSplitter\n",
|
||||
"from langchain_yugabytedb import YBEngine, YugabyteDBVectorStore\n",
|
||||
"from langchain_openai import OpenAIEmbeddings\n",
|
||||
"\n",
|
||||
"my_api_key = getpass.getpass(\"Enter your API Key: \")\n",
|
||||
"url = \"https://docs.yugabyte.com/preview/releases/ybdb-releases/v2.25/\"\n",
|
||||
"\n",
|
||||
"loader = WebBaseLoader(url)\n",
|
||||
"\n",
|
||||
"documents = loader.load()\n",
|
||||
"\n",
|
||||
"print(f\"Number of documents loaded: {len(documents)}\")\n",
|
||||
"\n",
|
||||
"# For very large HTML files, you'll want to split the text into smaller\n",
|
||||
"# chunks before sending them to an LLM, as LLMs have token limits.\n",
|
||||
"for i, doc in enumerate(documents):\n",
|
||||
" text_splitter = CharacterTextSplitter(\n",
|
||||
" separator=\"\\n\\n\", # Split by double newline (common paragraph separator)\n",
|
||||
" chunk_size=1000, # Each chunk will aim for 1000 characters\n",
|
||||
" chunk_overlap=200, # Allow 200 characters overlap between chunks\n",
|
||||
" length_function=len,\n",
|
||||
" is_separator_regex=False,\n",
|
||||
" )\n",
|
||||
"\n",
|
||||
" # Apply the splitter to the loaded documents\n",
|
||||
" chunks = text_splitter.split_documents(documents)\n",
|
||||
"\n",
|
||||
" print(f\"\\n--- After Splitting ({len(chunks)} chunks) ---\")\n",
|
||||
"\n",
|
||||
" CONNECTION_STRING = \"postgresql+psycopg://yugabyte:@localhost:5433/yugabyte\"\n",
|
||||
" TABLE_NAME = \"yb_relnotes_chunks\"\n",
|
||||
" VECTOR_SIZE = 1536\n",
|
||||
" engine = YBEngine.from_connection_string(url=CONNECTION_STRING)\n",
|
||||
" engine.init_vectorstore_table(\n",
|
||||
" table_name=TABLE_NAME,\n",
|
||||
" vector_size=VECTOR_SIZE,\n",
|
||||
" )\n",
|
||||
" embeddings = OpenAIEmbeddings(api_key=my_api_key)\n",
|
||||
"\n",
|
||||
" # The PGVector.from_documents method handles:\n",
|
||||
" # 1. Creating the table if it doesn't exist (with 'embedding' column).\n",
|
||||
" # 2. Generating embeddings for each chunk using the provided embeddings model.\n",
|
||||
" # 3. Inserting the chunk text, metadata, and embeddings into the table.\n",
|
||||
" vectorstore = YugabyteDBVectorStore.from_documents(\n",
|
||||
" engine=engine, table_name=TABLE_NAME, documents=chunks, embedding=embeddings\n",
|
||||
" )\n",
|
||||
"\n",
|
||||
" print(f\"Successfully stored {len(chunks)} chunks in PostgreSQL table: {TABLE_NAME}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "e6483d89",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"#### Configure the YugabyteDB retriever"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "18a84445",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from langchain_core.prompts import ChatPromptTemplate\n",
|
||||
"from langchain_core.runnables import RunnablePassthrough\n",
|
||||
"from langchain_core.output_parsers import StrOutputParser\n",
|
||||
"from langchain_openai import ChatOpenAI\n",
|
||||
"\n",
|
||||
"retriever = vectorstore.as_retriever(search_kwargs={\"k\": 3})\n",
|
||||
"print(\n",
|
||||
" f\"Retriever created, set to retrieve top {retriever.search_kwargs['k']} documents.\"\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"# Initialize the Chat Model (e.g., OpenAI's GPT-3.5 Turbo)\n",
|
||||
"llm = ChatOpenAI(model=\"gpt-3.5-turbo\", temperature=0, api_key=my_api_key)\n",
|
||||
"\n",
|
||||
"# Define the RAG prompt template\n",
|
||||
"prompt = ChatPromptTemplate.from_messages(\n",
|
||||
" [\n",
|
||||
" (\n",
|
||||
" \"system\",\n",
|
||||
" \"You are a helpful and friendly assistant named 'YugaAI'. You love to answer questions about YugabyteDB and distributed sql.\",\n",
|
||||
" ),\n",
|
||||
" (\"human\", \"Context: {context}\\nQuestion: {question}\"),\n",
|
||||
" ]\n",
|
||||
")\n",
|
||||
"# Build the RAG chain\n",
|
||||
"# 1. Take the input question.\n",
|
||||
"# 2. Pass it to the retriever to get relevant documents.\n",
|
||||
"# 3. Format the documents into a string for the context.\n",
|
||||
"# 4. Pass the context and question to the prompt template.\n",
|
||||
"# 5. Send the prompt to the LLM.\n",
|
||||
"# 6. Parse the LLM's string output.\n",
|
||||
"rag_chain = (\n",
|
||||
" {\"context\": retriever, \"question\": RunnablePassthrough()}\n",
|
||||
" | prompt\n",
|
||||
" | llm\n",
|
||||
" | StrOutputParser()\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "04d12dc5",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Now, let's try asking the same question `Can you tell me the current preview release version of YugabyteDB?` again to the LLM"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "846e9963",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Invoke the RAG chain with a question\n",
|
||||
"rag_query = \"Can you tell me the current preview release version of YugabyteDB?\"\n",
|
||||
"print(f\"\\nQuerying RAG chain: '{rag_query}'\")\n",
|
||||
"rag_response = rag_chain.invoke(rag_query)\n",
|
||||
"print(\"\\n--- RAG Chain Response ---\")\n",
|
||||
"print(rag_response)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "c6fc3e7f",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Querying RAG chain: 'Can you tell me the current preview release version of YugabyteDB?'\n",
|
||||
"\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "efc2e0c7",
|
||||
"metadata": {
|
||||
"vscode": {
|
||||
"languageId": "log"
|
||||
}
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"--- RAG Chain Response ---\n",
|
||||
"The current preview release version of YugabyteDB is v2.25.2.0."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "af388f24",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## API reference\n",
|
||||
" \n",
|
||||
"For detailed information of all YugabyteDBVectorStore features and configurations head to the langchain-yugabytedb github repo: https://github.com/yugabyte/langchain-yugabytedb\""
|
||||
]
|
||||
}
|
||||
],
|
||||
"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.10.12"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
Reference in New Issue
Block a user