feat: add batch append API

README.md

@@ -67,6 +67,22 @@ ctx.add("assistant", image)
 
 print("Current version:", ctx.version())
 
+# Batch append source chunks in one storage operation
+ctx.add_many([
+    {
+        "role": "source",
+        "content": "Chunk 1 from a runbook",
+        "content_type": "text/markdown",
+        "session_id": "runbook-import",
+    },
+    {
+        "role": "source",
+        "content": "Chunk 2 from the same runbook",
+        "content_type": "text/markdown",
+        "session_id": "runbook-import",
+    },
+])
+
 # Time-travel to prior state
 first_version = ctx.version()
 ctx.add("assistant", "Let me fetch suggestions…")

python/python/lance_context/api.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import warnings
+from collections.abc import Iterable, Mapping
 from datetime import datetime
 from io import BytesIO
 from typing import Any
@@ -368,6 +369,46 @@ def add(
             external_id,
         )
 
+    def add_many(self, records: Iterable[Mapping[str, Any]]) -> None:
+        """Append multiple records in one storage operation.
+
+        Each record accepts the same fields as :meth:`add`: ``role``,
+        ``content``, optional ``content_type``/``data_type``, ``embedding``,
+        ``bot_id``, ``session_id``, and ``external_id``.
+        """
+        normalized: list[dict[str, Any]] = []
+        for index, record in enumerate(records):
+            if not isinstance(record, Mapping):
+                raise TypeError(f"records[{index}] must be a mapping")
+            if "role" not in record:
+                raise ValueError(f"records[{index}] is missing required key 'role'")
+            if "content" not in record:
+                raise ValueError(f"records[{index}] is missing required key 'content'")
+
+            content_type = record.get("content_type")
+            data_type = record.get("data_type")
+            if content_type is not None and data_type is not None:
+                raise ValueError(
+                    f"records[{index}] specifies both content_type and data_type"
+                )
+            if content_type is None:
+                content_type = data_type
+
+            payload, resolved_type = _normalize_content(record["content"], content_type)
+            normalized.append(
+                {
+                    "role": record["role"],
+                    "content": payload,
+                    "data_type": resolved_type,
+                    "embedding": record.get("embedding"),
+                    "bot_id": record.get("bot_id"),
+                    "session_id": record.get("session_id"),
+                    "external_id": record.get("external_id"),
+                }
+            )
+
+        self._inner.add_many(normalized)
+
     def snapshot(self, label: str | None = None) -> str:
         return self._inner.snapshot(label)
 

python/src/lib.rs

@@ -2,7 +2,7 @@ use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 
 use chrono::{SecondsFormat, Utc};
-use pyo3::exceptions::PyRuntimeError;
+use pyo3::exceptions::{PyRuntimeError, PyTypeError};
 use pyo3::prelude::*;
 use pyo3::types::{PyBytes, PyDict, PyType};
 use pyo3::IntoPyObject;
