Skip to main content

engram_core/
import_handler.rs

1use std::sync::Arc;
2
3use serde::Deserialize;
4use serde_json::{Value, json};
5
6use engram_storage::Memory;
7
8use crate::error::CoreError;
9use crate::lock_helpers;
10use crate::server::ServerState;
11use crate::tags_normalize::{TagsInput, normalize_tags};
12
13#[derive(Deserialize)]
14struct ImportParams {
15    version: u64,
16    memories: Vec<ImportedMemory>,
17}
18
19#[derive(Deserialize)]
20struct ImportedMemory {
21    id: String,
22    memory_type: String,
23    context: String,
24    action: String,
25    result: String,
26    score: f32,
27    tags: Option<String>,
28    project: Option<String>,
29    parent_id: Option<String>,
30    source_ids: Option<String>,
31    insight_type: Option<String>,
32    created_at: String,
33    updated_at: String,
34    used_count: i64,
35    last_used_at: Option<String>,
36}
37
38pub async fn handle(state: &Arc<ServerState>, params: Value) -> Result<Value, CoreError> {
39    let parsed: ImportParams = serde_json::from_value(params)
40        .map_err(|error| CoreError::DispatchError(error.to_string()))?;
41    if parsed.version != 1 {
42        return Err(CoreError::ImportVersionMismatch(parsed.version));
43    }
44    let state_clone = Arc::clone(state);
45    tokio::task::spawn_blocking(move || {
46        let database = lock_helpers::lock_db(&state_clone);
47        let (imported, skipped) = insert_non_duplicate_memories(&database, &parsed.memories)?;
48        Ok::<Value, CoreError>(json!({ "imported": imported, "skipped": skipped }))
49    })
50    .await
51    .map_err(|error| CoreError::SocketError(error.to_string()))?
52}
53
54fn insert_non_duplicate_memories(
55    database: &engram_storage::Database,
56    memories: &[ImportedMemory],
57) -> Result<(usize, usize), CoreError> {
58    let mut imported = 0;
59    let mut skipped = 0;
60    for entry in memories {
61        if memory_exists(database, &entry.id)? {
62            skipped += 1;
63            continue;
64        }
65        let memory = imported_memory_to_storage(entry)?;
66        database
67            .insert_memory(&memory)
68            .map_err(|error| CoreError::ImportFailed(error.to_string()))?;
69        imported += 1;
70    }
71    Ok((imported, skipped))
72}
73
74fn memory_exists(database: &engram_storage::Database, id: &str) -> Result<bool, CoreError> {
75    match database.get_memory(id) {
76        Ok(_) => Ok(true),
77        Err(engram_storage::StorageError::NotFound(_)) => Ok(false),
78        Err(error) => Err(CoreError::ImportFailed(error.to_string())),
79    }
80}
81
82fn imported_memory_to_storage(entry: &ImportedMemory) -> Result<Memory, CoreError> {
83    let tags = normalize_tags(entry.tags.clone().map(TagsInput::Encoded));
84    Ok(Memory {
85        id: entry.id.clone(),
86        memory_type: entry.memory_type.clone(),
87        context: entry.context.clone(),
88        action: entry.action.clone(),
89        result: entry.result.clone(),
90        score: entry.score,
91        embedding_context: None,
92        embedding_action: None,
93        embedding_result: None,
94        indexed: false,
95        tags,
96        project: entry.project.clone(),
97        parent_id: entry.parent_id.clone(),
98        source_ids: entry.source_ids.clone(),
99        insight_type: entry.insight_type.clone(),
100        created_at: entry.created_at.clone(),
101        updated_at: entry.updated_at.clone(),
102        used_count: entry.used_count,
103        last_used_at: entry.last_used_at.clone(),
104        superseded_by: None,
105    })
106}