diff --git a/crates/lance-context-core/src/api_impl.rs b/crates/lance-context-core/src/api_impl.rs index d8cc71e5..ca3a6de6 100644 --- a/crates/lance-context-core/src/api_impl.rs +++ b/crates/lance-context-core/src/api_impl.rs @@ -515,48 +515,56 @@ impl RolloutStoreApi for RolloutStore { } pub fn rollout_record_from_add_request(r: &AddRolloutRequest) -> RolloutRecord { - RolloutRecord { - id: r.id.clone(), - rollout_id: r.rollout_id.clone(), - problem_id: r.problem_id.clone().unwrap_or_else(|| r.rollout_id.clone()), - dataset: r.dataset.clone(), - sequence_order: r.sequence_order, - role: r.role.clone(), - created_at: r.created_at.unwrap_or_else(Utc::now), - content: r.content.clone(), - content_type: r.content_type.clone(), - model_input_string: r.model_input_string.clone(), - model_output_string: r.model_output_string.clone(), - rationale: r.rationale.clone(), - problem_text: r.problem_text.clone(), - user_metadata: r.user_metadata.clone(), - input_tokens: r.input_tokens.clone(), - output_tokens: r.output_tokens.clone(), - num_input_tokens: r.num_input_tokens, - num_output_tokens: r.num_output_tokens, - output_logprobs: r.output_logprobs.clone(), - input_logprobs: r.input_logprobs.clone(), - ref_logprobs: r.ref_logprobs.clone(), - loss_mask: r.loss_mask.clone(), - advantage: r.advantage, - reward: r.reward, - raw_reward: r.raw_reward, - grader_id: r.grader_id.clone(), - score: r.score, - include_in_training: r.include_in_training, - exclude_reason: r.exclude_reason.clone(), - policy_version: r.policy_version.clone(), - relationships: r - .relationships - .iter() - .cloned() - .map(dto_to_relationship) - .collect(), - binary_payload: r.binary_payload.clone(), - payload_size: r.payload_size, - payload_checksum: r.payload_checksum.clone(), - artifact_type: r.artifact_type.clone(), - metadata: r.metadata.clone(), + r.clone().into() +} + +/// Move request buffers into the core record. HTTP ingress already owns these +/// allocations; borrowing the request would clone large blobs and token arrays. +impl From for RolloutRecord { + fn from(r: AddRolloutRequest) -> Self { + let problem_id = r.problem_id.unwrap_or_else(|| r.rollout_id.clone()); + Self { + id: r.id, + rollout_id: r.rollout_id, + problem_id, + dataset: r.dataset, + sequence_order: r.sequence_order, + role: r.role, + created_at: r.created_at.unwrap_or_else(Utc::now), + content: r.content, + content_type: r.content_type, + model_input_string: r.model_input_string, + model_output_string: r.model_output_string, + rationale: r.rationale, + problem_text: r.problem_text, + user_metadata: r.user_metadata, + input_tokens: r.input_tokens, + output_tokens: r.output_tokens, + num_input_tokens: r.num_input_tokens, + num_output_tokens: r.num_output_tokens, + output_logprobs: r.output_logprobs, + input_logprobs: r.input_logprobs, + ref_logprobs: r.ref_logprobs, + loss_mask: r.loss_mask, + advantage: r.advantage, + reward: r.reward, + raw_reward: r.raw_reward, + grader_id: r.grader_id, + score: r.score, + include_in_training: r.include_in_training, + exclude_reason: r.exclude_reason, + policy_version: r.policy_version, + relationships: r + .relationships + .into_iter() + .map(dto_to_relationship) + .collect(), + binary_payload: r.binary_payload, + payload_size: r.payload_size, + payload_checksum: r.payload_checksum, + artifact_type: r.artifact_type, + metadata: r.metadata, + } } } @@ -1139,3 +1147,6 @@ fn failure_to_dto(failure: &DatagenFailure) -> DatagenFailureDto { traceback: failure.error.traceback.clone(), } } + +#[cfg(test)] +mod tests; diff --git a/crates/lance-context-core/src/api_impl/tests.rs b/crates/lance-context-core/src/api_impl/tests.rs new file mode 100644 index 00000000..9e6dfe2a --- /dev/null +++ b/crates/lance-context-core/src/api_impl/tests.rs @@ -0,0 +1,54 @@ +use super::*; + +#[test] +fn owned_rollout_conversion_moves_buffers_and_preserves_every_field() { + let mut request: AddRolloutRequest = serde_json::from_value(serde_json::json!({ + "id":"row-1", "rollout_id":"session-1", "problem_id":"problem-1", + "dataset":"synthetic", "sequence_order":7, "role":"artifact", + "created_at":"2026-01-01T00:00:00Z", "content":"text", + "content_type":"application/octet-stream", "model_input_string":"input", + "model_output_string":"output", "rationale":"reason", "problem_text":"problem", + "user_metadata":"user", "input_tokens":[1,2], "output_tokens":[3,4], + "num_input_tokens":2, "num_output_tokens":2, + "output_logprobs":[-0.5,-0.25], "input_logprobs":[-1.0,-0.5], + "ref_logprobs":[-2.0,-1.0], "loss_mask":[1,0], + "advantage":0.5, "reward":1.0, "raw_reward":2.0, "grader_id":"grader", + "score":0.25, "include_in_training":false, "exclude_reason":"test", + "policy_version":"v1", "relationships":[{"target_id":"parent","relation":"parent","weight":0.5}], + "payload_size":1048576, "payload_checksum":"synthetic-checksum", + "artifact_type":"test", "metadata":{"nested":{"key":"value"}} + })).unwrap(); + request.binary_payload = Some(vec![91; 1024 * 1024]); + let expected = serde_json::to_value(&request).unwrap(); + let blob = request.binary_payload.as_ref().unwrap().as_ptr(); + let tokens = request.input_tokens.as_ref().unwrap().as_ptr(); + let probabilities = request.output_logprobs.as_ref().unwrap().as_ptr(); + let text = request.model_input_string.as_ref().unwrap().as_ptr(); + let record = RolloutRecord::from(request); + assert_eq!(record.binary_payload.as_ref().unwrap().as_ptr(), blob); + assert_eq!(record.input_tokens.as_ref().unwrap().as_ptr(), tokens); + assert_eq!( + record.output_logprobs.as_ref().unwrap().as_ptr(), + probabilities + ); + assert_eq!(record.model_input_string.as_ref().unwrap().as_ptr(), text); + assert_eq!( + serde_json::to_value(rollout_record_to_dto(record)).unwrap(), + expected + ); +} + +#[test] +fn owned_rollout_conversion_preserves_defaults_and_explicit_empty_problem() { + for (problem, expected) in [(None, "session"), (Some(""), ""), (Some("p"), "p")] { + let before = Utc::now(); + let record = RolloutRecord::from(AddRolloutRequest { + rollout_id: "session".to_owned(), + problem_id: problem.map(str::to_owned), + ..Default::default() + }); + assert_eq!(record.problem_id, expected); + assert!(record.created_at >= before && record.created_at <= Utc::now()); + assert!(record.binary_payload.is_none()); + } +} diff --git a/crates/lance-context-server/src/routes/rollouts.rs b/crates/lance-context-server/src/routes/rollouts.rs index 2b39f7ed..5e3684aa 100644 --- a/crates/lance-context-server/src/routes/rollouts.rs +++ b/crates/lance-context-server/src/routes/rollouts.rs @@ -11,7 +11,9 @@ use lance_context_api::{ CompactResponse, CompactStatsResponse, CreateRolloutStoreRequest, GetRolloutResponse, ListRolloutStoresResponse, ListRolloutsResponse, RolloutStoreInfo, VersionResponse, }; -use lance_context_core::{rollout_record_from_add_request, rollout_record_to_dto}; +#[cfg(test)] +use lance_context_core::rollout_record_from_add_request; +use lance_context_core::rollout_record_to_dto; use lance_context_core::{ CompactionConfig, RolloutFilters, RolloutRecord, RolloutStore, RolloutStoreOptions, }; @@ -356,10 +358,7 @@ pub async fn add_rollouts( let store_lock = state.get_or_open_rollout_store(&name).await?; let ids: Vec = records.iter().map(|r| r.id.clone()).collect(); - let core_records: Vec = records - .iter() - .map(rollout_record_from_add_request) - .collect(); + let core_records: Vec = records.into_iter().map(RolloutRecord::from).collect(); let count = core_records.len(); // A read lock: `add` is `&self` and MemWAL appends are internally