@@ -17,6 +17,13 @@ use lance_context::{
 const DEFAULT_BINARY_CONTENT_TYPE: &str = "application/octet-stream";
 const BINARY_PLACEHOLDER: &str = "[binary]";
 
+struct PreparedRecord {
+    record: ContextRecord,
+    role: String,
+    inner_content: String,
+    data_type: Option<String>,
+}
+
 #[pyfunction]
 fn version() -> &'static str {
     env!("CARGO_PKG_VERSION")
@@ -184,47 +191,55 @@ impl Context {
         session_id: Option<String>,
         external_id: Option<String>,
     ) -> PyResult<()> {
-        let (content_type, text_payload, binary_payload, inner_content) =
-            match content.extract::<&[u8]>() {
-                Ok(bytes) => (
-                    data_type.unwrap_or(DEFAULT_BINARY_CONTENT_TYPE).to_string(),
-                    None,
-                    Some(bytes.to_vec()),
-                    BINARY_PLACEHOLDER.to_string(),
-                ),
-                Err(_) => {
-                    let content_str = content.str()?.to_string();
-                    (
-                        data_type.unwrap_or(CONTENT_TYPE_TEXT).to_string(),
-                        Some(content_str.clone()),
-                        None,
-                        content_str,
-                    )
-                }
-            };
-
-        let record_id = format!("{}-{}", self.run_id, self.inner.entries() + 1);
-        let record = ContextRecord {
-            id: record_id,
-            external_id,
-            run_id: self.run_id.clone(),
+        let prepared = self.prepare_record(
+            role.to_string(),
+            content,
+            data_type.map(str::to_string),
+            embedding,
             bot_id,
             session_id,
-            created_at: Utc::now(),
-            role: role.to_string(),
-            state_metadata: None,
-            content_type,
-            text_payload,
-            binary_payload,
-            embedding,
-        };
+            external_id,
+            1,
+        )?;
 
         let add_res = py.allow_threads(|| {
             self.runtime
-                .block_on(self.store.add(std::slice::from_ref(&record)))
+                .block_on(self.store.add(std::slice::from_ref(&prepared.record)))
         });
         add_res.map_err(to_py_err)?;
-        self.inner.add(role, &inner_content, data_type);
+        self.inner.add(
+            &prepared.role,
+            &prepared.inner_content,
+            prepared.data_type.as_deref(),
+        );
+        Ok(())
+    }
+
+    #[pyo3(signature = (records))]
+    fn add_many(&mut self, py: Python<'_>, records: &Bound<'_, PyAny>) -> PyResult<()> {
+        let mut prepared = Vec::new();
+        for (index, item) in records.try_iter()?.enumerate() {
+            let item = item?;
+            let dict = item.downcast::<PyDict>().map_err(|_| {
+                PyTypeError::new_err(format!("records[{index}] must be a dict"))
+            })?;
+            prepared.push(self.prepare_record_from_dict(dict, index)?);
+        }
+
+        if prepared.is_empty() {
+            return Ok(());
+        }
+
+        let context_records: Vec<ContextRecord> =
+            prepared.iter().map(|item| item.record.clone()).collect();
+        let add_res =
+            py.allow_threads(|| self.runtime.block_on(self.store.add(&context_records)));
+        add_res.map_err(to_py_err)?;
+
+        for item in prepared {
+            self.inner
+                .add(&item.role, &item.inner_content, item.data_type.as_deref());
+        }
         Ok(())
     }
 
@@ -359,6 +374,110 @@ impl Context {
     }
 }
 
