Update vector_index for export() to take list of VectorIndexDef. (#304)

Author: badmonster0

README.md

@@ -80,7 +80,10 @@ def text_embedding_flow(flow_builder: cocoindex.FlowBuilder, data_scope: cocoind
         "doc_embeddings",
         cocoindex.storages.Postgres(),
         primary_key_fields=["filename", "location"],
-        vector_index=[("embedding", cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY)])
+        vector_indexes=[
+            cocoindex.VectorIndexDef(
+                field_name="embedding",
+                metric=cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY)])
 ```
 
 It defines an index flow like this:

docs/docs/core/flow_def.mdx

@@ -259,8 +259,10 @@ Export must happen at the top level of a flow, i.e. not within any child scopes
 
 *   `name`: the name to identify the export target.
 *   `target_spec`: the storage spec as the export target.
-*   `primary_key_fields` (optional): the fields to be used as primary key. Types of the fields must be supported as key fields. See [Key Types](data_types#key-types) for more details.
-*   `vector_index` (optional): the fields to create vector index. Each item is a tuple of a field name and a similarity metric. See [Vector Type](data_types#vector-type) for more details about supported similarity metrics.
+*   `primary_key_fields` (`Sequence[str]`): the fields to be used as primary key. Types of the fields must be supported as key fields. See [Key Types](data_types#key-types) for more details.
+*   `vector_indexes` (`Sequence[VectorIndexDef]`, optional): the fields to create vector index. `VectorIndexDef` has the following fields:
+    *   `field_name`: the field to create vector index.
+    *   `metric`: the similarity metric to use. See [Vector Type](data_types#vector-type) for more details about supported similarity metrics.
 *   `setup_by_user` (optional):
      whether the export target is setup by user.
      By default, CocoIndex is managing the target setup (surfaced by the `cocoindex setup` CLI subcommand), e.g. create related tables/collections/etc. with compatible schema, and update them upon change.

docs/docs/getting_started/quickstart.md

@@ -97,7 +97,10 @@ def text_embedding_flow(flow_builder: cocoindex.FlowBuilder, data_scope: cocoind
         "doc_embeddings",
         cocoindex.storages.Postgres(),
         primary_key_fields=["filename", "location"],
-        vector_index=[("embedding", cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY)])
+        vector_indexes=[
+            cocoindex.VectorIndexDef(
+                field_name="embedding",
+                metric=cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY)])
 ```
 
 Notes:

examples/code_embedding/main.py

@@ -41,7 +41,10 @@ def code_embedding_flow(flow_builder: cocoindex.FlowBuilder, data_scope: cocoind
         "code_embeddings",
         cocoindex.storages.Postgres(),
         primary_key_fields=["filename", "location"],
-        vector_index=[("embedding", cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY)])
+        vector_indexes=[
+            cocoindex.VectorIndexDef(
+                field_name="embedding",
+                metric=cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY)])
 
 
 query_handler = cocoindex.query.SimpleSemanticsQueryHandler(

examples/docs_to_kg/main.py

@@ -93,15 +93,13 @@ def docs_to_kg_flow(flow_builder: cocoindex.FlowBuilder, data_scope: cocoindex.D
             ),
             nodes={
                 "Entity": cocoindex.storages.Neo4jRelationshipNodeSpec(
-                    index_options=cocoindex.IndexOptions(
-                        primary_key_fields=["value"],
-                        vector_index_defs=[
-                            cocoindex.VectorIndexDef(
-                                field_name="embedding",
-                                metric=cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY,
-                            ),
-                        ],
-                    ),
+                    primary_key_fields=["value"],
+                    vector_indexes=[
+                        cocoindex.VectorIndexDef(
+                            field_name="embedding",
+                            metric=cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY,
+                        ),
+                    ],
                 ),
             },
         ),

examples/gdrive_text_embedding/main.py

