Skip to main content

locus_mcp/tools/
run_embedding_migration.rs

1use locus_core_rs::{EmbeddingMigrationFilter, EmbeddingMigrationRunRequest};
2use serde_json::json;
3use tracing::error;
4
5use crate::{
6    RunEmbeddingMigrationRequest, SttpMcpServer, mode_to_string, normalize_tiers,
7    parse_migration_mode, parse_utc_optional, to_json_string, tool_error, validate_batch_size,
8    validate_max_nodes,
9};
10
11pub(crate) async fn execute(
12    server: &SttpMcpServer,
13    request: RunEmbeddingMigrationRequest,
14) -> String {
15    let from_utc = match parse_utc_optional(request.from_utc.as_deref(), "from_utc") {
16        Ok(value) => value,
17        Err(message) => return tool_error("InvalidDate", &message),
18    };
19    let to_utc = match parse_utc_optional(request.to_utc.as_deref(), "to_utc") {
20        Ok(value) => value,
21        Err(message) => return tool_error("InvalidDate", &message),
22    };
23    let tiers = request
24        .tiers
25        .as_ref()
26        .map(|values| normalize_tiers(values.as_slice()));
27    let batch_size = match validate_batch_size(request.batch_size) {
28        Ok(value) => value,
29        Err(message) => return tool_error("InvalidArgument", &message),
30    };
31    let max_nodes = match validate_max_nodes(request.max_nodes) {
32        Ok(value) => value,
33        Err(message) => return tool_error("InvalidArgument", &message),
34    };
35    let mode = match parse_migration_mode(request.mode.as_deref()) {
36        Ok(value) => value,
37        Err(message) => return tool_error("InvalidArgument", &message),
38    };
39
40    let filter = EmbeddingMigrationFilter {
41        session_id: request.session_id,
42        from_utc,
43        to_utc,
44        tiers,
45        has_embedding: request.has_embedding,
46        embedding_model: request.embedding_model,
47        sync_keys: request.sync_keys,
48    };
49
50    match server
51        .embedding_migration
52        .run_async(EmbeddingMigrationRunRequest {
53            filter,
54            mode,
55            dry_run: request.dry_run,
56            batch_size,
57            max_nodes,
58        })
59        .await
60    {
61        Ok(result) => to_json_string(json!({
62            "scanned": result.scanned,
63            "selected": result.selected,
64            "updated": result.updated,
65            "skipped": result.skipped,
66            "failed": result.failed,
67            "duplicate": result.duplicate,
68            "started_at": result.started_at.to_rfc3339(),
69            "completed_at": result.completed_at.to_rfc3339(),
70            "provider_model": result.provider_model,
71            "dry_run": request.dry_run,
72            "mode": mode_to_string(mode),
73            "failure_reasons": result.failure_reasons,
74        })),
75        Err(err) => {
76            error!(error = %err, "run_embedding_migration failed");
77            tool_error("MigrationRunFailure", &err.to_string())
78        }
79    }
80}