Skip to main content

dstack/
cmd_memory.rs

1use crate::config::Config;
2use dstack_memory::{Field, MemoryProvider};
3
4pub async fn provider_from_config(cfg: &Config) -> anyhow::Result<Box<dyn MemoryProvider>> {
5    match cfg.memory.backend.as_str() {
6        "eruka" => {
7            let key = cfg
8                .eruka_service_key()
9                .unwrap_or_default();
10            Ok(Box::new(dstack_memory::eruka::ErukaProvider::new(
11                &cfg.memory.eruka.url,
12                &key,
13            )))
14        }
15        _ => {
16            let path = cfg.memory_path();
17            std::fs::create_dir_all(&path)?;
18            Ok(Box::new(dstack_memory::file::FileProvider::new(path)))
19        }
20    }
21}
22
23pub async fn load(cfg: &Config, project: Option<&str>) -> anyhow::Result<()> {
24    let provider = provider_from_config(cfg).await?;
25    let path = project.unwrap_or("");
26    let fields = provider.load(path).await?;
27    if fields.is_empty() {
28        eprintln!(
29            "No memory fields found for: {}",
30            if path.is_empty() { "(all)" } else { path }
31        );
32        return Ok(());
33    }
34    for f in &fields {
35        println!("[{:.0}%] {} = {}", f.confidence * 100.0, f.path, f.value);
36    }
37    eprintln!("\n{} field(s) loaded.", fields.len());
38    Ok(())
39}
40
41pub async fn save(cfg: &Config, key: &str, value: &str) -> anyhow::Result<()> {
42    let provider = provider_from_config(cfg).await?;
43    let field = Field::new(key, value, "user").with_confidence(0.9);
44    provider.write(&field).await?;
45    eprintln!("Saved: {} = {}", key, value);
46    Ok(())
47}
48
49pub async fn query(cfg: &Config, pattern: &str) -> anyhow::Result<()> {
50    let provider = provider_from_config(cfg).await?;
51    let results = provider.search(pattern).await?;
52    if results.is_empty() {
53        eprintln!("No matches for: {}", pattern);
54        return Ok(());
55    }
56    for f in &results {
57        println!("[{:.0}%] {} = {}", f.confidence * 100.0, f.path, f.value);
58    }
59    eprintln!("\n{} match(es).", results.len());
60    Ok(())
61}
62
63pub async fn export(cfg: &Config) -> anyhow::Result<()> {
64    let provider = provider_from_config(cfg).await?;
65    let all = provider.export_all().await?;
66    println!("{}", serde_json::to_string_pretty(&all)?);
67    Ok(())
68}