feat: add SpannerVectorStore for orchestrating and providing utility functions for a Spanner vector store

PiperOrigin-RevId: 854392465

Author: google-genai-bot

src/google/adk/features/_feature_registry.py

@@ -41,6 +41,7 @@ class FeatureName(str, Enum):
   PUBSUB_TOOLSET = "PUBSUB_TOOLSET"
   SPANNER_TOOLSET = "SPANNER_TOOLSET"
   SPANNER_TOOL_SETTINGS = "SPANNER_TOOL_SETTINGS"
+  SPANNER_VECTOR_STORE = "SPANNER_VECTOR_STORE"
   TOOL_CONFIG = "TOOL_CONFIG"
   TOOL_CONFIRMATION = "TOOL_CONFIRMATION"
 
@@ -120,6 +121,9 @@ class FeatureConfig:
     FeatureName.SPANNER_TOOL_SETTINGS: FeatureConfig(
         FeatureStage.EXPERIMENTAL, default_on=True
     ),
+    FeatureName.SPANNER_VECTOR_STORE: FeatureConfig(
+        FeatureStage.EXPERIMENTAL, default_on=True
+    ),
     FeatureName.TOOL_CONFIG: FeatureConfig(
         FeatureStage.EXPERIMENTAL, default_on=True
     ),

src/google/adk/tools/spanner/settings.py

@@ -55,6 +55,74 @@ class QueryResultMode(Enum):
   """
 
 
+class TableColumn(BaseModel):
+  """Represents column configuration, to be used as part of create DDL statement for a new vector store table set up."""
+
+  name: str
+  """Required. The name of the column."""
+
+  type: str
+  """Required. The type of the column.
+
+  For example,
+
+  - GoogleSQL: 'STRING(MAX)', 'INT64', 'FLOAT64', 'BOOL', etc.
+  - PostgreSQL: 'text', 'int8', 'float8', 'boolean', etc.
+  """
+
+  is_nullable: bool = True
+  """Optional. Whether the column is nullable. By default, the column is nullable."""
+
+
+class VectorSearchIndexSettings(BaseModel):
+  """Settings for the index for use with Approximate Nearest Neighbor (ANN) vector similarity search."""
+
+  index_name: str
+  """Required. The name of the vector similarity search index."""
+
+  additional_key_columns: Optional[list[str]] = None
+  """Optional. The list of the additional key column names in the vector similarity search index.
+
+  To further speed up filtering for highly selective filtering columns, organize
+  them as additional keys in the vector index after the embedding column.
+  For example: `category` as additional key column.
+  `CREATE VECTOR INDEX ON documents(embedding, category);`
+  """
+
+  additional_storing_columns: Optional[list[str]] = None
+  """Optional. The list of the storing column names in the vector similarity search index.
+
+  This enables filtering while walking the vector index, removing unqualified
+  rows early.
+  For example: `category` as storing column.
+  `CREATE VECTOR INDEX ON documents(embedding) STORING (category);`
+  """
+
+  tree_depth: int = 2
+  """Required. The tree depth (level). This value can be either 2 or 3.
+
+  A tree with 2 levels only has leaves (num_leaves) as nodes.
+  If the dataset has more than 100 million rows,
+  then you can use a tree with 3 levels and add branches (num_branches) to
+  further partition the dataset.
+  """
+
+  num_leaves: int = 1000
+  """Required. The number of leaves (i.e. potential partitions) for the vector data.
+
+  You can designate num_leaves for trees with 2 or 3 levels.
+  We recommend that the number of leaves is number_of_rows_in_dataset/1000.
+  """
+
+  num_branches: Optional[int] = None
+  """Optional. The number of branches to further parititon the vector data.
+
+  You can only designate num_branches for trees with 3 levels.
+  The number of branches must be fewer than the number of leaves
+  We recommend that the number of leaves is between 1000 and sqrt(number_of_rows_in_dataset).
+  """
+
+
 class SpannerVectorStoreSettings(BaseModel):
   """Settings for Spanner Vector Store.
 
@@ -86,27 +154,28 @@ class SpannerVectorStoreSettings(BaseModel):
 
   vertex_ai_embedding_model_name: str
   """Required. The Vertex AI embedding model name, which is used to generate embeddings for vector store and vector similarity search.
-    For example, 'text-embedding-005'.
 
-    Note: the output dimensionality of the embedding model should be the same as the value specified in the `vector_length` field.
-    Otherwise, a runtime error might be raised during a query.
+  For example, 'text-embedding-005'.
+
+  Note: the output dimensionality of the embedding model should be the same as the value specified in the `vector_length` field.
+  Otherwise, a runtime error might be raised during a query.
   """
 
