Implement semantic embeddings and duplicate review
This commit is contained in:
@@ -38,6 +38,9 @@ mlua = { workspace = true }
|
||||
url = { workspace = true }
|
||||
base64 = { workspace = true }
|
||||
dirs = { workspace = true }
|
||||
fastembed = { workspace = true }
|
||||
usearch = { workspace = true }
|
||||
ort = { workspace = true }
|
||||
|
||||
[dev-dependencies]
|
||||
tempfile = "3"
|
||||
|
||||
@@ -258,9 +258,12 @@ CREATE TABLE IF NOT EXISTS embedding_keys (
|
||||
post_id TEXT NOT NULL,
|
||||
project_id TEXT NOT NULL,
|
||||
content_hash TEXT NOT NULL,
|
||||
vector TEXT NOT NULL
|
||||
vector BLOB NOT NULL
|
||||
);
|
||||
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS embedding_keys_project_post_idx
|
||||
ON embedding_keys(project_id, post_id);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS dismissed_duplicate_pairs (
|
||||
id TEXT NOT NULL PRIMARY KEY,
|
||||
project_id TEXT NOT NULL REFERENCES projects(id),
|
||||
|
||||
@@ -0,0 +1,15 @@
|
||||
-- This file should undo anything in `up.sql`
|
||||
CREATE TABLE embedding_keys_old (
|
||||
label INTEGER NOT NULL PRIMARY KEY,
|
||||
post_id TEXT NOT NULL,
|
||||
project_id TEXT NOT NULL,
|
||||
content_hash TEXT NOT NULL,
|
||||
vector TEXT NOT NULL
|
||||
);
|
||||
|
||||
INSERT INTO embedding_keys_old (label, post_id, project_id, content_hash, vector)
|
||||
SELECT label, post_id, project_id, content_hash, CAST(vector AS TEXT)
|
||||
FROM embedding_keys;
|
||||
|
||||
DROP TABLE embedding_keys;
|
||||
ALTER TABLE embedding_keys_old RENAME TO embedding_keys;
|
||||
@@ -0,0 +1,18 @@
|
||||
-- Your SQL goes here
|
||||
CREATE TABLE embedding_keys_new (
|
||||
label INTEGER NOT NULL PRIMARY KEY,
|
||||
post_id TEXT NOT NULL,
|
||||
project_id TEXT NOT NULL,
|
||||
content_hash TEXT NOT NULL,
|
||||
vector BLOB NOT NULL
|
||||
);
|
||||
|
||||
INSERT INTO embedding_keys_new (label, post_id, project_id, content_hash, vector)
|
||||
SELECT label, post_id, project_id, content_hash, CAST(vector AS BLOB)
|
||||
FROM embedding_keys;
|
||||
|
||||
DROP TABLE embedding_keys;
|
||||
ALTER TABLE embedding_keys_new RENAME TO embedding_keys;
|
||||
|
||||
CREATE UNIQUE INDEX embedding_keys_project_post_idx
|
||||
ON embedding_keys(project_id, post_id);
|
||||
@@ -36,7 +36,7 @@ mod tests {
|
||||
let applied = db
|
||||
.conn()
|
||||
.with_migrations(|conn| conn.applied_migrations().unwrap().len());
|
||||
assert_eq!(applied, 6);
|
||||
assert_eq!(applied, 7);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
283
crates/bds-core/src/db/queries/embedding.rs
Normal file
283
crates/bds-core/src/db/queries/embedding.rs
Normal file
@@ -0,0 +1,283 @@
|
||||
use diesel::prelude::*;
|
||||
|
||||
use crate::db::DbConnection;
|
||||
use crate::db::schema::{dismissed_duplicate_pairs, embedding_keys};
|
||||
use crate::model::{DismissedDuplicatePair, EmbeddingKey};
|
||||
|
||||
pub fn get_key_for_post(
|
||||
conn: &DbConnection,
|
||||
project_id: &str,
|
||||
post_id: &str,
|
||||
) -> QueryResult<Option<EmbeddingKey>> {
|
||||
conn.with(|c| {
|
||||
embedding_keys::table
|
||||
.filter(embedding_keys::project_id.eq(project_id))
|
||||
.filter(embedding_keys::post_id.eq(post_id))
|
||||
.select(EmbeddingKey::as_select())
|
||||
.first(c)
|
||||
.optional()
|
||||
})
|
||||
}
|
||||
|
||||
pub fn list_keys(conn: &DbConnection, project_id: &str) -> QueryResult<Vec<EmbeddingKey>> {
|
||||
conn.with(|c| {
|
||||
embedding_keys::table
|
||||
.filter(embedding_keys::project_id.eq(project_id))
|
||||
.order(embedding_keys::label.asc())
|
||||
.select(EmbeddingKey::as_select())
|
||||
.load(c)
|
||||
})
|
||||
}
|
||||
|
||||
pub fn max_label(conn: &DbConnection) -> QueryResult<i64> {
|
||||
conn.with(|c| {
|
||||
embedding_keys::table
|
||||
.select(embedding_keys::label)
|
||||
.order(embedding_keys::label.desc())
|
||||
.first(c)
|
||||
.optional()
|
||||
.map(|label| label.unwrap_or(0))
|
||||
})
|
||||
}
|
||||
|
||||
pub fn upsert_key(conn: &DbConnection, key: &EmbeddingKey) -> QueryResult<()> {
|
||||
conn.with(|c| {
|
||||
diesel::insert_into(embedding_keys::table)
|
||||
.values(key)
|
||||
.on_conflict((embedding_keys::project_id, embedding_keys::post_id))
|
||||
.do_update()
|
||||
.set((
|
||||
embedding_keys::content_hash.eq(&key.content_hash),
|
||||
embedding_keys::vector.eq(&key.vector),
|
||||
))
|
||||
.execute(c)
|
||||
.map(|_| ())
|
||||
})
|
||||
}
|
||||
|
||||
pub fn delete_key_for_post(
|
||||
conn: &DbConnection,
|
||||
project_id: &str,
|
||||
post_id: &str,
|
||||
) -> QueryResult<()> {
|
||||
conn.with(|c| {
|
||||
diesel::delete(
|
||||
embedding_keys::table
|
||||
.filter(embedding_keys::project_id.eq(project_id))
|
||||
.filter(embedding_keys::post_id.eq(post_id)),
|
||||
)
|
||||
.execute(c)
|
||||
.map(|_| ())
|
||||
})
|
||||
}
|
||||
|
||||
pub fn delete_stale_keys(
|
||||
conn: &DbConnection,
|
||||
project_id: &str,
|
||||
live_post_ids: &[String],
|
||||
) -> QueryResult<usize> {
|
||||
conn.with(|c| {
|
||||
let query = embedding_keys::table.filter(embedding_keys::project_id.eq(project_id));
|
||||
if live_post_ids.is_empty() {
|
||||
diesel::delete(query).execute(c)
|
||||
} else {
|
||||
diesel::delete(query.filter(embedding_keys::post_id.ne_all(live_post_ids))).execute(c)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
pub fn insert_dismissed_pair(
|
||||
conn: &DbConnection,
|
||||
pair: &DismissedDuplicatePair,
|
||||
) -> QueryResult<()> {
|
||||
conn.with(|c| {
|
||||
diesel::insert_into(dismissed_duplicate_pairs::table)
|
||||
.values(pair)
|
||||
.on_conflict((
|
||||
dismissed_duplicate_pairs::project_id,
|
||||
dismissed_duplicate_pairs::post_id_a,
|
||||
dismissed_duplicate_pairs::post_id_b,
|
||||
))
|
||||
.do_nothing()
|
||||
.execute(c)
|
||||
.map(|_| ())
|
||||
})
|
||||
}
|
||||
|
||||
pub fn insert_dismissed_pairs(
|
||||
conn: &DbConnection,
|
||||
pairs: &[DismissedDuplicatePair],
|
||||
) -> QueryResult<usize> {
|
||||
if pairs.is_empty() {
|
||||
return Ok(0);
|
||||
}
|
||||
conn.with(|c| {
|
||||
diesel::insert_into(dismissed_duplicate_pairs::table)
|
||||
.values(pairs)
|
||||
.on_conflict((
|
||||
dismissed_duplicate_pairs::project_id,
|
||||
dismissed_duplicate_pairs::post_id_a,
|
||||
dismissed_duplicate_pairs::post_id_b,
|
||||
))
|
||||
.do_nothing()
|
||||
.execute(c)
|
||||
})
|
||||
}
|
||||
|
||||
pub fn list_dismissed_pairs(
|
||||
conn: &DbConnection,
|
||||
project_id: &str,
|
||||
) -> QueryResult<Vec<DismissedDuplicatePair>> {
|
||||
conn.with(|c| {
|
||||
dismissed_duplicate_pairs::table
|
||||
.filter(dismissed_duplicate_pairs::project_id.eq(project_id))
|
||||
.select(DismissedDuplicatePair::as_select())
|
||||
.load(c)
|
||||
})
|
||||
}
|
||||
|
||||
pub fn delete_orphan_dismissals(
|
||||
conn: &DbConnection,
|
||||
project_id: &str,
|
||||
live_post_ids: &[String],
|
||||
) -> QueryResult<usize> {
|
||||
conn.with(|c| {
|
||||
let query = dismissed_duplicate_pairs::table
|
||||
.filter(dismissed_duplicate_pairs::project_id.eq(project_id));
|
||||
if live_post_ids.is_empty() {
|
||||
diesel::delete(query).execute(c)
|
||||
} else {
|
||||
diesel::delete(
|
||||
query.filter(
|
||||
dismissed_duplicate_pairs::post_id_a
|
||||
.ne_all(live_post_ids)
|
||||
.or(dismissed_duplicate_pairs::post_id_b.ne_all(live_post_ids)),
|
||||
),
|
||||
)
|
||||
.execute(c)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
pub fn delete_dismissals_for_post(
|
||||
conn: &DbConnection,
|
||||
project_id: &str,
|
||||
post_id: &str,
|
||||
) -> QueryResult<usize> {
|
||||
conn.with(|c| {
|
||||
diesel::delete(
|
||||
dismissed_duplicate_pairs::table
|
||||
.filter(dismissed_duplicate_pairs::project_id.eq(project_id))
|
||||
.filter(
|
||||
dismissed_duplicate_pairs::post_id_a
|
||||
.eq(post_id)
|
||||
.or(dismissed_duplicate_pairs::post_id_b.eq(post_id)),
|
||||
),
|
||||
)
|
||||
.execute(c)
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::db::Database;
|
||||
use crate::model::{Post, PostStatus, Project};
|
||||
|
||||
fn seeded() -> (Database, String, String) {
|
||||
let db = Database::open_in_memory().unwrap();
|
||||
db.migrate().unwrap();
|
||||
let project_id = "embedding-project".to_string();
|
||||
crate::db::queries::project::insert_project(
|
||||
db.conn(),
|
||||
&Project {
|
||||
id: project_id.clone(),
|
||||
name: "Embedding".into(),
|
||||
slug: "embedding".into(),
|
||||
description: None,
|
||||
data_path: Some("/tmp/embedding".into()),
|
||||
is_active: true,
|
||||
created_at: 1,
|
||||
updated_at: 1,
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
let post_id = "embedding-post".to_string();
|
||||
crate::db::queries::post::insert_post(
|
||||
db.conn(),
|
||||
&Post {
|
||||
id: post_id.clone(),
|
||||
project_id: project_id.clone(),
|
||||
title: "Post".into(),
|
||||
slug: "post".into(),
|
||||
excerpt: None,
|
||||
content: Some("Body".into()),
|
||||
status: PostStatus::Draft,
|
||||
author: None,
|
||||
language: Some("en".into()),
|
||||
do_not_translate: false,
|
||||
template_slug: None,
|
||||
file_path: "posts/post.md".into(),
|
||||
checksum: None,
|
||||
tags: vec![],
|
||||
categories: vec![],
|
||||
published_title: None,
|
||||
published_content: None,
|
||||
published_tags: None,
|
||||
published_categories: None,
|
||||
published_excerpt: None,
|
||||
created_at: 1,
|
||||
updated_at: 1,
|
||||
published_at: None,
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
(db, project_id, post_id)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn embedding_vector_round_trips_as_blob_and_dismissals_are_canonical() {
|
||||
let (db, project_id, post_id) = seeded();
|
||||
let key = EmbeddingKey {
|
||||
label: 1,
|
||||
post_id: post_id.clone(),
|
||||
project_id: project_id.clone(),
|
||||
content_hash: "hash".into(),
|
||||
vector: vec![0, 1, 255],
|
||||
};
|
||||
upsert_key(db.conn(), &key).unwrap();
|
||||
assert_eq!(
|
||||
get_key_for_post(db.conn(), &project_id, &post_id).unwrap(),
|
||||
Some(key)
|
||||
);
|
||||
|
||||
let replacement = EmbeddingKey {
|
||||
label: 2,
|
||||
post_id: post_id.clone(),
|
||||
project_id: project_id.clone(),
|
||||
content_hash: "new-hash".into(),
|
||||
vector: vec![3, 2, 1],
|
||||
};
|
||||
upsert_key(db.conn(), &replacement).unwrap();
|
||||
let updated = get_key_for_post(db.conn(), &project_id, &post_id)
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(updated.label, 1, "a post keeps its stable HNSW label");
|
||||
assert_eq!(updated.content_hash, "new-hash");
|
||||
assert_eq!(updated.vector, vec![3, 2, 1]);
|
||||
|
||||
let pair = DismissedDuplicatePair {
|
||||
id: "dismissal".into(),
|
||||
project_id: project_id.clone(),
|
||||
post_id_a: "a".into(),
|
||||
post_id_b: "b".into(),
|
||||
dismissed_at: 1,
|
||||
};
|
||||
insert_dismissed_pair(db.conn(), &pair).unwrap();
|
||||
insert_dismissed_pair(db.conn(), &pair).unwrap();
|
||||
assert_eq!(
|
||||
list_dismissed_pairs(db.conn(), &project_id).unwrap(),
|
||||
vec![pair]
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,5 +1,6 @@
|
||||
pub mod chat;
|
||||
pub mod db_notification;
|
||||
pub mod embedding;
|
||||
pub mod generated_file_hash;
|
||||
pub mod import_definition;
|
||||
pub mod mcp_proposal;
|
||||
|
||||
@@ -115,7 +115,7 @@ diesel::table! {
|
||||
post_id -> Text,
|
||||
project_id -> Text,
|
||||
content_hash -> Text,
|
||||
vector -> Text,
|
||||
vector -> Binary,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
1382
crates/bds-core/src/engine/embedding.rs
Normal file
1382
crates/bds-core/src/engine/embedding.rs
Normal file
File diff suppressed because it is too large
Load Diff
@@ -5,6 +5,7 @@ use std::path::Path;
|
||||
use crate::db::DbConnection as Connection;
|
||||
use walkdir::WalkDir;
|
||||
|
||||
use crate::db::queries::embedding as qe;
|
||||
use crate::db::queries::media as qm;
|
||||
use crate::db::queries::media_translation as qmt;
|
||||
use crate::db::queries::post as qp;
|
||||
@@ -159,6 +160,49 @@ pub fn compute_metadata_diff(
|
||||
}
|
||||
|
||||
// 6. Detect orphans
|
||||
if crate::engine::meta::read_project_json(data_dir)
|
||||
.is_ok_and(|metadata| metadata.semantic_similarity_enabled)
|
||||
{
|
||||
let service = crate::engine::embedding::EmbeddingService::production(conn, data_dir);
|
||||
let keys = qe::list_keys(conn, project_id)?
|
||||
.into_iter()
|
||||
.map(|key| (key.post_id.clone(), key))
|
||||
.collect::<std::collections::HashMap<_, _>>();
|
||||
for post in &posts {
|
||||
let expected = service.content_hash_for_post(post)?;
|
||||
let key = keys.get(&post.id);
|
||||
let current_hash = key.map(|key| key.content_hash.as_str()).unwrap_or("");
|
||||
let vector_status = key
|
||||
.filter(|key| crate::engine::embedding::decode_vector(&key.vector).is_ok())
|
||||
.map(|_| "ready")
|
||||
.unwrap_or("missing");
|
||||
if current_hash != expected || vector_status != "ready" {
|
||||
let mut fields = Vec::new();
|
||||
if current_hash != expected {
|
||||
fields.push(DiffField {
|
||||
field_name: "content_hash".into(),
|
||||
db_value: current_hash.into(),
|
||||
file_value: expected,
|
||||
});
|
||||
}
|
||||
if vector_status != "ready" {
|
||||
fields.push(DiffField {
|
||||
field_name: "embedding".into(),
|
||||
db_value: vector_status.into(),
|
||||
file_value: "ready".into(),
|
||||
});
|
||||
}
|
||||
report.diffs.push(EntityDiff {
|
||||
entity_type: "embedding".into(),
|
||||
entity_id: post.id.clone(),
|
||||
file_path: format!("projects/{project_id}/embeddings.usearch"),
|
||||
fields,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 7. Detect orphans
|
||||
let orphans = detect_orphan_files(conn, data_dir, project_id)?;
|
||||
report.orphans = orphans;
|
||||
|
||||
@@ -204,6 +248,11 @@ pub fn repair_metadata_diff_item(
|
||||
conn, data_dir, project_id, &path,
|
||||
)?;
|
||||
}
|
||||
"embedding" => {
|
||||
let post = qp::get_post_by_id(conn, &item.entity_id)?;
|
||||
crate::engine::embedding::EmbeddingService::production(conn, data_dir)
|
||||
.sync_post(&post)?;
|
||||
}
|
||||
other => return unsupported_repair(other),
|
||||
}
|
||||
}
|
||||
@@ -222,6 +271,8 @@ pub fn repair_metadata_diff_item(
|
||||
)?,
|
||||
"script" => rewrite_script_from_database(conn, data_dir, &item.entity_id)?,
|
||||
"template" => rewrite_template_from_database(conn, data_dir, &item.entity_id)?,
|
||||
"embedding" => crate::engine::embedding::EmbeddingService::production(conn, data_dir)
|
||||
.flush_project(project_id)?,
|
||||
other => return unsupported_repair(other),
|
||||
},
|
||||
}
|
||||
|
||||
@@ -8,6 +8,7 @@ mod chat_tools;
|
||||
pub mod cli_launcher;
|
||||
pub mod cli_sync;
|
||||
pub mod domain_events;
|
||||
pub mod embedding;
|
||||
pub mod error;
|
||||
pub mod gallery_import;
|
||||
pub mod generation;
|
||||
|
||||
@@ -36,7 +36,7 @@ pub struct RebuildReport {
|
||||
)]
|
||||
pub fn create_post(
|
||||
conn: &Connection,
|
||||
_data_dir: &Path,
|
||||
data_dir: &Path,
|
||||
project_id: &str,
|
||||
title: &str,
|
||||
content: Option<&str>,
|
||||
@@ -91,6 +91,7 @@ pub fn create_post(
|
||||
fts_index_post(conn, &post)?;
|
||||
|
||||
emit_post(&post, NotificationAction::Created);
|
||||
crate::engine::embedding::sync_post_best_effort(conn, data_dir, &post);
|
||||
|
||||
Ok(post)
|
||||
}
|
||||
@@ -102,7 +103,7 @@ pub fn create_post(
|
||||
)]
|
||||
pub fn update_post(
|
||||
conn: &Connection,
|
||||
_data_dir: &Path,
|
||||
data_dir: &Path,
|
||||
post_id: &str,
|
||||
title: Option<&str>,
|
||||
slug: Option<&str>,
|
||||
@@ -169,7 +170,7 @@ pub fn update_post(
|
||||
if post.status == PostStatus::Published || post.status == PostStatus::Archived {
|
||||
// Reload content from filesystem if content field is NULL (published state)
|
||||
if post.content.is_none() && !post.file_path.is_empty() {
|
||||
let abs_path = _data_dir.join(&post.file_path);
|
||||
let abs_path = data_dir.join(&post.file_path);
|
||||
if abs_path.exists()
|
||||
&& let Ok(file_content) = fs::read_to_string(&abs_path)
|
||||
&& let Ok((_fm, body)) = read_post_file(&file_content)
|
||||
@@ -187,6 +188,7 @@ pub fn update_post(
|
||||
fts_index_post(conn, &post)?;
|
||||
|
||||
emit_post(&post, NotificationAction::Updated);
|
||||
crate::engine::embedding::sync_post_best_effort(conn, data_dir, &post);
|
||||
|
||||
Ok(post)
|
||||
}
|
||||
@@ -210,6 +212,7 @@ pub fn publish_post(conn: &Connection, data_dir: &Path, post_id: &str) -> Engine
|
||||
Ok(post) => {
|
||||
conn.release_savepoint()?;
|
||||
emit_post(&post, NotificationAction::Updated);
|
||||
crate::engine::embedding::sync_post_best_effort(conn, data_dir, &post);
|
||||
Ok(post)
|
||||
}
|
||||
Err(error) => {
|
||||
@@ -343,7 +346,9 @@ pub fn archive_post(conn: &Connection, data_dir: &Path, post_id: &str) -> Engine
|
||||
}
|
||||
let now = now_unix_ms();
|
||||
qp::update_post_status(conn, post_id, &PostStatus::Archived, now)?;
|
||||
post.status = PostStatus::Archived;
|
||||
emit_post(&post, NotificationAction::Updated);
|
||||
crate::engine::embedding::sync_post_best_effort(conn, data_dir, &post);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -455,6 +460,7 @@ pub fn discard_post_draft(conn: &Connection, data_dir: &Path, post_id: &str) ->
|
||||
Ok(post) => {
|
||||
conn.release_savepoint()?;
|
||||
emit_post(&post, NotificationAction::Updated);
|
||||
crate::engine::embedding::sync_post_best_effort(conn, data_dir, &post);
|
||||
Ok(post)
|
||||
}
|
||||
Err(error) => {
|
||||
@@ -503,6 +509,8 @@ pub fn delete_post(conn: &Connection, data_dir: &Path, post_id: &str) -> EngineR
|
||||
// Delete post from DB
|
||||
qp::delete_post(conn, post_id)?;
|
||||
|
||||
crate::engine::embedding::remove_post_best_effort(conn, data_dir, &post.project_id, post_id);
|
||||
|
||||
emit_post(&post, NotificationAction::Deleted);
|
||||
|
||||
Ok(())
|
||||
|
||||
@@ -238,6 +238,7 @@ pub fn delete_project(
|
||||
let is_custom_path = project.data_path.is_some();
|
||||
|
||||
q::delete_project(conn, project_id)?;
|
||||
crate::engine::embedding::EmbeddingService::forget_project(project_id);
|
||||
|
||||
// Clean up internal filesystem only (not custom external paths per spec)
|
||||
if !is_custom_path
|
||||
|
||||
@@ -170,15 +170,18 @@ fn rebuild_from_filesystem_inner(
|
||||
)));
|
||||
}
|
||||
|
||||
progress(0.98, "Refreshing semantic index...");
|
||||
crate::engine::embedding::EmbeddingService::production(conn, data_dir)
|
||||
.index_unindexed(project_id)?;
|
||||
|
||||
progress(1.0, "Rebuild complete");
|
||||
Ok(report)
|
||||
}
|
||||
|
||||
fn clear_project_rows(conn: &Connection, project_id: &str) -> EngineResult<()> {
|
||||
use crate::db::schema::{
|
||||
dismissed_duplicate_pairs, embedding_keys, generated_file_hashes, import_definitions,
|
||||
media, media_translations, post_links, post_media, post_translations, posts, scripts, tags,
|
||||
templates,
|
||||
generated_file_hashes, import_definitions, media, media_translations, post_links,
|
||||
post_media, post_translations, posts, scripts, tags, templates,
|
||||
};
|
||||
|
||||
let post_ids = crate::db::queries::post::list_posts_by_project(conn, project_id)?
|
||||
@@ -219,13 +222,6 @@ fn clear_project_rows(conn: &Connection, project_id: &str) -> EngineResult<()> {
|
||||
generated_file_hashes::table.filter(generated_file_hashes::project_id.eq(project_id)),
|
||||
)
|
||||
.execute(connection)?;
|
||||
diesel::delete(embedding_keys::table.filter(embedding_keys::project_id.eq(project_id)))
|
||||
.execute(connection)?;
|
||||
diesel::delete(
|
||||
dismissed_duplicate_pairs::table
|
||||
.filter(dismissed_duplicate_pairs::project_id.eq(project_id)),
|
||||
)
|
||||
.execute(connection)?;
|
||||
diesel::delete(
|
||||
import_definitions::table.filter(import_definitions::project_id.eq(project_id)),
|
||||
)
|
||||
|
||||
50
crates/bds-core/src/model/embedding.rs
Normal file
50
crates/bds-core/src/model/embedding.rs
Normal file
@@ -0,0 +1,50 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(
|
||||
Debug,
|
||||
Clone,
|
||||
PartialEq,
|
||||
Serialize,
|
||||
Deserialize,
|
||||
diesel::Queryable,
|
||||
diesel::Selectable,
|
||||
diesel::Insertable,
|
||||
diesel::AsChangeset,
|
||||
)]
|
||||
#[diesel(
|
||||
table_name = crate::db::schema::embedding_keys,
|
||||
check_for_backend(diesel::sqlite::Sqlite),
|
||||
treat_none_as_default_value = false
|
||||
)]
|
||||
pub struct EmbeddingKey {
|
||||
pub label: i64,
|
||||
pub post_id: String,
|
||||
pub project_id: String,
|
||||
pub content_hash: String,
|
||||
pub vector: Vec<u8>,
|
||||
}
|
||||
|
||||
#[derive(
|
||||
Debug,
|
||||
Clone,
|
||||
PartialEq,
|
||||
Eq,
|
||||
Serialize,
|
||||
Deserialize,
|
||||
diesel::Queryable,
|
||||
diesel::Selectable,
|
||||
diesel::Insertable,
|
||||
diesel::AsChangeset,
|
||||
)]
|
||||
#[diesel(
|
||||
table_name = crate::db::schema::dismissed_duplicate_pairs,
|
||||
check_for_backend(diesel::sqlite::Sqlite),
|
||||
treat_none_as_default_value = false
|
||||
)]
|
||||
pub struct DismissedDuplicatePair {
|
||||
pub id: String,
|
||||
pub project_id: String,
|
||||
pub post_id_a: String,
|
||||
pub post_id_b: String,
|
||||
pub dismissed_at: i64,
|
||||
}
|
||||
@@ -1,4 +1,5 @@
|
||||
mod chat;
|
||||
mod embedding;
|
||||
mod event;
|
||||
mod generation;
|
||||
mod import;
|
||||
@@ -12,6 +13,7 @@ mod tag;
|
||||
mod template;
|
||||
|
||||
pub use chat::{ChatConversation, ChatMessage, ChatRole, NewChatConversation, NewChatMessage};
|
||||
pub use embedding::{DismissedDuplicatePair, EmbeddingKey};
|
||||
pub use event::DomainEvent;
|
||||
pub use generation::{
|
||||
DbNotification, DomainEntity, GeneratedFileHash, NotificationAction, NotificationEntity,
|
||||
|
||||
@@ -1382,6 +1382,95 @@ impl CoreHost {
|
||||
(_, response) => one_shot_json(response),
|
||||
}
|
||||
}
|
||||
|
||||
fn embeddings(&self, method: &str, args: &[Value]) -> HostResult<Value> {
|
||||
let db = self.database()?;
|
||||
let service = engine::embedding::EmbeddingService::production(db.conn(), &self.data_dir);
|
||||
match method {
|
||||
"get_progress" => {
|
||||
let (indexed, total) = service.indexing_progress(&self.project_id)?;
|
||||
Ok(json!({"indexed": indexed, "total": total}))
|
||||
}
|
||||
"find_similar" => {
|
||||
let post_id = string_arg(args, 0)?;
|
||||
self.scoped(
|
||||
|conn| crate::db::queries::post::get_post_by_id(conn, post_id),
|
||||
|post| post.project_id.as_str(),
|
||||
)?;
|
||||
let limit = args.get(1).and_then(Value::as_u64).unwrap_or(5) as usize;
|
||||
Ok(Value::Array(
|
||||
service
|
||||
.find_similar(post_id, limit)?
|
||||
.into_iter()
|
||||
.map(|post| {
|
||||
json!({
|
||||
"post_id": post.post_id,
|
||||
"title": post.title,
|
||||
"score": post.similarity,
|
||||
})
|
||||
})
|
||||
.collect(),
|
||||
))
|
||||
}
|
||||
"compute_similarities" => {
|
||||
let post_id = string_arg(args, 0)?;
|
||||
self.scoped(
|
||||
|conn| crate::db::queries::post::get_post_by_id(conn, post_id),
|
||||
|post| post.project_id.as_str(),
|
||||
)?;
|
||||
let target_ids = string_array_arg(args, 1)?;
|
||||
json_value(service.compute_similarities(post_id, &target_ids))
|
||||
}
|
||||
"suggest_tags" => {
|
||||
let post_id = string_arg(args, 0)?;
|
||||
self.scoped(
|
||||
|conn| crate::db::queries::post::get_post_by_id(conn, post_id),
|
||||
|post| post.project_id.as_str(),
|
||||
)?;
|
||||
json_value(service.suggest_tags(post_id))
|
||||
}
|
||||
"find_duplicates" => {
|
||||
let mut page = 0;
|
||||
let pairs = loop {
|
||||
let result = service.find_duplicates(&self.project_id, page)?;
|
||||
if !result.has_more {
|
||||
break result.pairs;
|
||||
}
|
||||
page += 1;
|
||||
};
|
||||
Ok(Value::Array(
|
||||
pairs
|
||||
.into_iter()
|
||||
.map(|pair| {
|
||||
json!({
|
||||
"post_id_a": pair.post_id_a,
|
||||
"title_a": pair.title_a,
|
||||
"post_id_b": pair.post_id_b,
|
||||
"title_b": pair.title_b,
|
||||
"score": pair.similarity,
|
||||
"similarity": pair.similarity,
|
||||
"exact_match": pair.exact_match,
|
||||
})
|
||||
})
|
||||
.collect(),
|
||||
))
|
||||
}
|
||||
"dismiss_pair" => {
|
||||
let post_id_a = string_arg(args, 0)?;
|
||||
let post_id_b = string_arg(args, 1)?;
|
||||
for post_id in [post_id_a, post_id_b] {
|
||||
self.scoped(
|
||||
|conn| crate::db::queries::post::get_post_by_id(conn, post_id),
|
||||
|post| post.project_id.as_str(),
|
||||
)?;
|
||||
}
|
||||
service.dismiss_duplicate_pair(post_id_a, post_id_b)?;
|
||||
Ok(Value::Bool(true))
|
||||
}
|
||||
"index_unindexed_posts" => json_value(service.index_unindexed(&self.project_id)),
|
||||
_ => Err(format!("unknown embeddings capability: {method}").into()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl HostApi for CoreHost {
|
||||
@@ -1398,6 +1487,7 @@ impl HostApi for CoreHost {
|
||||
"tasks" => self.tasks(method, &arguments),
|
||||
"publish" => self.publish(method, &arguments),
|
||||
"chat" => self.chat(method, &arguments),
|
||||
"embeddings" => self.embeddings(method, &arguments),
|
||||
"bds" if method == "report_progress" => self.report_progress(&arguments),
|
||||
_ => Err(format!("unknown host capability: {namespace}.{method}").into()),
|
||||
};
|
||||
@@ -1537,6 +1627,19 @@ fn string_arg(args: &[Value], index: usize) -> HostResult<&str> {
|
||||
.ok_or_else(|| format!("argument {} must be a string", index + 1).into())
|
||||
}
|
||||
|
||||
fn string_array_arg(args: &[Value], index: usize) -> HostResult<Vec<String>> {
|
||||
args.get(index)
|
||||
.and_then(Value::as_array)
|
||||
.map(|items| {
|
||||
items
|
||||
.iter()
|
||||
.filter_map(Value::as_str)
|
||||
.map(str::to_owned)
|
||||
.collect()
|
||||
})
|
||||
.ok_or_else(|| format!("argument {} must be a table", index + 1).into())
|
||||
}
|
||||
|
||||
fn string_field<'a>(value: &'a Map<String, Value>, field: &str) -> HostResult<&'a str> {
|
||||
value
|
||||
.get(field)
|
||||
@@ -1821,6 +1924,9 @@ mod tests {
|
||||
upload = upload,
|
||||
timestamp = post.created_at,
|
||||
project_path = bds.app.get_default_project_path(),
|
||||
embedding_progress = bds.embeddings.get_progress(),
|
||||
embedding_backfill = bds.embeddings.index_unindexed_posts(),
|
||||
foreign_embedding = bds.embeddings.find_similar(input.foreign_post, 5),
|
||||
}
|
||||
end
|
||||
"#,
|
||||
@@ -1842,5 +1948,11 @@ mod tests {
|
||||
assert!(result.value["upload"].is_null());
|
||||
assert!(result.value["timestamp"].as_str().unwrap().contains('T'));
|
||||
assert_eq!(manager.progress(task_id), Some(0.5));
|
||||
assert_eq!(
|
||||
result.value["embedding_progress"],
|
||||
json!({"indexed": 0, "total": 1})
|
||||
);
|
||||
assert_eq!(result.value["embedding_backfill"], json!([]));
|
||||
assert!(result.value["foreign_embedding"].is_null());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -456,6 +456,7 @@ mod tests {
|
||||
[
|
||||
"app",
|
||||
"chat",
|
||||
"embeddings",
|
||||
"media",
|
||||
"meta",
|
||||
"posts",
|
||||
@@ -473,7 +474,7 @@ mod tests {
|
||||
function main()
|
||||
return {
|
||||
sync = bds.sync,
|
||||
embeddings = bds.embeddings,
|
||||
embeddings = type(bds.embeddings),
|
||||
report_progress = type(bds.report_progress),
|
||||
post_search = type(bds.posts.search),
|
||||
app_toast = type(bds.app.toast),
|
||||
@@ -492,6 +493,7 @@ mod tests {
|
||||
"report_progress": "function",
|
||||
"post_search": "function",
|
||||
"app_toast": "function",
|
||||
"embeddings": "table",
|
||||
})
|
||||
);
|
||||
}
|
||||
@@ -514,6 +516,7 @@ mod tests {
|
||||
bds.tasks.status_snapshot(),
|
||||
bds.publish.upload_site({}),
|
||||
bds.chat.detect_post_language("title", "body"),
|
||||
bds.embeddings.get_progress(),
|
||||
}
|
||||
end
|
||||
"#,
|
||||
@@ -525,8 +528,8 @@ mod tests {
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(execution.value.as_array().unwrap().len(), 11);
|
||||
assert_eq!(host.0.lock().unwrap().len(), 11);
|
||||
assert_eq!(execution.value.as_array().unwrap().len(), 12);
|
||||
assert_eq!(host.0.lock().unwrap().len(), 12);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
Reference in New Issue
Block a user