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