132bba8d8b
We believe this is causing the database to execute certain queries very slowly, and it's unclear whether we're actually benefitting from it. Release Notes: - N/A
87 lines
2.8 KiB
Rust
87 lines
2.8 KiB
Rust
use super::TestDb;
|
|
use crate::db::embedding;
|
|
use collections::HashMap;
|
|
use sea_orm::{ColumnTrait, EntityTrait, QueryFilter, sea_query::Expr};
|
|
use std::ops::Sub;
|
|
use time::{Duration, OffsetDateTime, PrimitiveDateTime};
|
|
|
|
// SQLite does not support array arguments, so we only test this against a real postgres instance
|
|
#[gpui::test]
|
|
async fn test_get_embeddings_postgres(cx: &mut gpui::TestAppContext) {
|
|
let test_db = TestDb::postgres(cx.executor().clone());
|
|
let db = test_db.db();
|
|
|
|
let provider = "test_model";
|
|
let digest1 = vec![1, 2, 3];
|
|
let digest2 = vec![4, 5, 6];
|
|
let embeddings = HashMap::from_iter([
|
|
(digest1.clone(), vec![0.1, 0.2, 0.3]),
|
|
(digest2.clone(), vec![0.4, 0.5, 0.6]),
|
|
]);
|
|
|
|
// Save embeddings
|
|
db.save_embeddings(provider, &embeddings).await.unwrap();
|
|
|
|
// Retrieve embeddings
|
|
let retrieved_embeddings = db
|
|
.get_embeddings(provider, &[digest1.clone(), digest2.clone()])
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(retrieved_embeddings.len(), 2);
|
|
assert!(retrieved_embeddings.contains_key(&digest1));
|
|
assert!(retrieved_embeddings.contains_key(&digest2));
|
|
|
|
// Check if the retrieved embeddings are correct
|
|
assert_eq!(retrieved_embeddings[&digest1], vec![0.1, 0.2, 0.3]);
|
|
assert_eq!(retrieved_embeddings[&digest2], vec![0.4, 0.5, 0.6]);
|
|
}
|
|
|
|
#[gpui::test]
|
|
async fn test_purge_old_embeddings(cx: &mut gpui::TestAppContext) {
|
|
let test_db = TestDb::postgres(cx.executor().clone());
|
|
let db = test_db.db();
|
|
|
|
let model = "test_model";
|
|
let digest = vec![7, 8, 9];
|
|
let embeddings = HashMap::from_iter([(digest.clone(), vec![0.7, 0.8, 0.9])]);
|
|
|
|
// Save old embeddings
|
|
db.save_embeddings(model, &embeddings).await.unwrap();
|
|
|
|
// Reach into the DB and change the retrieved at to be > 60 days
|
|
db.transaction(|tx| {
|
|
let digest = digest.clone();
|
|
async move {
|
|
let sixty_days_ago = OffsetDateTime::now_utc().sub(Duration::days(61));
|
|
let retrieved_at = PrimitiveDateTime::new(sixty_days_ago.date(), sixty_days_ago.time());
|
|
|
|
embedding::Entity::update_many()
|
|
.filter(
|
|
embedding::Column::Model
|
|
.eq(model)
|
|
.and(embedding::Column::Digest.eq(digest)),
|
|
)
|
|
.col_expr(embedding::Column::RetrievedAt, Expr::value(retrieved_at))
|
|
.exec(&*tx)
|
|
.await
|
|
.unwrap();
|
|
|
|
Ok(())
|
|
}
|
|
})
|
|
.await
|
|
.unwrap();
|
|
|
|
// Purge old embeddings
|
|
db.purge_old_embeddings().await.unwrap();
|
|
|
|
// Try to retrieve the purged embeddings
|
|
let retrieved_embeddings = db
|
|
.get_embeddings(model, std::slice::from_ref(&digest))
|
|
.await
|
|
.unwrap();
|
|
assert!(
|
|
retrieved_embeddings.is_empty(),
|
|
"Old embeddings should have been purged"
|
|
);
|
|
}
|