cocoindex 0.2.13__tar.gz → 0.2.15__tar.gz

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

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

{cocoindex-0.2.13 → cocoindex-0.2.15}/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.15"
 edition = "2024"
 rust-version = "1.89"
 license = "Apache-2.0"

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

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

{cocoindex-0.2.13 → cocoindex-0.2.15}/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.15</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.15}/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.15}/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.15}/python/cocoindex/__init__.py

@@ -2,6 +2,7 @@
 Cocoindex is a framework for building and running indexing pipelines.
 """
 
+from . import _engine  # type: ignore
 from . import functions, sources, targets, cli, utils
 
 from . import targets as storages  # Deprecated: Use targets instead
@@ -21,7 +22,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
@@ -36,6 +43,8 @@ from .typing import (
     Json,
 )
 
+_engine.init_pyo3_runtime()
+
 __all__ = [
     # Submodules
     "_engine",
@@ -82,6 +91,8 @@ __all__ = [
     "VectorSimilarityMetric",
     "VectorIndexDef",
     "IndexOptions",
+    "HnswVectorIndexMethod",
+    "IvfFlatVectorIndexMethod",
     # Settings
     "DatabaseConnectionSpec",
     "Settings",

{cocoindex-0.2.13 → cocoindex-0.2.15}/python/cocoindex/functions.py

@@ -2,13 +2,13 @@
 
 import dataclasses
 import functools
-from typing import Annotated, Any, Literal
+from typing import Any, Literal
 
 import numpy as np
 from numpy.typing import NDArray
 
 from . import llm, op
-from .typing import TypeAttr, Vector
+from .typing import Vector
 
 
 class ParseJson(op.FunctionSpec):
@@ -40,6 +40,24 @@ class SplitRecursively(op.FunctionSpec):
     custom_languages: list[CustomLanguageSpec] = dataclasses.field(default_factory=list)
 
 
+class SplitBySeparators(op.FunctionSpec):
+    """
+    Split text by specified regex separators only.
+    Output schema matches SplitRecursively for drop-in compatibility:
+        KTable rows with fields: location (Range), text (Str), start, end.
+    Args:
+        separators_regex: list[str]  # e.g., [r"\\n\\n+"]
+        keep_separator: Literal["NONE", "LEFT", "RIGHT"] = "NONE"
+        include_empty: bool = False
+        trim: bool = True
+    """
+
+    separators_regex: list[str] = dataclasses.field(default_factory=list)
+    keep_separator: Literal["NONE", "LEFT", "RIGHT"] = "NONE"
+    include_empty: bool = False
+    trim: bool = True
+
+
 class EmbedText(op.FunctionSpec):
     """Embed a text into a vector space."""
 

{cocoindex-0.2.13 → cocoindex-0.2.15}/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.15}/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.15}/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.15}/src/lib_context.rs

@@ -162,6 +162,13 @@ impl FlowContext {
 static TOKIO_RUNTIME: LazyLock<Runtime> = LazyLock::new(|| Runtime::new().unwrap());
 static AUTH_REGISTRY: LazyLock<Arc<AuthRegistry>> = LazyLock::new(|| Arc::new(AuthRegistry::new()));
 
+pub fn get_runtime() -> &'static Runtime {
+    &TOKIO_RUNTIME
+}
+pub fn get_auth_registry() -> &'static Arc<AuthRegistry> {
+    &AUTH_REGISTRY
+}
+
 type PoolKey = (String, Option<String>);
 type PoolValue = Arc<tokio::sync::OnceCell<PgPool>>;
 
@@ -271,21 +278,10 @@ impl LibContext {
     }
 }
 
-pub fn get_runtime() -> &'static Runtime {
-    &TOKIO_RUNTIME
-}
-
-pub fn get_auth_registry() -> &'static Arc<AuthRegistry> {
-    &AUTH_REGISTRY
-}
-
 static LIB_INIT: OnceLock<()> = OnceLock::new();
 pub async fn create_lib_context(settings: settings::Settings) -> Result<LibContext> {
     LIB_INIT.get_or_init(|| {
         let _ = env_logger::try_init();
-
-        pyo3_async_runtimes::tokio::init_with_runtime(get_runtime()).unwrap();
-
         let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
     });
 

{cocoindex-0.2.13 → cocoindex-0.2.15}/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.15}/src/ops/functions/mod.rs

@@ -1,6 +1,7 @@
 pub mod embed_text;
 pub mod extract_by_llm;
 pub mod parse_json;
+pub mod split_by_separators;
 pub mod split_recursively;
 
 #[cfg(test)]

cocoindex-0.2.15/src/ops/functions/split_by_separators.rs

@@ -0,0 +1,254 @@
+use anyhow::{Context, Result};
+use regex::Regex;
+use std::sync::Arc;
+
+use crate::ops::registry::ExecutorFactoryRegistry;
+use crate::ops::shared::split::{Position, make_common_chunk_schema, set_output_positions};
+use crate::{fields_value, ops::sdk::*};
+
+#[derive(Serialize, Deserialize, Clone, Copy, PartialEq, Eq)]
+#[serde(rename_all = "UPPERCASE")]
+enum KeepSep {
+    Left,
+    Right,
+}
+
+#[derive(Serialize, Deserialize)]
+struct Spec {
+    // Python SDK provides defaults/values.
+    separators_regex: Vec<String>,
+    keep_separator: Option<KeepSep>,
+    include_empty: bool,
+    trim: bool,
+}
+
+struct Args {
+    text: ResolvedOpArg,
+}
+
+struct Executor {
+    spec: Spec,
+    regex: Option<Regex>,
+    args: Args,
+}
+
+impl Executor {
+    fn new(args: Args, spec: Spec) -> Result<Self> {
+        let regex = if spec.separators_regex.is_empty() {
+            None
+        } else {
+            // OR-join all separators, multiline
+            let pattern = format!(
+                "(?m){}",
+                spec.separators_regex
+                    .iter()
+                    .map(|s| format!("(?:{s})"))
+                    .collect::<Vec<_>>()
+                    .join("|")
+            );
+            Some(Regex::new(&pattern).context("failed to compile separators_regex")?)
+        };
+        Ok(Self { args, spec, regex })
+    }
+}
+
+struct ChunkOutput<'s> {
+    start_pos: Position,
+    end_pos: Position,
+    text: &'s str,
+}
+
+#[async_trait]
+impl SimpleFunctionExecutor for Executor {
+    async fn evaluate(&self, input: Vec<Value>) -> Result<Value> {
+        let full_text = self.args.text.value(&input)?.as_str()?;
+        let bytes = full_text.as_bytes();
+
+        // add_range applies trim/include_empty and records the text slice
+        let mut chunks: Vec<ChunkOutput<'_>> = Vec::new();
+        let mut add_range = |mut s: usize, mut e: usize| {
+            if self.spec.trim {
+                while s < e && bytes[s].is_ascii_whitespace() {
+                    s += 1;
+                }
+                while e > s && bytes[e - 1].is_ascii_whitespace() {
+                    e -= 1;
+                }
+            }
+            if self.spec.include_empty || e > s {
+                chunks.push(ChunkOutput {
+                    start_pos: Position::new(s),
+                    end_pos: Position::new(e),
+                    text: &full_text[s..e],
+                });
+            }
+        };
+
+        if let Some(re) = &self.regex {
+            let mut start = 0usize;
+            for m in re.find_iter(full_text) {
+                let end = match self.spec.keep_separator {
+                    Some(KeepSep::Left) => m.end(),
+                    Some(KeepSep::Right) | None => m.start(),
+                };
+                add_range(start, end);
+                start = match self.spec.keep_separator {
+                    Some(KeepSep::Right) => m.start(),
+                    _ => m.end(),
+                };
+            }
+            add_range(start, full_text.len());
+        } else {
+            // No separators: emit whole text
+            add_range(0, full_text.len());
+        }
+
+        set_output_positions(
+            full_text,
+            chunks.iter_mut().flat_map(|c| {
+                std::iter::once(&mut c.start_pos).chain(std::iter::once(&mut c.end_pos))
+            }),
+        );
+
+        let table = chunks
+            .into_iter()
+            .map(|c| {
+                let s = c.start_pos.output.unwrap();
+                let e = c.end_pos.output.unwrap();
+                (
+                    KeyValue::from_single_part(RangeValue::new(s.char_offset, e.char_offset)),
+                    fields_value!(Arc::<str>::from(c.text), s.into_output(), e.into_output())
+                        .into(),
+                )
+            })
+            .collect();
+
+        Ok(Value::KTable(table))
+    }
+}
+
+struct Factory;
+
+#[async_trait]
+impl SimpleFunctionFactoryBase for Factory {
+    type Spec = Spec;
+    type ResolvedArgs = Args;
+
+    fn name(&self) -> &str {
+        "SplitBySeparators"
+    }
+
+    async fn resolve_schema<'a>(
+        &'a self,
+        _spec: &'a Spec,
+        args_resolver: &mut OpArgsResolver<'a>,
+        _context: &FlowInstanceContext,
+    ) -> Result<(Args, EnrichedValueType)> {
+        // one required arg: text: Str
+        let args = Args {
+            text: args_resolver
+                .next_arg("text")?
+                .expect_type(&ValueType::Basic(BasicValueType::Str))?
+                .required()?,
+        };
+
+        let output_schema = make_common_chunk_schema(args_resolver, &args.text)?;
+        Ok((args, output_schema))
+    }
+
+    async fn build_executor(
+        self: Arc<Self>,
+        spec: Spec,
+        args: Args,
+        _context: Arc<FlowInstanceContext>,
+    ) -> Result<impl SimpleFunctionExecutor> {
+        Executor::new(args, spec)
+    }
+}
+
+pub fn register(registry: &mut ExecutorFactoryRegistry) -> Result<()> {
+    Factory.register(registry)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::ops::functions::test_utils::test_flow_function;
+
+    #[tokio::test]
+    async fn test_split_by_separators_paragraphs() {
+        let spec = Spec {
+            separators_regex: vec![r"\n\n+".to_string()],
+            keep_separator: None,
+            include_empty: false,
+            trim: true,
+        };
+        let factory = Arc::new(Factory);
+        let text = "Para1\n\nPara2\n\n\nPara3";
+
+        let input_arg_schemas = &[(
+            Some("text"),
+            make_output_type(BasicValueType::Str).with_nullable(true),
+        )];
+
+        let result = test_flow_function(
+            &factory,
+            &spec,
+            input_arg_schemas,
+            vec![text.to_string().into()],
+        )
+        .await
+        .unwrap();
+
+        match result {
+            Value::KTable(table) => {
+                // Expected ranges after trimming whitespace:
+                let expected = vec![
+                    (RangeValue::new(0, 5), "Para1"),
+                    (RangeValue::new(7, 12), "Para2"),
+                    (RangeValue::new(15, 20), "Para3"),
+                ];
+                for (range, expected_text) in expected {
+                    let key = KeyValue::from_single_part(range);
+                    let row = table.get(&key).unwrap();
+                    let chunk_text = row.0.fields[0].as_str().unwrap();
+                    assert_eq!(**chunk_text, *expected_text);
+                }
+            }
+            other => panic!("Expected KTable, got {other:?}"),
+        }
+    }
+
+    #[tokio::test]
+    async fn test_split_by_separators_keep_right() {
+        let spec = Spec {
+            separators_regex: vec![r"\.".to_string()],
+            keep_separator: Some(KeepSep::Right),
+            include_empty: false,
+            trim: true,
+        };
+        let factory = Arc::new(Factory);
+        let text = "A. B. C.";
+
+        let input_arg_schemas = &[(
+            Some("text"),
+            make_output_type(BasicValueType::Str).with_nullable(true),
+        )];
+
+        let result = test_flow_function(
+            &factory,
+            &spec,
+            input_arg_schemas,
+            vec![text.to_string().into()],
+        )
+        .await
+        .unwrap();
+
+        match result {
+            Value::KTable(table) => {
+                assert!(table.len() >= 3);
+            }
+            _ => panic!("KTable expected"),
+        }
+    }
+}