diff --git a/conftest.py b/conftest.py new file mode 100644 index 00000000..640cc46e --- /dev/null +++ b/conftest.py @@ -0,0 +1,62 @@ +"""Global pytest configuration for langchain-google-alloydb-pg-python.""" + +import os +import pytest + +def pytest_collection_modifyitems(config, items): + """Skip tests that require GCP environment variables if PROJECT_ID is not set.""" + if os.environ.get("PROJECT_ID"): + return + + skip_gcp = pytest.mark.skip(reason="Missing required GCP environment variables") + + # List of test files/modules that require live GCP / AlloyDB connection + gcp_test_files = { + "test_async_chatmessagehistory.py", + "test_async_checkpoint.py", + "test_async_loader.py", + "test_async_vectorstore.py", + "test_async_vectorstore_from_methods.py", + "test_async_vectorstore_index.py", + "test_async_vectorstore_search.py", + "test_chatmessagehistory.py", + "test_checkpoint.py", + "test_embeddings.py", + "test_engine.py", + "test_loader.py", + "test_model_manager.py", + "test_pgvector_migrator.py", + "test_standard_test_suite.py", + "test_vectorstore.py", + "test_vectorstore_embeddings.py", + "test_vectorstore_from_methods.py", + "test_vectorstore_index.py", + "test_vectorstore_search.py", + } + + for item in items: + # Check if the test is from one of the GCP test files (supporting pytest 8/9 item.path) + if hasattr(item, "path"): + fspath_name = item.path.name + else: + fspath_name = getattr(getattr(item, "fspath", None), "basename", "") + if fspath_name in gcp_test_files: + item.add_marker(skip_gcp) + # Also skip tests with 'integration' or 'live' in their name/nodeid, or requiring 'engine' fixture + nodeid = getattr(item, "nodeid", "") + if ( + "integration" in nodeid.lower() + or "live" in nodeid.lower() + ): + item.add_marker(skip_gcp) + elif "engine" in getattr(item, "fixturenames", []): + fixtureinfo = getattr(item, "_fixtureinfo", None) + is_local_fixture = False + if fixtureinfo: + fixturedefs = fixtureinfo.name2fixturedefs.get("engine", []) + if fixturedefs and "conftest" not in fixturedefs[-1].func.__module__: + is_local_fixture = True + + if not is_local_fixture: + item.add_marker(skip_gcp) + diff --git a/docs/vector_store.ipynb b/docs/vector_store.ipynb index dde2bae4..17f493d2 100644 --- a/docs/vector_store.ipynb +++ b/docs/vector_store.ipynb @@ -2,7 +2,6 @@ "cells": [ { "cell_type": "markdown", - "metadata": {}, "source": [ "# Google AlloyDB for PostgreSQL\n", "\n", @@ -13,11 +12,13 @@ "Learn more about the package on [GitHub](https://github.com/googleapis/langchain-google-alloydb-pg-python/).\n", "\n", "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/googleapis/langchain-google-alloydb-pg-python/blob/main/docs/vector_store.ipynb)" - ] + ], + "metadata": { + "id": "9587bfed" + } }, { "cell_type": "markdown", - "metadata": {}, "source": [ "## Before you begin\n", "\n", @@ -28,21 +29,26 @@ " * [Create a AlloyDB cluster and instance.](https://cloud.google.com/alloydb/docs/cluster-create)\n", " * [Create a AlloyDB database.](https://cloud.google.com/alloydb/docs/quickstart/create-and-connect)\n", " * [Add a User to the database.](https://cloud.google.com/alloydb/docs/database-users/about)" - ] + ], + "metadata": { + "id": "83aa04a4" + } }, { "cell_type": "markdown", - "metadata": { - "id": "IR54BmgvdHT_" - }, "source": [ "### 🦜🔗 Library Installation\n", "Install the integration library, `langchain-google-alloydb-pg`, and the library for the embedding service, `langchain-google-vertexai`." - ] + ], + "metadata": { + "id": "IR54BmgvdHT_" + } }, { "cell_type": "code", - "execution_count": null, + "source": [ + "%pip install --upgrade --quiet langchain-google-alloydb-pg langchain-google-vertexai" + ], "metadata": { "colab": { "base_uri": "https://localhost:8080/", @@ -51,70 +57,61 @@ "id": "0ZITIDE160OD", "outputId": "e184bc0d-6541-4e0a-82d2-1e216db00a2d" }, - "outputs": [], - "source": [ - "%pip install --upgrade --quiet langchain-google-alloydb-pg langchain-google-vertexai" - ] + "execution_count": null, + "outputs": [] }, { "cell_type": "markdown", - "metadata": { - "id": "v40bB_GMcr9f" - }, "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." - ] + ], + "metadata": { + "id": "v40bB_GMcr9f" + } }, { "cell_type": "code", - "execution_count": null, - "id": "v6jBDnYnNM08", - "metadata": { - "id": "v6jBDnYnNM08" - }, - "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)" - ] + ], + "metadata": { + "id": "v6jBDnYnNM08" + }, + "execution_count": null, + "outputs": [] }, { "cell_type": "markdown", - "id": "yygMe6rPWxHS", - "metadata": { - "id": "yygMe6rPWxHS" - }, "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)." - ] + ], + "metadata": { + "id": "yygMe6rPWxHS" + } }, { "cell_type": "code", - "execution_count": 1, - "id": "PTXN1_DSXj2b", - "metadata": { - "id": "PTXN1_DSXj2b" - }, - "outputs": [], "source": [ "from google.colab import auth\n", "\n", "auth.authenticate_user()" - ] + ], + "metadata": { + "id": "PTXN1_DSXj2b" + }, + "execution_count": 1, + "outputs": [] }, { "cell_type": "markdown", - "id": "NEvB9BoLEulY", - "metadata": { - "id": "NEvB9BoLEulY" - }, "source": [ "### ☁ Set Your Google Cloud Project\n", "Set your Google Cloud project so that you can leverage Google Cloud resources within this notebook.\n", @@ -124,17 +121,13 @@ "* 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)." - ] + ], + "metadata": { + "id": "NEvB9BoLEulY" + } }, { "cell_type": "code", - "execution_count": null, - "id": "gfkS3yVRE4_W", - "metadata": { - "cellView": "form", - "id": "gfkS3yVRE4_W" - }, - "outputs": [], "source": [ "# @markdown Please fill in the value below with your Google Cloud project ID and then run the cell.\n", "\n", @@ -142,37 +135,35 @@ "\n", "# Set the project id\n", "!gcloud config set project {PROJECT_ID}" - ] + ], + "metadata": { + "cellView": "form", + "id": "gfkS3yVRE4_W" + }, + "execution_count": null, + "outputs": [] }, { "cell_type": "markdown", - "id": "f8f2830ee9ca1e01", - "metadata": { - "id": "f8f2830ee9ca1e01" - }, "source": [ "## Basic Usage" - ] + ], + "metadata": { + "id": "f8f2830ee9ca1e01" + } }, { "cell_type": "markdown", - "id": "OMvzMWRrR6n7", - "metadata": { - "id": "OMvzMWRrR6n7" - }, "source": [ "### Set AlloyDB database values\n", "Find your database values, in the [AlloyDB Instances page](https://console.cloud.google.com/alloydb/clusters)." - ] + ], + "metadata": { + "id": "OMvzMWRrR6n7" + } }, { "cell_type": "code", - "execution_count": 4, - "id": "irl7eMFnSPZr", - "metadata": { - "id": "irl7eMFnSPZr" - }, - "outputs": [], "source": [ "# @title Set Your Values Here { display-mode: \"form\" }\n", "REGION = \"us-central1\" # @param {type: \"string\"}\n", @@ -180,14 +171,15 @@ "INSTANCE = \"my-primary\" # @param {type: \"string\"}\n", "DATABASE = \"my-database\" # @param {type: \"string\"}\n", "TABLE_NAME = \"vector_store\" # @param {type: \"string\"}" - ] + ], + "metadata": { + "id": "irl7eMFnSPZr" + }, + "execution_count": 4, + "outputs": [] }, { "cell_type": "markdown", - "id": "QuQigs4UoFQ2", - "metadata": { - "id": "QuQigs4UoFQ2" - }, "source": [ "### AlloyDBEngine Connection Pool\n", "\n", @@ -210,20 +202,22 @@ "\n", "To connect to your AlloyDB instance from this notebook, you will need to enable public IP on your instance. Alternatively, you can follow [these instructions](https://cloud.google.com/alloydb/docs/connect-external) to connect to an AlloyDB for PostgreSQL instance with Private IP from outside your VPC.\n", "Learn more about [specifying IP types](https://github.com/GoogleCloudPlatform/alloydb-python-connector?tab=readme-ov-file#specifying-ip-address-type).\n" - ] + ], + "metadata": { + "id": "QuQigs4UoFQ2" + } }, { "cell_type": "markdown", - "metadata": {}, "source": [ "**Note:** This tutorial demonstrates the async interface. All async methods have corresponding sync methods." - ] + ], + "metadata": { + "id": "9c0533e8" + } }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], "source": [ "from langchain_google_alloydb_pg import AlloyDBEngine\n", "from google.cloud.alloydb.connector import IPTypes\n", @@ -236,54 +230,62 @@ " database=DATABASE,\n", " ip_type=IPTypes.PUBLIC,\n", ")" - ] + ], + "metadata": { + "id": "89ae517b" + }, + "execution_count": null, + "outputs": [] }, { "cell_type": "markdown", - "metadata": {}, "source": [ "### AlloyDBEngine for AlloyDB Omni\n", "To create an `AlloyDBEngine` for AlloyDB Omni, you will need a connection url. `AlloyDBEngine.from_engine_args` first creates an async engine and then turns it into an `AlloyDBEngine`. Here is an example connection with the `asyncpg` driver:" - ] + ], + "metadata": { + "id": "b9118293" + } }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], "source": [ "# Replace with your own AlloyDB Omni info\n", "connstring = f\"postgresql+asyncpg://{OMNI_USER}:{OMNI_PASSWORD}@{OMNI_HOST}:{OMNI_PORT}/{OMNI_DATABASE}\"\n", "engine = AlloyDBEngine.from_engine_args(connstring)" - ] + ], + "metadata": { + "id": "63e0bc49" + }, + "execution_count": null, + "outputs": [] }, { "cell_type": "markdown", - "metadata": { - "id": "D9Xs2qhm6X56" - }, "source": [ "### Initialize a table\n", "The `AlloyDBVectorStore` class requires a database table. The `AlloyDBEngine` engine has a helper method `init_vectorstore_table()` that can be used to create a table with the proper schema for you." - ] + ], + "metadata": { + "id": "D9Xs2qhm6X56" + } }, { "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "avlyHEMn6gzU" - }, - "outputs": [], "source": [ "await engine.ainit_vectorstore_table(\n", " table_name=TABLE_NAME,\n", " vector_size=768, # Vector size for VertexAI model(textembedding-gecko@latest)\n", ")" - ] + ], + "metadata": { + "id": "avlyHEMn6gzU" + }, + "execution_count": null, + "outputs": [] }, { "cell_type": "markdown", - "metadata": {}, "source": [ "#### Optional Tip: 💡\n", "You can also specify a schema name by passing `schema_name` wherever you pass `table_name`. Eg:\n", @@ -297,34 +299,42 @@ " schema_name=SCHEMA_NAME, # Default: \"public\"\n", ")\n", "```" - ] + ], + "metadata": { + "id": "cc21f455" + } }, { "cell_type": "markdown", - "metadata": {}, "source": [ "### Create an embedding class instance\n", "\n", "You can use any [LangChain embeddings model](https://python.langchain.com/docs/integrations/text_embedding/).\n", "You may need to enable Vertex AI API to use `VertexAIEmbeddings`. We recommend setting the embedding model's version for production, learn more about the [Text embeddings models](https://cloud.google.com/vertex-ai/docs/generative-ai/model-reference/text-embeddings)." - ] + ], + "metadata": { + "id": "35a7e376" + } }, { "cell_type": "code", - "execution_count": null, - "id": "5utKIdq7KYi5", - "metadata": { - "id": "5utKIdq7KYi5" - }, - "outputs": [], "source": [ "# enable Vertex AI API\n", "!gcloud services enable aiplatform.googleapis.com" - ] + ], + "metadata": { + "id": "5utKIdq7KYi5" + }, + "execution_count": null, + "outputs": [] }, { "cell_type": "code", - "execution_count": null, + "source": [ + "from langchain_google_vertexai import VertexAIEmbeddings\n", + "\n", + "embedding = VertexAIEmbeddings(model_name=\"text-embedding-005\", project=PROJECT_ID)" + ], "metadata": { "colab": { "base_uri": "https://localhost:8080/" @@ -332,29 +342,20 @@ "id": "Vb2RJocV9_LQ", "outputId": "37f5dc74-2512-47b2-c135-f34c10afdcf4" }, - "outputs": [], - "source": [ - "from langchain_google_vertexai import VertexAIEmbeddings\n", - "\n", - "embedding = VertexAIEmbeddings(model_name=\"text-embedding-005\", project=PROJECT_ID)" - ] + "execution_count": null, + "outputs": [] }, { "cell_type": "markdown", - "metadata": { - "id": "e1tl0aNx7SWy" - }, "source": [ "### Initialize a default AlloyDBVectorStore" - ] + ], + "metadata": { + "id": "e1tl0aNx7SWy" + } }, { "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "z-AZyzAQ7bsf" - }, - "outputs": [], "source": [ "from langchain_google_alloydb_pg import AlloyDBVectorStore\n", "\n", @@ -364,22 +365,26 @@ " # schema_name=SCHEMA_NAME,\n", " embedding_service=embedding,\n", ")" - ] + ], + "metadata": { + "id": "z-AZyzAQ7bsf" + }, + "execution_count": null, + "outputs": [] }, { "cell_type": "markdown", - "metadata": {}, "source": [ "#### Initialize Vector Store with documents\n", "\n", "This is a great way to get started quickly. However, the default method is recommended for most applications to avoid accidentally adding duplicate documents." - ] + ], + "metadata": { + "id": "8c877daf" + } }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], "source": [ "from langchain_core.documents import Document\n", "import uuid\n", @@ -408,20 +413,24 @@ " # schema_name=SCHEMA_NAME,\n", " embedding=embedding,\n", ")" - ] + ], + "metadata": { + "id": "7ffa103d" + }, + "execution_count": null, + "outputs": [] }, { "cell_type": "markdown", - "metadata": {}, "source": [ "### Add texts" - ] + ], + "metadata": { + "id": "24aaa4f7" + } }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], "source": [ "import uuid\n", "\n", @@ -430,22 +439,26 @@ "ids = [str(uuid.uuid4()) for _ in all_texts]\n", "\n", "await store.aadd_texts(all_texts, metadatas=metadatas, ids=ids)" - ] + ], + "metadata": { + "id": "bad6c191" + }, + "execution_count": null, + "outputs": [] }, { "cell_type": "markdown", - "metadata": {}, "source": [ "### Get document\n", "\n", "Get documents from the vectorstore using filters and parameters." - ] + ], + "metadata": { + "id": "4c4efb0e" + } }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], "source": [ "documents_with_apple = await store.aget(\n", " where_document={\"$ilike\": \"%apple%\"}, include=\"documents\"\n", @@ -454,128 +467,156 @@ "\n", "print(documents_with_apple[\"documents\"])\n", "print(paginated_ids[\"ids\"])" - ] + ], + "metadata": { + "id": "b39f5906" + }, + "execution_count": null, + "outputs": [] }, { "cell_type": "markdown", - "metadata": {}, "source": [ "### Delete documents\n", "\n", "Documents can be deleted using IDs or metadata filters." - ] + ], + "metadata": { + "id": "8f6a40e3" + } }, { "cell_type": "markdown", - "metadata": {}, "source": [ "#### Delete by IDs" - ] + ], + "metadata": { + "id": "94f0a89b" + } }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], "source": [ "await store.adelete([ids[1]])" - ] + ], + "metadata": { + "id": "51b88cf6" + }, + "execution_count": null, + "outputs": [] }, { "cell_type": "markdown", - "metadata": {}, "source": [ "#### Delete by metadata filter\n", "You can delete documents based on metadata filters. This is useful for bulk deletion operations." - ] + ], + "metadata": { + "id": "f382a101" + } }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], "source": [ "# Delete all documents with a specific metadata value\n", "await store.adelete(filter={\"source\": \"documentation\"})" - ] + ], + "metadata": { + "id": "8e243fae" + }, + "execution_count": null, + "outputs": [] }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], "source": [ "# Delete documents matching complex filter criteria\n", "await store.adelete(\n", " filter={\"$and\": [{\"category\": \"obsolete\"}, {\"year\": {\"$lt\": 2020}}]}\n", ")" - ] + ], + "metadata": { + "id": "a427decd" + }, + "execution_count": null, + "outputs": [] }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], "source": [ "# Delete by both IDs and filter (must match both criteria)\n", "await store.adelete(ids=[\"id1\", \"id2\"], filter={\"status\": \"archived\"})" - ] + ], + "metadata": { + "id": "28738d1b" + }, + "execution_count": null, + "outputs": [] }, { "cell_type": "markdown", - "metadata": {}, "source": [ "### Search for documents" - ] + ], + "metadata": { + "id": "e50cb3ff" + } }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], "source": [ "query = \"I'd like a fruit.\"\n", "docs = await store.asimilarity_search(query)\n", "print(docs)" - ] + ], + "metadata": { + "id": "5afe8831" + }, + "execution_count": null, + "outputs": [] }, { "cell_type": "markdown", - "metadata": {}, "source": [ "### Search for documents by vector" - ] + ], + "metadata": { + "id": "c4b21e78" + } }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], "source": [ "query_vector = embedding.embed_query(query)\n", "docs = await store.asimilarity_search_by_vector(query_vector, k=2)\n", "print(docs)" - ] + ], + "metadata": { + "id": "80f17947" + }, + "execution_count": null, + "outputs": [] }, { "cell_type": "markdown", - "metadata": {}, "source": [ "### Multi-modal Vector Store" - ] + ], + "metadata": { + "id": "b55763df" + } }, { "cell_type": "markdown", - "metadata": {}, "source": [ "You can also add local images using the `add_images()` or `aadd_images()` method by passing in a list of image URIs. If you are using `VertexAIEmbeddings` as your embedding service, you will also be able to pass in GCS and web URIs. Use `similarity_search_image()` or `asimilarity_search_image()` to perform similarity search with an input image, or search by text query using regular similarity search APIs. The result returned will be base64 encoded images." - ] + ], + "metadata": { + "id": "3d9e809c" + } }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], "source": [ "local_image_uris = [\"local_image_1.jpg\", \"local_image_2.jpg\", \"local_image_3.jpg\"]\n", "\n", @@ -585,97 +626,115 @@ "# similarity search\n", "image_uri = \"local_image_1.jpg\"\n", "result = await store.asimilarity_search_image(image_uri=image_uri)" - ] + ], + "metadata": { + "id": "0d2dc985" + }, + "execution_count": null, + "outputs": [] }, { "cell_type": "markdown", - "metadata": {}, "source": [ "## Add a Index\n", "Speed up vector search queries by applying a vector index. Learn more about [vector indexes](https://cloud.google.com/blog/products/databases/faster-similarity-search-performance-with-pgvector-indexes)." - ] + ], + "metadata": { + "id": "60ab9982" + } }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], "source": [ "from langchain_google_alloydb_pg.indexes import IVFFlatIndex\n", "\n", "index = IVFFlatIndex()\n", "await store.aapply_vector_index(index)" - ] + ], + "metadata": { + "id": "b2aaaef7" + }, + "execution_count": null, + "outputs": [] }, { "cell_type": "markdown", - "metadata": {}, "source": [ - "The `ScaNN` index creation (only available in AlloyDB Omni) requires sufficient maintenance work memory. You need to set the database flag `maintenance_work_mem` by calling `set_maintenance_work_mem` before applying the index." - ] + "The `ScaNN` index creation (only available in AlloyDB Omni) automatically manages `maintenance_work_mem` during index creation in `aapply_vector_index()`." + ], + "metadata": { + "id": "51ef50b5" + } }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], "source": [ "from langchain_google_alloydb_pg.indexes import ScaNNIndex\n", "\n", - "VECTOR_SIZE = 768 # Replace with the vector size of your embedding model\n", "index = ScaNNIndex()\n", - "await store.aset_maintenance_work_mem(index.num_leaves, VECTOR_SIZE)\n", "await store.aapply_vector_index(index)" - ] + ], + "metadata": { + "id": "5daa50d3" + }, + "execution_count": null, + "outputs": [] }, { "cell_type": "markdown", - "metadata": {}, "source": [ "### Re-index" - ] + ], + "metadata": { + "id": "21180692" + } }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], "source": [ "await store.areindex() # Re-index using default index name" - ] + ], + "metadata": { + "id": "bfa9eb19" + }, + "execution_count": null, + "outputs": [] }, { "cell_type": "markdown", - "metadata": {}, "source": [ "### Remove an index" - ] + ], + "metadata": { + "id": "49ec4d9a" + } }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], "source": [ "await store.adrop_vector_index() # Delete index using default name" - ] + ], + "metadata": { + "id": "cdd1ef52" + }, + "execution_count": null, + "outputs": [] }, { "cell_type": "markdown", - "metadata": {}, "source": [ "## Create a custom Vector Store\n", "A Vector Store can take advantage of relational data to filter similarity searches.\n", "\n", "Create a new table with custom metadata columns.\n", "You can also re-use an existing table which already has custom columns for a Document's id, content, embedding, and/or metadata." - ] + ], + "metadata": { + "id": "d39faef5" + } }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], "source": [ "from langchain_google_alloydb_pg import Column\n", "\n", @@ -699,20 +758,24 @@ " embedding_service=embedding,\n", " metadata_columns=[\"len\"],\n", ")" - ] + ], + "metadata": { + "id": "795cd221" + }, + "execution_count": null, + "outputs": [] }, { "cell_type": "markdown", - "metadata": {}, "source": [ "### Search for documents with metadata filter" - ] + ], + "metadata": { + "id": "2cfaeaa7" + } }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], "source": [ "import uuid\n", "\n", @@ -721,53 +784,65 @@ "metadatas = [{\"len\": len(t)} for t in all_texts]\n", "ids = [str(uuid.uuid4()) for _ in all_texts]\n", "await custom_store.aadd_texts(all_texts, metadatas=metadatas, ids=ids)" - ] + ], + "metadata": { + "id": "4f456e30" + }, + "execution_count": null, + "outputs": [] }, { "cell_type": "markdown", - "metadata": {}, "source": [ "#### For v0.13.0+\n", "\n", "**Important Update:** Support for string filters has been deprecated. Please use dictionaries to add filters." - ] + ], + "metadata": { + "id": "fbfe5eca" + } }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], "source": [ "# Use filter on search\n", "docs = await custom_store.asimilarity_search(query, filter={\"len\": {\"$gte\": 6}})\n", "\n", "print(docs)" - ] + ], + "metadata": { + "id": "7cb813dc" + }, + "execution_count": null, + "outputs": [] }, { "cell_type": "markdown", - "metadata": {}, "source": [ "#### For v0.12.0 and under\n", "\n", "You can make use of the string filters to filter on metadata" - ] + ], + "metadata": { + "id": "a7a08b39" + } }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], "source": [ "# Use filter on search\n", "docs = await custom_store.asimilarity_search(query, filter=\"len >= 6\")\n", "\n", "print(docs)" - ] + ], + "metadata": { + "id": "21ee08d8" + }, + "execution_count": null, + "outputs": [] }, { "cell_type": "markdown", - "metadata": {}, "source": [ "## Create a Vector Store using existing table\n", "\n", @@ -825,13 +900,13 @@ "- **`metadata_columns=[\"name\", \"category\", \"price_usd\", \"quantity\", \"sku\", \"image_url\"]`**: These columns are treated as metadata for each product. Metadata provides additional information about a product, such as its name, category, price, quantity available, SKU (Stock Keeping Unit), and an image URL. This information is useful for displaying product details in search results or for filtering and categorization.\n", "\n", "- **`metadata_json_column=\"metadata\"`**: The `metadata` column can store any additional information about the products in a flexible JSON format. This allows for storing varied and complex data that doesn't fit into the standard columns. Note that filtering on fields within the JSON but not in `metadata_columns` will be less efficient.\n" - ] + ], + "metadata": { + "id": "c9c51e9b" + } }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], "source": [ "# Set an existing table name\n", "TABLE_NAME = \"products\"\n", @@ -850,11 +925,15 @@ " metadata_columns=[\"name\", \"category\", \"price_usd\", \"quantity\", \"sku\", \"image_url\"],\n", " metadata_json_column=\"metadata\",\n", ")" - ] + ], + "metadata": { + "id": "e5cc2e2f" + }, + "execution_count": null, + "outputs": [] }, { "cell_type": "markdown", - "metadata": {}, "source": [ "Note: \n", "\n", @@ -863,77 +942,91 @@ " `ALTER TABLE products ADD COLUMN embed vector(768) DEFAULT NULL`\n", "\n", "1. For new records, added via `VectorStore` embeddings are automatically generated." - ] + ], + "metadata": { + "id": "0b01e05e" + } }, { "cell_type": "markdown", - "metadata": {}, "source": [ "### Search for documents with metadata filter\n", "Since price_usd is one of the metadata_columns, we can use price filter while searching" - ] + ], + "metadata": { + "id": "c7fe9e61" + } }, { "cell_type": "markdown", - "metadata": {}, "source": [ "#### For v0.13.0+\n", "\n", "**Important Update:** Support for string filters has been deprecated. Please use dictionaries to add filters." - ] + ], + "metadata": { + "id": "e6fabd51" + } }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], "source": [ "docs = await custom_store.asimilarity_search(query, filter={\"price_usd\": {\"$gte\": 100}})\n", "\n", "print(docs)" - ] + ], + "metadata": { + "id": "57dcc94b" + }, + "execution_count": null, + "outputs": [] }, { "cell_type": "markdown", - "metadata": {}, "source": [ "#### For v0.12.0 and under\n", "\n", "You can make use of the string filters to filter on metadata" - ] + ], + "metadata": { + "id": "82c05bc5" + } }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], "source": [ "docs = await custom_store.asimilarity_search(query, filter=\"price_usd > 100\")\n", "\n", "print(docs)" - ] + ], + "metadata": { + "id": "ef1b49d2" + }, + "execution_count": null, + "outputs": [] }, { "cell_type": "markdown", - "metadata": {}, "source": [ "### Search for documents with json filter\n" - ] + ], + "metadata": { + "id": "513d98f9" + } }, { "cell_type": "markdown", - "metadata": {}, "source": [ "#### For v0.15.0+\n", "\n", "Metadata filtering on the `metadata_json_column` is now supported in the `AlloyDBVectorStore`." - ] + ], + "metadata": { + "id": "4d9799b7" + } }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], "source": [ "TABLE_NAME = \"products\"\n", "# SCHEMA_NAME = \"my_schema\"\n", @@ -964,11 +1057,15 @@ ")\n", "\n", "print(docs)" - ] + ], + "metadata": { + "id": "72f039e4" + }, + "execution_count": null, + "outputs": [] }, { "cell_type": "markdown", - "metadata": {}, "source": [ "#### For v0.13.0 to v0.14.0\n", "\n", @@ -981,20 +1078,22 @@ "SET\n", " category_from_json = metadata ->> 'category';\n", "```" - ] + ], + "metadata": { + "id": "64429198" + } }, { "cell_type": "markdown", - "metadata": {}, "source": [ "Now that you've added the new column, you must update the Vectorstore instance to recognize it. After which the new column is available for filtering operations." - ] + ], + "metadata": { + "id": "6af5a0f6" + } }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], "source": [ "TABLE_NAME = \"products\"\n", "# SCHEMA_NAME = \"my_schema\"\n", @@ -1026,44 +1125,52 @@ ")\n", "\n", "print(docs)" - ] + ], + "metadata": { + "id": "e88ef294" + }, + "execution_count": null, + "outputs": [] }, { "cell_type": "markdown", - "metadata": {}, "source": [ "\n", "#### For v0.12.0 and under\n", "\n", "Since category is added in json metadata, we can use filter on JSON fields using string filters while searching." - ] + ], + "metadata": { + "id": "fa8e337b" + } }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], "source": [ "docs = await custom_store.asimilarity_search(\n", " query, filter=\"metadata->>'category' = 'Electronics'\"\n", ")\n", "\n", "print(docs)" - ] + ], + "metadata": { + "id": "7aaa4e2c" + }, + "execution_count": null, + "outputs": [] }, { "cell_type": "markdown", - "metadata": {}, "source": [ "### Search for documents without a Vector Store\n", "You may want to search documents based on Document metadata as a tool or as a part of an exploratory workflow. The Document Loader can be used to customize the search and load data in the form of Documents from your database. Learn how to ['Load Documents using a SQL query'](https://github.com/googleapis/langchain-google-alloydb-pg-python/blob/main/docs/document_loader.ipynb)\n" - ] + ], + "metadata": { + "id": "688b2665" + } }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], "source": [ "from langchain_google_alloydb_pg import AlloyDBLoader\n", "\n", @@ -1080,11 +1187,15 @@ "\n", "docs = await loader.aload()\n", "print(docs)" - ] + ], + "metadata": { + "id": "e4827315" + }, + "execution_count": null, + "outputs": [] }, { "cell_type": "markdown", - "metadata": {}, "source": [ "# Hybrid Search with AlloyDBVectorStore\n", "\n", @@ -1093,22 +1204,26 @@ "By integrating both semantic and lexical capabilities, hybrid search helps overcome the limitations of each individual method:\n", "* **Semantic Search**: Excellent for understanding the meaning of a query, even if the exact keywords aren't present. However, it can sometimes miss highly relevant documents that contain the precise keywords but have a slightly different semantic context.\n", "* **Keyword Search**: Highly effective for finding documents with exact keyword matches and is generally fast. Its weakness lies in its inability to understand synonyms, misspellings, or conceptual relationships." - ] + ], + "metadata": { + "id": "a8a63168" + } }, { "cell_type": "markdown", - "metadata": {}, "source": [ "## Hybrid Search Config\n", "\n", "You can take advantage of hybrid search with AlloyDBVectorStore using the `HybridSearchConfig`.\n", "\n", "With a `HybridSearchConfig` provided, the `AlloyDBVectorStore` class can efficiently manage a hybrid search vector store using AlloyDB as the backend, automatically handling the creation and population of the necessary TSV columns when possible." - ] + ], + "metadata": { + "id": "607384a1" + } }, { "cell_type": "markdown", - "metadata": {}, "source": [ "### Building the config\n", "\n", @@ -1122,20 +1237,22 @@ "* **secondary_top_k:** Max results fetched for secondary retrieval. Default: `4`\n", "* **index_name:** Name of the index built on the `tsv_column`\n", "* **index_type:** GIN or GIST. Default: `GIN`" - ] + ], + "metadata": { + "id": "ad9c5037" + } }, { "cell_type": "markdown", - "metadata": {}, "source": [ "Here is an example `HybridSearchConfig`" - ] + ], + "metadata": { + "id": "83755682" + } }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], "source": [ "from langchain_google_alloydb_pg import (\n", " HybridSearchConfig,\n", @@ -1151,11 +1268,15 @@ " \"fetch_top_k\": 10,\n", " },\n", ")" - ] + ], + "metadata": { + "id": "77746034" + }, + "execution_count": null, + "outputs": [] }, { "cell_type": "markdown", - "metadata": {}, "source": [ "**Note:** In this case, we have mentioned the fusion function to be a `reciprocal rank fusion` but you can also use the `weighted_sum_ranking`.\n", "\n", @@ -1169,32 +1290,36 @@ "* primary_results_weight: The weight for the primary source's scores. Defaults to 0.5\n", "* secondary_results_weight: The weight for the secondary source's scores. Defaults to 0.5\n", "* fetch_top_k: The number of documents to fetch after merging the results. Defaults to 4\n" - ] + ], + "metadata": { + "id": "1efbbf90" + } }, { "cell_type": "markdown", - "metadata": {}, "source": [ "## Usage\n", "\n", "Let's assume we are using the previously mentioned table [`products`](#create-a-vector-store-using-existing-table), which stores product details for an eComm venture.\n" - ] + ], + "metadata": { + "id": "4481ffbf" + } }, { "cell_type": "markdown", - "metadata": {}, "source": [ "### With a new hybrid search table\n", "To create a new AlloyDB table with the tsv column, specify the hybrid search config during the initialization of the vector store.\n", "\n", "In this case, all the similarity searches will make use of hybrid search." - ] + ], + "metadata": { + "id": "206bff25" + } }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], "source": [ "from langchain_google_alloydb_pg import Column\n", "\n", @@ -1242,24 +1367,28 @@ "# Use hybrid search\n", "hybrid_docs = await vs_hybrid.asimilarity_search(\"products\", k=5)\n", "print(hybrid_docs)" - ] + ], + "metadata": { + "id": "57066b09" + }, + "execution_count": null, + "outputs": [] }, { "cell_type": "markdown", - "metadata": {}, "source": [ "### With a pre-existing table\n", "\n", "If a hybrid search config is **NOT** provided during `init_vectorstore_table` while creating a table, the table will not contain a tsv_column. In this case you can still take advantage of hybrid search using the `HybridSearchConfig`.\n", "\n", "The specified TSV column is not present but the TSV vectors are created dynamically on-the-go for hybrid search." - ] + ], + "metadata": { + "id": "2a7306b3" + } }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], "source": [ "# Set the existing table name\n", "TABLE_NAME = \"products\"\n", @@ -1292,54 +1421,69 @@ "# Use hybrid search\n", "hybrid_docs = await custom_hybrid_store.asimilarity_search(\"products\", k=5)\n", "print(hybrid_docs)" - ] + ], + "metadata": { + "id": "0eff9861" + }, + "execution_count": null, + "outputs": [] }, { "cell_type": "markdown", - "metadata": {}, "source": [ "In this case, all the similarity searches will make use of hybrid search." - ] + ], + "metadata": { + "id": "df98fc43" + } }, { "cell_type": "markdown", - "metadata": {}, "source": [ "### Applying Hybrid Search to Specific Queries\n", "\n", "To use hybrid search only for certain queries, omit the configuration during initialization and pass it directly to the search method when needed." - ] + ], + "metadata": { + "id": "93de7022" + } }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], "source": [ "# Use hybrid search\n", "hybrid_docs = await custom_store.asimilarity_search(\n", " \"products\", k=5, hybrid_search_config=hybrid_search_config\n", ")\n", "print(hybrid_docs)" - ] + ], + "metadata": { + "id": "a8f7d70f" + }, + "execution_count": null, + "outputs": [] }, { "cell_type": "markdown", - "metadata": {}, "source": [ "## Hybrid Search Index\n", "\n", "Optionally, if you have created an AlloyDB table with a tsv_column, you can create an index." - ] + ], + "metadata": { + "id": "1056b4a8" + } }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], "source": [ "await vs_hybrid.aapply_hybrid_search_index()" - ] + ], + "metadata": { + "id": "350ccf29" + }, + "execution_count": null, + "outputs": [] } ], "metadata": { @@ -1365,6 +1509,6 @@ "version": "3.13.7" } }, - "nbformat": 4, - "nbformat_minor": 0 + "nbformat_minor": 0, + "nbformat": 4 } diff --git a/src/langchain_google_alloydb_pg/async_vectorstore.py b/src/langchain_google_alloydb_pg/async_vectorstore.py index 7437f5a3..c0475ede 100644 --- a/src/langchain_google_alloydb_pg/async_vectorstore.py +++ b/src/langchain_google_alloydb_pg/async_vectorstore.py @@ -16,6 +16,7 @@ from __future__ import annotations import base64 +import logging import re from typing import Any, Optional @@ -26,6 +27,13 @@ from langchain_postgres.v2.async_vectorstore import AsyncPGVectorStore from sqlalchemy import text +logger = logging.getLogger(__name__) + + +def _quote_ident(ident: str) -> str: + """Quote a PostgreSQL identifier to prevent SQL injection and syntax errors.""" + return '"' + ident.replace('"', '""') + '"' + class AsyncAlloyDBVectorStore(AsyncPGVectorStore): """Google AlloyDB Vector Store class""" @@ -33,6 +41,11 @@ class AsyncAlloyDBVectorStore(AsyncPGVectorStore): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) + @property + def _pool_engine(self) -> Any: + """Helper to access the underlying pool from AlloyDBEngine or PGEngine.""" + return getattr(self.engine, "_pool", self.engine) + def _encode_image(self, uri: str) -> str: """Get base64 string from a image URI.""" gcs_uri = re.match("gs://(.*?)/(.*)", uri) @@ -134,17 +147,380 @@ async def asimilarity_search_image( embedding=embedding, k=k, filter=filter, **kwargs ) - async def set_maintenance_work_mem(self, num_leaves: int, vector_size: int) -> None: - """Set database maintenance work memory (for ScaNN index creation).""" - # Required index memory in MB - buffer = 1 - index_memory_required = ( - round(50 * num_leaves * vector_size * 4 / 1024 / 1024) + buffer - ) # Convert bytes to MB - query = f"SET maintenance_work_mem TO '{index_memory_required} MB';" - async with self.engine.connect() as conn: - await conn.execute(text(query)) - await conn.commit() + async def aset_maintenance_work_mem( + self, num_leaves: Optional[int], vector_size: int + ) -> None: + """Deprecated: maintenance_work_mem is now automatically managed during aapply_vector_index.""" + import warnings + + warnings.warn( + "aset_maintenance_work_mem is deprecated and has no effect. " + "aapply_vector_index automatically calculates and sets maintenance_work_mem.", + DeprecationWarning, + stacklevel=2, + ) + + set_maintenance_work_mem = aset_maintenance_work_mem + + async def aapply_vector_index( + self, + index: Any, + name: Optional[str] = None, + *, + concurrently: bool = False, + ) -> None: + """Create index in the vector store table with ScaNN memory management.""" + from langchain_postgres.v2.indexes import ( + DEFAULT_INDEX_NAME_SUFFIX, + ExactNearestNeighbor, + ) + + from .indexes import ScaNNIndex + + if isinstance(index, ExactNearestNeighbor): + await self.adrop_vector_index() + return + + # Note: CREATE EXTENSION is omitted here as it requires SUPERUSER privileges. + # Extensions should be created during database setup by an administrator. + + function = index.get_index_function() + + filter = f"WHERE ({index.partial_indexes})" if index.partial_indexes else "" + params = "WITH " + index.index_options() + if name is None: + if index.name is None: + index.name = self.table_name + DEFAULT_INDEX_NAME_SUFFIX + name = index.name + + schema = getattr(self, "schema_name", None) + table_identifier = ( + f"{_quote_ident(schema)}.{_quote_ident(self.table_name)}" + if schema + else _quote_ident(self.table_name) + ) + stmt = f'CREATE INDEX {"CONCURRENTLY" if concurrently else ""} {_quote_ident(name)} ON {table_identifier} USING {index.index_type} ({_quote_ident(self.embedding_column)} {function}) {params} {filter};' + + mem_query = None + if isinstance(index, ScaNNIndex): + # For mode="AUTO", num_leaves is None. Use a default estimate of 1000 for memory calculation. + num_leaves: int = index.num_leaves if index.num_leaves is not None else 1000 + + # Resolve vector_size with proper precedence and no blocking/deadlocking I/O: + # 1. self.vector_size (explicitly set on store) + # 2. embedding_service.embedding_size (if present) + # 3. Fallback to 768 + vector_size: int = 768 + if hasattr(self, "vector_size") and self.vector_size is not None: + vector_size = getattr(self, "vector_size", 768) or 768 + elif ( + hasattr(self, "embedding_service") + and self.embedding_service is not None + and hasattr(self.embedding_service, "embedding_size") + ): + vector_size = ( + getattr(self.embedding_service, "embedding_size", 768) or 768 + ) + + # Calculate required memory in MB, capping at PostgreSQL's maximum limit of 2,097,151 MB (2 GB - 1 kB) + mem_mb = min( + 2_097_151, + max( + 10, + round(50 * num_leaves * vector_size * 4 / 1024 / 1024) + 1, + ), + ) + mem_query = f"SET maintenance_work_mem TO '{mem_mb} MB';" + + if concurrently: + async with self._pool_engine.connect() as conn: + autocommit_conn = await conn.execution_options( + isolation_level="AUTOCOMMIT" + ) + if mem_query: + await autocommit_conn.execute(text(mem_query)) + try: + await autocommit_conn.execute(text(stmt)) + finally: + if mem_query: + try: + await autocommit_conn.execute( + text("RESET maintenance_work_mem;") + ) + except Exception: + # Preserve the original CREATE INDEX exception if the connection is broken + pass + else: + async with self._pool_engine.begin() as conn: + if mem_query: + # SET LOCAL is automatically scoped to the transaction block + await conn.execute( + text(f"SET LOCAL maintenance_work_mem TO '{mem_mb} MB';") + ) + await conn.execute(text(stmt)) + + async def ainitialize_auto_vector_embeddings( + self, + model_id: str, + content_column: Optional[str] = None, + embedding_column: Optional[str] = None, + schema_name: Optional[str] = None, + ) -> None: + """Asynchronously initialize auto vector embeddings. + + Args: + model_id: The ID of the model to use for embeddings. + content_column: Optional name of the content column. Defaults to self.content_column. + embedding_column: Optional name of the embedding column. Defaults to self.embedding_column. + schema_name: Optional name of the database schema. Defaults to self.schema_name. + """ + content_col = content_column or self.content_column + embedding_col = embedding_column or self.embedding_column + schema = schema_name or getattr(self, "schema_name", "public") + + if not content_col: + raise ValueError( + "content_column must be provided or configured on the vector store." + ) + if not embedding_col: + raise ValueError( + "embedding_column must be provided or configured on the vector store." + ) + + def _quote_ident(ident: str) -> str: + return '"' + ident.replace('"', '""') + '"' + + table_identifier = ( + f"{_quote_ident(schema)}.{_quote_ident(self.table_name)}" + if schema + else _quote_ident(self.table_name) + ) + query = "CALL ai.initialize_embeddings(:model_id, :table_name, :content_column, :embedding_column)" + try: + async with self._pool_engine.connect() as conn: + await conn.execute( + text(query), + { + "model_id": model_id, + "table_name": table_identifier, + "content_column": content_col, + "embedding_column": embedding_col, + }, + ) + await conn.commit() + except Exception as e: + if ( + "ai.initialize_embeddings" in str(e) + or "UndefinedProcedureError" in type(e).__name__ + ): + raise RuntimeError( + "AlloyDB AI extension is not installed or enabled. " + "Please execute 'CREATE EXTENSION IF NOT EXISTS alloydb_ai CASCADE;' on your database." + ) from e + raise + + async def aenable_columnar_engine( + self, + columns: Optional[list[str]] = None, + ) -> None: + """Asynchronously add the table and its columns to the columnar engine. + + Args: + columns: Optional list of column names to add to the columnar engine. + """ + schema = getattr(self, "schema_name", None) + if schema: + table_identifier = f"{_quote_ident(schema)}.{_quote_ident(self.table_name)}" + schema_clause = "table_schema = :schema" + schema_param = schema + else: + table_identifier = _quote_ident(self.table_name) + schema_clause = "table_schema = CURRENT_SCHEMA()" + schema_param = None + + if columns: + columns_str = ",".join(_quote_ident(c) for c in columns) + query = "SELECT google_columnar_engine_add(relation => :table_name, columns => :columns)" + params = {"table_name": table_identifier, "columns": columns_str} + else: + # When columns is None, we should exclude vector columns to avoid wasting columnar memory + # Query table columns excluding the embedding_column + query = "SELECT google_columnar_engine_add(relation => :table_name, columns => :columns)" + # Fetch all columns except the vector embedding column + async with self._pool_engine.connect() as conn: + col_query = ( + "SELECT column_name FROM information_schema.columns " + f"WHERE table_name = :table_name AND {schema_clause} AND column_name != :embed_col" + ) + col_params: dict[str, Any] = { + "table_name": self.table_name, + "embed_col": self.embedding_column, + } + if schema_param: + col_params["schema"] = schema_param + col_result = await conn.execute(text(col_query), col_params) + col_names = [row[0] for row in col_result.fetchall()] + if col_names: + columns_str = ",".join(_quote_ident(c) for c in col_names) + params = {"table_name": table_identifier, "columns": columns_str} + else: + query = "SELECT google_columnar_engine_add(:table_name)" + params = {"table_name": table_identifier} + + try: + async with self._pool_engine.connect() as conn: + await conn.execute(text(query), params) + await conn.commit() + except Exception as e: + if ( + "google_columnar_engine" in str(e) + or "UndefinedFunctionError" in type(e).__name__ + ): + raise RuntimeError( + "AlloyDB Columnar Engine is not installed or enabled on this instance. " + "Please ensure 'google_columnar_engine' is in shared_preload_libraries " + "and 'google_columnar_engine.enabled = on' is set in instance flags." + ) from e + raise + + async def aenable_auto_columnarization(self) -> None: + """Asynchronously trigger auto-columnarization recommendations.""" + query = "SELECT google_columnar_engine_recommend('AUTO_COLUMNARIZATION')" + try: + async with self._pool_engine.connect() as conn: + await conn.execute(text(query)) + await conn.commit() + except Exception as e: + if ( + "google_columnar_engine" in str(e) + or "UndefinedFunctionError" in type(e).__name__ + ): + raise RuntimeError( + "AlloyDB Columnar Engine is not installed or enabled on this instance. " + "Please ensure 'google_columnar_engine' is in shared_preload_libraries " + "and 'google_columnar_engine.enabled = on' is set in instance flags." + ) from e + raise + + async def adefine_vector_assist_spec(self) -> list[dict]: + """Asynchronously define a Vector Assist spec for the current table.""" + schema = getattr(self, "schema_name", "public") or "public" + table_identifier = ( + f"{_quote_ident(schema)}.{_quote_ident(self.table_name)}" + if schema + else _quote_ident(self.table_name) + ) + query = "SELECT * FROM vector_assist.define_spec(table_name => :table_name, vector_column_name => :embedding_column)" + params = { + "table_name": table_identifier, + "embedding_column": self.embedding_column, + } + try: + async with self._pool_engine.connect() as conn: + result = await conn.execute(text(query), params) + rows = [dict(row) for row in result.mappings()] + await conn.commit() + return rows + except Exception as e: + if ( + "vector_assist" in str(e) + or "UndefinedSchemaError" in type(e).__name__ + or "UndefinedFunctionError" in type(e).__name__ + ): + raise RuntimeError( + "AlloyDB Vector Assist extension is not installed on this database. " + "Please execute 'CREATE EXTENSION IF NOT EXISTS vector_assist CASCADE;' as a superuser." + ) from e + raise + + async def aapply_vector_assist_spec( + self, spec_id: Optional[str] = None + ) -> list[dict]: + """Asynchronously apply the Vector Assist spec for the current table.""" + schema = getattr(self, "schema_name", "public") or "public" + table_identifier = ( + f"{_quote_ident(schema)}.{_quote_ident(self.table_name)}" + if schema + else _quote_ident(self.table_name) + ) + if spec_id is not None: + query = "SELECT * FROM vector_assist.apply_spec(spec_id => :spec_id)" + params = {"spec_id": spec_id} + else: + query = "SELECT * FROM vector_assist.apply_spec(table_name => :table_name, vector_column_name => :embedding_column)" + params = { + "table_name": table_identifier, + "embedding_column": self.embedding_column, + } + try: + async with self._pool_engine.connect() as conn: + result = await conn.execute(text(query), params) + rows = [dict(row) for row in result.mappings()] + await conn.commit() + return rows + except Exception as e: + if ( + "vector_assist" in str(e) + or "UndefinedSchemaError" in type(e).__name__ + or "UndefinedFunctionError" in type(e).__name__ + ): + raise RuntimeError( + "AlloyDB Vector Assist extension is not installed on this database. " + "Please execute 'CREATE EXTENSION IF NOT EXISTS vector_assist CASCADE;' as a superuser." + ) from e + raise + + async def aget_vector_assist_recommendations(self) -> list[dict]: + """Asynchronously get Vector Assist recommendations for the current table.""" + schema = getattr(self, "schema_name", "public") or "public" + table_identifier = ( + f"{_quote_ident(schema)}.{_quote_ident(self.table_name)}" + if schema + else _quote_ident(self.table_name) + ) + + # Query existing spec_id from vector_assist.specs instead of defining a new spec (avoids side-effects) + query_spec = ( + "SELECT spec_id FROM vector_assist.specs " + "WHERE table_name = :table_name AND vector_column_name = :embedding_column " + "ORDER BY created_at DESC LIMIT 1" + ) + try: + async with self._pool_engine.connect() as conn: + spec_result = await conn.execute( + text(query_spec), + { + "table_name": table_identifier, + "embedding_column": self.embedding_column, + }, + ) + spec_row = spec_result.mappings().first() + if not spec_row: + logger.warning( + "No vector assist spec found for table '%s'. " + "Call adefine_vector_assist_spec() first to create a spec.", + table_identifier, + ) + return [] + + spec_id = spec_row.get("spec_id") + query = "SELECT * FROM vector_assist.get_recommendations(spec_id => :spec_id)" + result = await conn.execute( + text(query), + {"spec_id": str(spec_id) if spec_id is not None else None}, + ) + return [dict(row) for row in result.mappings()] + except Exception as e: + if ( + "vector_assist" in str(e) + or "UndefinedSchemaError" in type(e).__name__ + or "UndefinedFunctionError" in type(e).__name__ + ): + raise RuntimeError( + "AlloyDB Vector Assist extension is not installed on this database. " + "Please execute 'CREATE EXTENSION IF NOT EXISTS vector_assist CASCADE;' as a superuser." + ) from e + raise def add_images( self, diff --git a/src/langchain_google_alloydb_pg/engine.py b/src/langchain_google_alloydb_pg/engine.py index b712abd6..bc7abf0b 100644 --- a/src/langchain_google_alloydb_pg/engine.py +++ b/src/langchain_google_alloydb_pg/engine.py @@ -106,7 +106,7 @@ def __start_background_loop( password: Optional[str] = None, ip_type: Union[str, IPTypes] = IPTypes.PUBLIC, iam_account_email: Optional[str] = None, - engine_args: Mapping = {}, + engine_args: Optional[Mapping[str, Any]] = None, ) -> Future: # Running a loop in a background thread allows us to support # async methods from non-async environments @@ -144,7 +144,7 @@ def from_instance( password: Optional[str] = None, ip_type: Union[str, IPTypes] = IPTypes.PUBLIC, iam_account_email: Optional[str] = None, - engine_args: Mapping = {}, + engine_args: Optional[Mapping[str, Any]] = None, ) -> AlloyDBEngine: """Create an AlloyDBEngine from an AlloyDB instance. @@ -158,9 +158,9 @@ def from_instance( password (Optional[str]): Cloud AlloyDB user password. Defaults to None. ip_type (Union[str, IPTypes], optional): IP address type. Defaults to IPTypes.PUBLIC. iam_account_email (Optional[str], optional): IAM service account email. Defaults to None. - engine_args (Mapping): Additional arguments that are passed directly to - :func:`~sqlalchemy.ext.asyncio.mymodule.MyClass.create_async_engine`. This can be - used to specify additional parameters to the underlying pool during it's creation. + engine_args (Optional[Mapping[str, Any]]): Additional arguments that are passed directly to + :func:`~sqlalchemy.ext.asyncio.create_async_engine`. This can be + used to specify additional parameters to the underlying pool during its creation. Returns: AlloyDBEngine: A newly created AlloyDBEngine instance. @@ -193,7 +193,7 @@ async def _create( loop: Optional[asyncio.AbstractEventLoop] = None, thread: Optional[Thread] = None, iam_account_email: Optional[str] = None, - engine_args: Mapping = {}, + engine_args: Optional[Mapping[str, Any]] = None, ) -> AlloyDBEngine: """Create an AlloyDBEngine from an AlloyDB instance. @@ -209,9 +209,9 @@ async def _create( loop (Optional[asyncio.AbstractEventLoop]): Async event loop used to create the engine. thread (Optional[Thread]): Thread used to create the engine async. iam_account_email (Optional[str]): IAM service account email. - engine_args (Mapping): Additional arguments that are passed directly to - :func:`~sqlalchemy.ext.asyncio.mymodule.MyClass.create_async_engine`. This can be - used to specify additional parameters to the underlying pool during it's creation. + engine_args (Optional[Mapping[str, Any]]): Additional arguments that are passed directly to + :func:`~sqlalchemy.ext.asyncio.create_async_engine`. This can be + used to specify additional parameters to the underlying pool during its creation. Raises: ValueError: Raises error if only one of 'user' or 'password' is specified. @@ -261,10 +261,11 @@ async def getconn() -> asyncpg.Connection: ) return conn + engine_kwargs = dict(engine_args) if engine_args is not None else {} engine = create_async_engine( "postgresql+asyncpg://", async_creator=getconn, - **engine_args, + **engine_kwargs, ) return cls(PGEngine._PGEngine__create_key, engine, loop, thread) # type: ignore @@ -280,7 +281,7 @@ async def afrom_instance( password: Optional[str] = None, ip_type: Union[str, IPTypes] = IPTypes.PUBLIC, iam_account_email: Optional[str] = None, - engine_args: Mapping = {}, + engine_args: Optional[Mapping[str, Any]] = None, ) -> AlloyDBEngine: """Create an AlloyDBEngine from an AlloyDB instance. @@ -294,9 +295,9 @@ async def afrom_instance( password (Optional[str], optional): Cloud AlloyDB user password. Defaults to None. ip_type (Union[str, IPTypes], optional): IP address type. Defaults to IPTypes.PUBLIC. iam_account_email (Optional[str], optional): IAM service account email. Defaults to None. - engine_args (Mapping): Additional arguments that are passed directly to - :func:`~sqlalchemy.ext.asyncio.mymodule.MyClass.create_async_engine`. This can be - used to specify additional parameters to the underlying pool during it's creation. + engine_args (Optional[Mapping[str, Any]]): Additional arguments that are passed directly to + :func:`~sqlalchemy.ext.asyncio.create_async_engine`. This can be + used to specify additional parameters to the underlying pool during its creation. Returns: AlloyDBEngine: A newly created AlloyDBEngine instance. @@ -321,15 +322,17 @@ def from_connection_string( url: str | URL, **kwargs: Any, ) -> AlloyDBEngine: - """Create an AlloyDBEngine instance from arguments + """Create an AlloyDBEngine instance from arguments. + Args: - url (Optional[str]): the URL used to connect to a database. Use url or set other arguments. + url (str | URL): The URL used to connect to a database. + Raises: - ValueError: If not all database url arguments are specified + ValueError: If not all database url arguments are specified. + Returns: - AlloyDBEngine + AlloyDBEngine: A newly created AlloyDBEngine instance. """ - return AlloyDBEngine.from_engine_args(url=url, **kwargs) @classmethod @@ -338,16 +341,16 @@ def from_engine_args( url: str | URL, **kwargs: Any, ) -> AlloyDBEngine: - """Create an AlloyDBEngine instance from arguments + """Create an AlloyDBEngine instance from arguments. Args: - url (Optional[str]): the URL used to connect to a database. Use url or set other arguments. + url (str | URL): The URL used to connect to a database. Raises: - ValueError: If not all database url arguments are specified + ValueError: If not all database url arguments are specified. Returns: - AlloyDBEngine + AlloyDBEngine: A newly created AlloyDBEngine instance. """ # Running a loop in a background thread allows us to support # async methods from non-async environments @@ -367,6 +370,82 @@ def from_engine_args( engine = create_async_engine(url, **kwargs) return cls(PGEngine._PGEngine__create_key, engine, cls._default_loop, cls._default_thread) # type: ignore + @classmethod + def from_engine( + cls: type[AlloyDBEngine], + engine: Any, + loop: Optional[asyncio.AbstractEventLoop] = None, + ) -> AlloyDBEngine: + """Create an AlloyDBEngine instance from an AsyncEngine.""" + return cls(PGEngine._PGEngine__create_key, engine, loop, None) # type: ignore + + def close(self) -> None: # type: ignore[override] + """Synchronously dispose of the connection pool.""" + if self._loop and self._loop.is_running(): + future = asyncio.run_coroutine_threadsafe(self._pool.dispose(), self._loop) + future.result() + elif self._default_loop and self._default_loop.is_running(): + future = asyncio.run_coroutine_threadsafe( + self._pool.dispose(), self._default_loop + ) + future.result() + else: + asyncio.run(self._pool.dispose()) + + async def aclose(self) -> None: + """Asynchronously dispose of the connection pool.""" + await self._run_as_async(self._pool.dispose()) + + async def __aenter__(self) -> AlloyDBEngine: + return self + + async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: + await self.aclose() + + def __enter__(self) -> AlloyDBEngine: + return self + + def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: + self.close() + + def __del__(self) -> None: + try: + self.close() + except Exception: + pass + + async def adrop_table( + self, + table_name: str, + *, + schema_name: str = "public", + ) -> None: + """Asynchronously drop a table from the database. + + Args: + table_name (str): The name of the table to drop. + schema_name (str): The schema name of the table. Default: "public". + """ + await self._run_as_async( + self._adrop_table(table_name=table_name, schema_name=schema_name) + ) + + def drop_table( + self, + table_name: str, + *, + schema_name: str = "public", + ) -> None: + """Synchronously drop a table from the database. + + Args: + table_name (str): The name of the table to drop. + schema_name (str): The schema name of the table. Default: "public". + """ + self._run_as_sync( + self._adrop_table(table_name=table_name, schema_name=schema_name) + ) + async def _ainit_chat_history_table( self, table_name: str, schema_name: str = "public" ) -> None: @@ -414,7 +493,7 @@ async def ainit_chat_history_table( def init_chat_history_table( self, table_name: str, schema_name: str = "public" ) -> None: - """Create a Cloud SQL table to store chat history. + """Create an AlloyDB table to store chat history. Args: table_name (str): Table name to store chat history. @@ -610,9 +689,10 @@ async def ainit_checkpoint_table( def init_checkpoint_table( self, table_name: str = CHECKPOINTS_TABLE, schema_name: str = "public" ) -> None: - """Create Cloud SQL tables to store checkpoints. + """Create AlloyDB tables to store checkpoints. Args: + table_name (str): The checkpoint table name. Default: "checkpoints". schema_name (str): The schema name to store checkpoint tables. Default: "public". @@ -621,6 +701,178 @@ def init_checkpoint_table( """ self._run_as_sync(self._ainit_checkpoint_table(table_name, schema_name)) + async def _aforecast( + self, + model_id: str, + source_table: str, + timestamp_col: str, + data_col: str, + horizon: int, + source_query: Optional[str] = None, + conf_level: Optional[float] = None, + ) -> list[dict]: + if not model_id or not isinstance(model_id, str) or not model_id.strip(): + raise ValueError("model_id must be a non-empty string.") + if ( + not source_table + or not isinstance(source_table, str) + or not source_table.strip() + ): + raise ValueError("source_table must be a non-empty string.") + if ( + not timestamp_col + or not isinstance(timestamp_col, str) + or not timestamp_col.strip() + ): + raise ValueError("timestamp_col must be a non-empty string.") + if not data_col or not isinstance(data_col, str) or not data_col.strip(): + raise ValueError("data_col must be a non-empty string.") + + # Validate horizon + import math + import operator + + try: + if isinstance(horizon, bool): + raise TypeError + horizon_val = operator.index(horizon) + if horizon_val <= 0: + raise ValueError + except (TypeError, ValueError): + raise ValueError("horizon must be a positive integer.") from None + + if horizon_val > 2_147_483_647: + raise ValueError( + "horizon exceeds maximum 32-bit integer limit (2,147,483,647)." + ) + + # Validate conf_level + if conf_level is not None: + if not isinstance(conf_level, (int, float)) or isinstance(conf_level, bool): + raise TypeError("conf_level must be a float between 0 and 1.") + if ( + not (0 < conf_level < 1) + or math.isnan(conf_level) + or math.isinf(conf_level) + ): + raise ValueError("conf_level must be a float strictly between 0 and 1.") + + # Clean source_query + if source_query is not None: + if not isinstance(source_query, str): + raise TypeError("source_query must be a string.") + source_query = source_query.strip() or None + + args = [ + "model_id => :model_id", + "source_table => :source_table", + "timestamp_col => :timestamp_col", + "data_col => :data_col", + "horizon => :horizon", + ] + params: dict[str, Any] = { + "model_id": model_id.strip(), + "source_table": source_table.strip(), + "timestamp_col": timestamp_col.strip(), + "data_col": data_col.strip(), + "horizon": horizon_val, + } + if source_query is not None: + args.append("source_query => :source_query") + params["source_query"] = source_query + if conf_level is not None: + args.append("conf_level => :conf_level") + params["conf_level"] = conf_level + + query = f"SELECT * FROM google_ml.forecast({', '.join(args)})" + try: + async with self._pool.connect() as conn: + result = await conn.execute(text(query), params) + return [dict(row) for row in result.mappings()] + except Exception as e: + if ( + "google_ml" in str(e) + or "UndefinedFunctionError" in type(e).__name__ + or "UndefinedSchemaError" in type(e).__name__ + ): + raise RuntimeError( + "AlloyDB AI google_ml extension is not installed or enabled. " + "Please execute 'CREATE EXTENSION IF NOT EXISTS google_ml CASCADE;' on your database." + ) from e + raise + + async def aforecast( + self, + model_id: str, + source_table: str, + timestamp_col: str, + data_col: str, + horizon: int, + source_query: Optional[str] = None, + conf_level: Optional[float] = None, + ) -> list[dict]: + """Asynchronously get forecasting from AlloyDB AI. + + Args: + model_id: The ID of the time series forecasting model. + source_table: The table to read historical time series data from. + timestamp_col: The column containing the timestamp. + data_col: The column containing the data to forecast. + horizon: Number of future time steps to forecast. + source_query: Optional query to filter historical data. + conf_level: Optional confidence level for prediction intervals. + + Returns: + A list of dictionaries with forecast_timestamp, forecast_value, and intervals. + """ + return await self._run_as_async( + self._aforecast( + model_id, + source_table, + timestamp_col, + data_col, + horizon, + source_query, + conf_level, + ) + ) + + def forecast( + self, + model_id: str, + source_table: str, + timestamp_col: str, + data_col: str, + horizon: int, + source_query: Optional[str] = None, + conf_level: Optional[float] = None, + ) -> list[dict]: + """Synchronously get forecasting from AlloyDB AI. + + Args: + model_id: The ID of the time series forecasting model. + source_table: The table to read historical time series data from. + timestamp_col: The column containing the timestamp. + data_col: The column containing the data to forecast. + horizon: Number of future time steps to forecast. + source_query: Optional query to filter historical data. + conf_level: Optional confidence level for prediction intervals. + + Returns: + A list of dictionaries with forecast_timestamp, forecast_value, and intervals. + """ + return self._run_as_sync( + self._aforecast( + model_id, + source_table, + timestamp_col, + data_col, + horizon, + source_query, + conf_level, + ) + ) + async def _aload_table_schema( self, table_name: str, schema_name: str = "public" ) -> Table: diff --git a/src/langchain_google_alloydb_pg/vectorstore.py b/src/langchain_google_alloydb_pg/vectorstore.py index 09fba583..6f63d132 100644 --- a/src/langchain_google_alloydb_pg/vectorstore.py +++ b/src/langchain_google_alloydb_pg/vectorstore.py @@ -167,6 +167,48 @@ def create_sync( vs = engine._run_as_sync(coro) return cls(cls._PGVectorStore__create_key, engine, vs) # type: ignore + async def ainitialize_auto_vector_embeddings( + self, + model_id: str, + content_column: Optional[str] = None, + embedding_column: Optional[str] = None, + schema_name: Optional[str] = None, + ) -> None: + """Generate and manage auto vector embeddings for large tables. + + Args: + model_id (str): The model id used for generating embeddings. + content_column (Optional[str]): Name of the content column. + embedding_column (Optional[str]): Name of the embedding column. + schema_name (Optional[str]): Name of the database schema. + """ + await self._engine._run_as_async( + self._PGVectorStore__vs.ainitialize_auto_vector_embeddings( # type: ignore + model_id, content_column, embedding_column, schema_name + ) + ) + + def initialize_auto_vector_embeddings( + self, + model_id: str, + content_column: Optional[str] = None, + embedding_column: Optional[str] = None, + schema_name: Optional[str] = None, + ) -> None: + """Generate and manage auto vector embeddings for large tables. + + Args: + model_id (str): The model id used for generating embeddings. + content_column (Optional[str]): Name of the content column. + embedding_column (Optional[str]): Name of the embedding column. + schema_name (Optional[str]): Name of the database schema. + """ + self._engine._run_as_sync( + self._PGVectorStore__vs.ainitialize_auto_vector_embeddings( # type: ignore + model_id, content_column, embedding_column, schema_name + ) + ) + async def aadd_images( self, uris: list[str], @@ -222,15 +264,91 @@ async def asimilarity_search_image( ) async def aset_maintenance_work_mem( - self, num_leaves: int, vector_size: int + self, num_leaves: Optional[int], vector_size: int + ) -> None: + """Deprecated: maintenance_work_mem is now automatically managed during aapply_vector_index.""" + await self._PGVectorStore__vs.aset_maintenance_work_mem(num_leaves, vector_size) # type: ignore + + def set_maintenance_work_mem( + self, num_leaves: Optional[int], vector_size: int + ) -> None: + """Deprecated: maintenance_work_mem is now automatically managed during aapply_vector_index.""" + self._engine._run_as_sync( + self._PGVectorStore__vs.aset_maintenance_work_mem(num_leaves, vector_size) # type: ignore + ) + + async def aenable_columnar_engine( + self, + columns: Optional[list[str]] = None, ) -> None: - """Set database maintenance work memory (for ScaNN index creation).""" + """Asynchronously add the table and its columns to the columnar engine. + + Args: + columns: Optional list of column names to add to the columnar engine. + """ await self._engine._run_as_async( - self._PGVectorStore__vs.set_maintenance_work_mem(num_leaves, vector_size) # type: ignore + self._PGVectorStore__vs.aenable_columnar_engine(columns) # type: ignore ) - def set_maintenance_work_mem(self, num_leaves: int, vector_size: int) -> None: - """Set database maintenance work memory (for ScaNN index creation).""" + def enable_columnar_engine( + self, + columns: Optional[list[str]] = None, + ) -> None: + """Add the table and its columns to the columnar engine. + + Args: + columns: Optional list of column names to add to the columnar engine. + """ + self._engine._run_as_sync( + self._PGVectorStore__vs.aenable_columnar_engine(columns) # type: ignore + ) + + async def aenable_auto_columnarization(self) -> None: + """Asynchronously trigger auto-columnarization recommendations.""" + await self._engine._run_as_async( + self._PGVectorStore__vs.aenable_auto_columnarization() # type: ignore + ) + + def enable_auto_columnarization(self) -> None: + """Trigger auto-columnarization recommendations.""" self._engine._run_as_sync( - self._PGVectorStore__vs.set_maintenance_work_mem(num_leaves, vector_size) # type: ignore + self._PGVectorStore__vs.aenable_auto_columnarization() # type: ignore + ) + + async def adefine_vector_assist_spec(self) -> list[dict]: + """Asynchronously define a Vector Assist spec for the current table.""" + return await self._engine._run_as_async( + self._PGVectorStore__vs.adefine_vector_assist_spec() # type: ignore + ) + + def define_vector_assist_spec(self) -> list[dict]: + """Define a Vector Assist spec for the current table.""" + return self._engine._run_as_sync( + self._PGVectorStore__vs.adefine_vector_assist_spec() # type: ignore + ) + + async def aapply_vector_assist_spec( + self, spec_id: Optional[str] = None + ) -> list[dict]: + """Asynchronously apply the Vector Assist spec for the current table.""" + return await self._engine._run_as_async( + self._PGVectorStore__vs.aapply_vector_assist_spec(spec_id=spec_id) # type: ignore + ) + + def apply_vector_assist_spec(self, spec_id: Optional[str] = None) -> list[dict]: + """Apply the Vector Assist spec for the current table.""" + return self._engine._run_as_sync( + self._PGVectorStore__vs.aapply_vector_assist_spec(spec_id=spec_id) # type: ignore + ) + + async def aget_vector_assist_recommendations(self) -> list[dict]: + """Asynchronously get Vector Assist recommendations for the current table.""" + return await self._engine._run_as_async( + self._PGVectorStore__vs.aget_vector_assist_recommendations() # type: ignore + ) + + def get_vector_assist_recommendations(self) -> list[dict]: + """Get Vector Assist recommendations for the current table.""" + return self._engine._run_as_sync( + self._PGVectorStore__vs.aget_vector_assist_recommendations() # type: ignore ) diff --git a/tests/test_async_vectorstore.py b/tests/test_async_vectorstore.py index 8dc93762..50c454ed 100644 --- a/tests/test_async_vectorstore.py +++ b/tests/test_async_vectorstore.py @@ -16,6 +16,7 @@ import os import uuid from typing import Sequence +from unittest.mock import AsyncMock, MagicMock, patch import pytest import pytest_asyncio @@ -27,6 +28,10 @@ from langchain_google_alloydb_pg import AlloyDBEngine, Column from langchain_google_alloydb_pg.async_vectorstore import AsyncAlloyDBVectorStore +from langchain_google_alloydb_pg.indexes import ( + DistanceStrategy, + ScaNNIndex, +) DEFAULT_TABLE = "test_table" + str(uuid.uuid4()) DEFAULT_TABLE_SYNC = "test_table_sync" + str(uuid.uuid4()) @@ -100,13 +105,23 @@ def db_name(self) -> str: @pytest_asyncio.fixture(scope="class") async def engine(self, db_project, db_region, db_cluster, db_instance, db_name): - engine = await AlloyDBEngine.afrom_instance( - project_id=db_project, - instance=db_instance, - cluster=db_cluster, - region=db_region, - database=db_name, - ) + host = os.environ.get("OMNI_HOST") or os.environ.get("IP_ADDRESS") + user = os.environ.get("OMNI_USER") or os.environ.get("DB_USER", "postgres") + password = os.environ.get("OMNI_PASSWORD") or os.environ.get("DB_PASSWORD") + if host and password: + import sqlalchemy.ext.asyncio + + connstring = f"postgresql+asyncpg://{user}:{password}@{host}:5432/{db_name}" + async_engine = sqlalchemy.ext.asyncio.create_async_engine(connstring) + engine = AlloyDBEngine.from_engine(async_engine) + else: + engine = await AlloyDBEngine.afrom_instance( + project_id=db_project, + instance=db_instance, + cluster=db_cluster, + region=db_region, + database=db_name, + ) yield engine await aexecute(engine, f'DROP TABLE IF EXISTS "{DEFAULT_TABLE}"') @@ -473,3 +488,308 @@ async def test_create_vectorstore_with_init(self, engine): embedding_column="myembedding", metadata_columns=["random_column"], # invalid metadata column ) + + async def test_live_columnar_engine(self, vs): + """Test enabling columnar engine against live AlloyDB instance.""" + await vs.aenable_columnar_engine(["content"]) + await vs.aenable_columnar_engine() + + # Assert functional similarity search still works on columnarized table + await vs.aadd_texts(["Columnar engine test document"]) + results = await vs.asimilarity_search("Columnar test", k=1) + assert len(results) > 0 + assert "Columnar" in results[0].page_content + + async def test_live_auto_columnarization(self, vs): + """Test triggering auto columnarization recommendations against live AlloyDB instance.""" + try: + await vs.aenable_auto_columnarization() + except Exception as e: + if "google_columnar_engine.enabled" in str( + e + ) or "shared_preload_libraries" in str(e): + pytest.skip(f"Columnar engine flag not enabled on instance: {e}") + raise + + # Assert functional similarity search still works after auto columnarization + await vs.aadd_texts(["Auto columnarization test document"]) + results = await vs.asimilarity_search("Auto columnarization", k=1) + assert len(results) > 0 + assert "Auto columnarization" in results[0].page_content + + async def test_live_vector_assist(self, engine): + """Test vector assist spec definition, application, and recommendations against live AlloyDB instance.""" + table_name = "va_live_table_" + str(uuid.uuid4()).replace("-", "_") + await aexecute(engine, f'DROP TABLE IF EXISTS "{table_name}" CASCADE;') + await aexecute( + engine, + f""" + CREATE TABLE "{table_name}" ( + langchain_id uuid PRIMARY KEY, + content text, + embedding vector({VECTOR_SIZE}), + meta jsonb + ); + """, + ) + await aexecute( + engine, + f""" + INSERT INTO "{table_name}" (langchain_id, content, embedding, meta) + SELECT + gen_random_uuid(), + 'Content ' || i, + (SELECT array_agg((random() * 2 - 1)::float4)::vector({VECTOR_SIZE}) FROM generate_series(1, {VECTOR_SIZE})), + '{{"page": 1}}'::jsonb + FROM generate_series(1, 100) AS i; + """, + ) + vs = await AsyncAlloyDBVectorStore.create( + engine, + embedding_service=embeddings_service, + table_name=table_name, + metadata_json_column="meta", + ) + specs = await vs.adefine_vector_assist_spec() + assert isinstance(specs, list) + assert len(specs) > 0 + apply_res = await vs.aapply_vector_assist_spec() + assert isinstance(apply_res, list) + recs = await vs.aget_vector_assist_recommendations() + assert isinstance(recs, list) + await aexecute(engine, f'DROP TABLE IF EXISTS "{table_name}" CASCADE;') + + +@pytest.mark.asyncio +class TestAsyncVectorStoreUnit: + @pytest.fixture + def vs(self): + vs = AsyncAlloyDBVectorStore.__new__(AsyncAlloyDBVectorStore) + vs.engine = MagicMock() + vs.schema_name = "public" + vs.table_name = "test_table" + vs.content_column = "content" + vs.embedding_column = "embedding" + return vs + + async def test_aenable_columnar_engine(self, vs): + """Test enabling columnar engine with a specific column list.""" + # 1. Mock the database connection + with patch.object(vs.engine._pool, "connect") as mock_connect: + mock_conn = AsyncMock() + mock_connect.return_value.__aenter__.return_value = mock_conn + + # 2. Call aenable_columnar_engine with a specific column list + await vs.aenable_columnar_engine(["content"]) + + # 3. Assert exact SQL signature and parameters sent to the database + call_args = mock_conn.execute.call_args + assert ( + str(call_args[0][0]) + == "SELECT google_columnar_engine_add(relation => :table_name, columns => :columns)" + ) + assert call_args[0][1] == { + "table_name": '"public"."test_table"', + "columns": '"content"', + } + + async def test_aenable_columnar_engine_without_columns(self, vs): + """Test enabling columnar engine without specifying columns (entire table).""" + # 1. Mock the database connection + with patch.object(vs.engine._pool, "connect") as mock_connect: + mock_conn = AsyncMock() + mock_result = MagicMock() + mock_result.fetchall.return_value = [("content",), ("langchain_id",)] + mock_conn.execute.return_value = mock_result + mock_connect.return_value.__aenter__.return_value = mock_conn + + # 2. Call aenable_columnar_engine without column arguments + await vs.aenable_columnar_engine() + + # 3. Assert default single-argument query is executed + call_args = mock_conn.execute.call_args + assert ( + str(call_args[0][0]) + == "SELECT google_columnar_engine_add(relation => :table_name, columns => :columns)" + ) + assert call_args[0][1] == { + "table_name": '"public"."test_table"', + "columns": '"content","langchain_id"', + } + + async def test_aenable_auto_columnarization(self, vs): + """Test enabling auto columnarization executes queries on engine.""" + # 1. Mock the database connection + with patch.object(vs.engine._pool, "connect") as mock_connect: + mock_conn = AsyncMock() + mock_connect.return_value.__aenter__.return_value = mock_conn + + # 2. Trigger auto columnarization recommendations + await vs.aenable_auto_columnarization() + + # 3. Assert recommendation query executed on engine + call_args = mock_conn.execute.call_args + assert ( + str(call_args[0][0]) + == "SELECT google_columnar_engine_recommend('AUTO_COLUMNARIZATION')" + ) + + async def test_adefine_vector_assist_spec(self, vs): + """Test definition of vector assist specification.""" + # 1. Mock database returning a vector assist spec row + with patch.object(vs.engine._pool, "connect") as mock_connect: + mock_conn = AsyncMock() + mock_result = MagicMock() + mock_result.mappings.return_value = [{"spec": "ok"}] + mock_conn.execute.return_value = mock_result + mock_connect.return_value.__aenter__.return_value = mock_conn + + # 2. Call adefine_vector_assist_spec + res = await vs.adefine_vector_assist_spec() + + # 3. Assert returned spec list and query parameters + assert res == [{"spec": "ok"}] + call_args = mock_conn.execute.call_args + assert ( + str(call_args[0][0]) + == "SELECT * FROM vector_assist.define_spec(table_name => :table_name, vector_column_name => :embedding_column)" + ) + assert call_args[0][1] == { + "table_name": '"public"."test_table"', + "embedding_column": "embedding", + } + + async def test_aapply_vector_assist_spec(self, vs): + """Test applying vector assist specifications.""" + # 1. Mock database applying vector assist spec + with patch.object(vs.engine._pool, "connect") as mock_connect: + mock_conn = AsyncMock() + mock_result = MagicMock() + mock_result.mappings.return_value = [{"apply": "ok"}] + mock_conn.execute.return_value = mock_result + mock_connect.return_value.__aenter__.return_value = mock_conn + + # 2. Apply spec + res = await vs.aapply_vector_assist_spec() + + # 3. Assert results and query parameters + assert res == [{"apply": "ok"}] + call_args = mock_conn.execute.call_args + assert ( + str(call_args[0][0]) + == "SELECT * FROM vector_assist.apply_spec(table_name => :table_name, vector_column_name => :embedding_column)" + ) + assert call_args[0][1] == { + "table_name": '"public"."test_table"', + "embedding_column": "embedding", + } + + async def test_aget_vector_assist_recommendations(self, vs): + """Test retrieving vector assist recommendations with a valid spec ID.""" + with patch.object(vs.engine._pool, "connect") as mock_connect: + mock_conn = AsyncMock() + # First query returns spec_id, second query returns recommendations + mock_spec_result = MagicMock() + mock_spec_result.mappings.return_value.first.return_value = { + "spec_id": "spec123" + } + mock_rec_result = MagicMock() + mock_rec_result.mappings.return_value = [{"rec": "ok"}] + mock_conn.execute.side_effect = [mock_spec_result, mock_rec_result] + mock_connect.return_value.__aenter__.return_value = mock_conn + + # 2. Retrieve recommendations + res = await vs.aget_vector_assist_recommendations() + + # 3. Assert recommendations and query calls + assert res == [{"rec": "ok"}] + assert mock_conn.execute.call_count == 2 + + async def test_aget_vector_assist_recommendations_empty_specs(self, vs): + """Test retrieving vector assist recommendations when no specs exist in vector_assist.specs.""" + with patch.object(vs.engine._pool, "connect") as mock_connect: + mock_conn = AsyncMock() + mock_spec_result = MagicMock() + mock_spec_result.mappings.return_value.first.return_value = None + mock_conn.execute.return_value = mock_spec_result + mock_connect.return_value.__aenter__.return_value = mock_conn + + res = await vs.aget_vector_assist_recommendations() + assert res == [] + + async def test_ainitialize_auto_vector_embeddings(self, vs): + """Test initializing auto vector embeddings asynchronously with default columns.""" + # 1. Mock the database connection + with patch.object(vs.engine._pool, "connect") as mock_connect: + mock_conn = AsyncMock() + mock_connect.return_value.__aenter__.return_value = mock_conn + + # 2. Call auto vector embedding initialization + await vs.ainitialize_auto_vector_embeddings( + model_id="test-model", + ) + + # 3. Assert exact procedure call and parameters + call_args = mock_conn.execute.call_args + assert ( + str(call_args[0][0]) + == "CALL ai.initialize_embeddings(:model_id, :table_name, :content_column, :embedding_column)" + ) + assert call_args[0][1] == { + "model_id": "test-model", + "table_name": '"public"."test_table"', + "content_column": "content", + "embedding_column": "embedding", + } + + async def test_ainitialize_auto_vector_embeddings_custom_columns(self, vs): + """Test initializing auto vector embeddings with custom columns and schema.""" + # 1. Mock the database connection + with patch.object(vs.engine._pool, "connect") as mock_connect: + mock_conn = AsyncMock() + mock_connect.return_value.__aenter__.return_value = mock_conn + + # 2. Call with custom content column, embedding column, and schema + await vs.ainitialize_auto_vector_embeddings( + model_id="test-model", + content_column="custom_content", + embedding_column="custom_embedding", + schema_name="myschema", + ) + + # 3. Assert custom parameters and quoted schema identifier + call_args = mock_conn.execute.call_args + assert ( + str(call_args[0][0]) + == "CALL ai.initialize_embeddings(:model_id, :table_name, :content_column, :embedding_column)" + ) + assert call_args[0][1] == { + "model_id": "test-model", + "table_name": '"myschema"."test_table"', + "content_column": "custom_content", + "embedding_column": "custom_embedding", + } + + async def test_ainitialize_auto_vector_embeddings_missing_columns(self, vs): + """Test error raised when required content column name is missing.""" + # 1. Clear content_column on vector store + vs.content_column = None + + # 2. Assert ValueError is raised when calling without content_column + with pytest.raises( + ValueError, match="content_column must be provided or configured" + ): + await vs.ainitialize_auto_vector_embeddings(model_id="test-model") + + async def test_ainitialize_auto_vector_embeddings_missing_embedding_column( + self, vs + ): + """Test error raised when required embedding_column name is missing.""" + # 1. Clear embedding_column on vector store + vs.embedding_column = None + + # 2. Assert ValueError is raised when calling without embedding_column + with pytest.raises( + ValueError, match="embedding_column must be provided or configured" + ): + await vs.ainitialize_auto_vector_embeddings(model_id="test-model") diff --git a/tests/test_vectorstore.py b/tests/test_vectorstore.py index 0ec411ff..a4fa547d 100644 --- a/tests/test_vectorstore.py +++ b/tests/test_vectorstore.py @@ -18,6 +18,7 @@ import uuid from threading import Thread from typing import Sequence +from unittest.mock import AsyncMock, MagicMock, patch import pytest import pytest_asyncio @@ -30,6 +31,10 @@ from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine from langchain_google_alloydb_pg import AlloyDBEngine, AlloyDBVectorStore, Column +from langchain_google_alloydb_pg.indexes import ( + DistanceStrategy, + ScaNNIndex, +) DEFAULT_TABLE = "test_table" + str(uuid.uuid4()) DEFAULT_TABLE_SYNC = "test_table_sync" + str(uuid.uuid4()) @@ -39,7 +44,7 @@ VECTOR_SIZE = 768 embeddings_service = DeterministicFakeEmbedding(size=VECTOR_SIZE) -host = os.environ["IP_ADDRESS"] +host = os.environ.get("IP_ADDRESS", "127.0.0.1") texts = ["foo", "bar", "baz"] metadatas = [{"page": str(i), "source": "google.com"} for i in range(len(texts))] @@ -121,13 +126,20 @@ def password(self) -> str: @pytest_asyncio.fixture(scope="class") async def engine(self, db_project, db_region, db_cluster, db_instance, db_name): - engine = await AlloyDBEngine.afrom_instance( - project_id=db_project, - cluster=db_cluster, - instance=db_instance, - region=db_region, - database=db_name, - ) + host = os.environ.get("OMNI_HOST") or os.environ.get("IP_ADDRESS") + user = os.environ.get("OMNI_USER") or os.environ.get("DB_USER", "postgres") + password = os.environ.get("OMNI_PASSWORD") or os.environ.get("DB_PASSWORD") + if host and password: + connstring = f"postgresql+asyncpg://{user}:{password}@{host}:5432/{db_name}" + engine = AlloyDBEngine.from_connection_string(connstring) + else: + engine = await AlloyDBEngine.afrom_instance( + project_id=db_project, + cluster=db_cluster, + instance=db_instance, + region=db_region, + database=db_name, + ) yield engine await aexecute(engine, f'DROP TABLE IF EXISTS "{DEFAULT_TABLE}"') @@ -745,3 +757,470 @@ async def test_from_engine_loop( def test_get_table_name(self, vs): assert vs.get_table_name() == DEFAULT_TABLE + + def test_live_columnar_engine(self, vs): + """Test enabling columnar engine against live AlloyDB instance.""" + vs.enable_columnar_engine(["content"]) + vs.enable_columnar_engine() + + # Assert functional similarity search still works on columnarized table + vs.add_texts(["Columnar engine test document"]) + results = vs.similarity_search("Columnar test", k=1) + assert len(results) > 0 + assert "Columnar" in results[0].page_content + + def test_live_auto_columnarization(self, vs): + """Test triggering auto columnarization recommendations against live AlloyDB instance.""" + try: + vs.enable_auto_columnarization() + except Exception as e: + if "google_columnar_engine.enabled" in str( + e + ) or "shared_preload_libraries" in str(e): + pytest.skip(f"Columnar engine flag not enabled on instance: {e}") + raise + + # Assert functional similarity search still works after auto columnarization + vs.add_texts(["Auto columnarization test document"]) + results = vs.similarity_search("Auto columnarization", k=1) + assert len(results) > 0 + assert "Auto columnarization" in results[0].page_content + + def test_live_vector_assist(self, engine): + """Test vector assist spec definition, application, and recommendations against live AlloyDB instance.""" + table_name = "va_live_sync_table_" + str(uuid.uuid4()).replace("-", "_") + engine._run_as_sync( + aexecute(engine, f'DROP TABLE IF EXISTS "{table_name}" CASCADE;') + ) + engine._run_as_sync( + aexecute( + engine, + f""" + CREATE TABLE "{table_name}" ( + langchain_id uuid PRIMARY KEY, + content text, + embedding vector({VECTOR_SIZE}), + meta jsonb + ); + """, + ) + ) + engine._run_as_sync( + aexecute( + engine, + f""" + INSERT INTO "{table_name}" (langchain_id, content, embedding, meta) + SELECT + gen_random_uuid(), + 'Content ' || i, + (SELECT array_agg((random() * 2 - 1)::float4)::vector({VECTOR_SIZE}) FROM generate_series(1, {VECTOR_SIZE})), + '{{"page": 1}}'::jsonb + FROM generate_series(1, 100) AS i; + """, + ) + ) + vs = AlloyDBVectorStore.create_sync( + engine, + embedding_service=embeddings_service, + table_name=table_name, + metadata_json_column="meta", + ) + specs = vs.define_vector_assist_spec() + assert isinstance(specs, list) + assert len(specs) > 0 + apply_res = vs.apply_vector_assist_spec() + assert isinstance(apply_res, list) + recs = vs.get_vector_assist_recommendations() + assert isinstance(recs, list) + engine._run_as_sync( + aexecute(engine, f'DROP TABLE IF EXISTS "{table_name}" CASCADE;') + ) + + +class TestVectorStoreUnit: + @pytest.fixture + def vs(self): + vs = AlloyDBVectorStore.__new__(AlloyDBVectorStore) + vs._engine = MagicMock() + mock_vs = MagicMock() + vs._PGVectorStore__vs = mock_vs + vs._AlloyDBVectorStore__vs = mock_vs + + def mock_sync(coro): + if hasattr(coro, "close"): + coro.close() + return getattr(vs._engine._run_as_sync, "return_value", None) + + async def mock_async(coro): + if hasattr(coro, "close"): + coro.close() + return getattr(vs._engine._run_as_async, "return_value", None) + + vs._engine._run_as_sync = MagicMock(side_effect=mock_sync) + vs._engine._run_as_async = AsyncMock(side_effect=mock_async) + return vs + + def test_enable_columnar_engine(self, vs): + """Test enabling the columnar engine triggers the appropriate sync method on the underlying store.""" + # 1. Call sync facade with specific column + vs.enable_columnar_engine(["content"]) + # 2. Assert delegation to underlying async implementation + vs._PGVectorStore__vs.aenable_columnar_engine.assert_called_once_with( + ["content"] + ) + + def test_enable_columnar_engine_without_columns(self, vs): + """Test enabling columnar engine without columns triggers default on underlying store.""" + # 1. Call sync facade with default columns (None) + vs.enable_columnar_engine() + # 2. Assert delegation with None + vs._PGVectorStore__vs.aenable_columnar_engine.assert_called_once_with(None) + + @pytest.mark.asyncio + async def test_aenable_columnar_engine(self, vs): + """Test enabling the columnar engine triggers the appropriate async method on the underlying store.""" + # 1. Call async facade with specific column + await vs.aenable_columnar_engine(["content"]) + # 2. Assert delegation to underlying async implementation + vs._PGVectorStore__vs.aenable_columnar_engine.assert_called_once_with( + ["content"] + ) + + @pytest.mark.asyncio + async def test_aenable_columnar_engine_without_columns(self, vs): + """Test enabling columnar engine without columns asynchronously.""" + # 1. Call async facade with default columns (None) + await vs.aenable_columnar_engine() + # 2. Assert delegation with None + vs._PGVectorStore__vs.aenable_columnar_engine.assert_called_once_with(None) + + def test_enable_auto_columnarization(self, vs): + """Test enabling auto columnarization triggers the sync engine wrapper.""" + # 1. Call sync facade + vs.enable_auto_columnarization() + # 2. Assert delegation + vs._PGVectorStore__vs.aenable_auto_columnarization.assert_called_once_with() + + @pytest.mark.asyncio + async def test_aenable_auto_columnarization(self, vs): + """Test enabling auto columnarization triggers the async engine wrapper.""" + # 1. Call async facade + await vs.aenable_auto_columnarization() + # 2. Assert delegation + vs._PGVectorStore__vs.aenable_auto_columnarization.assert_called_once_with() + + def test_define_vector_assist_spec(self, vs): + """Test definition of vector assist specification.""" + # 1. Mock engine sync return value + expected = [{"spec": "ok"}] + vs._engine._run_as_sync.return_value = expected + # 2. Execute define_vector_assist_spec + res = vs.define_vector_assist_spec() + # 3. Assert return value and underlying call + assert res == expected + vs._PGVectorStore__vs.adefine_vector_assist_spec.assert_called_once_with() + + @pytest.mark.asyncio + async def test_adefine_vector_assist_spec(self, vs): + """Test definition of vector assist specification asynchronously.""" + # 1. Mock engine async return value + expected = [{"spec": "ok"}] + vs._engine._run_as_async.return_value = expected + # 2. Execute adefine_vector_assist_spec + res = await vs.adefine_vector_assist_spec() + # 3. Assert return value and underlying call + assert res == expected + vs._PGVectorStore__vs.adefine_vector_assist_spec.assert_called_once_with() + + def test_apply_vector_assist_spec(self, vs): + """Test applying vector assist specifications.""" + # 1. Mock engine sync return value + expected = [{"apply": "ok"}] + vs._engine._run_as_sync.return_value = expected + # 2. Execute apply_vector_assist_spec + res = vs.apply_vector_assist_spec() + # 3. Assert return value and underlying call + assert res == expected + vs._PGVectorStore__vs.aapply_vector_assist_spec.assert_called_once_with( + spec_id=None + ) + + @pytest.mark.asyncio + async def test_aapply_vector_assist_spec(self, vs): + """Test applying vector assist specifications asynchronously.""" + # 1. Mock engine async return value + expected = [{"apply": "ok"}] + vs._engine._run_as_async.return_value = expected + # 2. Execute aapply_vector_assist_spec + res = await vs.aapply_vector_assist_spec() + # 3. Assert return value and underlying call + assert res == expected + vs._PGVectorStore__vs.aapply_vector_assist_spec.assert_called_once_with( + spec_id=None + ) + + def test_apply_vector_assist_spec_id_passthrough(self, vs): + """Test spec_id pass-through in apply_vector_assist_spec and aapply_vector_assist_spec.""" + # Test with a valid spec_id + vs.apply_vector_assist_spec(spec_id="spec_123") + vs._PGVectorStore__vs.aapply_vector_assist_spec.assert_called_with( + spec_id="spec_123" + ) + + # Test with empty string spec_id to verify if it is passed through or treated as None + vs.apply_vector_assist_spec(spec_id="") + vs._PGVectorStore__vs.aapply_vector_assist_spec.assert_called_with(spec_id="") + + @pytest.mark.asyncio + async def test_aapply_vector_assist_spec_id_passthrough(self, vs): + """Test spec_id pass-through in aapply_vector_assist_spec asynchronously.""" + await vs.aapply_vector_assist_spec(spec_id="spec_456") + vs._PGVectorStore__vs.aapply_vector_assist_spec.assert_called_with( + spec_id="spec_456" + ) + + await vs.aapply_vector_assist_spec(spec_id="") + vs._PGVectorStore__vs.aapply_vector_assist_spec.assert_called_with(spec_id="") + + @pytest.mark.asyncio + async def test_quote_ident_adversarial_payloads(self): + """Test _quote_ident in aenable_columnar_engine, adefine_vector_assist_spec, aapply_vector_assist_spec, aget_vector_assist_recommendations with adversarial payloads.""" + from langchain_google_alloydb_pg.async_vectorstore import ( + AsyncAlloyDBVectorStore, + ) + + store = AsyncAlloyDBVectorStore.__new__(AsyncAlloyDBVectorStore) + store.table_name = 'table"; DROP TABLE users;--' + store.schema_name = 'public"; DROP TABLE users;--' + store.embedding_column = 'embedding"; DROP TABLE users;--' + store.engine = MagicMock() + + mock_conn = AsyncMock() + mock_conn.execute = AsyncMock() + mock_conn.commit = AsyncMock() + mock_conn.fetchall = MagicMock(return_value=[]) + + mock_ctx = AsyncMock() + mock_ctx.__aenter__ = AsyncMock(return_value=mock_conn) + mock_ctx.__aexit__ = AsyncMock(return_value=None) + + store.engine._pool = MagicMock() + store.engine._pool.connect = MagicMock(return_value=mock_ctx) + + # 1. Test aenable_columnar_engine + await store.aenable_columnar_engine(columns=['col"; DROP TABLE--']) + # Verify that the table_identifier in params has properly escaped quotes + call_args = mock_conn.execute.call_args + assert call_args is not None + params = call_args[0][1] + assert ( + params["table_name"] + == '"public""; DROP TABLE users;--"."table""; DROP TABLE users;--"' + ) + assert params["columns"] == '"col""; DROP TABLE--"' + + # 2. Test adefine_vector_assist_spec + mock_conn.execute.reset_mock() + mock_result = MagicMock() + mock_result.mappings.return_value = [] + mock_conn.execute.return_value = mock_result + await store.adefine_vector_assist_spec() + call_args = mock_conn.execute.call_args + assert call_args is not None + params = call_args[0][1] + assert ( + params["table_name"] + == '"public""; DROP TABLE users;--"."table""; DROP TABLE users;--"' + ) + + # 3. Test aapply_vector_assist_spec + mock_conn.execute.reset_mock() + mock_conn.execute.return_value = mock_result + await store.aapply_vector_assist_spec() + call_args = mock_conn.execute.call_args + assert call_args is not None + params = call_args[0][1] + assert ( + params["table_name"] + == '"public""; DROP TABLE users;--"."table""; DROP TABLE users;--"' + ) + + # 4. Test aget_vector_assist_recommendations + mock_conn.execute.reset_mock() + mock_mappings = MagicMock() + mock_mappings.first.return_value = None + mock_result.mappings.return_value = mock_mappings + mock_conn.execute.return_value = mock_result + await store.aget_vector_assist_recommendations() + call_args = mock_conn.execute.call_args + assert call_args is not None + params = call_args[0][1] + assert ( + params["table_name"] + == '"public""; DROP TABLE users;--"."table""; DROP TABLE users;--"' + ) + + @pytest.mark.asyncio + async def test_enable_columnar_engine_schema_handling(self): + """Test aenable_columnar_engine with schema="" and schema=None to verify if it uses "public" or CURRENT_SCHEMA().""" + from langchain_google_alloydb_pg.async_vectorstore import ( + AsyncAlloyDBVectorStore, + ) + + store = AsyncAlloyDBVectorStore.__new__(AsyncAlloyDBVectorStore) + store.table_name = "test_table" + store.embedding_column = "embedding" + store.engine = MagicMock() + + mock_conn = AsyncMock() + mock_conn.execute = AsyncMock() + mock_conn.commit = AsyncMock() + + mock_ctx = AsyncMock() + mock_ctx.__aenter__ = AsyncMock(return_value=mock_conn) + mock_ctx.__aexit__ = AsyncMock(return_value=None) + + store.engine._pool = MagicMock() + store.engine._pool.connect = MagicMock(return_value=mock_ctx) + + # Test with schema_name = "" -> should NOT hardcode "public", should use CURRENT_SCHEMA() + store.schema_name = "" + await store.aenable_columnar_engine(columns=["content"]) + call_args = mock_conn.execute.call_args + params = call_args[0][1] + assert params["table_name"] == '"test_table"' + + # Test with schema_name = None -> should NOT hardcode "public", should use CURRENT_SCHEMA() + store.schema_name = None + await store.aenable_columnar_engine(columns=["content"]) + call_args = mock_conn.execute.call_args + params = call_args[0][1] + assert params["table_name"] == '"test_table"' + + @pytest.mark.asyncio + async def test_apply_vector_index_schema_handling(self): + """Test aapply_vector_index with schema_name=None and schema_name="" to verify it does not produce "None"."table" or ""."table".""" + from langchain_google_alloydb_pg.async_vectorstore import ( + AsyncAlloyDBVectorStore, + ) + from langchain_google_alloydb_pg.indexes import HNSWIndex + + store = AsyncAlloyDBVectorStore.__new__(AsyncAlloyDBVectorStore) + store.table_name = "test_table" + store.embedding_column = "embedding" + store.engine = MagicMock() + + mock_conn = AsyncMock() + mock_conn.execute = AsyncMock() + + mock_ctx = AsyncMock() + mock_ctx.__aenter__ = AsyncMock(return_value=mock_conn) + mock_ctx.__aexit__ = AsyncMock(return_value=None) + + store.engine._pool = MagicMock() + store.engine._pool.begin = MagicMock(return_value=mock_ctx) + + index = HNSWIndex(name="idx_test") + + # Test with schema_name = None -> should produce ON "test_table" + store.schema_name = None + await store.aapply_vector_index(index) + call_args = mock_conn.execute.call_args + stmt = str(call_args[0][0]) + assert 'ON "test_table"' in stmt + assert '"None"' not in stmt + + # Test with schema_name = "" -> should produce ON "test_table" + store.schema_name = "" + await store.aapply_vector_index(index) + call_args = mock_conn.execute.call_args + stmt = str(call_args[0][0]) + assert 'ON "test_table"' in stmt + assert '""."test_table"' not in stmt + + @pytest.mark.asyncio + async def test_set_maintenance_work_mem_retention(self): + """Test aset_maintenance_work_mem to verify it is deprecated and issues a DeprecationWarning.""" + from langchain_google_alloydb_pg.async_vectorstore import ( + AsyncAlloyDBVectorStore, + ) + + store = AsyncAlloyDBVectorStore.__new__(AsyncAlloyDBVectorStore) + + with pytest.deprecated_call(): + await store.aset_maintenance_work_mem(num_leaves=10, vector_size=768) + + def test_get_vector_assist_recommendations(self, vs): + """Test retrieving vector assist recommendations.""" + # 1. Mock engine sync return value + expected = [{"rec": "ok"}] + vs._engine._run_as_sync.return_value = expected + # 2. Execute get_vector_assist_recommendations + res = vs.get_vector_assist_recommendations() + # 3. Assert return value and underlying call + assert res == expected + vs._PGVectorStore__vs.aget_vector_assist_recommendations.assert_called_once_with() + + @pytest.mark.asyncio + async def test_aget_vector_assist_recommendations(self, vs): + """Test retrieving vector assist recommendations asynchronously.""" + # 1. Mock engine async return value + expected = [{"rec": "ok"}] + vs._engine._run_as_async.return_value = expected + # 2. Execute aget_vector_assist_recommendations + res = await vs.aget_vector_assist_recommendations() + # 3. Assert return value and underlying call + assert res == expected + vs._PGVectorStore__vs.aget_vector_assist_recommendations.assert_called_once_with() + + def test_initialize_auto_vector_embeddings(self, vs): + """Test initializing auto vector embeddings with default arguments.""" + # 1. Execute initialize_auto_vector_embeddings + vs.initialize_auto_vector_embeddings( + model_id="test-model", + ) + # 2. Assert delegation to underlying async method + vs._PGVectorStore__vs.ainitialize_auto_vector_embeddings.assert_called_once_with( + "test-model", None, None, None + ) + + def test_initialize_auto_vector_embeddings_with_columns(self, vs): + """Test initializing auto vector embeddings with custom columns.""" + # 1. Execute initialize_auto_vector_embeddings with custom column parameters + vs.initialize_auto_vector_embeddings( + model_id="test-model", + content_column="custom_content", + embedding_column="custom_embedding", + schema_name="myschema", + ) + # 2. Assert custom parameters passed to underlying async method + vs._PGVectorStore__vs.ainitialize_auto_vector_embeddings.assert_called_once_with( + "test-model", "custom_content", "custom_embedding", "myschema" + ) + + @pytest.mark.asyncio + async def test_ainitialize_auto_vector_embeddings(self, vs): + """Test initializing auto vector embeddings asynchronously.""" + # 1. Execute ainitialize_auto_vector_embeddings asynchronously + await vs.ainitialize_auto_vector_embeddings( + model_id="test-model", + ) + # 2. Assert delegation + vs._PGVectorStore__vs.ainitialize_auto_vector_embeddings.assert_called_once_with( + "test-model", None, None, None + ) + + @pytest.mark.asyncio + async def test_ainitialize_auto_vector_embeddings_with_columns(self, vs): + """Test initializing auto vector embeddings with custom columns asynchronously.""" + # 1. Execute ainitialize_auto_vector_embeddings with custom parameters + await vs.ainitialize_auto_vector_embeddings( + model_id="test-model", + content_column="custom_content", + embedding_column="custom_embedding", + schema_name="myschema", + ) + # 2. Assert custom parameters passed to underlying async method + vs._PGVectorStore__vs.ainitialize_auto_vector_embeddings.assert_called_once_with( + "test-model", "custom_content", "custom_embedding", "myschema" + )