@@ -37,7 +37,10 @@ def gdrive_text_embedding_flow(flow_builder: cocoindex.FlowBuilder, data_scope:
         "doc_embeddings",
         cocoindex.storages.Postgres(),
         primary_key_fields=["filename", "location"],
-        vector_index=[("embedding", cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY)])
+        vector_indexes=[
+            cocoindex.VectorIndexDef(
+                field_name="embedding",
+                metric=cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY)])
 
 query_handler = cocoindex.query.SimpleSemanticsQueryHandler(
     name="SemanticsSearch",

examples/pdf_embedding/main.py

@@ -63,7 +63,10 @@ def pdf_embedding_flow(flow_builder: cocoindex.FlowBuilder, data_scope: cocoinde
         "doc_embeddings",
         cocoindex.storages.Postgres(),
         primary_key_fields=["id"],
-        vector_index=[("embedding", cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY)])
+        vector_indexes=[
+            cocoindex.VectorIndexDef(
+                field_name="embedding",
+                metric=cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY)])
 
 query_handler = cocoindex.query.SimpleSemanticsQueryHandler(
     name="SemanticsSearch",

examples/text_embedding/main.py

@@ -35,7 +35,10 @@ def text_embedding_flow(flow_builder: cocoindex.FlowBuilder, data_scope: cocoind
         "doc_embeddings",
         cocoindex.storages.Postgres(),
         primary_key_fields=["filename", "location"],
-        vector_index=[("embedding", cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY)])
+        vector_indexes=[
+            cocoindex.VectorIndexDef(
+                field_name="embedding",
+                metric=cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY)])
 
 query_handler = cocoindex.query.SimpleSemanticsQueryHandler(
     name="SemanticsSearch",

python/cocoindex/flow.py

@@ -267,21 +267,27 @@ def collect(self, **kwargs):
             self._engine_data_collector, regular_kwargs, auto_uuid_field)
 
     def export(self, name: str, target_spec: op.StorageSpec, /, *,
-              primary_key_fields: Sequence[str] | None = None,
+              primary_key_fields: Sequence[str],
+              vector_indexes: Sequence[index.VectorIndexDef] = (),
               vector_index: Sequence[tuple[str, index.VectorSimilarityMetric]] = (),
               setup_by_user: bool = False):
         """
         Export the collected data to the specified target.
+
+        `vector_index` is for backward compatibility only. Please use `vector_indexes` instead.
         """
-        index_options: dict[str, Any] = {}
-        if primary_key_fields is not None:
-            index_options["primary_key_fields"] = primary_key_fields
-        index_options["vector_index_defs"] = [
-            {"field_name": field_name, "metric": metric.value}
-            for field_name, metric in vector_index]
+        # For backward compatibility only.
+        if len(vector_indexes) == 0 and len(vector_index) > 0:
+            vector_indexes = [index.VectorIndexDef(field_name=field_name, metric=metric)
+                             for field_name, metric in vector_index]
+
+        index_options = index.IndexOptions(
+            primary_key_fields=primary_key_fields,
+            vector_indexes=vector_indexes,
+        )
         self._flow_builder_state.engine_flow_builder.export(
             name, _spec_kind(target_spec), dump_engine_object(target_spec),
-            index_options, self._engine_data_collector, setup_by_user)
+            dump_engine_object(index_options), self._engine_data_collector, setup_by_user)
 
 
 _flow_name_builder = _NameBuilder()

python/cocoindex/index.py

@@ -1,6 +1,6 @@
 from enum import Enum
 from dataclasses import dataclass
-
+from typing import Sequence
 class VectorSimilarityMetric(Enum):
     COSINE_SIMILARITY = "CosineSimilarity"
     L2_DISTANCE = "L2Distance"
@@ -19,5 +19,5 @@ class IndexOptions:
     """
     Options for an index.
     """
-    primary_key_fields: list[str] | None = None
-    vector_index_defs: list[VectorIndexDef] | None = None
+    primary_key_fields: Sequence[str]
+    vector_indexes: Sequence[VectorIndexDef] = ()