cocoindex 0.2.14__tar.gz → 0.2.15__tar.gz

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

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

{cocoindex-0.2.14 → 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.14"
+version = "0.2.15"
 edition = "2024"
 rust-version = "1.89"
 license = "Apache-2.0"

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

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

{cocoindex-0.2.14 → 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.14</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.14 → 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
@@ -42,6 +43,8 @@ from .typing import (
     Json,
 )
 
+_engine.init_pyo3_runtime()
+
 __all__ = [
     # Submodules
     "_engine",

{cocoindex-0.2.14 → 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.14 → 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.14 → 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"),
+        }
+    }
+}

{cocoindex-0.2.14 → cocoindex-0.2.15}/src/ops/functions/split_recursively.rs

@@ -6,8 +6,7 @@ use std::sync::LazyLock;
 use std::{collections::HashMap, sync::Arc};
 use unicase::UniCase;
 
-use crate::base::field_attrs;
-use crate::ops::registry::ExecutorFactoryRegistry;
+use crate::ops::shared::split::{Position, set_output_positions};
 use crate::{fields_value, ops::sdk::*};
 
 #[derive(Serialize, Deserialize)]
@@ -479,36 +478,6 @@ impl<'s> AtomChunksCollector<'s> {
     }
 }
 
-#[derive(Debug, Clone, PartialEq, Eq)]
-struct OutputPosition {
-    char_offset: usize,
-    line: u32,
-    column: u32,
-}
-
-impl OutputPosition {
-    fn into_output(self) -> value::Value {
-        value::Value::Struct(fields_value!(
-            self.char_offset as i64,
-            self.line as i64,
-            self.column as i64
-        ))
-    }
-}
-struct Position {
-    byte_offset: usize,
-    output: Option<OutputPosition>,
-}
-
-impl Position {
-    fn new(byte_offset: usize) -> Self {
-        Self {
-            byte_offset,
-            output: None,
-        }
-    }
-}
-
 struct ChunkOutput<'s> {
     start_pos: Position,
     end_pos: Position,
@@ -826,55 +795,6 @@ impl Executor {
     }
 }
 
-fn set_output_positions<'a>(text: &str, positions: impl Iterator<Item = &'a mut Position>) {
-    let mut positions = positions.collect::<Vec<_>>();
-    positions.sort_by_key(|o| o.byte_offset);
-
-    let mut positions_iter = positions.iter_mut();
-    let Some(mut next_position) = positions_iter.next() else {
-        return;
-    };
-
-    let mut char_offset = 0;
-    let mut line = 1;
-    let mut column = 1;
-    for (byte_offset, ch) in text.char_indices() {
-        while next_position.byte_offset == byte_offset {
-            next_position.output = Some(OutputPosition {
-                char_offset,
-                line,
-                column,
-            });
-            if let Some(position) = positions_iter.next() {
-                next_position = position;
-            } else {
-                return;
-            }
-        }
-        char_offset += 1;
-        if ch == '\n' {
-            line += 1;
-            column = 1;
-        } else {
-            column += 1;
-        }
-    }
-
-    // Offsets after the last char.
-    loop {
-        next_position.output = Some(OutputPosition {
-            char_offset,
-            line,
-            column,
-        });
-        if let Some(position) = positions_iter.next() {
-            next_position = position;
-        } else {
-            return;
-        }
-    }
-}
-
 #[async_trait]
 impl SimpleFunctionExecutor for Executor {
     async fn evaluate(&self, input: Vec<Value>) -> Result<Value> {
@@ -997,49 +917,8 @@ impl SimpleFunctionFactoryBase for Factory {
                 .optional(),
         };
 
-        let pos_struct = schema::ValueType::Struct(schema::StructSchema {
-            fields: Arc::new(vec![
-                schema::FieldSchema::new("offset", make_output_type(BasicValueType::Int64)),
-                schema::FieldSchema::new("line", make_output_type(BasicValueType::Int64)),
-                schema::FieldSchema::new("column", make_output_type(BasicValueType::Int64)),
-            ]),
-            description: None,
-        });
-
-        let mut struct_schema = StructSchema::default();
-        let mut schema_builder = StructSchemaBuilder::new(&mut struct_schema);
-        schema_builder.add_field(FieldSchema::new(
-            "location",
-            make_output_type(BasicValueType::Range),
-        ));
-        schema_builder.add_field(FieldSchema::new(
-            "text",
-            make_output_type(BasicValueType::Str),
-        ));
-        schema_builder.add_field(FieldSchema::new(
-            "start",
-            schema::EnrichedValueType {
-                typ: pos_struct.clone(),
-                nullable: false,
-                attrs: Default::default(),
-            },
-        ));
-        schema_builder.add_field(FieldSchema::new(
-            "end",
-            schema::EnrichedValueType {
-                typ: pos_struct,
-                nullable: false,
-                attrs: Default::default(),
-            },
-        ));
-        let output_schema = make_output_type(TableSchema::new(
-            TableKind::KTable(KTableInfo { num_key_parts: 1 }),
-            struct_schema,
-        ))
-        .with_attr(
-            field_attrs::CHUNK_BASE_TEXT,
-            serde_json::to_value(args_resolver.get_analyze_value(&args.text))?,
-        );
+        let output_schema =
+            crate::ops::shared::split::make_common_chunk_schema(args_resolver, &args.text)?;
         Ok((args, output_schema))
     }
 
@@ -1060,7 +939,7 @@ pub fn register(registry: &mut ExecutorFactoryRegistry) -> Result<()> {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::ops::functions::test_utils::test_flow_function;
+    use crate::ops::{functions::test_utils::test_flow_function, shared::split::OutputPosition};
 
     // Helper function to assert chunk text and its consistency with the range within the original text.
     fn assert_chunk_text_consistency(

{cocoindex-0.2.14 → cocoindex-0.2.15}/src/ops/registration.rs

@@ -18,6 +18,7 @@ fn register_executor_factories(registry: &mut ExecutorFactoryRegistry) -> Result
     functions::split_recursively::register(registry)?;
     functions::extract_by_llm::Factory.register(registry)?;
     functions::embed_text::register(registry)?;
+    functions::split_by_separators::register(registry)?;
 
     targets::postgres::Factory::default().register(registry)?;
     targets::qdrant::register(registry)?;

{cocoindex-0.2.14 → cocoindex-0.2.15}/src/ops/shared/mod.rs

@@ -1 +1,2 @@
 pub mod postgres;
+pub mod split;