-  selected_columns: List[str] = []
+  selected_columns: list[str] = []
   """Required. The vector store table columns to return in the vector similarity search result.
 
-    By default, only the `content_column` value and the distance value are returned.
-    If sepecified, the list of selected columns and the distance value are returned.
-    For example, if `selected_columns` is ['col1', 'col2'], then the result will contain the values of 'col1' and 'col2' columns and the distance value.
+  By default, only the `content_column` value and the distance value are returned.
+  If sepecified, the list of selected columns and the distance value are returned.
+  For example, if `selected_columns` is ['col1', 'col2'], then the result will contain the values of 'col1' and 'col2' columns and the distance value.
   """
 
   nearest_neighbors_algorithm: NearestNeighborsAlgorithm = (
       "EXACT_NEAREST_NEIGHBORS"
   )
   """The algorithm used to perform vector similarity search. This value can be EXACT_NEAREST_NEIGHBORS or APPROXIMATE_NEAREST_NEIGHBORS.
 
-    For more details about EXACT_NEAREST_NEIGHBORS, see https://docs.cloud.google.com/spanner/docs/find-k-nearest-neighbors
-    For more details about APPROXIMATE_NEAREST_NEIGHBORS, see https://docs.cloud.google.com/spanner/docs/find-approximate-nearest-neighbors
+  For more details about EXACT_NEAREST_NEIGHBORS, see https://docs.cloud.google.com/spanner/docs/find-k-nearest-neighbors
+  For more details about APPROXIMATE_NEAREST_NEIGHBORS, see https://docs.cloud.google.com/spanner/docs/find-approximate-nearest-neighbors
   """
 
   top_k: int = 4
@@ -118,16 +187,41 @@ class SpannerVectorStoreSettings(BaseModel):
   num_leaves_to_search: Optional[int] = None
   """Optional. This option specifies how many leaf nodes of the index are searched.
 
-    Note: this option is only used when the nearest neighbors search algorithm (`nearest_neighbors_algorithm`) is APPROXIMATE_NEAREST_NEIGHBORS.
-    For more details, see https://docs.cloud.google.com/spanner/docs/vector-index-best-practices
+  Note: This option is only used when the nearest neighbors search algorithm (`nearest_neighbors_algorithm`) is APPROXIMATE_NEAREST_NEIGHBORS.
+  For more details, see https://docs.cloud.google.com/spanner/docs/vector-index-best-practices
   """
 
   additional_filter: Optional[str] = None
   """Optional. An optional filter to apply to the search query. If provided, this will be added to the WHERE clause of the final query."""
 
+  vector_search_index_settings: Optional[VectorSearchIndexSettings] = None
+  """Optional. Settings for the index for use with Approximate Nearest Neighbor (ANN) in the vector store.
+
+  Note: This option is only required when the nearest neighbors search algorithm (`nearest_neighbors_algorithm`) is APPROXIMATE_NEAREST_NEIGHBORS.
+  For more details, see https://docs.cloud.google.com/spanner/docs/vector-indexes
+  """
+
+  additional_columns_to_setup: Optional[list[TableColumn]] = None
+  """Optional. A list of supplemental columns to be created when initializing a new vector store table or inserting content rows.
+
+  Note: This configuration is only utilized during the initial table setup
+  or when inserting content rows.
+  """
+
+  primary_key_columns: Optional[list[str]] = None
+  """Optional. Specifies the column names to be used as the primary key for a new vector store table.
+
+  If provided, every column name listed here must be defined within
+  `additional_columns_to_setup`. If this field is omitted (set to `None`),
+  defaults to a single primary key column named `id` which automatically
+  generates UUIDs for each entry.
+
+  Note: This field is only used during the creation phase of a new vector store.
+  """
+
   @model_validator(mode="after")
   def __post_init__(self):
-    """Validate the embedding settings."""
+    """Validate the vector store settings."""
     if not self.vector_length or self.vector_length <= 0:
       raise ValueError(
           "Invalid vector length in the Spanner vector store settings."
@@ -136,6 +230,17 @@ def __post_init__(self):
     if not self.selected_columns:
       self.selected_columns = [self.content_column]
 
+    if self.primary_key_columns:
+      cols = {self.content_column, self.embedding_column}
+      if self.additional_columns_to_setup:
+        cols.update({c.name for c in self.additional_columns_to_setup})
+
+      for pk in self.primary_key_columns:
+        if pk not in cols:
+          raise ValueError(
+              f"Primary key column '{pk}' not found in column definitions."
+          )
+
     return self