cocoindex 0.2.13__tar.gz → 0.2.14__tar.gz

{cocoindex-0.2.13 → cocoindex-0.2.14}/Cargo.lock

@@ -1283,7 +1283,7 @@ dependencies = [
 
 [[package]]
 name = "cocoindex"
-version = "0.2.13"
+version = "0.2.14"
 dependencies = [
  "anyhow",
  "async-openai",

{cocoindex-0.2.13 → cocoindex-0.2.14}/Cargo.toml

@@ -2,7 +2,7 @@
 name = "cocoindex"
 # Version used for local development is always higher than others to take precedence.
 # Will be overridden for specific release versions.
-version = "0.2.13"
+version = "0.2.14"
 edition = "2024"
 rust-version = "1.89"
 license = "Apache-2.0"

{cocoindex-0.2.13 → cocoindex-0.2.14}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cocoindex
-Version: 0.2.13
+Version: 0.2.14
 Classifier: Development Status :: 3 - Alpha
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Operating System :: OS Independent

{cocoindex-0.2.13 → cocoindex-0.2.14}/THIRD_PARTY_NOTICES.html

@@ -2428,7 +2428,7 @@ Software.
                 <h3 id="Apache-2.0">Apache License 2.0</h3>
                 <h4>Used by:</h4>
                 <ul class="license-used-by">
-                    <li><a href=" https://crates.io/crates/cocoindex ">cocoindex 0.2.13</a></li>
+                    <li><a href=" https://crates.io/crates/cocoindex ">cocoindex 0.2.14</a></li>
                     <li><a href=" https://github.com/awesomized/crc-fast-rust ">crc-fast 1.3.0</a></li>
                     <li><a href=" https://github.com/qdrant/rust-client ">qdrant-client 1.15.0</a></li>
                 </ul>

{cocoindex-0.2.13 → cocoindex-0.2.14}/docs/docs/core/flow_def.mdx

@@ -313,6 +313,7 @@ Types of the fields must be key types. See [Key Types](data_types#key-types) for
 
     *   `field_name`: the field to create vector index.
     *   `metric`: the similarity metric to use.
+    *   `method` (optional): the index algorithm and optional tuning parameters. Leave unset to use the target default (HNSW for Postgres). Use `cocoindex.HnswVectorIndexMethod()` or `cocoindex.IvfFlatVectorIndexMethod()` to customize the method and its parameters.
 
 #### Similarity Metrics
 

{cocoindex-0.2.13 → cocoindex-0.2.14}/docs/docs/examples/examples/simple_vector_index.md

@@ -105,6 +105,16 @@ doc_embeddings.export(
 CocoIndex supports other vector databases as well, with 1-line switch.
 <DocumentationButton url="https://cocoindex.io/docs/ops/targets" text="Targets" />
 
+Need IVFFlat or custom HNSW parameters? Pass a method, for example:
+
+```python
+cocoindex.VectorIndexDef(
+    field_name="embedding",
+    metric=cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY,
+    method=cocoindex.IvfFlatVectorIndexMethod(lists=200),
+)
+```
+
 ## Query the index
 
 ### Define a shared flow for both indexing and querying

{cocoindex-0.2.13 → cocoindex-0.2.14}/python/cocoindex/__init__.py

@@ -21,7 +21,13 @@ from .flow import add_flow_def, remove_flow  # DEPRECATED
 from .flow import update_all_flows_async, setup_all_flows, drop_all_flows
 from .lib import settings, init, start_server, stop
 from .llm import LlmSpec, LlmApiType
-from .index import VectorSimilarityMetric, VectorIndexDef, IndexOptions
+from .index import (
+    VectorSimilarityMetric,
+    VectorIndexDef,
+    IndexOptions,
+    HnswVectorIndexMethod,
+    IvfFlatVectorIndexMethod,
+)
 from .setting import DatabaseConnectionSpec, Settings, ServerSettings
 from .setting import get_app_namespace
 from .query_handler import QueryHandlerResultFields, QueryInfo, QueryOutput
@@ -82,6 +88,8 @@ __all__ = [
     "VectorSimilarityMetric",
     "VectorIndexDef",
     "IndexOptions",
+    "HnswVectorIndexMethod",
+    "IvfFlatVectorIndexMethod",
     # Settings
     "DatabaseConnectionSpec",
     "Settings",

{cocoindex-0.2.13 → cocoindex-0.2.14}/python/cocoindex/index.py

@@ -1,6 +1,6 @@
 from enum import Enum
 from dataclasses import dataclass
-from typing import Sequence
+from typing import Sequence, Union
 
 
 class VectorSimilarityMetric(Enum):
@@ -9,6 +9,26 @@ class VectorSimilarityMetric(Enum):
     INNER_PRODUCT = "InnerProduct"
 
 
+@dataclass
+class HnswVectorIndexMethod:
+    """HNSW vector index parameters."""
+
+    kind: str = "Hnsw"
+    m: int | None = None
+    ef_construction: int | None = None
+
+
+@dataclass
+class IvfFlatVectorIndexMethod:
+    """IVFFlat vector index parameters."""
+
+    kind: str = "IvfFlat"
+    lists: int | None = None
+
+
+VectorIndexMethod = Union[HnswVectorIndexMethod, IvfFlatVectorIndexMethod]
+
+
 @dataclass
 class VectorIndexDef:
     """
@@ -17,6 +37,7 @@ class VectorIndexDef:
 
     field_name: str
     metric: VectorSimilarityMetric
+    method: VectorIndexMethod | None = None
 
 
 @dataclass

{cocoindex-0.2.13 → cocoindex-0.2.14}/python/cocoindex/targets/lancedb.py

@@ -296,6 +296,12 @@ class _Connector:
     ) -> _State:
         if len(key_fields_schema) != 1:
             raise ValueError("LanceDB only supports a single key field")
+        if index_options.vector_indexes is not None:
+            for vector_index in index_options.vector_indexes:
+                if vector_index.method is not None:
+                    raise ValueError(
+                        "Vector index method is not configurable for LanceDB yet"
+                    )
         return _State(
             key_field_schema=key_fields_schema[0],
             value_fields_schema=value_fields_schema,

{cocoindex-0.2.13 → cocoindex-0.2.14}/src/base/spec.rs

@@ -384,15 +384,72 @@ impl fmt::Display for VectorSimilarityMetric {
     }
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+#[serde(tag = "kind")]
+pub enum VectorIndexMethod {
+    Hnsw {
+        #[serde(default, skip_serializing_if = "Option::is_none")]
+        m: Option<u32>,
+        #[serde(default, skip_serializing_if = "Option::is_none")]
+        ef_construction: Option<u32>,
+    },
+    IvfFlat {
+        #[serde(default, skip_serializing_if = "Option::is_none")]
+        lists: Option<u32>,
+    },
+}
+
+impl VectorIndexMethod {
+    pub fn kind(&self) -> &'static str {
+        match self {
+            Self::Hnsw { .. } => "Hnsw",
+            Self::IvfFlat { .. } => "IvfFlat",
+        }
+    }
+}
+
+impl fmt::Display for VectorIndexMethod {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Self::Hnsw { m, ef_construction } => {
+                let mut parts = Vec::new();
+                if let Some(m) = m {
+                    parts.push(format!("m={}", m));
+                }
+                if let Some(ef) = ef_construction {
+                    parts.push(format!("ef_construction={}", ef));
+                }
+                if parts.is_empty() {
+                    write!(f, "Hnsw")
+                } else {
+                    write!(f, "Hnsw({})", parts.join(","))
+                }
+            }
+            Self::IvfFlat { lists } => {
+                if let Some(lists) = lists {
+                    write!(f, "IvfFlat(lists={lists})")
+                } else {
+                    write!(f, "IvfFlat")
+                }
+            }
+        }
+    }
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
 pub struct VectorIndexDef {
     pub field_name: FieldName,
     pub metric: VectorSimilarityMetric,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub method: Option<VectorIndexMethod>,
 }
 
 impl fmt::Display for VectorIndexDef {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        write!(f, "{}:{}", self.field_name, self.metric)
+        match &self.method {
+            None => write!(f, "{}:{}", self.field_name, self.metric),
+            Some(method) => write!(f, "{}:{}:{}", self.field_name, self.metric, method),
+        }
     }
 }
 

{cocoindex-0.2.13 → cocoindex-0.2.14}/src/llm/gemini.rs

@@ -74,6 +74,30 @@ impl AiStudioClient {
     }
 }
 
+fn build_embed_payload(
+    model: &str,
+    text: &str,
+    task_type: Option<&str>,
+    output_dimension: Option<u32>,
+) -> serde_json::Value {
+    let mut payload = serde_json::json!({
+        "model": model,
+        "content": { "parts": [{ "text": text }] },
+    });
+    if let Some(task_type) = task_type {
+        payload["taskType"] = serde_json::Value::String(task_type.to_string());
+    }
+    if let Some(output_dimension) = output_dimension {
+        payload["outputDimensionality"] = serde_json::json!(output_dimension);
+        if model.starts_with("gemini-embedding-") {
+            payload["config"] = serde_json::json!({
+                "outputDimensionality": output_dimension,
+            });
+        }
+    }
+    payload
+}
+
 #[async_trait]
 impl LlmGenerationClient for AiStudioClient {
     async fn generate<'req>(
@@ -174,16 +198,12 @@ impl LlmEmbeddingClient for AiStudioClient {
         request: super::LlmEmbeddingRequest<'req>,
     ) -> Result<super::LlmEmbeddingResponse> {
         let url = self.get_api_url(request.model, "embedContent");
-        let mut payload = serde_json::json!({
-            "model": request.model,
-            "content": { "parts": [{ "text": request.text }] },
-        });
-        if let Some(task_type) = request.task_type {
-            payload["taskType"] = serde_json::Value::String(task_type.into());
-        }
-        if let Some(output_dimension) = request.output_dimension {
-            payload["outputDimensionality"] = serde_json::Value::Number(output_dimension.into());
-        }
+        let payload = build_embed_payload(
+            request.model,
+            request.text.as_ref(),
+            request.task_type.as_deref(),
+            request.output_dimension,
+        );
         let resp = retryable::run(
             || async {
                 self.client

{cocoindex-0.2.13 → cocoindex-0.2.14}/src/ops/targets/kuzu.rs

@@ -772,6 +772,9 @@ impl TargetFactoryBase for Factory {
         let data_coll_outputs: Vec<TypedExportDataCollectionBuildOutput<Self>> =
             std::iter::zip(data_collections, analyzed_data_colls.into_iter())
                 .map(|(data_coll, analyzed)| {
+                    if !data_coll.index_options.vector_indexes.is_empty() {
+                        api_bail!("Vector indexes are not supported for Kuzu yet");
+                    }
                     fn to_dep_table(
                         field_mapping: &AnalyzedGraphElementFieldMapping,
                     ) -> Result<ReferencedNodeTable> {

{cocoindex-0.2.13 → cocoindex-0.2.14}/src/ops/targets/neo4j.rs

@@ -568,6 +568,9 @@ impl SetupState {
             .map(|f| (f.name.as_str(), &f.value_type.typ))
             .collect::<HashMap<_, _>>();
         for index_def in index_options.vector_indexes.iter() {
+            if index_def.method.is_some() {
+                api_bail!("Vector index method is not configurable for Neo4j yet");
+            }
             sub_components.push(ComponentState {
                 object_label: schema.elem_type.clone(),
                 index_def: IndexDef::from_vector_index_def(

{cocoindex-0.2.13 → cocoindex-0.2.14}/src/ops/targets/postgres.rs

@@ -461,21 +461,51 @@ fn to_vector_similarity_metric_sql(metric: VectorSimilarityMetric) -> &'static s
 }
 
 fn to_index_spec_sql(index_spec: &VectorIndexDef) -> Cow<'static, str> {
+    let (method, options) = match index_spec.method.as_ref() {
+        Some(spec::VectorIndexMethod::Hnsw { m, ef_construction }) => {
+            let mut opts = Vec::new();
+            if let Some(m) = m {
+                opts.push(format!("m = {}", m));
+            }
+            if let Some(ef) = ef_construction {
+                opts.push(format!("ef_construction = {}", ef));
+            }
+            ("hnsw", opts)
+        }
+        Some(spec::VectorIndexMethod::IvfFlat { lists }) => (
+            "ivfflat",
+            lists
+                .map(|lists| vec![format!("lists = {}", lists)])
+                .unwrap_or_default(),
+        ),
+        None => ("hnsw", Vec::new()),
+    };
+    let with_clause = if options.is_empty() {
+        String::new()
+    } else {
+        format!(" WITH ({})", options.join(", "))
+    };
     format!(
-        "USING hnsw ({} {})",
+        "USING {method} ({} {}){}",
         index_spec.field_name,
-        to_vector_similarity_metric_sql(index_spec.metric)
+        to_vector_similarity_metric_sql(index_spec.metric),
+        with_clause
     )
     .into()
 }
 
 fn to_vector_index_name(table_name: &str, vector_index_def: &spec::VectorIndexDef) -> String {
-    format!(
+    let mut name = format!(
         "{}__{}__{}",
         table_name,
         vector_index_def.field_name,
         to_vector_similarity_metric_sql(vector_index_def.metric)
-    )
+    );
+    if let Some(method) = vector_index_def.method.as_ref() {
+        name.push_str("__");
+        name.push_str(&method.kind().to_ascii_lowercase());
+    }
+    name
 }
 
 fn describe_index_spec(index_name: &str, index_spec: &VectorIndexDef) -> String {

{cocoindex-0.2.13 → cocoindex-0.2.14}/src/ops/targets/qdrant.rs

@@ -422,6 +422,9 @@ impl TargetFactoryBase for Factory {
                             } else {
                                 api_bail!("Field `{}` specified more than once in vector index definition", vector_index.field_name);
                             }
+                            if vector_index.method.is_some() {
+                                api_bail!("Vector index method is not configurable for Qdrant yet");
+                            }
                         }
                         None => {
                             if let Some(field) = d.value_fields_schema.iter().find(|f| f.name == vector_index.field_name) {