From 316cce432a6ce83e703a1a46f9cebe4a41063765 Mon Sep 17 00:00:00 2001 From: Allen Cheng Date: Mon, 8 Jun 2026 21:15:14 -0700 Subject: [PATCH 1/4] feat: support external record identifiers --- README.md | 8 +- crates/lance-context-core/src/record.rs | 1 + crates/lance-context-core/src/store.rs | 194 +++++++++++++++++++++--- python/python/lance_context/api.py | 23 ++- python/src/lib.rs | 39 ++++- python/tests/test_external_id.py | 50 ++++++ python/tests/test_search.py | 129 +++++++++++++--- 7 files changed, 405 insertions(+), 39 deletions(-) create mode 100644 python/tests/test_external_id.py diff --git a/README.md b/README.md index 71d4c1e..d6261d6 100644 --- a/README.md +++ b/README.md @@ -54,7 +54,12 @@ uri = Path("context.lance").as_posix() ctx = Context.create(uri) # Add multimodal entries -ctx.add("user", "Where should I travel in spring?") +ctx.add( + "user", + "Where should I travel in spring?", + external_id="conversation-2026-03-01#turn-1", +) +print(ctx.get(external_id="conversation-2026-03-01#turn-1")) from PIL import Image image = Image.new("RGB", (2, 2), color="teal") @@ -138,6 +143,7 @@ use chrono::Utc; let mut store = ContextStore::open("context.lance").await?; let record = ContextRecord { id: "run-1-1".into(), + external_id: None, run_id: "run-1".into(), created_at: Utc::now(), role: "user".into(), diff --git a/crates/lance-context-core/src/record.rs b/crates/lance-context-core/src/record.rs index e5729eb..95cf27b 100644 --- a/crates/lance-context-core/src/record.rs +++ b/crates/lance-context-core/src/record.rs @@ -13,6 +13,7 @@ pub struct StateMetadata { #[derive(Debug, Clone)] pub struct ContextRecord { pub id: String, + pub external_id: Option, pub run_id: String, pub bot_id: Option, pub session_id: Option, diff --git a/crates/lance-context-core/src/store.rs b/crates/lance-context-core/src/store.rs index 8f6d8f4..d13640d 100644 --- a/crates/lance-context-core/src/store.rs +++ b/crates/lance-context-core/src/store.rs @@ -203,6 +203,8 @@ impl ContextStore { return Ok(self.dataset.manifest.version); } + self.validate_unique_ids(entries).await?; + // Group entries by (bot_id, session_id) let mut groups: HashMap<(Option, Option), Vec> = HashMap::new(); @@ -250,6 +252,50 @@ impl ContextStore { Ok(self.dataset.manifest.version) } + async fn validate_unique_ids(&self, entries: &[ContextRecord]) -> LanceResult<()> { + let mut ids = HashSet::new(); + let mut external_ids = HashSet::new(); + for entry in entries { + if !ids.insert(entry.id.as_str()) { + return Err(ArrowError::InvalidArgumentError(format!( + "duplicate id '{}' in batch", + entry.id + )) + .into()); + } + if let Some(external_id) = &entry.external_id { + if !external_ids.insert(external_id.as_str()) { + return Err(ArrowError::InvalidArgumentError(format!( + "duplicate external_id '{}' in batch", + external_id + )) + .into()); + } + } + } + + for record in self.list(None, None).await? { + if ids.contains(record.id.as_str()) { + return Err(ArrowError::InvalidArgumentError(format!( + "id '{}' already exists", + record.id + )) + .into()); + } + if let Some(external_id) = record.external_id { + if external_ids.contains(external_id.as_str()) { + return Err(ArrowError::InvalidArgumentError(format!( + "external_id '{}' already exists", + external_id + )) + .into()); + } + } + } + + Ok(()) + } + fn derive_region_id(bot_id: &Option, session_id: &Option) -> Uuid { let mut input = String::new(); @@ -299,6 +345,27 @@ impl ContextStore { Ok(results) } + /// Find a record by its internal storage id. + pub async fn get_by_id(&self, id: &str) -> LanceResult> { + Ok(self + .list(None, None) + .await? + .into_iter() + .find(|record| record.id == id)) + } + + /// Find a record by its caller-supplied external id. + pub async fn get_by_external_id( + &self, + external_id: &str, + ) -> LanceResult> { + Ok(self + .list(None, None) + .await? + .into_iter() + .find(|record| record.external_id.as_deref() == Some(external_id))) + } + /// Perform a nearest-neighbor search over stored embeddings. pub async fn search( &self, @@ -583,6 +650,10 @@ impl ContextStore { /// Lance V1 blob encoding (out-of-line binary buffers). For `text_payload`, /// this also changes the Arrow type from `LargeUtf8` to `LargeBinary`. pub fn schema(blob_columns: &HashSet) -> Schema { + Self::schema_with_options(blob_columns, true) + } + + fn schema_with_options(blob_columns: &HashSet, include_external_id: bool) -> Schema { let mut id_metadata = HashMap::new(); id_metadata.insert( "lance-schema:unenforced-primary-key".to_string(), @@ -605,8 +676,13 @@ impl ContextStore { Field::new("binary_payload", DataType::LargeBinary, true) }; - Schema::new(vec![ + let mut fields = vec![ Field::new("id", DataType::Utf8, false).with_metadata(id_metadata), + ]; + if include_external_id { + fields.push(Field::new("external_id", DataType::Utf8, true)); + } + fields.extend([ Field::new("run_id", DataType::Utf8, false), Field::new("bot_id", DataType::Utf8, true), Field::new("session_id", DataType::Utf8, true), @@ -644,7 +720,9 @@ impl ContextStore { ), true, ), - ]) + ]); + + Schema::new(fields) } async fn load_with_options( @@ -692,7 +770,21 @@ impl ContextStore { } fn records_to_batch(&self, entries: &[ContextRecord]) -> LanceResult { + let include_external_id = self + .dataset + .schema() + .field_with_name("external_id") + .is_ok(); + if !include_external_id && entries.iter().any(|entry| entry.external_id.is_some()) { + return Err(ArrowError::InvalidArgumentError( + "external_id requires a context dataset created with external_id support" + .to_string(), + ) + .into()); + } + let mut id_builder = StringBuilder::new(); + let mut external_id_builder = StringBuilder::new(); let mut run_id_builder = StringBuilder::new(); let mut bot_id_builder = StringBuilder::new(); let mut session_id_builder = StringBuilder::new(); @@ -734,6 +826,7 @@ impl ContextStore { for entry in entries { id_builder.append_value(&entry.id); + external_id_builder.append_option(entry.external_id.as_deref()); run_id_builder.append_value(&entry.run_id); bot_id_builder.append_option(entry.bot_id.as_deref()); session_id_builder.append_option(entry.session_id.as_deref()); @@ -826,6 +919,7 @@ impl ContextStore { } let id_array: ArrayRef = Arc::new(id_builder.finish()); + let external_id_array: ArrayRef = Arc::new(external_id_builder.finish()); let run_id_array: ArrayRef = Arc::new(run_id_builder.finish()); let bot_id_array: ArrayRef = Arc::new(bot_id_builder.finish()); let session_id_array: ArrayRef = Arc::new(session_id_builder.finish()); @@ -841,23 +935,27 @@ impl ContextStore { let state_array: ArrayRef = Arc::new(state_builder.finish()); let embedding_array: ArrayRef = Arc::new(embedding_builder.finish()); - let schema = Arc::new(Self::schema(&self.blob_columns)); - let batch = RecordBatch::try_new( - schema, - vec![ - id_array, - run_id_array, - bot_id_array, - session_id_array, - created_at_array, - role_array, - state_array, - content_type_array, - text_array, - binary_array, - embedding_array, - ], - )?; + let schema = Arc::new(Self::schema_with_options( + &self.blob_columns, + include_external_id, + )); + let mut arrays = vec![id_array]; + if include_external_id { + arrays.push(external_id_array); + } + arrays.extend([ + run_id_array, + bot_id_array, + session_id_array, + created_at_array, + role_array, + state_array, + content_type_array, + text_array, + binary_array, + embedding_array, + ]); + let batch = RecordBatch::try_new(schema, arrays)?; Ok(batch) } @@ -877,6 +975,7 @@ impl Drop for ContextStore { /// Convert a record batch to context records. fn batch_to_records(batch: &RecordBatch) -> LanceResult> { let id_array = column_as::(batch, "id")?; + let external_id_array = column_as_optional::(batch, "external_id"); let run_id_array = column_as::(batch, "run_id")?; let bot_id_array = column_as_optional::(batch, "bot_id"); let session_id_array = column_as_optional::(batch, "session_id"); @@ -1046,6 +1145,13 @@ fn batch_to_records(batch: &RecordBatch) -> LanceResult> { results.push(ContextRecord { id: id_array.value(row).to_string(), + external_id: external_id_array.and_then(|arr| { + if arr.is_null(row) { + None + } else { + Some(arr.value(row).to_string()) + } + }), run_id: run_id_array.value(row).to_string(), bot_id, session_id, @@ -1134,6 +1240,7 @@ mod tests { fn text_record(id: &str, embedding_pivot: f32) -> ContextRecord { ContextRecord { id: id.to_string(), + external_id: None, run_id: format!("run-{id}"), bot_id: None, session_id: None, @@ -1192,6 +1299,55 @@ mod tests { }); } + #[test] + fn external_id_roundtrips_and_supports_lookup() { + let dir = TempDir::new().unwrap(); + let uri = dir.path().to_string_lossy().to_string(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let mut store = ContextStore::open(&uri).await.unwrap(); + let mut record = text_record("a", 0.0); + record.external_id = Some("doc-123#chunk-1".to_string()); + store.add(std::slice::from_ref(&record)).await.unwrap(); + + let by_external_id = store + .get_by_external_id("doc-123#chunk-1") + .await + .unwrap() + .unwrap(); + assert_eq!(by_external_id.id, record.id); + assert_eq!(by_external_id.external_id, record.external_id); + + let by_id = store.get_by_id(&record.id).await.unwrap().unwrap(); + assert_eq!(by_id.external_id.as_deref(), Some("doc-123#chunk-1")); + + let missing = store.get_by_external_id("missing").await.unwrap(); + assert!(missing.is_none()); + }); + } + + #[test] + fn add_rejects_duplicate_external_id() { + let dir = TempDir::new().unwrap(); + let uri = dir.path().to_string_lossy().to_string(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let mut store = ContextStore::open(&uri).await.unwrap(); + let mut first = text_record("a", 0.0); + first.external_id = Some("doc-123#chunk-1".to_string()); + store.add(std::slice::from_ref(&first)).await.unwrap(); + + let mut duplicate = text_record("b", 0.0); + duplicate.external_id = first.external_id.clone(); + let err = store.add(&[duplicate]).await.unwrap_err(); + let message = err.to_string(); + assert!( + message.contains("external_id") && message.contains("already exists"), + "unexpected error message: {message}" + ); + }); + } + #[test] fn test_region_id_derivation_explicit() { let bot_id = Some("bot-123".to_string()); diff --git a/python/python/lance_context/api.py b/python/python/lance_context/api.py index 5e7161b..d431def 100644 --- a/python/python/lance_context/api.py +++ b/python/python/lance_context/api.py @@ -115,6 +115,7 @@ def _normalize_record(raw: dict[str, Any]) -> dict[str, Any]: created_at = datetime.fromisoformat(created_at.replace("Z", "+00:00")) return { "id": raw.get("id"), + "external_id": raw.get("external_id"), "run_id": raw.get("run_id"), "bot_id": raw.get("bot_id"), "session_id": raw.get("session_id"), @@ -350,13 +351,22 @@ def add( embedding: list[float] | None = None, bot_id: str | None = None, session_id: str | None = None, + external_id: str | None = None, ) -> None: if content_type is not None and data_type is not None: raise ValueError("Specify only one of content_type or data_type") if content_type is None: content_type = data_type payload, resolved_type = _normalize_content(content, content_type) - self._inner.add(role, payload, resolved_type, embedding, bot_id, session_id) + self._inner.add( + role, + payload, + resolved_type, + embedding, + bot_id, + session_id, + external_id, + ) def snapshot(self, label: str | None = None) -> str: return self._inner.snapshot(label) @@ -389,6 +399,17 @@ def list( results = self._inner.list(limit, offset) return [_normalize_record(item) for item in results] + def get( + self, *, id: str | None = None, external_id: str | None = None + ) -> dict[str, Any] | None: + """Return one entry by internal id or caller-supplied external id.""" + if (id is None) == (external_id is None): + raise ValueError("Specify exactly one of id or external_id") + result = self._inner.get(id, external_id) + if result is None: + return None + return _normalize_record(result) + def compact( self, *, diff --git a/python/src/lib.rs b/python/src/lib.rs index 2e2f855..f40a460 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -172,7 +172,7 @@ impl Context { } #[allow(clippy::too_many_arguments)] - #[pyo3(signature = (role, content, data_type = None, embedding = None, bot_id = None, session_id = None))] + #[pyo3(signature = (role, content, data_type = None, embedding = None, bot_id = None, session_id = None, external_id = None))] fn add( &mut self, py: Python<'_>, @@ -182,6 +182,7 @@ impl Context { embedding: Option>, bot_id: Option, session_id: Option, + external_id: Option, ) -> PyResult<()> { let (content_type, text_payload, binary_payload, inner_content) = match content.extract::<&[u8]>() { @@ -205,6 +206,7 @@ impl Context { 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(), bot_id, session_id, @@ -281,6 +283,39 @@ impl Context { .collect() } + #[pyo3(signature = (id = None, external_id = None))] + fn get( + &self, + py: Python<'_>, + id: Option, + external_id: Option, + ) -> PyResult> { + let record = match (id, external_id) { + (Some(id), None) => py.allow_threads(|| { + self.runtime + .block_on(self.store.get_by_id(&id)) + .map_err(to_py_err) + })?, + (None, Some(external_id)) => py.allow_threads(|| { + self.runtime + .block_on(self.store.get_by_external_id(&external_id)) + .map_err(to_py_err) + })?, + (None, None) => { + return Err(PyRuntimeError::new_err( + "get() requires either id or external_id", + )); + } + (Some(_), Some(_)) => { + return Err(PyRuntimeError::new_err( + "get() accepts only one of id or external_id", + )); + } + }; + + record.map(|record| record_to_py(py, record)).transpose() + } + #[pyo3(signature = (target_rows_per_fragment=None, materialize_deletions=None))] fn compact( &mut self, @@ -367,6 +402,7 @@ fn search_hit_to_py(py: Python<'_>, hit: SearchResult) -> PyResult { fn record_to_py(py: Python<'_>, record: ContextRecord) -> PyResult { let ContextRecord { id, + external_id, run_id, bot_id, session_id, @@ -381,6 +417,7 @@ fn record_to_py(py: Python<'_>, record: ContextRecord) -> PyResult { let dict = PyDict::new(py); dict.set_item("id", id)?; + dict.set_item("external_id", external_id)?; dict.set_item("run_id", run_id)?; dict.set_item("bot_id", bot_id)?; dict.set_item("session_id", session_id)?; diff --git a/python/tests/test_external_id.py b/python/tests/test_external_id.py new file mode 100644 index 0000000..caab842 --- /dev/null +++ b/python/tests/test_external_id.py @@ -0,0 +1,50 @@ +"""Tests for caller-supplied external record identifiers.""" + +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_external_id_roundtrip_and_lookup(tmp_path: Path) -> None: + uri = str(tmp_path / "context.lance") + ctx = Context.create(uri) + + ctx.add("user", "stable memory", external_id="doc-123#chunk-1") + + entries = ctx.list() + assert entries[0]["external_id"] == "doc-123#chunk-1" + + by_external_id = ctx.get(external_id="doc-123#chunk-1") + assert by_external_id is not None + assert by_external_id["text"] == "stable memory" + assert by_external_id["external_id"] == "doc-123#chunk-1" + + by_id = ctx.get(id=by_external_id["id"]) + assert by_id is not None + assert by_id["external_id"] == "doc-123#chunk-1" + + +def test_external_id_missing_lookup_returns_none(tmp_path: Path) -> None: + uri = str(tmp_path / "context.lance") + ctx = Context.create(uri) + + ctx.add("user", "stable memory", external_id="doc-123#chunk-1") + + assert ctx.get(external_id="missing") is None + + +def test_duplicate_external_id_is_rejected(tmp_path: Path) -> None: + uri = str(tmp_path / "context.lance") + ctx = Context.create(uri) + + ctx.add("user", "first", external_id="doc-123#chunk-1") + + with pytest.raises(RuntimeError, match="external_id.*already exists"): + ctx.add("user", "duplicate", external_id="doc-123#chunk-1") diff --git a/python/tests/test_search.py b/python/tests/test_search.py index 32eae26..51dbe1e 100644 --- a/python/tests/test_search.py +++ b/python/tests/test_search.py @@ -14,8 +14,17 @@ class DummyInner: def __init__(self) -> None: self.search_calls: list[tuple[list[float], int | None]] = [] self.list_calls: list[tuple[int | None, int | None]] = [] + self.get_calls: list[tuple[str | None, str | None]] = [] self.add_calls: list[ - tuple[str, Any, str | None, list[float] | None, str | None, str | None] + tuple[ + str, + Any, + str | None, + list[float] | None, + str | None, + str | None, + str | None, + ] ] = [] def add( @@ -26,14 +35,24 @@ def add( embedding: list[float] | None, bot_id: str | None, session_id: str | None, + external_id: str | None, ): - self.add_calls.append((role, content, data_type, embedding, bot_id, session_id)) + self.add_calls.append( + (role, content, data_type, embedding, bot_id, session_id, external_id) + ) + + def get(self, id: str | None, external_id: str | None): + self.get_calls.append((id, external_id)) + if id == "rec-1" or external_id == "source-1": + return self.list(None, None)[0] + return None def search(self, vector: list[float], limit: int | None): self.search_calls.append((vector, limit)) return [ { "id": "rec-1", + "external_id": "source-1", "run_id": "run-1", "bot_id": "support_bot", "session_id": None, @@ -53,6 +72,7 @@ def list(self, limit: int | None, offset: int | None): return [ { "id": "rec-1", + "external_id": "source-1", "run_id": "run-1", "bot_id": "support_bot", "session_id": "user_1", @@ -66,6 +86,7 @@ def list(self, limit: int | None, offset: int | None): }, { "id": "rec-2", + "external_id": None, "run_id": "run-1", "bot_id": None, "session_id": None, @@ -80,6 +101,11 @@ def list(self, limit: int | None, offset: int | None): ] +def _only_add_call(dummy: DummyInner): + assert len(dummy.add_calls) == 1 + return dummy.add_calls[0] + + def test_coerce_vector_from_list(): assert _coerce_vector([1, 2.5]) == [1.0, 2.5] @@ -93,6 +119,7 @@ def test_normalize_search_hit_converts_timestamp(): result = _normalize_search_hit( { "id": "rec-2", + "external_id": None, "created_at": "2024-01-01T00:00:00Z", "content_type": "text/plain", "text_payload": None, @@ -125,6 +152,7 @@ def test_normalize_record_without_distance(): result = _normalize_record( { "id": "rec-1", + "external_id": "source-1", "created_at": "2024-01-01T00:00:00Z", "content_type": "text/plain", "text_payload": "hello", @@ -158,6 +186,50 @@ def test_context_list_returns_entries(): assert isinstance(entries[0]["created_at"], datetime) +def test_context_get_by_external_id(): + ctx = Context.__new__(Context) + dummy = DummyInner() + ctx._inner = dummy # type: ignore[attr-defined] + + entry = ctx.get(external_id="source-1") + + assert dummy.get_calls == [(None, "source-1")] + assert entry is not None + assert entry["id"] == "rec-1" + assert entry["external_id"] == "source-1" + + +def test_context_get_by_id(): + ctx = Context.__new__(Context) + dummy = DummyInner() + ctx._inner = dummy # type: ignore[attr-defined] + + entry = ctx.get(id="rec-1") + + assert dummy.get_calls == [("rec-1", None)] + assert entry is not None + assert entry["id"] == "rec-1" + + +def test_context_get_missing_returns_none(): + ctx = Context.__new__(Context) + dummy = DummyInner() + ctx._inner = dummy # type: ignore[attr-defined] + + assert ctx.get(external_id="missing") is None + + +def test_context_get_requires_exactly_one_identifier(): + ctx = Context.__new__(Context) + dummy = DummyInner() + ctx._inner = dummy # type: ignore[attr-defined] + + with pytest.raises(ValueError, match="exactly one"): + ctx.get() + with pytest.raises(ValueError, match="exactly one"): + ctx.get(id="rec-1", external_id="source-1") + + def test_context_list_default_args(): ctx = Context.__new__(Context) dummy = DummyInner() @@ -176,14 +248,16 @@ def test_context_add_with_embedding(): embedding = [0.1, 0.2, 0.3] ctx.add("user", "hello", embedding=embedding) - assert len(dummy.add_calls) == 1 - role, content, data_type, passed_embedding, bot_id, session_id = dummy.add_calls[0] + role, content, data_type, passed_embedding, bot_id, session_id, external_id = ( + _only_add_call(dummy) + ) assert role == "user" assert content == "hello" assert data_type is None assert passed_embedding == [0.1, 0.2, 0.3] assert bot_id is None assert session_id is None + assert external_id is None def test_context_add_without_embedding(): @@ -193,13 +267,15 @@ def test_context_add_without_embedding(): ctx.add("assistant", "world") - assert len(dummy.add_calls) == 1 - role, content, data_type, passed_embedding, bot_id, session_id = dummy.add_calls[0] + role, content, data_type, passed_embedding, bot_id, session_id, external_id = ( + _only_add_call(dummy) + ) assert role == "assistant" assert content == "world" assert passed_embedding is None assert bot_id is None assert session_id is None + assert external_id is None def test_context_add_with_content_type_and_embedding(): @@ -210,13 +286,15 @@ def test_context_add_with_content_type_and_embedding(): embedding = [0.5, 0.6] ctx.add("system", "prompt", content_type="text/markdown", embedding=embedding) - assert len(dummy.add_calls) == 1 - role, content, data_type, passed_embedding, bot_id, session_id = dummy.add_calls[0] + role, content, data_type, passed_embedding, bot_id, session_id, external_id = ( + _only_add_call(dummy) + ) assert role == "system" assert data_type == "text/markdown" assert passed_embedding == [0.5, 0.6] assert bot_id is None assert session_id is None + assert external_id is None def test_context_add_with_bot_id(): @@ -226,12 +304,14 @@ def test_context_add_with_bot_id(): ctx.add("user", "hello", bot_id="support_bot") - assert len(dummy.add_calls) == 1 - role, content, data_type, passed_embedding, bot_id, session_id = dummy.add_calls[0] + role, content, data_type, passed_embedding, bot_id, session_id, external_id = ( + _only_add_call(dummy) + ) assert role == "user" assert content == "hello" assert bot_id == "support_bot" assert session_id is None + assert external_id is None def test_context_add_with_session_id(): @@ -241,12 +321,14 @@ def test_context_add_with_session_id(): ctx.add("user", "hello", session_id="user_123") - assert len(dummy.add_calls) == 1 - role, content, data_type, passed_embedding, bot_id, session_id = dummy.add_calls[0] + role, content, data_type, passed_embedding, bot_id, session_id, external_id = ( + _only_add_call(dummy) + ) assert role == "user" assert content == "hello" assert bot_id is None assert session_id == "user_123" + assert external_id is None def test_context_add_with_agent_and_session_id(): @@ -256,11 +338,13 @@ def test_context_add_with_agent_and_session_id(): ctx.add("user", "hello", bot_id="sales_bot", session_id="conv_456") - assert len(dummy.add_calls) == 1 - role, content, data_type, passed_embedding, bot_id, session_id = dummy.add_calls[0] + role, content, data_type, passed_embedding, bot_id, session_id, external_id = ( + _only_add_call(dummy) + ) assert role == "user" assert bot_id == "sales_bot" assert session_id == "conv_456" + assert external_id is None def test_context_add_with_all_options(): @@ -269,20 +353,30 @@ def test_context_add_with_all_options(): ctx._inner = dummy # type: ignore[attr-defined] embedding = [0.1, 0.2] - ctx.add("user", "hello", embedding=embedding, bot_id="bot", session_id="sess") + ctx.add( + "user", + "hello", + embedding=embedding, + bot_id="bot", + session_id="sess", + external_id="doc-1#chunk-1", + ) - assert len(dummy.add_calls) == 1 - role, content, data_type, passed_embedding, bot_id, session_id = dummy.add_calls[0] + role, content, data_type, passed_embedding, bot_id, session_id, external_id = ( + _only_add_call(dummy) + ) assert role == "user" assert passed_embedding == [0.1, 0.2] assert bot_id == "bot" assert session_id == "sess" + assert external_id == "doc-1#chunk-1" def test_normalize_record_with_agent_and_session_id(): result = _normalize_record( { "id": "rec-1", + "external_id": "source-1", "created_at": "2024-01-01T00:00:00Z", "content_type": "text/plain", "text_payload": "hello", @@ -297,3 +391,4 @@ def test_normalize_record_with_agent_and_session_id(): ) assert result["bot_id"] == "support_bot" assert result["session_id"] == "user_88" + assert result["external_id"] == "source-1" From 3e767e2f828dbb094ea3a3f60a7d9372047fa9a8 Mon Sep 17 00:00:00 2001 From: Allen Cheng Date: Mon, 8 Jun 2026 21:19:07 -0700 Subject: [PATCH 2/4] style: apply rustfmt to external id changes --- crates/lance-context-core/src/store.rs | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/crates/lance-context-core/src/store.rs b/crates/lance-context-core/src/store.rs index d13640d..cbd2b23 100644 --- a/crates/lance-context-core/src/store.rs +++ b/crates/lance-context-core/src/store.rs @@ -676,9 +676,7 @@ impl ContextStore { Field::new("binary_payload", DataType::LargeBinary, true) }; - let mut fields = vec![ - Field::new("id", DataType::Utf8, false).with_metadata(id_metadata), - ]; + let mut fields = vec![Field::new("id", DataType::Utf8, false).with_metadata(id_metadata)]; if include_external_id { fields.push(Field::new("external_id", DataType::Utf8, true)); } @@ -770,11 +768,7 @@ impl ContextStore { } fn records_to_batch(&self, entries: &[ContextRecord]) -> LanceResult { - let include_external_id = self - .dataset - .schema() - .field_with_name("external_id") - .is_ok(); + let include_external_id = self.dataset.schema().field_with_name("external_id").is_ok(); if !include_external_id && entries.iter().any(|entry| entry.external_id.is_some()) { return Err(ArrowError::InvalidArgumentError( "external_id requires a context dataset created with external_id support" From cbab85f76c7680bf7c7b9daa38234ae00fd98750 Mon Sep 17 00:00:00 2001 From: Allen Cheng Date: Mon, 8 Jun 2026 21:39:02 -0700 Subject: [PATCH 3/4] fix: use Lance schema field_path for external id check --- crates/lance-context-core/src/store.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/lance-context-core/src/store.rs b/crates/lance-context-core/src/store.rs index cbd2b23..c90c823 100644 --- a/crates/lance-context-core/src/store.rs +++ b/crates/lance-context-core/src/store.rs @@ -768,7 +768,7 @@ impl ContextStore { } fn records_to_batch(&self, entries: &[ContextRecord]) -> LanceResult { - let include_external_id = self.dataset.schema().field_with_name("external_id").is_ok(); + let include_external_id = self.dataset.schema().field_path("external_id").is_ok(); if !include_external_id && entries.iter().any(|entry| entry.external_id.is_some()) { return Err(ArrowError::InvalidArgumentError( "external_id requires a context dataset created with external_id support" From 239f30af0fe9f1d36977abfabe87a5787abf5f1a Mon Sep 17 00:00:00 2001 From: Allen Cheng Date: Mon, 8 Jun 2026 21:53:26 -0700 Subject: [PATCH 4/4] fix: detect external id schema field --- crates/lance-context-core/src/store.rs | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/crates/lance-context-core/src/store.rs b/crates/lance-context-core/src/store.rs index c90c823..5f1f6a1 100644 --- a/crates/lance-context-core/src/store.rs +++ b/crates/lance-context-core/src/store.rs @@ -768,7 +768,12 @@ impl ContextStore { } fn records_to_batch(&self, entries: &[ContextRecord]) -> LanceResult { - let include_external_id = self.dataset.schema().field_path("external_id").is_ok(); + let include_external_id = self + .dataset + .schema() + .field_paths() + .iter() + .any(|path| path == "external_id"); if !include_external_id && entries.iter().any(|entry| entry.external_id.is_some()) { return Err(ArrowError::InvalidArgumentError( "external_id requires a context dataset created with external_id support"