rs_adk/artifacts/
in_memory.rs1use async_trait::async_trait;
4use dashmap::DashMap;
5
6use super::{now_secs, Artifact, ArtifactError, ArtifactMetadata, ArtifactService};
7
8pub struct InMemoryArtifactService {
13 store: DashMap<String, DashMap<String, Vec<Artifact>>>,
15}
16
17impl InMemoryArtifactService {
18 pub fn new() -> Self {
20 Self {
21 store: DashMap::new(),
22 }
23 }
24}
25
26impl Default for InMemoryArtifactService {
27 fn default() -> Self {
28 Self::new()
29 }
30}
31
32#[async_trait]
33impl ArtifactService for InMemoryArtifactService {
34 async fn save(
35 &self,
36 session_id: &str,
37 mut artifact: Artifact,
38 ) -> Result<ArtifactMetadata, ArtifactError> {
39 let session = self.store.entry(session_id.to_string()).or_default();
40
41 let mut versions = session.entry(artifact.metadata.name.clone()).or_default();
42
43 let version = versions.len() as u32 + 1;
44 artifact.metadata.version = version;
45 artifact.metadata.updated_at = now_secs();
46 if version == 1 {
47 artifact.metadata.created_at = artifact.metadata.updated_at;
48 }
49
50 let metadata = artifact.metadata.clone();
51 versions.push(artifact);
52 Ok(metadata)
53 }
54
55 async fn load(&self, session_id: &str, name: &str) -> Result<Option<Artifact>, ArtifactError> {
56 let result = self.store.get(session_id).and_then(|session| {
57 session
58 .get(name)
59 .and_then(|versions| versions.last().cloned())
60 });
61 Ok(result)
62 }
63
64 async fn load_version(
65 &self,
66 session_id: &str,
67 name: &str,
68 version: u32,
69 ) -> Result<Option<Artifact>, ArtifactError> {
70 let result = self.store.get(session_id).and_then(|session| {
71 session.get(name).and_then(|versions| {
72 if version == 0 || version as usize > versions.len() {
73 None
74 } else {
75 Some(versions[(version - 1) as usize].clone())
76 }
77 })
78 });
79 Ok(result)
80 }
81
82 async fn list(&self, session_id: &str) -> Result<Vec<ArtifactMetadata>, ArtifactError> {
83 let result = self
84 .store
85 .get(session_id)
86 .map(|session| {
87 session
88 .iter()
89 .filter_map(|entry| entry.value().last().map(|a| a.metadata.clone()))
90 .collect()
91 })
92 .unwrap_or_default();
93 Ok(result)
94 }
95
96 async fn delete(&self, session_id: &str, name: &str) -> Result<(), ArtifactError> {
97 if let Some(session) = self.store.get(session_id) {
98 session.remove(name);
99 }
100 Ok(())
101 }
102}