use crate::client::QdrantHandle;
use crate::embedder::Embedder;
use crate::error::store_err;
use async_trait::async_trait;
use klieo_core::error::MemoryError;
use klieo_core::ids::FactId;
use klieo_core::memory::{Fact, LongTermMemory, Scope};
use klieo_memory_graph::FilterableLongTermMemory;
use qdrant_client::qdrant::{
Condition, DeletePointsBuilder, Filter, PointStruct, ScoredPoint, SearchPointsBuilder,
UpsertPointsBuilder, Value,
};
use std::collections::HashMap;
use std::sync::Arc;
fn is_missing_collection<E: std::fmt::Display>(error: &E) -> bool {
let message = error.to_string().to_ascii_lowercase();
let mentions_collection = message.contains("collection");
let missing_shape = message.contains("doesn't exist")
|| message.contains("does not exist")
|| message.contains("not found")
|| message.contains("not_found");
mentions_collection && missing_shape
}
fn decode_point_to_fact(point: ScoredPoint) -> Fact {
let text = point
.payload
.get("text")
.and_then(|v| v.as_str())
.map(|s| s.to_string())
.unwrap_or_default();
let metadata = point
.payload
.get("metadata")
.and_then(|v| v.as_str())
.map(|s| serde_json::from_str::<serde_json::Value>(s).unwrap_or(serde_json::Value::Null))
.unwrap_or(serde_json::Value::Null);
Fact { text, metadata }
}
pub struct QdrantLongTerm {
handle: QdrantHandle,
embedder: Arc<dyn Embedder>,
embedder_id: String,
}
impl QdrantLongTerm {
pub(crate) fn new(
handle: QdrantHandle,
embedder: Arc<dyn Embedder>,
embedder_id: String,
) -> Self {
Self {
handle,
embedder,
embedder_id,
}
}
fn collection_name(&self) -> String {
format!("{}_v1", self.handle.collection_prefix)
}
fn payload_for(scope: &Scope, fact: &Fact, fact_id: &str) -> HashMap<String, Value> {
let (kind, value) = match scope {
Scope::Workspace(s) => ("workspace", s.clone()),
Scope::Agent(s) => ("agent", s.clone()),
Scope::Global => ("global", String::new()),
};
let mut p = HashMap::new();
p.insert("fact_id".into(), Value::from(fact_id.to_string()));
p.insert("text".into(), Value::from(fact.text.clone()));
p.insert("scope_kind".into(), Value::from(kind));
p.insert("scope_value".into(), Value::from(value));
let metadata_str =
serde_json::to_string(&fact.metadata).unwrap_or_else(|_| "null".to_string());
p.insert("metadata".into(), Value::from(metadata_str));
p
}
fn scope_filter_conditions(scope: &Scope) -> Vec<Condition> {
let (kind, value) = match scope {
Scope::Workspace(s) => ("workspace", s.clone()),
Scope::Agent(s) => ("agent", s.clone()),
Scope::Global => ("global", String::new()),
};
vec![
Condition::matches("scope_kind", kind.to_string()),
Condition::matches("scope_value", value),
]
}
fn scope_filter(scope: &Scope) -> Filter {
Filter::must(Self::scope_filter_conditions(scope))
}
}
#[async_trait]
impl LongTermMemory for QdrantLongTerm {
async fn remember(&self, scope: Scope, fact: Fact) -> Result<FactId, MemoryError> {
let collection = self.collection_name();
let dim = self.embedder.dimension() as u64;
self.handle.ensure_collection(&collection, dim).await?;
let vectors = self
.embedder
.embed(std::slice::from_ref(&fact.text))
.await?;
let vector = vectors
.into_iter()
.next()
.ok_or_else(|| MemoryError::Embedding("embedder returned empty vec".into()))?;
if vector.len() != dim as usize {
return Err(MemoryError::Embedding(format!(
"embedder produced {}-dim vector, expected {}",
vector.len(),
dim
)));
}
let id_string = uuid::Uuid::new_v4().to_string();
let payload = Self::payload_for(&scope, &fact, &id_string);
let point = PointStruct::new(id_string.clone(), vector, payload);
self.handle
.client
.upsert_points(UpsertPointsBuilder::new(&collection, vec![point]).wait(true))
.await
.map_err(store_err)?;
Ok(FactId::new(id_string))
}
async fn recall(&self, scope: Scope, query: &str, k: usize) -> Result<Vec<Fact>, MemoryError> {
if k == 0 {
return Ok(Vec::new());
}
let collection = self.collection_name();
let dim = self.embedder.dimension() as u64;
let query_vec = self
.embedder
.embed(&[query.to_string()])
.await?
.into_iter()
.next()
.ok_or_else(|| MemoryError::Embedding("embedder returned empty vec".into()))?;
if query_vec.len() != dim as usize {
return Err(MemoryError::Embedding(format!(
"embedder produced {}-dim vector, expected {}",
query_vec.len(),
dim
)));
}
let req = SearchPointsBuilder::new(&collection, query_vec, k as u64)
.filter(Self::scope_filter(&scope))
.with_payload(true);
let search_response = match self.handle.client.search_points(req).await {
Ok(r) => r,
Err(e) if is_missing_collection(&e) => return Ok(Vec::new()),
Err(e) => return Err(store_err(e)),
};
Ok(search_response
.result
.into_iter()
.map(decode_point_to_fact)
.collect())
}
async fn forget(&self, id: FactId) -> Result<(), MemoryError> {
let collection = self.collection_name();
let filter = Filter::must([Condition::matches("fact_id", id.to_string())]);
let res = self
.handle
.client
.delete_points(
DeletePointsBuilder::new(&collection)
.points(filter)
.wait(true),
)
.await;
match res {
Ok(_) => Ok(()),
Err(e) if is_missing_collection(&e) => Ok(()),
Err(e) => Err(store_err(e)),
}
}
}
#[async_trait]
impl FilterableLongTermMemory for QdrantLongTerm {
async fn recall_filtered(
&self,
scope: Scope,
query: &str,
k: usize,
candidate_ids: &[FactId],
) -> Result<Vec<Fact>, MemoryError> {
if k == 0 || candidate_ids.is_empty() {
return Ok(Vec::new());
}
let collection = self.collection_name();
let dim = self.embedder.dimension() as u64;
let query_vec = self
.embedder
.embed(&[query.to_string()])
.await?
.into_iter()
.next()
.ok_or_else(|| MemoryError::Embedding("embedder returned empty vec".into()))?;
if query_vec.len() != dim as usize {
return Err(MemoryError::Embedding(format!(
"embedder produced {}-dim vector, expected {}",
query_vec.len(),
dim
)));
}
let id_strings: Vec<String> = candidate_ids.iter().map(|f| f.to_string()).collect();
let mut conditions = Self::scope_filter_conditions(&scope);
conditions.push(Condition::matches("fact_id", id_strings));
let req = SearchPointsBuilder::new(&collection, query_vec, k as u64)
.filter(Filter::must(conditions))
.with_payload(true);
let response = match self.handle.client.search_points(req).await {
Ok(r) => r,
Err(e) if is_missing_collection(&e) => return Ok(Vec::new()),
Err(e) => return Err(store_err(e)),
};
Ok(response
.result
.into_iter()
.map(decode_point_to_fact)
.collect())
}
fn embedder_id(&self) -> &str {
&self.embedder_id
}
}
impl QdrantLongTerm {
pub async fn recall_filtered_checked(
&self,
caller_embedder_id: &str,
scope: Scope,
query: &str,
k: usize,
candidate_ids: &[FactId],
) -> Result<Vec<Fact>, MemoryError> {
if caller_embedder_id != self.embedder_id {
return Err(MemoryError::Embedding(format!(
"embedder mismatch: store has '{}', caller uses '{}'; \
re-index required before mixing embedder versions",
self.embedder_id, caller_embedder_id
)));
}
self.recall_filtered(scope, query, k, candidate_ids).await
}
}
#[cfg(test)]
mod tests {
use super::*;
#[allow(dead_code)]
fn _trait_object_safe(_: &dyn LongTermMemory) {}
#[test]
fn payload_carries_scope_text_and_fact_id() {
let scope = Scope::Workspace("ws".into());
let fact = Fact {
text: "hi".into(),
metadata: serde_json::Value::Null,
};
let p = QdrantLongTerm::payload_for(&scope, &fact, "fact-abc");
assert_eq!(p.get("text").unwrap().as_str().unwrap(), "hi");
assert_eq!(p.get("scope_kind").unwrap().as_str().unwrap(), "workspace");
assert_eq!(p.get("scope_value").unwrap().as_str().unwrap(), "ws");
assert_eq!(p.get("fact_id").unwrap().as_str().unwrap(), "fact-abc");
assert_eq!(p.get("metadata").unwrap().as_str().unwrap(), "null");
}
#[test]
fn payload_for_global_scope_uses_empty_string_value() {
let fact = Fact {
text: "g".into(),
metadata: serde_json::Value::Null,
};
let p = QdrantLongTerm::payload_for(&Scope::Global, &fact, "fact-1");
assert_eq!(p.get("scope_kind").unwrap().as_str().unwrap(), "global");
assert_eq!(p.get("scope_value").unwrap().as_str().unwrap(), "");
}
struct FakeErr(&'static str);
impl std::fmt::Display for FakeErr {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(self.0)
}
}
#[test]
fn is_missing_collection_matches_doesnt_exist() {
assert!(is_missing_collection(&FakeErr(
"Status { code: NotFound, message: \"Collection 'klieo_facts_v1' doesn't exist\" }"
)));
}
#[test]
fn is_missing_collection_matches_does_not_exist() {
assert!(is_missing_collection(&FakeErr(
"collection 'klieo_facts_v1' does not exist"
)));
}
#[test]
fn is_missing_collection_matches_not_found_case_insensitive() {
assert!(is_missing_collection(&FakeErr(
"Status { code: NOT_FOUND, message: \"Collection 'x' missing\" }"
)));
assert!(is_missing_collection(&FakeErr(
"collection resource Not Found"
)));
}
#[test]
fn is_missing_collection_rejects_unrelated_errors() {
assert!(!is_missing_collection(&FakeErr("dimension mismatch")));
assert!(!is_missing_collection(&FakeErr("connection refused")));
assert!(!is_missing_collection(&FakeErr(
"permission denied for collection"
)));
}
#[test]
fn is_missing_collection_rejects_vector_or_index_not_found() {
assert!(!is_missing_collection(&FakeErr(
"vector index not found for field bar"
)));
assert!(!is_missing_collection(&FakeErr("payload field not_found")));
assert!(!is_missing_collection(&FakeErr(
"Status { code: NOT_FOUND, message: \"vector named X doesn't exist\" }"
)));
}
}