+ As of January 1, 2020 this library no longer supports Python 2 on the latest released version.
+ Library versions released prior to that date will continue to be available. For more information please
+ visit Python 2 support on Google Cloud.
+
+{%- else %}
+{{ super() }}
+{%- endif %}
+{%- endblock %}
diff --git a/docs/changelog.md b/docs/changelog.md
new file mode 120000
index 0000000..04c99a5
--- /dev/null
+++ b/docs/changelog.md
@@ -0,0 +1 @@
+../CHANGELOG.md
\ No newline at end of file
diff --git a/docs/conf.py b/docs/conf.py
new file mode 100644
index 0000000..8c2efd6
--- /dev/null
+++ b/docs/conf.py
@@ -0,0 +1,384 @@
+# -*- coding: utf-8 -*-
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# llama-index-cloud-sql-pg documentation build configuration file
+#
+# This file is execfile()d with the current directory set to its
+# containing dir.
+#
+# Note that not all possible configuration values are present in this
+# autogenerated file.
+#
+# All configuration values have a default; values that are commented out
+# serve to show the default.
+
+import os
+import shlex
+import sys
+
+# If extensions (or modules to document with autodoc) are in another directory,
+# add these directories to sys.path here. If the directory is relative to the
+# documentation root, use os.path.abspath to make it absolute, like shown here.
+sys.path.insert(0, os.path.abspath(".."))
+
+# For plugins that can not read conf.py.
+# See also: https://github.com/docascode/sphinx-docfx-yaml/issues/85
+sys.path.insert(0, os.path.abspath("."))
+
+__version__ = ""
+
+# -- General configuration ------------------------------------------------
+
+# If your documentation needs a minimal Sphinx version, state it here.
+needs_sphinx = "1.5.5"
+
+# Add any Sphinx extension module names here, as strings. They can be
+# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
+# ones.
+extensions = [
+ "sphinx.ext.autodoc",
+ "sphinx.ext.autosummary",
+ "sphinx.ext.intersphinx",
+ "sphinx.ext.coverage",
+ "sphinx.ext.doctest",
+ "sphinx.ext.napoleon",
+ "sphinx.ext.todo",
+ "sphinx.ext.viewcode",
+ "recommonmark",
+]
+
+# autodoc/autosummary flags
+autoclass_content = "both"
+autodoc_default_options = {"members": True}
+autosummary_generate = True
+
+
+# Add any paths that contain templates here, relative to this directory.
+templates_path = ["_templates"]
+
+# The suffix(es) of source filenames.
+# You can specify multiple suffix as a list of string:
+# source_suffix = ['.rst', '.md']
+source_suffix = [".rst", ".md"]
+
+# The encoding of source files.
+# source_encoding = 'utf-8-sig'
+
+# The root toctree document.
+root_doc = "index"
+
+# General information about the project.
+project = "llama-index-cloud-sql-pg"
+copyright = "2024, Google"
+author = "Google APIs"
+
+# The version info for the project you're documenting, acts as replacement for
+# |version| and |release|, also used in various other places throughout the
+# built documents.
+#
+# The full version, including alpha/beta/rc tags.
+release = __version__
+# The short X.Y version.
+version = ".".join(release.split(".")[0:2])
+
+# The language for content autogenerated by Sphinx. Refer to documentation
+# for a list of supported languages.
+#
+# This is also used if you do content translation via gettext catalogs.
+# Usually you set "language" from the command line for these cases.
+language = None
+
+# There are two options for replacing |today|: either, you set today to some
+# non-false value, then it is used:
+# today = ''
+# Else, today_fmt is used as the format for a strftime call.
+# today_fmt = '%B %d, %Y'
+
+# List of patterns, relative to source directory, that match files and
+# directories to ignore when looking for source files.
+exclude_patterns = [
+ "_build",
+ "**/.nox/**/*",
+ "samples/AUTHORING_GUIDE.md",
+ "samples/CONTRIBUTING.md",
+ "samples/snippets/README.rst",
+]
+
+# The reST default role (used for this markup: `text`) to use for all
+# documents.
+# default_role = None
+
+# If true, '()' will be appended to :func: etc. cross-reference text.
+# add_function_parentheses = True
+
+# If true, the current module name will be prepended to all description
+# unit titles (such as .. function::).
+# add_module_names = True
+
+# If true, sectionauthor and moduleauthor directives will be shown in the
+# output. They are ignored by default.
+# show_authors = False
+
+# The name of the Pygments (syntax highlighting) style to use.
+pygments_style = "sphinx"
+
+# A list of ignored prefixes for module index sorting.
+# modindex_common_prefix = []
+
+# If true, keep warnings as "system message" paragraphs in the built documents.
+# keep_warnings = False
+
+# If true, `todo` and `todoList` produce output, else they produce nothing.
+todo_include_todos = True
+
+
+# -- Options for HTML output ----------------------------------------------
+
+# The theme to use for HTML and HTML Help pages. See the documentation for
+# a list of builtin themes.
+html_theme = "alabaster"
+
+# Theme options are theme-specific and customize the look and feel of a theme
+# further. For a list of options available for each theme, see the
+# documentation.
+html_theme_options = {
+ "description": "Google Cloud Client Libraries for llama-index-cloud-sql-pg",
+ "github_user": "googleapis",
+ "github_repo": "llama-index-cloud-sql-pg-python",
+ "github_banner": True,
+ "font_family": "'Roboto', Georgia, sans",
+ "head_font_family": "'Roboto', Georgia, serif",
+ "code_font_family": "'Roboto Mono', 'Consolas', monospace",
+}
+
+# Add any paths that contain custom themes here, relative to this directory.
+# html_theme_path = []
+
+# The name for this set of Sphinx documents. If None, it defaults to
+# " v documentation".
+# html_title = None
+
+# A shorter title for the navigation bar. Default is the same as html_title.
+# html_short_title = None
+
+# The name of an image file (relative to this directory) to place at the top
+# of the sidebar.
+# html_logo = None
+
+# The name of an image file (within the static path) to use as favicon of the
+# docs. This file should be a Windows icon file (.ico) being 16x16 or 32x32
+# pixels large.
+# html_favicon = None
+
+# Add any paths that contain custom static files (such as style sheets) here,
+# relative to this directory. They are copied after the builtin static files,
+# so a file named "default.css" will overwrite the builtin "default.css".
+html_static_path = ["_static"]
+
+# Add any extra paths that contain custom files (such as robots.txt or
+# .htaccess) here, relative to this directory. These files are copied
+# directly to the root of the documentation.
+# html_extra_path = []
+
+# If not '', a 'Last updated on:' timestamp is inserted at every page bottom,
+# using the given strftime format.
+# html_last_updated_fmt = '%b %d, %Y'
+
+# If true, SmartyPants will be used to convert quotes and dashes to
+# typographically correct entities.
+# html_use_smartypants = True
+
+# Custom sidebar templates, maps document names to template names.
+# html_sidebars = {}
+
+# Additional templates that should be rendered to pages, maps page names to
+# template names.
+# html_additional_pages = {}
+
+# If false, no module index is generated.
+# html_domain_indices = True
+
+# If false, no index is generated.
+# html_use_index = True
+
+# If true, the index is split into individual pages for each letter.
+# html_split_index = False
+
+# If true, links to the reST sources are added to the pages.
+# html_show_sourcelink = True
+
+# If true, "Created using Sphinx" is shown in the HTML footer. Default is True.
+# html_show_sphinx = True
+
+# If true, "(C) Copyright ..." is shown in the HTML footer. Default is True.
+# html_show_copyright = True
+
+# If true, an OpenSearch description file will be output, and all pages will
+# contain a tag referring to it. The value of this option must be the
+# base URL from which the finished HTML is served.
+# html_use_opensearch = ''
+
+# This is the file name suffix for HTML files (e.g. ".xhtml").
+# html_file_suffix = None
+
+# Language to be used for generating the HTML full-text search index.
+# Sphinx supports the following languages:
+# 'da', 'de', 'en', 'es', 'fi', 'fr', 'hu', 'it', 'ja'
+# 'nl', 'no', 'pt', 'ro', 'ru', 'sv', 'tr'
+# html_search_language = 'en'
+
+# A dictionary with options for the search language support, empty by default.
+# Now only 'ja' uses this config value
+# html_search_options = {'type': 'default'}
+
+# The name of a javascript file (relative to the configuration directory) that
+# implements a search results scorer. If empty, the default will be used.
+# html_search_scorer = 'scorer.js'
+
+# Output file base name for HTML help builder.
+htmlhelp_basename = "llama-index-cloud-sql-pg-doc"
+
+# -- Options for warnings ------------------------------------------------------
+
+
+suppress_warnings = [
+ # Temporarily suppress this to avoid "more than one target found for
+ # cross-reference" warning, which are intractable for us to avoid while in
+ # a mono-repo.
+ # See https://github.com/sphinx-doc/sphinx/blob
+ # /2a65ffeef5c107c19084fabdd706cdff3f52d93c/sphinx/domains/python.py#L843
+ "ref.python"
+]
+
+# -- Options for LaTeX output ---------------------------------------------
+
+latex_elements = {
+ # The paper size ('letterpaper' or 'a4paper').
+ #'papersize': 'letterpaper',
+ # The font size ('10pt', '11pt' or '12pt').
+ #'pointsize': '10pt',
+ # Additional stuff for the LaTeX preamble.
+ #'preamble': '',
+ # Latex figure (float) alignment
+ #'figure_align': 'htbp',
+}
+
+# Grouping the document tree into LaTeX files. List of tuples
+# (source start file, target name, title,
+# author, documentclass [howto, manual, or own class]).
+latex_documents = [
+ (
+ root_doc,
+ "llama-index-cloud-sql-pg.tex",
+ "llama-index-cloud-sql-pg Documentation",
+ author,
+ "manual",
+ )
+]
+
+# The name of an image file (relative to this directory) to place at the top of
+# the title page.
+# latex_logo = None
+
+# For "manual" documents, if this is true, then toplevel headings are parts,
+# not chapters.
+# latex_use_parts = False
+
+# If true, show page references after internal links.
+# latex_show_pagerefs = False
+
+# If true, show URL addresses after external links.
+# latex_show_urls = False
+
+# Documents to append as an appendix to all manuals.
+# latex_appendices = []
+
+# If false, no module index is generated.
+# latex_domain_indices = True
+
+
+# -- Options for manual page output ---------------------------------------
+
+# One entry per manual page. List of tuples
+# (source start file, name, description, authors, manual section).
+man_pages = [
+ (
+ root_doc,
+ "llama-index-cloud-sql-pg",
+ "llama-index-cloud-sql-pg Documentation",
+ [author],
+ 1,
+ )
+]
+
+# If true, show URL addresses after external links.
+# man_show_urls = False
+
+
+# -- Options for Texinfo output -------------------------------------------
+
+# Grouping the document tree into Texinfo files. List of tuples
+# (source start file, target name, title, author,
+# dir menu entry, description, category)
+texinfo_documents = [
+ (
+ root_doc,
+ "llama-index-cloud-sql-pg",
+ "llama-index-cloud-sql-pg Documentation",
+ author,
+ "llama-index-cloud-sql-pg",
+ "llama-index-cloud-sql-pg Library",
+ "APIs",
+ )
+]
+
+# Documents to append as an appendix to all manuals.
+# texinfo_appendices = []
+
+# If false, no module index is generated.
+# texinfo_domain_indices = True
+
+# How to display URL addresses: 'footnote', 'no', or 'inline'.
+# texinfo_show_urls = 'footnote'
+
+# If true, do not generate a @detailmenu in the "Top" node's menu.
+# texinfo_no_detailmenu = False
+
+
+# Example configuration for intersphinx: refer to the Python standard library.
+intersphinx_mapping = {
+ "python": ("https://python.readthedocs.org/en/latest/", None),
+ "google-auth": ("https://googleapis.dev/python/google-auth/latest/", None),
+ "google.api_core": (
+ "https://googleapis.dev/python/google-api-core/latest/",
+ None,
+ ),
+ "grpc": ("https://grpc.github.io/grpc/python/", None),
+ "proto-plus": ("https://proto-plus-python.readthedocs.io/en/latest/", None),
+ "protobuf": ("https://googleapis.dev/python/protobuf/latest/", None),
+}
+
+
+# Napoleon settings
+napoleon_google_docstring = True
+napoleon_numpy_docstring = True
+napoleon_include_private_with_doc = False
+napoleon_include_special_with_doc = True
+napoleon_use_admonition_for_examples = False
+napoleon_use_admonition_for_notes = False
+napoleon_use_admonition_for_references = False
+napoleon_use_ivar = False
+napoleon_use_param = True
+napoleon_use_rtype = True
diff --git a/docs/index.rst b/docs/index.rst
new file mode 100644
index 0000000..a4305df
--- /dev/null
+++ b/docs/index.rst
@@ -0,0 +1,20 @@
+.. include:: README.rst
+
+API Reference
+-------------
+.. toctree::
+ :maxdepth: 2
+
+ llama_index_cloud_sql_pg/chat_store
+ llama_index_cloud_sql_pg/document_store
+ llama_index_cloud_sql_pg/engine
+ llama_index_cloud_sql_pg/index_store
+ llama_index_cloud_sql_pg/reader
+ llama_index_cloud_sql_pg/vector_store
+
+Changelog
+---------
+.. toctree::
+ :maxdepth: 2
+
+ changelog
diff --git a/docs/llama_index_cloud_sql_pg/chat_store.rst b/docs/llama_index_cloud_sql_pg/chat_store.rst
new file mode 100644
index 0000000..433a624
--- /dev/null
+++ b/docs/llama_index_cloud_sql_pg/chat_store.rst
@@ -0,0 +1,7 @@
+Chat Store
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. automodule:: llama_index_cloud_sql_pg.chat_store
+ :members:
+ :private-members:
+ :noindex:
diff --git a/docs/llama_index_cloud_sql_pg/document_store.rst b/docs/llama_index_cloud_sql_pg/document_store.rst
new file mode 100644
index 0000000..39ee4f9
--- /dev/null
+++ b/docs/llama_index_cloud_sql_pg/document_store.rst
@@ -0,0 +1,7 @@
+Document Store
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. automodule:: llama_index_cloud_sql_pg.document_store
+ :members:
+ :private-members:
+ :noindex:
diff --git a/docs/llama_index_cloud_sql_pg/engine.rst b/docs/llama_index_cloud_sql_pg/engine.rst
new file mode 100644
index 0000000..37bdb1e
--- /dev/null
+++ b/docs/llama_index_cloud_sql_pg/engine.rst
@@ -0,0 +1,7 @@
+Engine
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. automodule:: llama_index_cloud_sql_pg.engine
+ :members:
+ :private-members:
+ :noindex:
diff --git a/docs/llama_index_cloud_sql_pg/index_store.rst b/docs/llama_index_cloud_sql_pg/index_store.rst
new file mode 100644
index 0000000..bc983b9
--- /dev/null
+++ b/docs/llama_index_cloud_sql_pg/index_store.rst
@@ -0,0 +1,7 @@
+Index Store
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. automodule:: llama_index_cloud_sql_pg.index_store
+ :members:
+ :private-members:
+ :noindex:
diff --git a/docs/llama_index_cloud_sql_pg/reader.rst b/docs/llama_index_cloud_sql_pg/reader.rst
new file mode 100644
index 0000000..e082aaa
--- /dev/null
+++ b/docs/llama_index_cloud_sql_pg/reader.rst
@@ -0,0 +1,7 @@
+Document Reader
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. automodule:: llama_index_cloud_sql_pg.reader
+ :members:
+ :private-members:
+ :noindex:
diff --git a/docs/llama_index_cloud_sql_pg/vector_store.rst b/docs/llama_index_cloud_sql_pg/vector_store.rst
new file mode 100644
index 0000000..860b945
--- /dev/null
+++ b/docs/llama_index_cloud_sql_pg/vector_store.rst
@@ -0,0 +1,13 @@
+Vector Store
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. automodule:: llama_index_cloud_sql_pg.vector_store
+ :members:
+ :private-members:
+ :noindex:
+
+
+.. automodule:: llama_index_cloud_sql_pg.indexes
+ :members:
+ :private-members:
+ :noindex:
diff --git a/noxfile.py b/noxfile.py
new file mode 100644
index 0000000..bbbc40e
--- /dev/null
+++ b/noxfile.py
@@ -0,0 +1,117 @@
+# -*- coding: utf-8 -*-
+#
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import absolute_import
+
+import os
+import pathlib
+import shutil
+from pathlib import Path
+from typing import Optional
+
+import nox
+
+DEFAULT_PYTHON_VERSION = "3.10"
+CURRENT_DIRECTORY = pathlib.Path(__file__).parent.absolute()
+
+nox.options.sessions = [
+ "docs",
+ "docfx",
+]
+
+# Error if a python version is missing
+nox.options.error_on_missing_interpreters = True
+
+
+@nox.session(python="3.10")
+def docs(session):
+ """Build the docs for this library."""
+
+ session.install("-e", ".")
+ session.install(
+ # We need to pin to specific versions of the `sphinxcontrib-*` packages
+ # which still support sphinx 4.x.
+ # See https://github.com/googleapis/sphinx-docfx-yaml/issues/344
+ # and https://github.com/googleapis/sphinx-docfx-yaml/issues/345.
+ "sphinxcontrib-applehelp==1.0.4",
+ "sphinxcontrib-devhelp==1.0.2",
+ "sphinxcontrib-htmlhelp==2.0.1",
+ "sphinxcontrib-qthelp==1.0.3",
+ "sphinxcontrib-serializinghtml==1.1.5",
+ "sphinx==4.5.0",
+ "alabaster",
+ "recommonmark",
+ )
+
+ shutil.rmtree(os.path.join("docs", "_build"), ignore_errors=True)
+ session.run(
+ "sphinx-build",
+ "-W", # warnings as errors
+ "-T", # show full traceback on exception
+ "-N", # no colors
+ "-b",
+ "html",
+ "-d",
+ os.path.join("docs", "_build", "doctrees", ""),
+ os.path.join("docs", ""),
+ os.path.join("docs", "_build", "html", ""),
+ )
+
+
+@nox.session(python="3.10")
+def docfx(session):
+ """Build the docfx yaml files for this library."""
+
+ session.install("-e", ".")
+ session.install(
+ # We need to pin to specific versions of the `sphinxcontrib-*` packages
+ # which still support sphinx 4.x.
+ # See https://github.com/googleapis/sphinx-docfx-yaml/issues/344
+ # and https://github.com/googleapis/sphinx-docfx-yaml/issues/345.
+ "sphinxcontrib-applehelp==1.0.4",
+ "sphinxcontrib-devhelp==1.0.2",
+ "sphinxcontrib-htmlhelp==2.0.1",
+ "sphinxcontrib-qthelp==1.0.3",
+ "sphinxcontrib-serializinghtml==1.1.5",
+ "gcp-sphinx-docfx-yaml",
+ "alabaster",
+ "recommonmark",
+ )
+
+ shutil.rmtree(os.path.join("docs", "_build"), ignore_errors=True)
+ session.run(
+ "sphinx-build",
+ "-T", # show full traceback on exception
+ "-N", # no colors
+ "-D",
+ (
+ "extensions=sphinx.ext.autodoc,"
+ "sphinx.ext.autosummary,"
+ "docfx_yaml.extension,"
+ "sphinx.ext.intersphinx,"
+ "sphinx.ext.coverage,"
+ "sphinx.ext.napoleon,"
+ "sphinx.ext.todo,"
+ "sphinx.ext.viewcode,"
+ "recommonmark"
+ ),
+ "-b",
+ "html",
+ "-d",
+ os.path.join("docs", "_build", "doctrees", ""),
+ os.path.join("docs", ""),
+ os.path.join("docs", "_build", "html", ""),
+ )
diff --git a/pyproject.toml b/pyproject.toml
index d59c095..c51222d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -40,8 +40,8 @@ Changelog = "https://github.com/googleapis/llama-index-cloud-sql-pg-python/blob/
test = [
"black[jupyter]==25.1.0",
"isort==6.0.0",
- "mypy==1.14.1",
- "pytest-asyncio==0.25.2",
+ "mypy==1.15.0",
+ "pytest-asyncio==0.25.3",
"pytest==8.3.4",
"pytest-cov==6.0.0"
]
diff --git a/requirements.txt b/requirements.txt
index d8f19ff..22cef8e 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,4 +1,4 @@
-cloud-sql-python-connector[asyncpg]==1.16.0
-llama-index-core==0.12.12
+cloud-sql-python-connector[asyncpg]==1.17.0
+llama-index-core==0.12.17
pgvector==0.3.6
-SQLAlchemy[asyncio]==2.0.37
+SQLAlchemy[asyncio]==2.0.38
diff --git a/samples/llama_index_chat_store.ipynb b/samples/llama_index_chat_store.ipynb
new file mode 100644
index 0000000..c3073b3
--- /dev/null
+++ b/samples/llama_index_chat_store.ipynb
@@ -0,0 +1,396 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# Google Cloud SQL for PostgreSQL - `PostgresChatStore`\n",
+ "\n",
+ "> [Cloud SQL](https://cloud.google.com/sql) is a fully managed relational database service that offers high performance, seamless integration, and impressive scalability. It offers MySQL, PostgreSQL, and SQL Server database engines. Extend your database application to build AI-powered experiences leveraging Cloud SQL's LlamaIndex integrations.\n",
+ "\n",
+ "This notebook goes over how to use `Cloud SQL for PostgreSQL` to store chat history with `PostgresChatStore` class.\n",
+ "\n",
+ "Learn more about the package on [GitHub](https://github.com/googleapis/llama-index-cloud-sql-pg-python/).\n",
+ "\n",
+ "[](https://colab.research.google.com/github/googleapis/llama-index-cloud-sql-pg-python/blob/main/samples/llama_index_chat_store.ipynb)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## Before you begin\n",
+ "\n",
+ "To run this notebook, you will need to do the following:\n",
+ "\n",
+ " * [Create a Google Cloud Project](https://developers.google.com/workspace/guides/create-project)\n",
+ " * [Enable the Cloud SQL Admin API.](https://console.cloud.google.com/flows/enableapi?apiid=sqladmin.googleapis.com)\n",
+ " * [Create a Cloud SQL instance.](https://cloud.google.com/sql/docs/postgres/connect-instance-auth-proxy#create-instance)\n",
+ " * [Create a Cloud SQL database.](https://cloud.google.com/sql/docs/postgres/create-manage-databases)\n",
+ " * [Add a User to the database.](https://cloud.google.com/sql/docs/postgres/create-manage-users)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### 🦙 Library Installation\n",
+ "Install the integration library, `llama-index-cloud-sql-pg`, and the library for the embedding service, `llama-index-embeddings-vertex`."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "%pip install --upgrade --quiet llama-index-cloud-sql-pg llama-index-llms-vertex llama-index"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "**Colab only:** Uncomment the following cell to restart the kernel or use the button to restart the kernel. For Vertex AI Workbench you can restart the terminal using the button on top."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# # Automatically restart kernel after installs so that your environment can access the new packages\n",
+ "# import IPython\n",
+ "\n",
+ "# app = IPython.Application.instance()\n",
+ "# app.kernel.do_shutdown(True)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### 🔐 Authentication\n",
+ "Authenticate to Google Cloud as the IAM user logged into this notebook in order to access your Google Cloud Project.\n",
+ "\n",
+ "* If you are using Colab to run this notebook, use the cell below and continue.\n",
+ "* If you are using Vertex AI Workbench, check out the setup instructions [here](https://github.com/GoogleCloudPlatform/generative-ai/tree/main/setup-env)."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from google.colab import auth\n",
+ "\n",
+ "auth.authenticate_user()"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### ☁ Set Your Google Cloud Project\n",
+ "Set your Google Cloud project so that you can leverage Google Cloud resources within this notebook.\n",
+ "\n",
+ "If you don't know your project ID, try the following:\n",
+ "\n",
+ "* Run `gcloud config list`.\n",
+ "* Run `gcloud projects list`.\n",
+ "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# @markdown Please fill in the value below with your Google Cloud project ID and then run the cell.\n",
+ "\n",
+ "PROJECT_ID = \"my-project-id\" # @param {type:\"string\"}\n",
+ "\n",
+ "# Set the project id\n",
+ "!gcloud config set project {PROJECT_ID}"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## Basic Usage"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Set Cloud SQL database values\n",
+ "Find your database values, in the [Cloud SQL Instances page](https://console.cloud.google.com/sql?_ga=2.223735448.2062268965.1707700487-2088871159.1707257687)."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# @title Set Your Values Here { display-mode: \"form\" }\n",
+ "REGION = \"us-central1\" # @param {type: \"string\"}\n",
+ "INSTANCE = \"my-primary\" # @param {type: \"string\"}\n",
+ "DATABASE = \"my-database\" # @param {type: \"string\"}\n",
+ "TABLE_NAME = \"chat_store\" # @param {type: \"string\"}\n",
+ "USER = \"postgres\" # @param {type: \"string\"}\n",
+ "PASSWORD = \"my-password\" # @param {type: \"string\"}"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### PostgresEngine Connection Pool\n",
+ "\n",
+ "One of the requirements and arguments to establish Cloud SQL as a chat store is a `PostgresEngine` object. The `PostgresEngine` configures a connection pool to your Cloud SQL database, enabling successful connections from your application and following industry best practices.\n",
+ "\n",
+ "To create a `PostgresEngine` using `PostgresEngine.from_instance()` you need to provide only 4 things:\n",
+ "\n",
+ "1. `project_id` : Project ID of the Google Cloud Project where the Cloud SQL instance is located.\n",
+ "1. `region` : Region where the Cloud SQL instance is located.\n",
+ "1. `instance` : The name of the Cloud SQL instance.\n",
+ "1. `database` : The name of the database to connect to on the Cloud SQL instance.\n",
+ "\n",
+ "By default, [IAM database authentication](https://cloud.google.com/sql/docs/postgres/iam-authentication#iam-db-auth) will be used as the method of database authentication. This library uses the IAM principal belonging to the [Application Default Credentials (ADC)](https://cloud.google.com/docs/authentication/application-default-credentials) sourced from the envionment.\n",
+ "\n",
+ "For more informatin on IAM database authentication please see:\n",
+ "\n",
+ "* [Configure an instance for IAM database authentication](https://cloud.google.com/sql/docs/postgres/create-edit-iam-instances)\n",
+ "* [Manage users with IAM database authentication](https://cloud.google.com/sql/docs/postgres/add-manage-iam-users)\n",
+ "\n",
+ "Optionally, [built-in database authentication](https://cloud.google.com/sql/docs/postgres/built-in-authentication) using a username and password to access the Cloud SQL database can also be used. Just provide the optional `user` and `password` arguments to `PostgresEngine.from_instance()`:\n",
+ "\n",
+ "* `user` : Database user to use for built-in database authentication and login\n",
+ "* `password` : Database password to use for built-in database authentication and login.\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "**Note:** This tutorial demonstrates the async interface. All async methods have corresponding sync methods."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from llama_index_cloud_sql_pg import PostgresEngine\n",
+ "\n",
+ "engine = await PostgresEngine.afrom_instance(\n",
+ " project_id=PROJECT_ID,\n",
+ " region=REGION,\n",
+ " instance=INSTANCE,\n",
+ " database=DATABASE,\n",
+ " user=USER,\n",
+ " password=PASSWORD,\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Initialize a table\n",
+ "The `PostgresChatStore` class requires a database table. The `PostgresEngine` engine has a helper method `ainit_chat_store_table()` that can be used to create a table with the proper schema for you."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "await engine.ainit_chat_store_table(table_name=TABLE_NAME)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "#### Optional Tip: 💡\n",
+ "You can also specify a schema name by passing `schema_name` wherever you pass `table_name`."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "SCHEMA_NAME = \"my_schema\"\n",
+ "\n",
+ "await engine.ainit_chat_store_table(\n",
+ " table_name=TABLE_NAME,\n",
+ " schema_name=SCHEMA_NAME,\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Initialize a default PostgresChatStore"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from llama_index_cloud_sql_pg import PostgresChatStore\n",
+ "\n",
+ "chat_store = await PostgresChatStore.create(\n",
+ " engine=engine,\n",
+ " table_name=TABLE_NAME,\n",
+ " # schema_name=SCHEMA_NAME\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Create a ChatMemoryBuffer\n",
+ "The `ChatMemoryBuffer` stores a history of recent chat messages, enabling the LLM to access relevant context from prior interactions.\n",
+ "\n",
+ "By passing our chat store into the `ChatMemoryBuffer`, it can automatically retrieve and update messages associated with a specific session ID or `chat_store_key`."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from llama_index.core.memory import ChatMemoryBuffer\n",
+ "\n",
+ "memory = ChatMemoryBuffer.from_defaults(\n",
+ " token_limit=3000,\n",
+ " chat_store=chat_store,\n",
+ " chat_store_key=\"user1\",\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Create an LLM class instance\n",
+ "\n",
+ "You can use any of the [LLMs compatible with LlamaIndex](https://docs.llamaindex.ai/en/stable/module_guides/models/llms/modules/).\n",
+ "You may need to enable Vertex AI API to use `Vertex`."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from llama_index.llms.vertex import Vertex\n",
+ "\n",
+ "llm = Vertex(model=\"gemini-1.5-flash-002\", project=PROJECT_ID)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Use the PostgresChatStore without a storage context"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "#### Create a Simple Chat Engine"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from llama_index.core.chat_engine import SimpleChatEngine\n",
+ "\n",
+ "chat_engine = SimpleChatEngine(memory=memory, llm=llm, prefix_messages=[])\n",
+ "\n",
+ "response = chat_engine.chat(\"Hello\")\n",
+ "\n",
+ "print(response)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Use the PostgresChatStore with a storage context"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "#### Create a LlamaIndex `Index`\n",
+ "\n",
+ "An `Index` is allows us to quickly retrieve relevant context for a user query.\n",
+ "They are used to build `QueryEngines` and `ChatEngines`.\n",
+ "For a list of indexes that can be built in LlamaIndex, see [Index Guide](https://docs.llamaindex.ai/en/stable/module_guides/indexing/index_guide/).\n",
+ "\n",
+ "A `VectorStoreIndex`, can be built using the `PostgresVectorStore`. See the detailed guide on how to use the `PostgresVectorStore` to build an index [here](https://github.com/googleapis/llama-index-cloud-sql-pg-python/blob/main/samples/llama_index_vector_store.ipynb).\n",
+ "\n",
+ "You can also use the `PostgresDocumentStore` and `PostgresIndexStore` to persist documents and index metadata.\n",
+ "These modules can be used to build other `Indexes`.\n",
+ "For a detailed python notebook on this, see [LlamaIndex Doc Store Guide](https://github.com/googleapis/llama-index-cloud-sql-pg-python/blob/main/samples/llama_index_doc_store.ipynb)."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "#### Create and use the Chat Engine"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Create an `index` here\n",
+ "\n",
+ "chat_engine = index.as_chat_engine(llm=llm, chat_mode=\"context\", memory=memory) # type: ignore\n",
+ "response = chat_engine.chat(\"What did the author do?\")"
+ ]
+ }
+ ],
+ "metadata": {
+ "kernelspec": {
+ "display_name": "senseAIenv",
+ "language": "python",
+ "name": "python3"
+ },
+ "language_info": {
+ "name": "python",
+ "version": "3.11.9"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/samples/llama_index_quick_start.ipynb b/samples/llama_index_quick_start.ipynb
new file mode 100644
index 0000000..4f3cebd
--- /dev/null
+++ b/samples/llama_index_quick_start.ipynb
@@ -0,0 +1,1087 @@
+{
+ "cells": [
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "upi2EY4L9ei3"
+ },
+ "outputs": [],
+ "source": [
+ "# Copyright 2025 Google LLC\n",
+ "#\n",
+ "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+ "# you may not use this file except in compliance with the License.\n",
+ "# You may obtain a copy of the License at\n",
+ "#\n",
+ "# https://www.apache.org/licenses/LICENSE-2.0\n",
+ "#\n",
+ "# Unless required by applicable law or agreed to in writing, software\n",
+ "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+ "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+ "# See the License for the specific language governing permissions and\n",
+ "# limitations under the License."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "mbF2F2miAT4a"
+ },
+ "source": [
+ "[](https://colab.research.google.com/github/googleapis/llama-index-cloud-sql-pg-python/blob/main/samples/llama_index_quick_start.ipynb)\n",
+ "\n",
+ "---\n",
+ "# Introduction\n",
+ "\n",
+ "In this codelab, you'll learn how to create a powerful interactive generative AI application using Retrieval Augmented Generation powered by [Cloud SQL for PostgreSQL](https://cloud.google.com/sql/docs/postgres) and [LlamaIndex](https://www.llamaindex.ai/). We will be creating an application grounded in a [Netflix Movie dataset](https://www.kaggle.com/datasets/shivamb/netflix-shows), allowing you to interact with movie data in exciting new ways."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "Ma6pEng3ypbA"
+ },
+ "source": [
+ "## Prerequisites\n",
+ "\n",
+ "* A basic understanding of the Google Cloud Console\n",
+ "* Basic skills in command line interface and Google Cloud shell\n",
+ "* Basic python knowledge"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "DzDgqJHgysy1"
+ },
+ "source": [
+ "## What you'll learn\n",
+ "\n",
+ "* How to deploy a Cloud SQL for PostgreSQL instance\n",
+ "* How to use Cloud SQL for PostgreSQL as a Document Reader\n",
+ "* How to use Cloud SQL for PostgreSQL as a Vector Store\n",
+ "* How to use Cloud SQL for PostgreSQL as a Document Store with a Index Store\n",
+ "* How to use Cloud SQL for PostgreSQL as a Chat Store"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "FbcZUjT1yvTq"
+ },
+ "source": [
+ "## What you'll need\n",
+ "\n",
+ "* A Google Cloud Account and Google Cloud Project\n",
+ "* A web browser such as [Chrome](https://www.google.com/chrome/)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "vHdR4fF3vLWA"
+ },
+ "source": [
+ "# Setup and Requirements\n",
+ "\n",
+ "In the following instructions you will learn to:\n",
+ "\n",
+ "1. Install required dependencies for our application\n",
+ "2. Set up authentication for our project\n",
+ "3. Set up a Cloud SQL for PostgreSQL Instance\n",
+ "4. Import the data used by our application"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "uy9KqgPQ4GBi"
+ },
+ "source": [
+ "## Install dependencies\n",
+ "First you will need to install the dependencies needed to run this demo app."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "M_ppDxYf4Gqs"
+ },
+ "outputs": [],
+ "source": [
+ "%pip install llama-index-cloud-sql-pg llama-index llama-index-embeddings-vertex llama-index-llms-vertex cloud-sql-python-connector[pg8000]"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "**Colab only:** Uncomment the following cell to restart the kernel or use the button to restart the kernel. For Vertex AI Workbench you can restart the terminal using the button on top."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Automatically restart kernel after installs so that your environment can access the new packages\n",
+ "import IPython\n",
+ "\n",
+ "app = IPython.Application.instance()\n",
+ "app.kernel.do_shutdown(True)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "DeUbHclxw7_l"
+ },
+ "source": [
+ "## Authenticate to Google Cloud within Colab\n",
+ "In order to access your Google Cloud Project from this notebook, you will need to Authenticate as an IAM user."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "_Q9hyqdyEx6l"
+ },
+ "outputs": [],
+ "source": [
+ "from google.colab import auth\n",
+ "\n",
+ "auth.authenticate_user()"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "UCiNGP1Qxd6x"
+ },
+ "source": [
+ "## Connect Your Google Cloud Project"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "SLUGlG6UE2CK"
+ },
+ "outputs": [],
+ "source": [
+ "# @markdown Please fill in the value below with your GCP project ID and then run the cell.\n",
+ "\n",
+ "# Please fill in these values.\n",
+ "project_id = \"\" # @param {type:\"string\"}\n",
+ "\n",
+ "# Quick input validations.\n",
+ "assert project_id, \"⚠️ Please provide a Google Cloud project ID\"\n",
+ "\n",
+ "# Configure gcloud.\n",
+ "!gcloud config set project {project_id}"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "O-oqMC5Ox-ZM"
+ },
+ "source": [
+ "## Configure Your Google Cloud Project\n",
+ "\n",
+ "Configure the following in your Google Cloud Project."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "X-bzfFb4A-xK"
+ },
+ "source": [
+ "You will need to enable these APIs in order to use `VertexTextEmbeddings` as an embeddings service!"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "1. IAM principal (user, service account, etc.) with the [Cloud SQL Client][client-role] role. The user logged into this notebook will be used as the IAM principal and will be granted the Cloud SQL Client role.\n",
+ "\n",
+ "[client-role]: https://cloud.google.com/sql/docs/mysql/roles-and-permissions"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "current_user = !gcloud auth list --filter=status:ACTIVE --format=\"value(account)\"\n",
+ "!gcloud projects add-iam-policy-binding {project_id} \\\n",
+ " --member=user:{current_user[0]} \\\n",
+ " --role=\"roles/cloudsql.client\""
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "2. Enable the APIs for Cloud SQL and Vertex AI within your project."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "CKWrwyfzyTwH"
+ },
+ "outputs": [],
+ "source": [
+ "# Enable GCP services\n",
+ "!gcloud services enable sqladmin.googleapis.com aiplatform.googleapis.com"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "Gn8g7-wCyZU6"
+ },
+ "source": [
+ "## Set up Cloud SQL\n",
+ "You will need a **Postgres** Cloud SQL instance for the following stages of this notebook. Please set the following variables."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# @markdown Please fill in the both the Google Cloud region and name of your Cloud SQL instance. Once filled in, run the cell.\n",
+ "\n",
+ "# Please fill in these values.\n",
+ "region = \"us-central1\" # @param {type:\"string\"}\n",
+ "instance_name = \"llamaindex-quickstart-instance\" # @param {type:\"string\"}\n",
+ "database_name = \"llamaindex-quickstart-db\" # @param {type:\"string\"}\n",
+ "user = \"postgres\" # @param {type:\"string\"}\n",
+ "password = input(\"Please provide a password to be used for 'postgres' database user: \")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "T616pEOUygYQ"
+ },
+ "source": [
+ "### Create a Postgres Instance\n",
+ "Running the below cell will verify the existence of the Cloud SQL instance and or create a new instance and database if one does not exist.\n",
+ "\n",
+ "> ⏳ - Creating a Cloud SQL instance may take a few minutes."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "XXI1uUu3y8gc"
+ },
+ "outputs": [],
+ "source": [
+ "# check if Cloud SQL instance exists in the provided region\n",
+ "database_version = !gcloud sql instances describe {instance_name} --format=\"value(databaseVersion)\"\n",
+ "if database_version[0].startswith(\"POSTGRES\"):\n",
+ " print(\"Found existing Postgres Cloud SQL Instance!\")\n",
+ "else:\n",
+ " print(\"Creating new Cloud SQL instance...\")\n",
+ " !gcloud sql instances create {instance_name} --database-version=POSTGRES_15 \\\n",
+ " --region={region} --cpu=1 --memory=4GB --root-password={password} \\\n",
+ " --database-flags=cloudsql.iam_authentication=On"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Create a Database\n",
+ "\n",
+ "Next you will create database to store the data for this application."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "databases = !gcloud sql databases list --instance={instance_name} --format=\"value(name)\"\n",
+ "if database_name not in databases:\n",
+ " !gcloud sql databases create {database_name} --instance={instance_name}\n",
+ "else:\n",
+ " print(\"Found existing database!\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "#### Connect to our New Database\n",
+ "\n",
+ "Now you will use `PostgresEngine` that connects to your new database!"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from google.cloud.sql.connector import Connector\n",
+ "import sqlalchemy\n",
+ "\n",
+ "# initialize Connector object\n",
+ "connector = Connector()\n",
+ "\n",
+ "\n",
+ "# function to return the database connection\n",
+ "def getconn():\n",
+ " conn = connector.connect(\n",
+ " f\"{project_id}:{region}:{instance_name}\",\n",
+ " \"pg8000\",\n",
+ " user=user,\n",
+ " password=password,\n",
+ " db=database_name,\n",
+ " )\n",
+ " return conn\n",
+ "\n",
+ "\n",
+ "# create connection pool\n",
+ "pool = sqlalchemy.create_engine(\n",
+ " \"postgresql+pg8000://\",\n",
+ " creator=getconn,\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "HdolCWyatZmG"
+ },
+ "source": [
+ "## Import data to your database\n",
+ "\n",
+ "Now that you have your database, you will need to import data! We will be using a [Netflix Dataset from Kaggle](https://www.kaggle.com/datasets/shivamb/netflix-shows). Here is what the data looks like:"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "36-FBKzJ-tLa"
+ },
+ "source": [
+ "| show_id | type | title | director | cast | country | date_added | release_year | rating | duration | listed_in | description |\n",
+ "|---------|---------|----------------------|------------------|--------------------------------------------------------------------------------------------------------------------------------------------------------|---------------|-------------------|--------------|--------|-----------|----------------------------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|\n",
+ "| s1 | Movie | Dick Johnson Is Dead | Kirsten Johnson | | United States | September 25, 2021 | 2020 | PG-13 | 90 min | Documentaries | As her father nears the end of his life, filmmaker Kirsten Johnson stages his death in inventive and comical ways to help them both face the inevitable. |\n",
+ "| s2 | TV Show | Blood & Water | | Ama Qamata, Khosi Ngema, Gail Mabalane, Thabang Molaba, Dillon Windvogel, Natasha Thahane, Arno Greeff, Xolile Tshabalala, Getmore Sithole, Cindy Mahlangu, Ryle De Morny, Greteli Fincham, Sello Maake Ka-Ncube, Odwa Gwanya, Mekaila Mathys, Sandi Schultz, Duane Williams, Shamilla Miller, Patrick Mofokeng | South Africa | September 24, 2021 | 2021 | TV-MA | 2 Seasons | International TV Shows, TV Dramas, TV Mysteries | After crossing paths at a party, a Cape Town teen sets out to prove whether a private-school swimming star is her sister who was abducted at birth. |\n",
+ "| s3 | TV Show | Ganglands | Julien Leclercq | Sami Bouajila, Tracy Gotoas, Samuel Jouy, Nabiha Akkari, Sofia Lesaffre, Salim Kechiouche, Noureddine Farihi, Geert Van Rampelberg, Bakary Diombera | | September 24, 2021 | 2021 | TV-MA | 1 Season | Crime TV Shows, International TV Shows, TV Action & Adventure | To protect his family from a powerful drug lord, skilled thief Mehdi and his expert team of robbers are pulled into a violent and deadly turf war. |\n",
+ "| s4 | TV Show | Jailbirds New Orleans | | | | September 24, 2021 | 2021 | TV-MA | 1 Season | Docuseries, Reality TV | Feuds, flirtations and toilet talk go down among the incarcerated women at the Orleans Justice Center in New Orleans on this gritty reality series. |\n",
+ "| s5 | TV Show | Kota Factory | | Mayur More, Jitendra Kumar, Ranjan Raj, Alam Khan, Ahsaas Channa, Revathi Pillai, Urvi Singh, Arun Kumar | India | September 24, 2021 | 2021 | TV-MA | 2 Seasons | International TV Shows, Romantic TV Shows, TV Comedies | In a city of coaching centers known to train India’s finest collegiate minds, an earnest but unexceptional student and his friends navigate campus life. |\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "kQ2KWsYI_Msa"
+ },
+ "source": [
+ "The following code has been prepared code to help insert the CSV data into your Cloud SQL for PostgreSQL database."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "Dzr-2VZIkvtY"
+ },
+ "source": [
+ "Download the CSV file:"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "5KkIQ2zSvQkN"
+ },
+ "outputs": [],
+ "source": [
+ "!gsutil cp gs://cloud-samples-data/llamaindex/netflix_titles.csv ."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "oFU13dCBlYHh"
+ },
+ "source": [
+ "The download can be verified by the following command or using the \"Files\" tab."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "nQBs10I8vShh"
+ },
+ "outputs": [],
+ "source": [
+ "!ls"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "2H7rorG9Ivur"
+ },
+ "source": [
+ "In this next step you will:\n",
+ "\n",
+ "1. Create the table into store data\n",
+ "2. And insert the data from the CSV into the database table\n",
+ "\n",
+ "> To avoid costs, the following code uses only 100 rows as an example"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "qCsM2KXbdYiv"
+ },
+ "outputs": [],
+ "source": [
+ "import pandas as pd\n",
+ "import sqlalchemy\n",
+ "\n",
+ "create_table_cmd = sqlalchemy.text(\n",
+ " 'CREATE TABLE netflix_titles ( \\\n",
+ " show_id VARCHAR, \\\n",
+ " type VARCHAR, \\\n",
+ " title VARCHAR, \\\n",
+ " director VARCHAR, \\\n",
+ " \"cast\" VARCHAR, \\\n",
+ " country VARCHAR, \\\n",
+ " date_added VARCHAR, \\\n",
+ " release_year INTEGER, \\\n",
+ " rating VARCHAR, \\\n",
+ " duration VARCHAR, \\\n",
+ " listed_in VARCHAR, \\\n",
+ " description TEXT \\\n",
+ " )',\n",
+ ")\n",
+ "\n",
+ "netflix_data = \"/content/netflix_titles.csv\"\n",
+ "\n",
+ "df = pd.read_csv(netflix_data)\n",
+ "insert_data_cmd = sqlalchemy.text(\n",
+ " \"\"\"\n",
+ " INSERT INTO netflix_titles VALUES (:show_id, :type, :title, :director,\n",
+ " :cast, :country, :date_added, :release_year, :rating,\n",
+ " :duration, :listed_in, :description)\n",
+ " \"\"\"\n",
+ ")\n",
+ "\n",
+ "parameter_map = [\n",
+ " {\n",
+ " \"show_id\": row[\"show_id\"],\n",
+ " \"type\": row[\"type\"],\n",
+ " \"title\": row[\"title\"],\n",
+ " \"director\": row[\"director\"],\n",
+ " \"cast\": row[\"cast\"],\n",
+ " \"country\": row[\"country\"],\n",
+ " \"date_added\": row[\"date_added\"],\n",
+ " \"release_year\": row[\"release_year\"],\n",
+ " \"rating\": row[\"rating\"],\n",
+ " \"duration\": row[\"duration\"],\n",
+ " \"listed_in\": row[\"listed_in\"],\n",
+ " \"description\": row[\"description\"],\n",
+ " }\n",
+ " for index, row in df.head(100).iterrows() # limit to 100 rows\n",
+ "]\n",
+ "\n",
+ "with pool.connect() as db_conn:\n",
+ " db_conn.execute(create_table_cmd)\n",
+ " db_conn.execute(\n",
+ " insert_data_cmd,\n",
+ " parameter_map,\n",
+ " )\n",
+ " db_conn.commit()\n",
+ "connector.close()"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "#### Set LLM and embedding models globally for Llama Index components"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from llama_index.core import Settings\n",
+ "from llama_index.embeddings.vertex import VertexTextEmbedding\n",
+ "from llama_index.llms.vertex import Vertex\n",
+ "import google.auth\n",
+ "\n",
+ "credentials, _ = google.auth.default()\n",
+ "\n",
+ "Settings.embed_model = VertexTextEmbedding(\n",
+ " model_name=\"text-embedding-005\", project=project_id, credentials=credentials\n",
+ ")\n",
+ "Settings.llm = Vertex(model=\"gemini-2.0-flash-001\", project=project_id)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "SsGS80H04bDN"
+ },
+ "source": [
+ "## Use case 1: Cloud SQL for Postgres as a Document Reader\n",
+ "\n",
+ "---\n",
+ "\n",
+ "\n",
+ "\n",
+ "Now that you have data in your database, you are ready to use Cloud SQL for PostgreSQL as a [Document Reader](https://docs.llamaindex.ai/en/stable/module_guides/loading/connector/). This means you will pull data from the database and load it into memory as documents. These documents can be used to create a Vector Store."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "-CQgPON8dwSK"
+ },
+ "source": [
+ "First, create a connection to your Cloud SQL for PostgreSQL instance using the `PostgresEngine` class."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "zrwTsWHMkQ_v"
+ },
+ "outputs": [],
+ "source": [
+ "from llama_index_cloud_sql_pg import PostgresEngine\n",
+ "\n",
+ "engine = PostgresEngine.from_instance(\n",
+ " project_id=project_id,\n",
+ " instance=instance_name,\n",
+ " region=region,\n",
+ " database=database_name,\n",
+ " user=user,\n",
+ " password=password,\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "8s-C0P-Oee69"
+ },
+ "source": [
+ "The `PostgresReader` requires an `PostgresEngine` object to define the database connection and a `table_name` to define which data is to be retrieved. The `content_columns` argument can be used to define the columns that will be used as \"content\" in the document object we will later construct. The rest of the columns in that table will become the \"metadata\" associated with the documents."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "2SdFJT6Vece1"
+ },
+ "outputs": [],
+ "source": [
+ "from llama_index_cloud_sql_pg import PostgresReader\n",
+ "\n",
+ "\n",
+ "table_name = \"netflix_titles\"\n",
+ "content_columns = [\"title\", \"director\", \"cast\", \"description\"]\n",
+ "metadata_columns = [\n",
+ " \"show_id\",\n",
+ " \"type\",\n",
+ " \"country\",\n",
+ " \"date_added\",\n",
+ " \"release_year\",\n",
+ " \"rating\",\n",
+ " \"duration\",\n",
+ " \"listed_in\",\n",
+ "]\n",
+ "reader = PostgresReader.create_sync(\n",
+ " engine=engine,\n",
+ " table_name=table_name,\n",
+ " content_columns=content_columns,\n",
+ " metadata_columns=metadata_columns,\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "dsL-KFrmfuS1"
+ },
+ "source": [
+ "Use method `load_data()` to pull documents from out database. You can see the documents from the database here."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "t4zTx-HLfwmW"
+ },
+ "outputs": [],
+ "source": [
+ "documents = reader.load_data()\n",
+ "print(f\"Loaded {len(documents)} from the database. Examples:\")\n",
+ "for doc in documents[:5]:\n",
+ " print(doc)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Nice, you just used Cloud SQL for Postgres as a Document Reader!"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "z9uLV3bs4noo"
+ },
+ "source": [
+ "## Use case 2: Cloud SQL for PostgreSQL as Vector Store\n",
+ "\n",
+ "---\n",
+ "\n",
+ "\n",
+ "Now, you will learn how to put all of the documents into a [Vector Store](https://docs.llamaindex.ai/en/stable/module_guides/storing/vector_stores/) so that you perform a vector search."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "jfH8oQJ945Ko"
+ },
+ "source": [
+ "### Create Your Vector Store table\n",
+ "\n",
+ "Create a Vector Store table that can preserve the Document's metadata by using the method `init_vector_store_table` and defining specific metadata columns. The vector size is required. The example shows the vector size, `768`, that corresponds with the length of the vectors computed by the model our embeddings service uses, Vertex AI's `text-embedding-005`. "
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "e_rmjywG47pv"
+ },
+ "outputs": [],
+ "source": [
+ "from llama_index_cloud_sql_pg import Column\n",
+ "\n",
+ "sample_vector_table_name = \"movie_vector_table_samples\"\n",
+ "engine.init_vector_store_table(\n",
+ " sample_vector_table_name,\n",
+ " vector_size=768,\n",
+ " metadata_columns=[\n",
+ " Column(\"show_id\", \"VARCHAR\", nullable=True),\n",
+ " Column(\"type\", \"VARCHAR\", nullable=True),\n",
+ " Column(\"country\", \"VARCHAR\", nullable=True),\n",
+ " Column(\"date_added\", \"VARCHAR\", nullable=True),\n",
+ " Column(\"release_year\", \"INTEGER\", nullable=True),\n",
+ " Column(\"rating\", \"VARCHAR\", nullable=True),\n",
+ " Column(\"duration\", \"VARCHAR\", nullable=True),\n",
+ " Column(\"listed_in\", \"VARCHAR\", nullable=True),\n",
+ " ],\n",
+ " overwrite_existing=True, # Enabling this will delete and recreate the table if exists.\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "KG6rwEuJLNIo"
+ },
+ "source": [
+ "### Create the Vector Store instance\n",
+ "\n",
+ "Next, you will create a `PostgresVectorStore` object that connects to the new Cloud SQL database table to store the data from the documents."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "Wo4-7EYCIFF9"
+ },
+ "outputs": [],
+ "source": [
+ "from llama_index_cloud_sql_pg import PostgresVectorStore\n",
+ "\n",
+ "vector_store = PostgresVectorStore.create_sync(\n",
+ " engine=engine,\n",
+ " table_name=sample_vector_table_name,\n",
+ " metadata_columns=[\n",
+ " \"show_id\",\n",
+ " \"type\",\n",
+ " \"country\",\n",
+ " \"date_added\",\n",
+ " \"release_year\",\n",
+ " \"rating\",\n",
+ " \"duration\",\n",
+ " \"listed_in\",\n",
+ " ],\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "fr1rP6KQ-8ag"
+ },
+ "source": [
+ "#### Create a LlamaIndex `Index`\n",
+ "\n",
+ "An `Index` allows us to quickly retrieve relevant context for a user query. They are used to build `QueryEngines` and `ChatEngines` over which a user can get answers to their queries.\n",
+ "For a list of indexes that can be built in LlamaIndex, see [Index Guide](https://docs.llamaindex.ai/en/stable/module_guides/indexing/index_guide/).\n",
+ "\n",
+ "A `VectorStoreIndex`, can be built using the `PostgresVectorStore`. You can also use the `PostgresDocumentStore` and `PostgresIndexStore` to persist documents and index metadata. These modules can be used to build other `Indexes`.\n",
+ "For a detailed python notebook on this, see [LlamaIndex Document Store Guide](https://github.com/googleapis/llama-index-cloud-sql-pg-python/blob/main/samples/llama_index_doc_store.ipynb).\n",
+ "\n",
+ "Now, add the documents data into the vector table. Here is a code example to load the first 100 documents in the list."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "CTks8Cy--93B"
+ },
+ "outputs": [],
+ "source": [
+ "from llama_index.core import StorageContext, VectorStoreIndex\n",
+ "\n",
+ "storage_context = StorageContext.from_defaults(vector_store=vector_store)\n",
+ "index = VectorStoreIndex.from_documents(\n",
+ " documents, storage_context=storage_context, show_progress=True\n",
+ ")\n",
+ "\n",
+ "# If you have more documents, you can still add them to the database Vector Store table using the embedding service to create embeddings for each record\n",
+ "# vector_store.add_documents(docs_to_load)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "#### Query over indexed data\n",
+ "\n",
+ "A query engine takes in a natural language query and returns a rich response. It is built on one of the indexes, see [Query Engine Guide](https://docs.llamaindex.ai/en/stable/module_guides/deploying/query_engine/).\n",
+ "\n",
+ "You can compose multiple query engines to achieve more advanced querying, see [Query Engine usage patterns](https://docs.llamaindex.ai/en/stable/module_guides/deploying/query_engine/usage_pattern/)."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "query_engine = index.as_query_engine()\n",
+ "query = query_engine.query(\"List shows that are about teenagers\")\n",
+ "query.response"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# Use case 3: Cloud SQL for PostgreSQL as a Document Store with a Index Store\n",
+ "\n",
+ "---\n",
+ "\n",
+ "\n",
+ "Llama Index breaks down documents into smaller units called nodes, storing them in a [Document Store](https://docs.llamaindex.ai/en/stable/module_guides/storing/docstores/). The Document Store can be used with multiple indexes where each index is stored in its own [Index Store](https://docs.llamaindex.ai/en/stable/module_guides/storing/index_stores/) and uses the same underlying nodes but can provide a different search capability."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Set up a Document Store"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from llama_index_cloud_sql_pg import PostgresDocumentStore\n",
+ "\n",
+ "document_store_table_name = \"document_store\"\n",
+ "engine.init_doc_store_table(table_name=document_store_table_name)\n",
+ "doc_store = PostgresDocumentStore.create_sync(\n",
+ " engine=engine,\n",
+ " table_name=document_store_table_name,\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "#### Parse documents into nodes\n",
+ "\n",
+ "Using a `TokenTextSplitter`, you can split the movie names on whitespace characters for building a keyword based search index."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from llama_index.core.node_parser import TokenTextSplitter\n",
+ "\n",
+ "splitter = TokenTextSplitter(\n",
+ " chunk_size=1024,\n",
+ " chunk_overlap=20,\n",
+ " separator=\" \",\n",
+ ")\n",
+ "nodes = splitter.get_nodes_from_documents(documents)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "#### Add nodes to Document Store"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from llama_index.core import StorageContext\n",
+ "\n",
+ "storage_context = StorageContext.from_defaults(docstore=doc_store)\n",
+ "storage_context.docstore.add_documents(nodes)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Set up an Index Store"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from llama_index_cloud_sql_pg import PostgresIndexStore\n",
+ "\n",
+ "index_store_table_name = \"index_store\"\n",
+ "engine.init_index_store_table(\n",
+ " table_name=index_store_table_name,\n",
+ ")\n",
+ "\n",
+ "index_store = PostgresIndexStore.create_sync(\n",
+ " engine=engine,\n",
+ " table_name=index_store_table_name,\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "#### Create a Storage Context"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from llama_index.core import StorageContext\n",
+ "\n",
+ "storage_context = StorageContext.from_defaults(\n",
+ " docstore=doc_store, index_store=index_store\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Create Indexes\n",
+ "\n",
+ "The Document Store can be used with multiple indexes. Each index uses the same underlying nodes. For example, the keyword table index extracts keywords from each Node and builds a mapping to all nodes containing that keyword. Let's use this to build a keyword search."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from llama_index.core import SimpleKeywordTableIndex\n",
+ "\n",
+ "keyword_table_index = SimpleKeywordTableIndex(nodes, storage_context=storage_context)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Query the index"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "query_engine = keyword_table_index.as_query_engine()\n",
+ "response = query_engine.query(\"What tv shows resonate with crime?\")\n",
+ "print(response)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "ZM_OFzZrQEPs"
+ },
+ "source": [
+ "# Use case 4: Cloud SQL for PostgreSQL as Chat Store\n",
+ "\n",
+ "---\n",
+ "\n",
+ "\n",
+ "Next, create a [Chat Store](https://docs.llamaindex.ai/en/stable/module_guides/storing/chat_stores/) so the LLM can retain context and information across multiple interactions, leading to more coherent and sophisticated conversations or text generation."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "vyYQILyoEAqg"
+ },
+ "outputs": [],
+ "source": [
+ "from llama_index_cloud_sql_pg import PostgresChatStore\n",
+ "\n",
+ "chat_store_table_name = \"chat_store\"\n",
+ "engine.init_chat_store_table(table_name=chat_store_table_name)\n",
+ "chat_store = PostgresChatStore.create_sync(\n",
+ " engine,\n",
+ " table_name=chat_store_table_name,\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Create a ChatMemoryBuffer\n",
+ "The `ChatMemoryBuffer` stores a history of recent chat messages, enabling the LLM to access relevant context from prior interactions.\n",
+ "\n",
+ "By providing our Chat Store into the `ChatMemoryBuffer`, it can automatically retrieve and update messages associated with a specific session ID or `chat_store_key`."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from llama_index.core.memory import ChatMemoryBuffer\n",
+ "\n",
+ "memory = ChatMemoryBuffer.from_defaults(\n",
+ " token_limit=3000,\n",
+ " chat_store=chat_store,\n",
+ " chat_store_key=\"user1\",\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Create a LlamaIndex `ChatEngine`\n",
+ "\n",
+ "You can re-use the `VectorStoreIndex` created above to create a [ChatEngine](https://docs.llamaindex.ai/en/stable/module_guides/deploying/chat_engines/) which includes a `ChatStore` to save the chats between the user and AI assistant."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "2yuXYLTCl2K1"
+ },
+ "source": [
+ "Here is an example of how you would add a chat message and fetch all messages from the Chat Store."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "qDVoTWZal0ZF"
+ },
+ "outputs": [],
+ "source": [
+ "chat_engine = index.as_chat_engine(chat_mode=\"context\", memory=memory)\n",
+ "\n",
+ "op = chat_engine.chat(\"What was the cast in Blood and Water?\")\n",
+ "op = chat_engine.chat(\"How many seasons are there for Kota Factory?\")\n",
+ "# Retrieve response for a chat message\n",
+ "print(op.response)\n",
+ "\n",
+ "# Retrieve all messages for a user / session\n",
+ "retrieved_messages = chat_store.get_messages(\"user1\")\n",
+ "for msg in retrieved_messages:\n",
+ " print(f\"{msg.role} -> {msg.content}\")"
+ ]
+ }
+ ],
+ "metadata": {
+ "colab": {
+ "provenance": []
+ },
+ "kernelspec": {
+ "display_name": "Python 3",
+ "name": "python3"
+ },
+ "language_info": {
+ "name": "python"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}
diff --git a/samples/llama_index_reader.ipynb b/samples/llama_index_reader.ipynb
new file mode 100644
index 0000000..68b6cf5
--- /dev/null
+++ b/samples/llama_index_reader.ipynb
@@ -0,0 +1,380 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# Google Cloud SQL for PostgreSQL - `PostgresReader`\n",
+ "\n",
+ "> [Cloud SQL](https://cloud.google.com/sql) is a fully managed relational database service that offers high performance, seamless integration, and impressive scalability. It offers MySQL, PostgreSQL, and SQL Server database engines. Extend your database application to build AI-powered experiences leveraging Cloud SQL's LlamaIndex integrations.\n",
+ "\n",
+ "This notebook goes over how to use `Cloud SQL for PostgreSQL` to retrieve data as documents with the `PostgresReader` class.\n",
+ "\n",
+ "Learn more about the package on [GitHub](https://github.com/googleapis/llama-index-cloud-sql-pg-python/).\n",
+ "\n",
+ "[](https://colab.research.google.com/github/googleapis/llama-index-cloud-sql-pg-python/blob/main/samples/llama_index_doc_store.ipynb)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## Before you begin\n",
+ "\n",
+ "To run this notebook, you will need to do the following:\n",
+ "\n",
+ " * [Create a Google Cloud Project](https://developers.google.com/workspace/guides/create-project)\n",
+ " * [Enable the Cloud SQL Admin API.](https://console.cloud.google.com/flows/enableapi?apiid=sqladmin.googleapis.com)\n",
+ " * [Create a Cloud SQL instance.](https://cloud.google.com/sql/docs/postgres/connect-instance-auth-proxy#create-instance)\n",
+ " * [Create a Cloud SQL database.](https://cloud.google.com/sql/docs/postgres/create-manage-databases)\n",
+ " * [Add a User to the database.](https://cloud.google.com/sql/docs/postgres/create-manage-users)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### 🦙 Library Installation\n",
+ "Install the integration library, `llama-index-cloud-sql-pg`."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "**Colab only:** Uncomment the following cell to restart the kernel or use the button to restart the kernel. For Vertex AI Workbench you can restart the terminal using the button on top."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# # Automatically restart kernel after installs so that your environment can access the new packages\n",
+ "# import IPython\n",
+ "\n",
+ "# app = IPython.Application.instance()\n",
+ "# app.kernel.do_shutdown(True)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### 🔐 Authentication\n",
+ "Authenticate to Google Cloud as the IAM user logged into this notebook in order to access your Google Cloud Project.\n",
+ "\n",
+ "* If you are using Colab to run this notebook, use the cell below and continue.\n",
+ "* If you are using Vertex AI Workbench, check out the setup instructions [here](https://github.com/GoogleCloudPlatform/generative-ai/tree/main/setup-env)."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from google.colab import auth\n",
+ "\n",
+ "auth.authenticate_user()"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### ☁ Set Your Google Cloud Project\n",
+ "Set your Google Cloud project so that you can leverage Google Cloud resources within this notebook.\n",
+ "\n",
+ "If you don't know your project ID, try the following:\n",
+ "\n",
+ "* Run `gcloud config list`.\n",
+ "* Run `gcloud projects list`.\n",
+ "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# @markdown Please fill in the value below with your Google Cloud project ID and then run the cell.\n",
+ "\n",
+ "PROJECT_ID = \"my-project-id\" # @param {type:\"string\"}\n",
+ "\n",
+ "# Set the project id\n",
+ "!gcloud config set project {PROJECT_ID}"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## Basic Usage"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Set Cloud SQL database values\n",
+ "Find your database values, in the [Cloud SQL Instances page](https://console.cloud.google.com/sql?_ga=2.223735448.2062268965.1707700487-2088871159.1707257687)."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# @title Set Your Values Here { display-mode: \"form\" }\n",
+ "REGION = \"us-central1\" # @param {type: \"string\"}\n",
+ "INSTANCE = \"my-primary\" # @param {type: \"string\"}\n",
+ "DATABASE = \"my-database\" # @param {type: \"string\"}\n",
+ "TABLE_NAME = \"reader_table\" # @param {type: \"string\"}\n",
+ "USER = \"postgres\" # @param {type: \"string\"}\n",
+ "PASSWORD = \"my-password\" # @param {type: \"string\"}"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### PostgresEngine Connection Pool\n",
+ "\n",
+ "One of the requirements and arguments to establish Cloud SQL as a reader is a `PostgresEngine` object. The `PostgresEngine` configures a connection pool to your Cloud SQL database, enabling successful connections from your application and following industry best practices.\n",
+ "\n",
+ "To create a `PostgresEngine` using `PostgresEngine.from_instance()` you need to provide only 4 things:\n",
+ "\n",
+ "1. `project_id` : Project ID of the Google Cloud Project where the Cloud SQL instance is located.\n",
+ "1. `region` : Region where the Cloud SQL instance is located.\n",
+ "1. `instance` : The name of the Cloud SQL instance.\n",
+ "1. `database` : The name of the database to connect to on the Cloud SQL instance.\n",
+ "\n",
+ "By default, [IAM database authentication](https://cloud.google.com/sql/docs/postgres/iam-authentication#iam-db-auth) will be used as the method of database authentication. This library uses the IAM principal belonging to the [Application Default Credentials (ADC)](https://cloud.google.com/docs/authentication/application-default-credentials) sourced from the envionment.\n",
+ "\n",
+ "For more informatin on IAM database authentication please see:\n",
+ "\n",
+ "* [Configure an instance for IAM database authentication](https://cloud.google.com/sql/docs/postgres/create-edit-iam-instances)\n",
+ "* [Manage users with IAM database authentication](https://cloud.google.com/sql/docs/postgres/add-manage-iam-users)\n",
+ "\n",
+ "Optionally, [built-in database authentication](https://cloud.google.com/sql/docs/postgres/built-in-authentication) using a username and password to access the Cloud SQL database can also be used. Just provide the optional `user` and `password` arguments to `PostgresEngine.from_instance()`:\n",
+ "\n",
+ "* `user` : Database user to use for built-in database authentication and login\n",
+ "* `password` : Database password to use for built-in database authentication and login.\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "**Note:** This tutorial demonstrates the async interface. All async methods have corresponding sync methods."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from llama_index_cloud_sql_pg import PostgresEngine\n",
+ "\n",
+ "engine = await PostgresEngine.afrom_instance(\n",
+ " project_id=PROJECT_ID,\n",
+ " region=REGION,\n",
+ " instance=INSTANCE,\n",
+ " database=DATABASE,\n",
+ " user=USER,\n",
+ " password=PASSWORD,\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Create PostgresReader"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "When creating an `PostgresReader` for fetching data from Cloud SQL Postgres, you have two main options to specify the data you want to load:\n",
+ "* using the table_name argument - When you specify the table_name argument, you're telling the reader to fetch all the data from the given table.\n",
+ "* using the query argument - When you specify the query argument, you can provide a custom SQL query to fetch the data. This allows you to have full control over the SQL query, including selecting specific columns, applying filters, sorting, joining tables, etc.\n",
+ "\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Load Documents using the `table_name` argument"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "#### Load Documents via default table\n",
+ "The reader returns a list of Documents from the table using the first column as text and all other columns as metadata. The default table will have the first column as\n",
+ "text and the second column as metadata (JSON). Each row becomes a document."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from llama_index_cloud_sql_pg import PostgresReader\n",
+ "\n",
+ "# Creating a basic PostgresReader object\n",
+ "reader = await PostgresReader.create(\n",
+ " engine,\n",
+ " table_name=TABLE_NAME,\n",
+ " # schema_name=SCHEMA_NAME,\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "#### Load documents via custom table/metadata or custom page content columns"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "reader = await PostgresReader.create(\n",
+ " engine,\n",
+ " table_name=TABLE_NAME,\n",
+ " # schema_name=SCHEMA_NAME,\n",
+ " content_columns=[\"product_name\"], # Optional\n",
+ " metadata_columns=[\"id\"], # Optional\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Load Documents using a SQL query\n",
+ "The query parameter allows users to specify a custom SQL query which can include filters to load specific documents from a database."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "table_name = \"products\"\n",
+ "content_columns = [\"product_name\", \"description\"]\n",
+ "metadata_columns = [\"id\", \"content\"]\n",
+ "\n",
+ "reader = PostgresReader.create(\n",
+ " engine=engine,\n",
+ " query=f\"SELECT * FROM {table_name};\",\n",
+ " content_columns=content_columns,\n",
+ " metadata_columns=metadata_columns,\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "**Note**: If the `content_columns` and `metadata_columns` are not specified, the reader will automatically treat the first returned column as the document’s `text` and all subsequent columns as `metadata`."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Set page content format\n",
+ "The reader returns a list of Documents, with one document per row, with page content in specified string format, i.e. text (space separated concatenation), JSON, YAML, CSV, etc. JSON and YAML formats include headers, while text and CSV do not include field headers."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "reader = await PostgresReader.create(\n",
+ " engine,\n",
+ " table_name=TABLE_NAME,\n",
+ " # schema_name=SCHEMA_NAME,\n",
+ " content_columns=[\"product_name\", \"description\"],\n",
+ " format=\"YAML\",\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Load the documents\n",
+ "\n",
+ "You can choose to load the documents in two ways:\n",
+ "* Load all the data at once\n",
+ "* Lazy load data"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "#### Load data all at once"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "docs = await reader.aload_data()\n",
+ "\n",
+ "print(docs)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "#### Lazy Load the data"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "docs_iterable = reader.alazy_load_data()\n",
+ "\n",
+ "docs = []\n",
+ "async for doc in docs_iterable:\n",
+ " docs.append(doc)\n",
+ "\n",
+ "print(docs)"
+ ]
+ }
+ ],
+ "metadata": {
+ "language_info": {
+ "name": "python"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/src/llama_index_cloud_sql_pg/async_chat_store.py b/src/llama_index_cloud_sql_pg/async_chat_store.py
index 8d80543..eb1b569 100644
--- a/src/llama_index_cloud_sql_pg/async_chat_store.py
+++ b/src/llama_index_cloud_sql_pg/async_chat_store.py
@@ -137,7 +137,7 @@ async def aset_messages(self, key: str, messages: List[ChatMessage]) -> None:
params = [
{
"key": key,
- "message": json.dumps(message.dict()),
+ "message": json.dumps(message.model_dump()),
}
for message in messages
]
@@ -175,7 +175,7 @@ async def async_add_message(self, key: str, message: ChatMessage) -> None:
insert_query = f"""
INSERT INTO "{self._schema_name}"."{self._table_name}" (key, message)
VALUES (:key, :message);"""
- params = {"key": key, "message": json.dumps(message.dict())}
+ params = {"key": key, "message": json.dumps(message.model_dump())}
await self.__aexecute_query(insert_query, params)
@@ -261,35 +261,35 @@ async def aget_keys(self) -> List[str]:
def set_messages(self, key: str, messages: List[ChatMessage]) -> None:
raise NotImplementedError(
- "Sync methods are not implemented for AsyncPostgresChatStore . Use PostgresChatStore interface instead."
+ "Sync methods are not implemented for AsyncPostgresChatStore. Use PostgresChatStore interface instead."
)
def get_messages(self, key: str) -> List[ChatMessage]:
raise NotImplementedError(
- "Sync methods are not implemented for AsyncPostgresChatStore . Use PostgresChatStore interface instead."
+ "Sync methods are not implemented for AsyncPostgresChatStore. Use PostgresChatStore interface instead."
)
def add_message(self, key: str, message: ChatMessage) -> None:
raise NotImplementedError(
- "Sync methods are not implemented for AsyncPostgresChatStore . Use PostgresChatStore interface instead."
+ "Sync methods are not implemented for AsyncPostgresChatStore. Use PostgresChatStore interface instead."
)
def delete_messages(self, key: str) -> Optional[List[ChatMessage]]:
raise NotImplementedError(
- "Sync methods are not implemented for AsyncPostgresChatStore . Use PostgresChatStore interface instead."
+ "Sync methods are not implemented for AsyncPostgresChatStore. Use PostgresChatStore interface instead."
)
def delete_message(self, key: str, idx: int) -> Optional[ChatMessage]:
raise NotImplementedError(
- "Sync methods are not implemented for AsyncPostgresChatStore . Use PostgresChatStore interface instead."
+ "Sync methods are not implemented for AsyncPostgresChatStore. Use PostgresChatStore interface instead."
)
def delete_last_message(self, key: str) -> Optional[ChatMessage]:
raise NotImplementedError(
- "Sync methods are not implemented for AsyncPostgresChatStore . Use PostgresChatStore interface instead."
+ "Sync methods are not implemented for AsyncPostgresChatStore. Use PostgresChatStore interface instead."
)
def get_keys(self) -> List[str]:
raise NotImplementedError(
- "Sync methods are not implemented for AsyncPostgresChatStore . Use PostgresChatStore interface instead."
+ "Sync methods are not implemented for AsyncPostgresChatStore. Use PostgresChatStore interface instead."
)
diff --git a/src/llama_index_cloud_sql_pg/async_index_store.py b/src/llama_index_cloud_sql_pg/async_index_store.py
index ee06c53..76a5917 100644
--- a/src/llama_index_cloud_sql_pg/async_index_store.py
+++ b/src/llama_index_cloud_sql_pg/async_index_store.py
@@ -191,22 +191,22 @@ async def aget_index_struct(
def index_structs(self) -> list[IndexStruct]:
raise NotImplementedError(
- "Sync methods are not implemented for AsyncPostgresIndexStore . Use PostgresIndexStore interface instead."
+ "Sync methods are not implemented for AsyncPostgresIndexStore. Use PostgresIndexStore interface instead."
)
def add_index_struct(self, index_struct: IndexStruct) -> None:
raise NotImplementedError(
- "Sync methods are not implemented for AsyncPostgresIndexStore . Use PostgresIndexStore interface instead."
+ "Sync methods are not implemented for AsyncPostgresIndexStore. Use PostgresIndexStore interface instead."
)
def delete_index_struct(self, key: str) -> None:
raise NotImplementedError(
- "Sync methods are not implemented for AsyncPostgresIndexStore . Use PostgresIndexStore interface instead."
+ "Sync methods are not implemented for AsyncPostgresIndexStore. Use PostgresIndexStore interface instead."
)
def get_index_struct(
self, struct_id: Optional[str] = None
) -> Optional[IndexStruct]:
raise NotImplementedError(
- "Sync methods are not implemented for AsyncPostgresIndexStore . Use PostgresIndexStore interface instead."
+ "Sync methods are not implemented for AsyncPostgresIndexStore. Use PostgresIndexStore interface instead."
)
diff --git a/src/llama_index_cloud_sql_pg/document_store.py b/src/llama_index_cloud_sql_pg/document_store.py
index f4ff3db..40aa83d 100644
--- a/src/llama_index_cloud_sql_pg/document_store.py
+++ b/src/llama_index_cloud_sql_pg/document_store.py
@@ -346,11 +346,7 @@ def get_document_hash(self, doc_id: str) -> Optional[str]:
async def aget_all_document_hashes(self) -> dict[str, str]:
"""Get the stored hash for all documents.
- Returns:
- dict[
- str, # doc_hash
- str # doc_id
- ]
+ Returns: A dictionary mapping document hashes to document IDs.
"""
return await self._engine._run_as_async(
self.__document_store.aget_all_document_hashes()
@@ -359,11 +355,7 @@ async def aget_all_document_hashes(self) -> dict[str, str]:
def get_all_document_hashes(self) -> dict[str, str]:
"""Get the stored hash for all documents.
- Returns:
- dict[
- str, # doc_hash
- str # doc_id
- ]
+ Returns: A dictionary mapping document hashes to document IDs.
"""
return self._engine._run_as_sync(
self.__document_store.aget_all_document_hashes()
@@ -372,13 +364,7 @@ def get_all_document_hashes(self) -> dict[str, str]:
async def aget_all_ref_doc_info(self) -> Optional[dict[str, RefDocInfo]]:
"""Get a mapping of ref_doc_id -> RefDocInfo for all ingested documents.
- Returns:
- Optional[
- dict[
- str, #Ref_doc_id
- RefDocInfo, #Ref_doc_info of the id
- ]
- ]
+ Returns: A dictionary mapping ref_doc_ids to `RefDocInfo` objects, or None if no documents have been ingested.
"""
return await self._engine._run_as_async(
self.__document_store.aget_all_ref_doc_info()
@@ -387,13 +373,7 @@ async def aget_all_ref_doc_info(self) -> Optional[dict[str, RefDocInfo]]:
def get_all_ref_doc_info(self) -> Optional[dict[str, RefDocInfo]]:
"""Get a mapping of ref_doc_id -> RefDocInfo for all ingested documents.
- Returns:
- Optional[
- dict[
- str, #Ref_doc_id
- RefDocInfo, #Ref_doc_info of the id
- ]
- ]
+ Returns: A dictionary mapping ref_doc_ids to `RefDocInfo` objects, or None if no documents have been ingested.
"""
return self._engine._run_as_sync(self.__document_store.aget_all_ref_doc_info())
diff --git a/src/llama_index_cloud_sql_pg/engine.py b/src/llama_index_cloud_sql_pg/engine.py
index 2faa943..0ec0cf8 100644
--- a/src/llama_index_cloud_sql_pg/engine.py
+++ b/src/llama_index_cloud_sql_pg/engine.py
@@ -762,16 +762,16 @@ async def _ainit_chat_store_table(
schema_name: str = "public",
overwrite_existing: bool = False,
) -> None:
- """
- Create an table to save chat store.
+ """Create a table to save chat store.
+
Args:
table_name (str): The table name to store chat history.
schema_name (str): The schema name to store the chat store table.
Default: "public".
overwrite_existing (bool): Whether to drop existing table.
Default: False.
- Returns:
- None
+
+ Returns: None
"""
if overwrite_existing:
async with self._pool.connect() as conn:
@@ -797,16 +797,15 @@ async def ainit_chat_store_table(
schema_name: str = "public",
overwrite_existing: bool = False,
) -> None:
- """
- Create an table to save chat store.
+ """Create a table to save chat store.
+
Args:
table_name (str): The table name to store chat store.
schema_name (str): The schema name to store the chat store table.
Default: "public".
overwrite_existing (bool): Whether to drop existing table.
Default: False.
- Returns:
- None
+ Returns: None
"""
await self._run_as_async(
self._ainit_chat_store_table(
@@ -822,16 +821,16 @@ def init_chat_store_table(
schema_name: str = "public",
overwrite_existing: bool = False,
) -> None:
- """
- Create an table to save chat store.
+ """Create a table to save chat store.
+
Args:
table_name (str): The table name to store chat store.
schema_name (str): The schema name to store the chat store table.
Default: "public".
overwrite_existing (bool): Whether to drop existing table.
Default: False.
- Returns:
- None
+
+ Returns: None
"""
self._run_as_sync(
self._ainit_chat_store_table(
diff --git a/src/llama_index_cloud_sql_pg/version.py b/src/llama_index_cloud_sql_pg/version.py
index 20c5861..fffa9d9 100644
--- a/src/llama_index_cloud_sql_pg/version.py
+++ b/src/llama_index_cloud_sql_pg/version.py
@@ -12,4 +12,4 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-__version__ = "0.2.0"
+__version__ = "0.2.1"
diff --git a/tests/test_async_chat_store.py b/tests/test_async_chat_store.py
index bdaf13f..ef34fd1 100644
--- a/tests/test_async_chat_store.py
+++ b/tests/test_async_chat_store.py
@@ -25,6 +25,7 @@
from llama_index_cloud_sql_pg.async_chat_store import AsyncPostgresChatStore
default_table_name_async = "chat_store_" + str(uuid.uuid4())
+sync_method_exception_str = "Sync methods are not implemented for AsyncPostgresChatStore. Use PostgresChatStore interface instead."
async def aexecute(engine: PostgresEngine, query: str) -> None:
@@ -122,7 +123,7 @@ async def test_async_add_message(self, async_engine, chat_store):
query = f"""select * from "public"."{default_table_name_async}" where key = '{key}';"""
results = await afetch(async_engine, query)
result = results[0]
- assert result["message"] == message.dict()
+ assert result["message"] == message.model_dump()
async def test_aset_and_aget_messages(self, chat_store):
message_1 = ChatMessage(content="First message", role="user")
@@ -160,7 +161,7 @@ async def test_adelete_message(self, async_engine, chat_store):
results = await afetch(async_engine, query)
assert len(results) == 1
- assert results[0]["message"] == message_1.dict()
+ assert results[0]["message"] == message_1.model_dump()
async def test_adelete_last_message(self, async_engine, chat_store):
message_1 = ChatMessage(content="Message 1", role="user")
@@ -175,8 +176,8 @@ async def test_adelete_last_message(self, async_engine, chat_store):
results = await afetch(async_engine, query)
assert len(results) == 2
- assert results[0]["message"] == message_1.dict()
- assert results[1]["message"] == message_2.dict()
+ assert results[0]["message"] == message_1.model_dump()
+ assert results[1]["message"] == message_2.model_dump()
async def test_aget_keys(self, async_engine, chat_store):
message_1 = [ChatMessage(content="First message", role="user")]
@@ -201,7 +202,7 @@ async def test_set_exisiting_key(self, async_engine, chat_store):
assert len(results) == 1
result = results[0]
- assert result["message"] == message_1[0].dict()
+ assert result["message"] == message_1[0].model_dump()
message_2 = ChatMessage(content="Second message", role="user")
message_3 = ChatMessage(content="Third message", role="user")
@@ -215,5 +216,33 @@ async def test_set_exisiting_key(self, async_engine, chat_store):
# Assert the previous messages are deleted and only the newest ones exist.
assert len(results) == 2
- assert results[0]["message"] == message_2.dict()
- assert results[1]["message"] == message_3.dict()
+ assert results[0]["message"] == message_2.model_dump()
+ assert results[1]["message"] == message_3.model_dump()
+
+ async def test_set_messages(self, chat_store):
+ with pytest.raises(Exception, match=sync_method_exception_str):
+ chat_store.set_messages("test_key", [])
+
+ async def test_get_messages(self, chat_store):
+ with pytest.raises(Exception, match=sync_method_exception_str):
+ chat_store.get_messages("test_key")
+
+ async def test_add_message(self, chat_store):
+ with pytest.raises(Exception, match=sync_method_exception_str):
+ chat_store.add_message("test_key", ChatMessage(content="test", role="user"))
+
+ async def test_delete_messages(self, chat_store):
+ with pytest.raises(Exception, match=sync_method_exception_str):
+ chat_store.delete_messages("test_key")
+
+ async def test_delete_message(self, chat_store):
+ with pytest.raises(Exception, match=sync_method_exception_str):
+ chat_store.delete_message("test_key", 0)
+
+ async def test_delete_last_message(self, chat_store):
+ with pytest.raises(Exception, match=sync_method_exception_str):
+ chat_store.delete_last_message("test_key")
+
+ async def test_get_keys(self, chat_store):
+ with pytest.raises(Exception, match=sync_method_exception_str):
+ chat_store.get_keys()
diff --git a/tests/test_async_index_store.py b/tests/test_async_index_store.py
index 0b7bbe8..a5cd11a 100644
--- a/tests/test_async_index_store.py
+++ b/tests/test_async_index_store.py
@@ -26,7 +26,7 @@
from llama_index_cloud_sql_pg.async_index_store import AsyncPostgresIndexStore
default_table_name_async = "index_store_" + str(uuid.uuid4())
-sync_method_exception_str = "Sync methods are not implemented for AsyncPostgresIndexStore . Use PostgresIndexStore interface instead."
+sync_method_exception_str = "Sync methods are not implemented for AsyncPostgresIndexStore. Use PostgresIndexStore interface instead."
async def aexecute(engine: PostgresEngine, query: str) -> None:
diff --git a/tests/test_chat_store.py b/tests/test_chat_store.py
index 694119b..896e066 100644
--- a/tests/test_chat_store.py
+++ b/tests/test_chat_store.py
@@ -124,7 +124,7 @@ async def test_async_add_message(self, async_engine, async_chat_store):
query = f"""select * from "public"."{default_table_name_async}" where key = '{key}';"""
results = await afetch(async_engine, query)
result = results[0]
- assert result["message"] == message.dict()
+ assert result["message"] == message.model_dump()
async def test_aset_and_aget_messages(self, async_chat_store):
message_1 = ChatMessage(content="First message", role="user")
@@ -162,7 +162,7 @@ async def test_adelete_message(self, async_engine, async_chat_store):
results = await afetch(async_engine, query)
assert len(results) == 1
- assert results[0]["message"] == message_1.dict()
+ assert results[0]["message"] == message_1.model_dump()
async def test_adelete_last_message(self, async_engine, async_chat_store):
message_1 = ChatMessage(content="Message 1", role="user")
@@ -177,8 +177,8 @@ async def test_adelete_last_message(self, async_engine, async_chat_store):
results = await afetch(async_engine, query)
assert len(results) == 2
- assert results[0]["message"] == message_1.dict()
- assert results[1]["message"] == message_2.dict()
+ assert results[0]["message"] == message_1.model_dump()
+ assert results[1]["message"] == message_2.model_dump()
async def test_aget_keys(self, async_engine, async_chat_store):
message_1 = [ChatMessage(content="First message", role="user")]
@@ -203,7 +203,7 @@ async def test_set_exisiting_key(self, async_engine, async_chat_store):
assert len(results) == 1
result = results[0]
- assert result["message"] == message_1[0].dict()
+ assert result["message"] == message_1[0].model_dump()
message_2 = ChatMessage(content="Second message", role="user")
message_3 = ChatMessage(content="Third message", role="user")
@@ -217,8 +217,8 @@ async def test_set_exisiting_key(self, async_engine, async_chat_store):
# Assert the previous messages are deleted and only the newest ones exist.
assert len(results) == 2
- assert results[0]["message"] == message_2.dict()
- assert results[1]["message"] == message_3.dict()
+ assert results[0]["message"] == message_2.model_dump()
+ assert results[1]["message"] == message_3.model_dump()
@pytest.mark.asyncio(loop_scope="class")
@@ -287,7 +287,7 @@ async def test_add_message(self, sync_engine, sync_chat_store):
query = f"""select * from "public"."{default_table_name_sync}" where key = '{key}';"""
results = await afetch(sync_engine, query)
result = results[0]
- assert result["message"] == message.dict()
+ assert result["message"] == message.model_dump()
async def test_set_and_get_messages(self, sync_chat_store):
message_1 = ChatMessage(content="First message", role="user")
@@ -325,7 +325,7 @@ async def test_delete_message(self, sync_engine, sync_chat_store):
results = await afetch(sync_engine, query)
assert len(results) == 1
- assert results[0]["message"] == message_1.dict()
+ assert results[0]["message"] == message_1.model_dump()
async def test_delete_last_message(self, sync_engine, sync_chat_store):
message_1 = ChatMessage(content="Message 1", role="user")
@@ -340,8 +340,8 @@ async def test_delete_last_message(self, sync_engine, sync_chat_store):
results = await afetch(sync_engine, query)
assert len(results) == 2
- assert results[0]["message"] == message_1.dict()
- assert results[1]["message"] == message_2.dict()
+ assert results[0]["message"] == message_1.model_dump()
+ assert results[1]["message"] == message_2.model_dump()
async def test_get_keys(self, sync_engine, sync_chat_store):
message_1 = [ChatMessage(content="First message", role="user")]
@@ -366,7 +366,7 @@ async def test_set_exisiting_key(self, sync_engine, sync_chat_store):
assert len(results) == 1
result = results[0]
- assert result["message"] == message_1[0].dict()
+ assert result["message"] == message_1[0].model_dump()
message_2 = ChatMessage(content="Second message", role="user")
message_3 = ChatMessage(content="Third message", role="user")
@@ -380,5 +380,5 @@ async def test_set_exisiting_key(self, sync_engine, sync_chat_store):
# Assert the previous messages are deleted and only the newest ones exist.
assert len(results) == 2
- assert results[0]["message"] == message_2.dict()
- assert results[1]["message"] == message_3.dict()
+ assert results[0]["message"] == message_2.model_dump()
+ assert results[1]["message"] == message_3.model_dump()