From 1f7fbf1dea0b51f1c526a9a3e7599686b40bce04 Mon Sep 17 00:00:00 2001 From: Allen Cheng Date: Mon, 8 Jun 2026 21:31:48 -0700 Subject: [PATCH 1/2] feat: add batch append API --- README.md | 16 +++ python/python/lance_context/api.py | 40 +++++++ python/src/lib.rs | 180 +++++++++++++++++++++++------ python/tests/test_add_many.py | 65 +++++++++++ python/tests/test_search.py | 79 +++++++++++++ 5 files changed, 346 insertions(+), 34 deletions(-) create mode 100644 python/tests/test_add_many.py diff --git a/README.md b/README.md index 71d4c1e..5799c6a 100644 --- a/README.md +++ b/README.md @@ -62,6 +62,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…") diff --git a/python/python/lance_context/api.py b/python/python/lance_context/api.py index 5e7161b..91aeda1 100644 --- a/python/python/lance_context/api.py +++ b/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 @@ -358,6 +359,45 @@ def add( payload, resolved_type = _normalize_content(content, content_type) self._inner.add(role, payload, resolved_type, embedding, bot_id, session_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``, and ``session_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"), + } + ) + + self._inner.add_many(normalized) + def snapshot(self, label: str | None = None) -> str: return self._inner.snapshot(label) diff --git a/python/src/lib.rs b/python/src/lib.rs index 2e2f855..cd23eb4 100644 --- a/python/src/lib.rs +++ b/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, +} + #[pyfunction] fn version() -> &'static str { env!("CARGO_PKG_VERSION") @@ -183,46 +190,55 @@ impl Context { bot_id: Option, session_id: Option, ) -> 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, - 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, - }; + 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.cast::().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 = 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(()) } @@ -324,6 +340,102 @@ impl Context { } } +impl Context { + fn prepare_record_from_dict( + &self, + dict: &Bound<'_, PyDict>, + index: usize, + ) -> PyResult { + let role = required_item(dict, "role", index)?.extract::()?; + let content = required_item(dict, "content", index)?; + let data_type = optional_item(dict, "data_type")?.map(|value| value.extract::()); + let embedding = optional_item(dict, "embedding")?.map(|value| value.extract::>()); + let bot_id = optional_item(dict, "bot_id")?.map(|value| value.extract::()); + let session_id = optional_item(dict, "session_id")?.map(|value| value.extract::()); + + self.prepare_record( + role, + &content, + data_type.transpose()?, + embedding.transpose()?, + bot_id.transpose()?, + session_id.transpose()?, + index as u64 + 1, + ) + } + + fn prepare_record( + &self, + role: String, + content: &Bound<'_, PyAny>, + data_type: Option, + embedding: Option>, + bot_id: Option, + session_id: Option, + offset: u64, + ) -> PyResult { + 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, + 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> { + 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>> { + Ok(dict.get_item(key)?.filter(|value| !value.is_none())) +} + fn compaction_metrics_to_py(py: Python<'_>, metrics: CompactionMetrics) -> PyResult { let dict = PyDict::new(py); dict.set_item("fragments_removed", metrics.fragments_removed)?; diff --git a/python/tests/test_add_many.py b/python/tests/test_add_many.py new file mode 100644 index 0000000..1728daa --- /dev/null +++ b/python/tests/test_add_many.py @@ -0,0 +1,65 @@ +"""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", + }, + { + "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[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() == [] diff --git a/python/tests/test_search.py b/python/tests/test_search.py index 32eae26..59edb90 100644 --- a/python/tests/test_search.py +++ b/python/tests/test_search.py @@ -17,6 +17,7 @@ def __init__(self) -> None: self.add_calls: list[ tuple[str, Any, str | None, list[float] | None, str | None, str | None] ] = [] + self.add_many_calls: list[list[dict[str, Any]]] = [] def add( self, @@ -29,6 +30,9 @@ def add( ): self.add_calls.append((role, content, data_type, embedding, bot_id, session_id)) + def add_many(self, records: list[dict[str, Any]]): + self.add_many_calls.append(records) + def search(self, vector: list[float], limit: int | None): self.search_calls.append((vector, limit)) return [ @@ -279,6 +283,81 @@ def test_context_add_with_all_options(): assert session_id == "sess" +def test_context_add_many_normalizes_records(): + ctx = Context.__new__(Context) + dummy = DummyInner() + ctx._inner = dummy # type: ignore[attr-defined] + + ctx.add_many( + [ + {"role": "user", "content": "hello"}, + { + "role": "assistant", + "content": "world", + "content_type": "text/markdown", + "embedding": [0.1, 0.2], + "bot_id": "bot", + "session_id": "sess", + }, + ] + ) + + assert dummy.add_many_calls == [ + [ + { + "role": "user", + "content": "hello", + "data_type": None, + "embedding": None, + "bot_id": None, + "session_id": None, + }, + { + "role": "assistant", + "content": "world", + "data_type": "text/markdown", + "embedding": [0.1, 0.2], + "bot_id": "bot", + "session_id": "sess", + }, + ] + ] + + +def test_context_add_many_accepts_data_type_alias(): + ctx = Context.__new__(Context) + dummy = DummyInner() + ctx._inner = dummy # type: ignore[attr-defined] + + ctx.add_many([{"role": "system", "content": "prompt", "data_type": "text/plain"}]) + + assert dummy.add_many_calls[0][0]["data_type"] == "text/plain" + + +def test_context_add_many_rejects_invalid_records(): + ctx = Context.__new__(Context) + dummy = DummyInner() + ctx._inner = dummy # type: ignore[attr-defined] + + with pytest.raises(TypeError, match="records\\[0\\]"): + ctx.add_many(["not-a-record"]) # type: ignore[list-item] + with pytest.raises(ValueError, match="missing required key 'role'"): + ctx.add_many([{"content": "hello"}]) + with pytest.raises(ValueError, match="missing required key 'content'"): + ctx.add_many([{"role": "user"}]) + with pytest.raises(ValueError, match="both content_type and data_type"): + ctx.add_many( + [ + { + "role": "user", + "content": "hello", + "content_type": "text/plain", + "data_type": "text/markdown", + } + ] + ) + + def test_normalize_record_with_agent_and_session_id(): result = _normalize_record( { From 94b3d78959eee9e68de34da8cf47927dcaf44c7f Mon Sep 17 00:00:00 2001 From: Allen Cheng Date: Mon, 8 Jun 2026 21:56:54 -0700 Subject: [PATCH 2/2] fix: use PyO3 downcast for add_many records --- python/src/lib.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/src/lib.rs b/python/src/lib.rs index cd23eb4..016e708 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -218,7 +218,7 @@ impl Context { let mut prepared = Vec::new(); for (index, item) in records.try_iter()?.enumerate() { let item = item?; - let dict = item.cast::().map_err(|_| { + let dict = item.downcast::().map_err(|_| { PyTypeError::new_err(format!("records[{index}] must be a dict")) })?; prepared.push(self.prepare_record_from_dict(dict, index)?);