Skip to main content

locus_mcp/tools/
preview_embedding_migration.rs

1use locus_core_rs::{EmbeddingMigrationFilter, EmbeddingMigrationPreviewRequest};
2use serde_json::json;
3use tracing::error;
4
5use crate::{
6    PreviewEmbeddingMigrationRequest, SttpMcpServer, normalize_tiers, parse_utc_optional,
7    to_json_string, tool_error, validate_limit, validate_max_nodes,
8};
9
10pub(crate) async fn execute(
11    server: &SttpMcpServer,
12    request: PreviewEmbeddingMigrationRequest,
13) -> String {
14    let from_utc = match parse_utc_optional(request.from_utc.as_deref(), "from_utc") {
15        Ok(value) => value,
16        Err(message) => return tool_error("InvalidDate", &message),
17    };
18    let to_utc = match parse_utc_optional(request.to_utc.as_deref(), "to_utc") {
19        Ok(value) => value,
20        Err(message) => return tool_error("InvalidDate", &message),
21    };
22    let tiers = request
23        .tiers
24        .as_ref()
25        .map(|values| normalize_tiers(values.as_slice()));
26    let sample_limit = match validate_limit(request.sample_limit, "sample_limit") {
27        Ok(value) => value,
28        Err(message) => return tool_error("InvalidArgument", &message),
29    };
30    let max_nodes = match validate_max_nodes(request.max_nodes) {
31        Ok(value) => value,
32        Err(message) => return tool_error("InvalidArgument", &message),
33    };
34
35    let filter = EmbeddingMigrationFilter {
36        session_id: request.session_id,
37        from_utc,
38        to_utc,
39        tiers,
40        has_embedding: request.has_embedding,
41        embedding_model: request.embedding_model,
42        sync_keys: request.sync_keys,
43    };
44
45    match server
46        .embedding_migration
47        .preview_async(EmbeddingMigrationPreviewRequest {
48            filter,
49            sample_limit,
50            max_nodes,
51        })
52        .await
53    {
54        Ok(result) => to_json_string(json!({
55            "total_candidates": result.total_candidates,
56            "provider_available": result.provider_available,
57            "provider_model": result.provider_model,
58            "sample": result
59                .sample
60                .iter()
61                .map(|sample| json!({
62                    "sync_key": sample.sync_key,
63                    "session_id": sample.session_id,
64                    "tier": sample.tier,
65                    "has_embedding": sample.has_embedding,
66                    "embedding_model": sample.embedding_model,
67                    "embedding_dimensions": sample.embedding_dimensions,
68                    "embedded_at": sample.embedded_at.map(|value| value.to_rfc3339()),
69                    "updated_at": sample.updated_at.to_rfc3339(),
70                    "context_summary": sample.context_summary,
71                }))
72                .collect::<Vec<_>>(),
73        })),
74        Err(err) => {
75            error!(error = %err, "preview_embedding_migration failed");
76            tool_error("MigrationPreviewFailure", &err.to_string())
77        }
78    }
79}