Fix the test and simplify types

This commit is contained in:
Mubelotix
2025-06-23 18:55:23 +02:00
parent 4925b30196
commit 2f82d94502
9 changed files with 87 additions and 51 deletions

View File

@ -29,7 +29,7 @@ pub struct ProcessingBatch {
pub uid: BatchId,
pub details: DetailsView,
pub stats: BatchStats,
pub embedder_stats: Option<Arc<EmbedderStats>>,
pub embedder_stats: Arc<EmbedderStats>,
pub statuses: HashSet<Status>,
pub kinds: HashSet<Kind>,
@ -47,11 +47,13 @@ impl ProcessingBatch {
let mut statuses = HashSet::default();
statuses.insert(Status::Processing);
println!("Processing batch created: {}", uid);
Self {
uid,
details: DetailsView::default(),
stats: BatchStats::default(),
embedder_stats: None,
embedder_stats: Default::default(),
statuses,
kinds: HashSet::default(),
@ -64,17 +66,6 @@ impl ProcessingBatch {
}
}
pub fn clone_embedder_stats(&mut self) -> Arc<EmbedderStats> {
match self.embedder_stats {
Some(ref stats) => stats.clone(),
None => {
let embedder_stats: Arc<EmbedderStats> = Default::default();
self.embedder_stats = Some(embedder_stats.clone());
embedder_stats
},
}
}
/// Update itself with the content of the task and update the batch id in the task.
pub fn processing<'a>(&mut self, tasks: impl IntoIterator<Item = &'a mut Task>) {
for task in tasks.into_iter() {
@ -113,11 +104,14 @@ impl ProcessingBatch {
}
pub fn reason(&mut self, reason: BatchStopReason) {
println!("batch stopped: {:?}", reason);
self.reason = reason;
}
/// Must be called once the batch has finished processing.
pub fn finished(&mut self) {
println!("Batch finished: {}", self.uid);
self.details = DetailsView::default();
self.stats = BatchStats::default();
self.finished_at = Some(OffsetDateTime::now_utc());
@ -132,6 +126,8 @@ impl ProcessingBatch {
/// Update the timestamp of the tasks and the inner structure of this structure.
pub fn update(&mut self, task: &mut Task) {
println!("Updating task: {} in batch: {}", task.uid, self.uid);
// We must re-set this value in case we're dealing with a task that has been added between
// the `processing` and `finished` state
// We must re-set this value in case we're dealing with a task that has been added between
@ -156,13 +152,13 @@ impl ProcessingBatch {
}
pub fn to_batch(&self) -> Batch {
println!("Converting to batch: {:?}", self.embedder_stats);
println!("Converting to batch: {:?} {:?}", self.uid, self.embedder_stats);
Batch {
uid: self.uid,
progress: None,
details: self.details.clone(),
stats: self.stats.clone(),
embedder_stats: self.embedder_stats.as_ref().map(|s| BatchEmbeddingStats::from(s.as_ref())),
embedder_stats: self.embedder_stats.as_ref().into(),
started_at: self.started_at,
finished_at: self.finished_at,
enqueued_at: self.enqueued_at,