+impl Context {
+    fn prepare_record_from_dict(
+        &self,
+        dict: &Bound<'_, PyDict>,
+        index: usize,
+    ) -> PyResult<PreparedRecord> {
+        let role = required_item(dict, "role", index)?.extract::<String>()?;
+        let content = required_item(dict, "content", index)?;
+        let data_type =
+            optional_item(dict, "data_type")?.map(|value| value.extract::<String>());
+        let embedding =
+            optional_item(dict, "embedding")?.map(|value| value.extract::<Vec<f32>>());
+        let bot_id = optional_item(dict, "bot_id")?.map(|value| value.extract::<String>());
+        let session_id =
+            optional_item(dict, "session_id")?.map(|value| value.extract::<String>());
+        let external_id =
+            optional_item(dict, "external_id")?.map(|value| value.extract::<String>());
+
+        self.prepare_record(
+            role,
+            &content,
+            data_type.transpose()?,
+            embedding.transpose()?,
+            bot_id.transpose()?,
+            session_id.transpose()?,
+            external_id.transpose()?,
+            index as u64 + 1,
+        )
+    }
+
+    fn prepare_record(
+        &self,
+        role: String,
+        content: &Bound<'_, PyAny>,
+        data_type: Option<String>,
+        embedding: Option<Vec<f32>>,
+        bot_id: Option<String>,
+        session_id: Option<String>,
+        external_id: Option<String>,
+        offset: u64,
+    ) -> PyResult<PreparedRecord> {
+        let (content_type, text_payload, binary_payload, inner_content) =
+            match content.extract::<&[u8]>() {
+                Ok(bytes) => (
+                    data_type
+                        .clone()
+                        .unwrap_or_else(|| DEFAULT_BINARY_CONTENT_TYPE.to_string()),
+                    None,
+                    Some(bytes.to_vec()),
+                    BINARY_PLACEHOLDER.to_string(),
+                ),
+                Err(_) => {
+                    let content_str = content.str()?.to_string();
+                    (
+                        data_type
+                            .clone()
+                            .unwrap_or_else(|| CONTENT_TYPE_TEXT.to_string()),
+                        Some(content_str.clone()),
+                        None,
+                        content_str,
+                    )
+                }
+            };
+
+        let record_id = format!("{}-{}", self.run_id, self.inner.entries() + offset);
+        Ok(PreparedRecord {
+            record: ContextRecord {
+                id: record_id,
+                external_id,
+                run_id: self.run_id.clone(),
+                bot_id,
+                session_id,
+                created_at: Utc::now(),
+                role: role.clone(),
+                state_metadata: None,
+                content_type,
+                text_payload,
+                binary_payload,
+                embedding,
+            },
+            role,
+            inner_content,
+            data_type,
+        })
+    }
+}
+
+fn required_item<'py>(
+    dict: &Bound<'py, PyDict>,
+    key: &str,
+    index: usize,
+) -> PyResult<Bound<'py, PyAny>> {
+    dict.get_item(key)?.ok_or_else(|| {
+        PyRuntimeError::new_err(format!("records[{index}] is missing required key '{key}'"))
+    })
+}
+
+fn optional_item<'py>(
+    dict: &Bound<'py, PyDict>,
+    key: &str,
+) -> PyResult<Option<Bound<'py, PyAny>>> {
+    Ok(dict.get_item(key)?.filter(|value| !value.is_none()))
+}
+
 fn compaction_metrics_to_py(py: Python<'_>, metrics: CompactionMetrics) -> PyResult<PyObject> {
     let dict = PyDict::new(py);
     dict.set_item("fragments_removed", metrics.fragments_removed)?;

python/tests/test_add_many.py

@@ -0,0 +1,67 @@
+"""Tests for batch append support."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import pytest
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+from lance_context.api import Context
+
+
+def test_add_many_appends_records_in_one_call(tmp_path: Path) -> None:
+    uri = str(tmp_path / "context.lance")
+    ctx = Context.create(uri)
+
+    ctx.add_many(
+        [
+            {"role": "user", "content": "first"},
+            {
+                "role": "assistant",
+                "content": "second",
+                "content_type": "text/markdown",
+                "bot_id": "bot",
+                "session_id": "session",
+                "external_id": "doc-1#chunk-2",
+            },
+            {
+                "role": "tool",
+                "content": b"\x01\x02",
+                "data_type": "application/octet-stream",
+            },
+        ]
+    )
+
+    assert ctx.entries() == 3
+    records = ctx.list()
+    assert [record["role"] for record in records] == ["user", "assistant", "tool"]
+    assert [record["text"] for record in records[:2]] == ["first", "second"]
+    assert records[1]["content_type"] == "text/markdown"
+    assert records[1]["bot_id"] == "bot"
+    assert records[1]["session_id"] == "session"
+    assert records[1]["external_id"] == "doc-1#chunk-2"
+    assert records[2]["binary"] == b"\x01\x02"
+
+
+def test_add_many_empty_batch_is_noop(tmp_path: Path) -> None:
+    uri = str(tmp_path / "context.lance")
+    ctx = Context.create(uri)
+
+    ctx.add_many([])
+
+    assert ctx.entries() == 0
+    assert ctx.list() == []
+
+
+def test_add_many_validates_records_before_write(tmp_path: Path) -> None:
+    uri = str(tmp_path / "context.lance")
+    ctx = Context.create(uri)
+
+    with pytest.raises(ValueError, match="missing required key 'content'"):
+        ctx.add_many([{"role": "user"}])
+
+    assert ctx.entries() == 0
+    assert ctx.list() == []