use super::embed_store::{ChunkPlanner, expected_chunk_count};
use crate::memory_core::embed::Embedder;
use crate::memory_core::palace::Drawer;
use crate::memory_core::store::vector::ChunkShape;
use parking_lot::RwLock;
use std::collections::{HashMap, HashSet};
use uuid::Uuid;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ChunkCheck {
Ran,
Unavailable(String),
}
impl ChunkCheck {
pub fn unavailable_reason(&self) -> Option<&str> {
match self {
ChunkCheck::Ran => None,
ChunkCheck::Unavailable(reason) => Some(reason),
}
}
}
pub(super) fn short_chunk_ids(
embedder: Option<&dyn Embedder>,
drawers: &RwLock<Vec<Drawer>>,
with_vector: &HashSet<Uuid>,
shapes: anyhow::Result<HashMap<Uuid, ChunkShape>>,
) -> (Vec<Uuid>, ChunkCheck) {
if with_vector.is_empty() {
return (Vec::new(), ChunkCheck::Ran);
}
let Some(embedder) = embedder else {
return unavailable("no embedder has initialised in this process".to_string());
};
let planner = match ChunkPlanner::for_embedder(embedder) {
Ok(planner) => planner,
Err(e) => return unavailable(format!("{e:#}")),
};
let shapes = match shapes {
Ok(shapes) => shapes,
Err(e) => return unavailable(format!("scan vector keys: {e:#}")),
};
let shape_of = |id: &Uuid| shapes.get(id).copied().unwrap_or_default();
let mut short = Vec::new();
let mut long: Vec<(Uuid, String, ChunkShape)> = Vec::new();
for drawer in drawers
.read()
.iter()
.filter(|d| with_vector.contains(&d.id))
{
let shape = shape_of(&drawer.id);
let fits = planner
.as_ref()
.is_none_or(|p| p.fits_one_chunk(drawer.content()));
if fits {
if !shape.is_complete(1) {
short.push(drawer.id);
}
} else {
long.push((drawer.id, drawer.content().to_string(), shape));
}
}
for (id, text, shape) in long {
match expected_chunk_count(embedder, &text) {
Ok(expected) if shape.is_complete(expected) => {}
Ok(_) => short.push(id),
Err(e) => {
if first_sighting(id) {
tracing::warn!(drawer = %id, "#9275: drawer text will not split: {e:#}");
} else {
tracing::debug!(drawer = %id, "#9275: drawer text will not split: {e:#}");
}
short.push(id);
}
}
}
(short, ChunkCheck::Ran)
}
fn unavailable(reason: String) -> (Vec<Uuid>, ChunkCheck) {
(Vec::new(), ChunkCheck::Unavailable(reason))
}
pub(super) const SIGHTING_CAP: usize = 4096;
pub(super) struct Sightings {
seen: HashSet<Uuid>,
cap: usize,
}
impl Sightings {
pub(super) fn new(cap: usize) -> Self {
Self {
seen: HashSet::new(),
cap,
}
}
pub(super) fn first(&mut self, id: Uuid) -> bool {
if self.seen.len() >= self.cap {
return false;
}
self.seen.insert(id)
}
}
pub(super) fn first_sighting(id: Uuid) -> bool {
static SEEN: std::sync::OnceLock<parking_lot::Mutex<Sightings>> = std::sync::OnceLock::new();
SEEN.get_or_init(|| parking_lot::Mutex::new(Sightings::new(SIGHTING_CAP)))
.lock()
.first(id)
}