diff --git a/crates/lance-context-api/src/lib.rs b/crates/lance-context-api/src/lib.rs index 78d13ea..d8e8881 100644 --- a/crates/lance-context-api/src/lib.rs +++ b/crates/lance-context-api/src/lib.rs @@ -137,6 +137,35 @@ pub trait RolloutStoreApi { filters: Option, ) -> impl Future>> + Send; + fn get_trajectory( + &self, + rollout_id: &str, + ) -> impl Future>> + Send + where + Self: Sync, + { + async move { + if rollout_id.is_empty() { + return Err(ContextError::InvalidRequest( + "rollout_id must not be empty".to_string(), + )); + } + let mut records = self + .list_filtered( + None, + None, + Some(serde_json::json!({"rollout_id": rollout_id})), + ) + .await?; + records.sort_by(|left, right| { + left.sequence_order + .cmp(&right.sequence_order) + .then_with(|| left.id.cmp(&right.id)) + }); + Ok(records) + } + } + fn get(&self, id: &str) -> impl Future>> + Send; diff --git a/crates/lance-context-core/src/api_impl.rs b/crates/lance-context-core/src/api_impl.rs index 6031569..9f26cb8 100644 --- a/crates/lance-context-core/src/api_impl.rs +++ b/crates/lance-context-core/src/api_impl.rs @@ -416,6 +416,13 @@ impl RolloutStoreApi for RolloutStore { Ok(records.into_iter().map(rollout_record_to_dto).collect()) } + async fn get_trajectory(&self, rollout_id: &str) -> ContextResult> { + let records = RolloutStore::get_trajectory(self, rollout_id) + .await + .map_err(to_ctx_err)?; + Ok(records.into_iter().map(rollout_record_to_dto).collect()) + } + async fn get(&self, id: &str) -> ContextResult> { let record = RolloutStore::get_by_id(self, id) .await diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index bae712d..838e655 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -965,6 +965,27 @@ impl RolloutStore { Ok(results) } + /// Return one complete trajectory in deterministic message order. + pub async fn get_trajectory(&self, rollout_id: &str) -> LanceResult> { + if rollout_id.is_empty() { + return Err(LanceError::from(ArrowError::InvalidArgumentError( + "rollout_id must not be empty".to_string(), + ))); + } + + let filters = RolloutFilters { + rollout_id: Some(rollout_id.to_string()), + ..Default::default() + }; + let mut records = self.list_with_filters(None, None, Some(&filters)).await?; + records.sort_by(|left, right| { + left.sequence_order + .cmp(&right.sequence_order) + .then_with(|| left.id.cmp(&right.id)) + }); + Ok(records) + } + /// Filter and page rollout rows in the LSM execution plan. /// /// Reads one row beyond the requested page to report `has_more`, avoiding @@ -2264,6 +2285,54 @@ mod tests { }); } + #[test] + fn trajectory_rows_are_filtered_and_sorted_across_fragments() { + 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 = RolloutStore::open_with_options( + &uri, + RolloutStoreOptions { + shard_id: Some("trajectory-test".to_string()), + merge_after_generations: Some(1), + ..Default::default() + }, + ) + .await + .unwrap(); + + let mut later = assistant_record("row-a"); + later.rollout_id = "target".to_string(); + later.sequence_order = 2; + let mut unrelated = assistant_record("other"); + unrelated.rollout_id = "other".to_string(); + store.add(&[later, unrelated]).await.unwrap(); + + let mut earlier = assistant_record("row-b"); + earlier.rollout_id = "target".to_string(); + earlier.sequence_order = 0; + store.add(&[earlier]).await.unwrap(); + + assert_eq!(store.observe().await.unwrap().fragment_count, 2); + let rows = store.get_trajectory("target").await.unwrap(); + let ids: Vec<&str> = rows.iter().map(|row| row.id.as_str()).collect(); + assert_eq!(ids, vec!["row-b", "row-a"]); + }); + } + + #[test] + fn trajectory_read_rejects_an_empty_rollout_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 store = RolloutStore::open(&uri).await.unwrap(); + let err = store.get_trajectory("").await.unwrap_err(); + assert!(err.to_string().contains("rollout_id must not be empty")); + }); + } + /// Read the number of un-merged flushed generations recorded for a store's /// own write shard. Used by merge tests to assert the manifest drains. async fn flushed_generation_count(store: &RolloutStore) -> usize { diff --git a/python/python/lance_context/api.py b/python/python/lance_context/api.py index fe2d2f6..a46cf66 100644 --- a/python/python/lance_context/api.py +++ b/python/python/lance_context/api.py @@ -2594,6 +2594,11 @@ def list( raw = self._sync.list(limit, offset, _json_dumps(filters, "filters")) return [_rollout_record_from_json(r) for r in json.loads(raw)] + def get_trajectory(self, rollout_id: str) -> list[dict[str, Any]]: + """Return one complete trajectory ordered by sequence and row id.""" + raw = self._sync.get_trajectory(rollout_id) + return [_rollout_record_from_json(r) for r in json.loads(raw)] + def get(self, id: str) -> dict[str, Any] | None: """Fetch a single rollout row by id, or ``None``.""" raw = self._sync.get(id) @@ -2681,6 +2686,14 @@ async def list( ) return [_rollout_record_from_json(r) for r in json.loads(raw)] + async def get_trajectory(self, rollout_id: str) -> list[dict[str, Any]]: + """Return one complete trajectory ordered by sequence and row id.""" + loop = asyncio.get_running_loop() + raw = await loop.run_in_executor( + None, lambda: self._sync.get_trajectory(rollout_id) + ) + return [_rollout_record_from_json(r) for r in json.loads(raw)] + async def get(self, id: str) -> dict[str, Any] | None: """Fetch a single rollout row by id, or ``None``.""" loop = asyncio.get_running_loop() diff --git a/python/src/lib.rs b/python/src/lib.rs index 9fa075d..6cbed07 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -2498,6 +2498,14 @@ impl RolloutStore { rollout_records_to_json(&records) } + /// Return one complete trajectory in deterministic message order. + fn get_trajectory(&self, py: Python<'_>, rollout_id: &str) -> PyResult { + let records = py + .allow_threads(|| self.runtime.block_on(self.store.get_trajectory(rollout_id))) + .map_err(to_py_err)?; + rollout_records_to_json(&records) + } + /// Fetch a single rollout row by id, or `None`. Returns a JSON object /// string (artifact bytes projected out). fn get(&self, py: Python<'_>, id: &str) -> PyResult> { diff --git a/python/tests/test_rollout.py b/python/tests/test_rollout.py index 175a31f..49afeb4 100644 --- a/python/tests/test_rollout.py +++ b/python/tests/test_rollout.py @@ -179,3 +179,17 @@ def test_list_rejects_invalid_filter_fields(store_uri): store = RolloutStore.open(store_uri) with pytest.raises(RuntimeError, match="unsupported rollout filter"): store.list(filters={"reward": 1.0}) + + +def test_get_trajectory_orders_rows(store_uri): + store = RolloutStore.open(store_uri) + store.add( + [ + {"id": "row-a", "rollout_id": "target", "sequence_order": 2}, + {"id": "other", "rollout_id": "other", "sequence_order": 0}, + {"id": "row-b", "rollout_id": "target", "sequence_order": 0}, + ] + ) + + rows = store.get_trajectory("target") + assert [row["id"] for row in rows] == ["row-b", "row-a"] diff --git a/python/tests/test_rollout_remote.py b/python/tests/test_rollout_remote.py index 8ee8576..9d0b620 100644 --- a/python/tests/test_rollout_remote.py +++ b/python/tests/test_rollout_remote.py @@ -159,3 +159,20 @@ async def run(): assert [row["id"] for row in rows] == ["row-7"] asyncio.run(run()) + + +def test_remote_get_trajectory_orders_rows(server): + async def run(): + store = await AsyncRolloutStore.connect_or_create(server, "rl-trajectory") + await store.add( + [ + {"id": "row-a", "rollout_id": "target", "sequence_order": 2}, + {"id": "other", "rollout_id": "other", "sequence_order": 0}, + {"id": "row-b", "rollout_id": "target", "sequence_order": 0}, + ] + ) + + rows = await store.get_trajectory("target") + assert [row["id"] for row in rows] == ["row-b", "row-a"] + + asyncio.run(run())