cocoindex 0.1.77__tar.gz → 0.1.79__tar.gz

{cocoindex-0.1.77 → cocoindex-0.1.79}/Cargo.lock

@@ -1297,7 +1297,7 @@ dependencies = [
 
 [[package]]
 name = "cocoindex"
-version = "0.1.77"
+version = "0.1.79"
 dependencies = [
  "anyhow",
  "async-openai",

{cocoindex-0.1.77 → cocoindex-0.1.79}/Cargo.toml

@@ -2,9 +2,10 @@
 name = "cocoindex"
 # Version used for local development is always higher than others to take precedence.
 # Will be overridden for specific release versions.
-version = "0.1.77"
+version = "0.1.79"
 edition = "2024"
 rust-version = "1.88"
+readme = "README.md"
 
 [profile.release]
 codegen-units = 1

{cocoindex-0.1.77 → cocoindex-0.1.79}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cocoindex
-Version: 0.1.77
+Version: 0.1.79
 Requires-Dist: click>=8.1.8
 Requires-Dist: rich>=14.0.0
 Requires-Dist: python-dotenv>=1.1.0

{cocoindex-0.1.77 → cocoindex-0.1.79}/docs/docs/ai/llm.mdx

@@ -75,7 +75,11 @@ CocoIndex integrates with various LLM APIs for these functions.
 To use the OpenAI LLM API, you need to set the environment variable `OPENAI_API_KEY`.
 You can generate the API key from [OpenAI Dashboard](https://platform.openai.com/api-keys).
 
-Currently we don't support custom address for OpenAI API.
+If you want to use a custom address, you can either provide the `address` parameter in `LlmSpec` / `EmbedText`, or set the environment variable `OPENAI_API_BASE`. The `address` parameter takes precedence over the environment variable.
+
+Spec for OpenAI takes additional `api_config` field, in type `cocoindex.llm.OpenAiConfig` with the following fields:
+-   `org_id` (type: `str`, optional): The organization ID of the OpenAI account.
+-   `project_id` (type: `str`, optional): The project ID of the OpenAI account.
 
 You can find the full list of models supported by OpenAI [here](https://platform.openai.com/docs/models).
 

{cocoindex-0.1.77 → cocoindex-0.1.79}/docs/docs/core/flow_def.mdx

@@ -360,21 +360,40 @@ It will use `Staging__doc_embeddings` as the collection name if the current app
 
 ### Control Processing Concurrency
 
-You can control the concurrency of the processing by setting the following options:
+CocoIndex processes data in parallel to maximize throughput, but unconstrained parallelism can overwhelm your system.
+Processing too many items simultaneously can lead to:
 
-*   `max_inflight_rows`: the maximum number of concurrent inflight requests for the processing.
-*   `max_inflight_bytes`: the maximum number of concurrent inflight bytes for the processing.
+- **Memory exhaustion**: Large datasets loaded concurrently can consume excessive RAM
+- **Resource contention**: Too many parallel operations competing for CPU, disk I/O, or network bandwidth
+- **System instability**: High concurrency can cause timeouts, crashes, or degraded performance
 
-These options can be passed in to the following APIs:
+To prevent these issues, CocoIndex provides concurrency controls that limit how many data items are processed simultaneously.
 
-*   [`FlowBuilder.add_source()`](#import-from-source): The options above control the processing concurrency of multiple rows from a source. New rows will not be loaded in memory if it'll be over the limit.
+#### Concurrency Options
 
-    Besides, global limits on overall processing concurrency of all sources from all flows can be specified by [`GlobalExecutionOptions`](/docs/core/settings#globalexecutionoptions) or corresponding [environment variables](/docs/core/settings#list-of-environment-variables).
-    If both global and per-source limits are specified, both need to be satisfied to admit additional source rows.
+You can control processing concurrency using these options:
 
-*   [`DataSlice.row()`](#for-each-row):  The options above provides a finer-grained control, to limit the processing concurrency of multiple rows within a table at any level.
+*   `max_inflight_rows`: Limits the maximum number of data rows being processed concurrently
+*   `max_inflight_bytes`: Limits the total memory footprint of data being processed concurrently (measured in bytes)
 
-`max_inflight_bytes` only counts the number of bytes already existing in the current row before any further processing.
+When these limits are reached, CocoIndex will pause loading new data until some of the current processing completes, ensuring your system remains stable.
+
+#### Where to Apply Concurrency Controls
+
+These concurrency options can be configured at different levels:
+
+*   **Source level** via [`FlowBuilder.add_source()`](#import-from-source): Controls how many rows from a data source are processed simultaneously. This prevents overwhelming your system when ingesting large datasets.
+
+    You can also set global limits across all sources and flows using [`GlobalExecutionOptions`](/docs/core/settings#globalexecutionoptions) or environment variables [`COCOINDEX_SOURCE_MAX_INFLIGHT_ROWS`](/docs/core/settings#list-of-environment-variables)/[`COCOINDEX_SOURCE_MAX_INFLIGHT_BYTES`](/docs/core/settings#list-of-environment-variables).
+    When both global and per-source limits are specified, both limits are enforced independently - a new row can only be processed if there's available capacity in both the global budget (shared across all sources) and the per-source budget (specific to that source).
+
+*   **Row iteration level** via [`DataSlice.row()`](#for-each-row): Provides fine-grained control over parallel processing within nested data structures, allowing you to tune concurrency at any level of your data hierarchy.
+
+:::note
+
+The `max_inflight_bytes` limit only counts the size of data that already exists in memory before any transformations are applied. It doesn't include the memory used by intermediate processing results.
+
+:::
 
 For example:
 

{cocoindex-0.1.77 → cocoindex-0.1.79}/docs/docs/core/settings.mdx

@@ -77,24 +77,33 @@ If not set, all flows are in a default unnamed namespace.
 
 `DatabaseConnectionSpec` configures the connection to a database. Only Postgres is supported for now. It has the following fields:
 
-*   `url` (type: `str`, required): The URL of the Postgres database to use as the internal storage, e.g. `postgres://cocoindex:cocoindex@localhost/cocoindex`.
+*   `url` (type: `str`): The URL of the Postgres database to use as the internal storage, e.g. `postgres://cocoindex:cocoindex@localhost/cocoindex`.
 
     *Environment variable* for `Settings.database.url`: `COCOINDEX_DATABASE_URL`
 
-*   `user` (type: `str`, optional): The username for the Postgres database. If not provided, username will come from `url`.
+*   `user` (type: `str | None`, default: `None`): The username for the Postgres database. If not provided, username will come from `url`.
 
     *Environment variable* for `Settings.database.user`: `COCOINDEX_DATABASE_USER`
 
-*   `password` (type: `str`, optional): The password for the Postgres database. If not provided, password will come from `url`.
+*   `password` (type: `str | None`, default: `None`): The password for the Postgres database. If not provided, password will come from `url`.
 
     *Environment variable* for `Settings.database.password`: `COCOINDEX_DATABASE_PASSWORD`
 
-:::tip
+    :::tip
 
-Please be careful that all values in `url` needs to be url-encoded if they contain special characters.
-For this reason, prefer to use the separated `user` and `password` fields for username and password.
+    Please be careful that all values in `url` needs to be url-encoded if they contain special characters.
+    For this reason, prefer to use the separated `user` and `password` fields for username and password.
+
+    :::
+
+*   `max_connections` (type: `int`, default: `64`): The maximum number of connections to keep in the pool.
+
+    *Environment variable* for `Settings.database.max_connections`: `COCOINDEX_DATABASE_MAX_CONNECTIONS`
+
+*   `min_connections` (type: `int`, default: `16`): The minimum number of connections to keep in the pool.
+
+    *Environment variable* for `Settings.database.min_connections`: `COCOINDEX_DATABASE_MIN_CONNECTIONS`
 
-:::
 
 :::info
 
@@ -109,10 +118,10 @@ If you use the Postgres database hosted by [Supabase](https://supabase.com/), pl
 
 `GlobalExecutionOptions` is used to configure the global execution options shared by all flows. It has the following fields:
 
-*   `source_max_inflight_rows` (type: `int`, optional): The maximum number of concurrent inflight requests for all source operations.
-*   `source_max_inflight_bytes` (type: `int`, optional): The maximum number of concurrent inflight bytes for all source operations.
+*   `source_max_inflight_rows` (type: `int | None`, default: `1024`): The maximum number of concurrent inflight rows for all source operations.
+*   `source_max_inflight_bytes` (type: `int | None`, default: `None`): The maximum number of concurrent inflight bytes for all source operations.
 
-See also [flow definition docs](/docs/core/flow_def#control-processing-concurrency) to control processing concurrency on per-source basis.
+See also [flow definition docs](/docs/core/flow_def#control-processing-concurrency) about why it's necessary to control processing concurrency, and how to configure it on per-source basis.
 If both global and per-source limits are specified, both need to be satisfied to admit additional source rows.
 
 ## List of Environment Variables
@@ -125,5 +134,7 @@ This is the list of environment variables, each of which has a corresponding fie
 | `COCOINDEX_DATABASE_URL` | `database.url` | Yes |
 | `COCOINDEX_DATABASE_USER` | `database.user` | No |
 | `COCOINDEX_DATABASE_PASSWORD` | `database.password` | No |
-| `COCOINDEX_SOURCE_MAX_INFLIGHT_ROWS` | `global_execution_options.source_max_inflight_rows` | No |
+| `COCOINDEX_DATABASE_MAX_CONNECTIONS` | `database.max_connections` | No (default: `64`) |
+| `COCOINDEX_DATABASE_MIN_CONNECTIONS` | `database.min_connections` | No (default: `16`) |
+| `COCOINDEX_SOURCE_MAX_INFLIGHT_ROWS` | `global_execution_options.source_max_inflight_rows` | No (default: `1024`) |
 | `COCOINDEX_SOURCE_MAX_INFLIGHT_BYTES` | `global_execution_options.source_max_inflight_bytes` | No |

{cocoindex-0.1.77 → cocoindex-0.1.79}/examples/image_search/colpali_main.py

@@ -27,10 +27,6 @@ COLPALI_MODEL_NAME = os.getenv("COLPALI_MODEL", "vidore/colpali-v1.2")
 print(f"📐 Using ColPali model {COLPALI_MODEL_NAME}")
 
 
-# Create ColPali embedding function using the class-based pattern
-colpali_embed = cocoindex.functions.ColPaliEmbedImage(model=COLPALI_MODEL_NAME)
-
-
 @cocoindex.transform_flow()
 def text_to_colpali_embedding(
     text: cocoindex.DataSlice[str],
@@ -56,38 +52,15 @@ def image_object_embedding_flow(
     )
     img_embeddings = data_scope.add_collector()
     with data_scope["images"].row() as img:
-        ollama_model_name = os.getenv("OLLAMA_MODEL")
-        if ollama_model_name is not None:
-            # If an Ollama model is specified, generate an image caption
-            img["caption"] = flow_builder.transform(
-                cocoindex.functions.ExtractByLlm(
-                    llm_spec=cocoindex.llm.LlmSpec(
-                        api_type=cocoindex.LlmApiType.OLLAMA, model=ollama_model_name
-                    ),
-                    instruction=(
-                        "Describe the image in one detailed sentence. "
-                        "Name all visible animal species, objects, and the main scene. "
-                        "Be specific about type, color, and notable features. "
-                        "Mention what each animal is doing."
-                    ),
-                    output_type=str,
-                ),
-                image=img["content"],
-            )
-        img["embedding"] = img["content"].transform(colpali_embed)
+        img["embedding"] = img["content"].transform(
+            cocoindex.functions.ColPaliEmbedImage(model=COLPALI_MODEL_NAME)
+        )
 
         collect_fields = {
             "id": cocoindex.GeneratedField.UUID,
             "filename": img["filename"],
             "embedding": img["embedding"],
         }
-
-        if ollama_model_name is not None:
-            print(f"Using Ollama model '{ollama_model_name}' for captioning.")
-            collect_fields["caption"] = img["caption"]
-        else:
-            print(f"No Ollama model '{ollama_model_name}' found — skipping captioning.")
-
         img_embeddings.collect(**collect_fields)
 
     img_embeddings.export(

{cocoindex-0.1.77 → cocoindex-0.1.79}/python/cocoindex/llm.py

@@ -26,6 +26,16 @@ class VertexAiConfig:
     region: str | None = None
 
 
+@dataclass
+class OpenAiConfig:
+    """A specification for a OpenAI LLM."""
+
+    kind = "OpenAi"
+
+    org_id: str | None = None
+    project_id: str | None = None
+
+
 @dataclass
 class LlmSpec:
     """A specification for a LLM."""
@@ -33,4 +43,4 @@ class LlmSpec:
     api_type: LlmApiType
     model: str
     address: str | None = None
-    api_config: VertexAiConfig | None = None
+    api_config: VertexAiConfig | OpenAiConfig | None = None

{cocoindex-0.1.77 → cocoindex-0.1.79}/python/cocoindex/setting.py

@@ -44,6 +44,8 @@ class DatabaseConnectionSpec:
     url: str
     user: str | None = None
     password: str | None = None
+    max_connections: int = 64
+    min_connections: int = 16
 
 
 @dataclass
@@ -51,7 +53,7 @@ class GlobalExecutionOptions:
     """Global execution options."""
 
     # The maximum number of concurrent inflight requests, shared among all sources from all flows.
-    source_max_inflight_rows: int | None = None
+    source_max_inflight_rows: int | None = 1024
     source_max_inflight_bytes: int | None = None
 
 
@@ -92,10 +94,22 @@ class Settings:
 
         database_url = os.getenv("COCOINDEX_DATABASE_URL")
         if database_url is not None:
-            db_kwargs: dict[str, str] = dict()
+            db_kwargs: dict[str, Any] = dict()
             _load_field(db_kwargs, "url", "COCOINDEX_DATABASE_URL", required=True)
             _load_field(db_kwargs, "user", "COCOINDEX_DATABASE_USER")
             _load_field(db_kwargs, "password", "COCOINDEX_DATABASE_PASSWORD")
+            _load_field(
+                db_kwargs,
+                "max_connections",
+                "COCOINDEX_DATABASE_MAX_CONNECTIONS",
+                parse=int,
+            )
+            _load_field(
+                db_kwargs,
+                "min_connections",
+                "COCOINDEX_DATABASE_MIN_CONNECTIONS",
+                parse=int,
+            )
             database = DatabaseConnectionSpec(**db_kwargs)
         else:
             database = None

{cocoindex-0.1.77 → cocoindex-0.1.79}/src/lib_context.rs

@@ -1,3 +1,5 @@
+use std::time::Duration;
+
 use crate::prelude::*;
 
 use crate::builder::AnalyzedFlow;
@@ -7,7 +9,7 @@ use crate::settings;
 use crate::setup::ObjectSetupStatus;
 use axum::http::StatusCode;
 use sqlx::PgPool;
-use sqlx::postgres::PgConnectOptions;
+use sqlx::postgres::{PgConnectOptions, PgPoolOptions};
 use tokio::runtime::Runtime;
 
 pub struct FlowExecutionContext {
@@ -176,7 +178,29 @@ impl DbPools {
                 if let Some(password) = &conn_spec.password {
                     pg_options = pg_options.password(password);
                 }
-                let pool = PgPool::connect_with(pg_options)
+
+                // Try to connect to the database with a low timeout first.
+                {
+                    let pool_options = PgPoolOptions::new()
+                        .max_connections(1)
+                        .min_connections(1)
+                        .acquire_timeout(Duration::from_secs(30));
+                    let pool = pool_options
+                        .connect_with(pg_options.clone())
+                        .await
+                        .context(format!("Failed to connect to database {}", conn_spec.url))?;
+                    let _ = pool.acquire().await?;
+                }
+
+                // Now create the actual pool.
+                let pool_options = PgPoolOptions::new()
+                    .max_connections(conn_spec.max_connections)
+                    .min_connections(conn_spec.min_connections)
+                    .acquire_timeout(Duration::from_secs(5 * 60))
+                    .idle_timeout(Duration::from_secs(10 * 60))
+                    .max_lifetime(Duration::from_secs(60 * 60));
+                let pool = pool_options
+                    .connect_with(pg_options)
                     .await
                     .context("Failed to connect to database")?;
                 anyhow::Ok(pool)
@@ -330,6 +354,8 @@ mod tests {
                 url: "postgresql://test".to_string(),
                 user: None,
                 password: None,
+                max_connections: 10,
+                min_connections: 1,
             }),
             ..Default::default()
         };

{cocoindex-0.1.77 → cocoindex-0.1.79}/src/llm/mod.rs

@@ -26,10 +26,17 @@ pub struct VertexAiConfig {
     pub region: Option<String>,
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize, Default)]
+pub struct OpenAiConfig {
+    pub org_id: Option<String>,
+    pub project_id: Option<String>,
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(tag = "kind")]
 pub enum LlmApiConfig {
     VertexAi(VertexAiConfig),
+    OpenAi(OpenAiConfig),
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -113,7 +120,7 @@ pub async fn new_llm_generation_client(
             Box::new(ollama::Client::new(address).await?) as Box<dyn LlmGenerationClient>
         }
         LlmApiType::OpenAi => {
-            Box::new(openai::Client::new(address)?) as Box<dyn LlmGenerationClient>
+            Box::new(openai::Client::new(address, api_config)?) as Box<dyn LlmGenerationClient>
         }
         LlmApiType::Gemini => {
             Box::new(gemini::AiStudioClient::new(address)?) as Box<dyn LlmGenerationClient>
@@ -151,7 +158,7 @@ pub async fn new_llm_embedding_client(
             Box::new(gemini::AiStudioClient::new(address)?) as Box<dyn LlmEmbeddingClient>
         }
         LlmApiType::OpenAi => {
-            Box::new(openai::Client::new(address)?) as Box<dyn LlmEmbeddingClient>
+            Box::new(openai::Client::new(address, api_config)?) as Box<dyn LlmEmbeddingClient>
         }
         LlmApiType::Voyage => {
             Box::new(voyage::Client::new(address)?) as Box<dyn LlmEmbeddingClient>

{cocoindex-0.1.77 → cocoindex-0.1.79}/src/llm/openai.rs

@@ -33,17 +33,31 @@ impl Client {
         Self { client }
     }
 
-    pub fn new(address: Option<String>) -> Result<Self> {
+    pub fn new(address: Option<String>, api_config: Option<super::LlmApiConfig>) -> Result<Self> {
+        let config = match api_config {
+            Some(super::LlmApiConfig::OpenAi(config)) => config,
+            Some(_) => api_bail!("unexpected config type, expected OpenAiConfig"),
+            None => super::OpenAiConfig::default(),
+        };
+
+        let mut openai_config = OpenAIConfig::new();
         if let Some(address) = address {
-            api_bail!("OpenAI doesn't support custom API address: {address}");
+            openai_config = openai_config.with_api_base(address);
+        }
+        if let Some(org_id) = config.org_id {
+            openai_config = openai_config.with_org_id(org_id);
         }
+        if let Some(project_id) = config.project_id {
+            openai_config = openai_config.with_project_id(project_id);
+        }
+
         // Verify API key is set
         if std::env::var("OPENAI_API_KEY").is_err() {
             api_bail!("OPENAI_API_KEY environment variable must be set");
         }
         Ok(Self {
-            // OpenAI client will use OPENAI_API_KEY env variable by default
-            client: OpenAIClient::new(),
+            // OpenAI client will use OPENAI_API_KEY and OPENAI_API_BASE env variables by default
+            client: OpenAIClient::with_config(openai_config),
         })
     }
 }

{cocoindex-0.1.77 → cocoindex-0.1.79}/src/settings.rs

@@ -5,6 +5,8 @@ pub struct DatabaseConnectionSpec {
     pub url: String,
     pub user: Option<String>,
     pub password: Option<String>,
+    pub max_connections: u32,
+    pub min_connections: u32,
 }
 
 #[derive(Deserialize, Debug, Default)]
@@ -34,7 +36,9 @@ mod tests {
             "database": {
                 "url": "postgresql://localhost:5432/test",
                 "user": "testuser",
-                "password": "testpass"
+                "password": "testpass",
+                "min_connections": 1,
+                "max_connections": 10
             },
             "app_namespace": "test_app"
         }"#;
@@ -46,6 +50,8 @@ mod tests {
         assert_eq!(db.url, "postgresql://localhost:5432/test");
         assert_eq!(db.user, Some("testuser".to_string()));
         assert_eq!(db.password, Some("testpass".to_string()));
+        assert_eq!(db.min_connections, 1);
+        assert_eq!(db.max_connections, 10);
         assert_eq!(settings.app_namespace, "test_app");
     }
 
@@ -75,7 +81,9 @@ mod tests {
     fn test_settings_deserialize_database_without_user_password() {
         let json = r#"{
             "database": {
-                "url": "postgresql://localhost:5432/test"
+                "url": "postgresql://localhost:5432/test",
+                "min_connections": 1,
+                "max_connections": 10
             }
         }"#;
 
@@ -86,6 +94,8 @@ mod tests {
         assert_eq!(db.url, "postgresql://localhost:5432/test");
         assert_eq!(db.user, None);
         assert_eq!(db.password, None);
+        assert_eq!(db.min_connections, 1);
+        assert_eq!(db.max_connections, 10);
         assert_eq!(settings.app_namespace, "");
     }
 
@@ -94,7 +104,9 @@ mod tests {
         let json = r#"{
             "url": "postgresql://localhost:5432/test",
             "user": "testuser",
-            "password": "testpass"
+            "password": "testpass",
+            "min_connections": 1,
+            "max_connections": 10
         }"#;
 
         let db_spec: DatabaseConnectionSpec = serde_json::from_str(json).unwrap();
@@ -102,5 +114,7 @@ mod tests {
         assert_eq!(db_spec.url, "postgresql://localhost:5432/test");
         assert_eq!(db_spec.user, Some("testuser".to_string()));
         assert_eq!(db_spec.password, Some("testpass".to_string()));
+        assert_eq!(db_spec.min_connections, 1);
+        assert_eq!(db_spec.max_connections, 10);
     }
 }