ai_agents_storage/
lib.rs

1//! Storage backends for AI Agents framework
2
3mod snapshot;
4mod storage;
5
6#[cfg(feature = "sqlite")]
7mod sqlite;
8
9#[cfg(feature = "redis-storage")]
10mod redis;
11
12pub use ai_agents_core::{AgentError, AgentSnapshot, AgentStorage, MemorySnapshot, Result};
13pub use snapshot::StateMachineSnapshot;
14pub use storage::FileStorage;
15
16#[cfg(feature = "sqlite")]
17pub use sqlite::{SessionInfo, SessionMetadata, SessionOrderBy, SessionQuery, SqliteStorage};
18
19#[cfg(feature = "redis-storage")]
20pub use redis::{RedisSessionMeta, RedisStorage};
21
22use serde::{Deserialize, Serialize};
23use std::sync::Arc;
24
25#[derive(Debug, Clone, Serialize, Deserialize)]
26#[serde(tag = "type")]
27pub enum StorageConfig {
28    #[serde(rename = "none")]
29    None,
30    #[serde(rename = "file")]
31    File { path: String },
32    #[serde(rename = "sqlite")]
33    Sqlite { path: String },
34    #[serde(rename = "redis")]
35    Redis {
36        url: String,
37        #[serde(default)]
38        prefix: Option<String>,
39        #[serde(default)]
40        ttl_seconds: Option<u64>,
41    },
42}
43
44impl Default for StorageConfig {
45    fn default() -> Self {
46        StorageConfig::None
47    }
48}
49
50pub async fn create_storage(config: &StorageConfig) -> Result<Option<Arc<dyn AgentStorage>>> {
51    match config {
52        StorageConfig::None => Ok(None),
53        StorageConfig::File { path } => Ok(Some(Arc::new(FileStorage::new(path)))),
54
55        #[cfg(feature = "sqlite")]
56        StorageConfig::Sqlite { path } => {
57            let storage = SqliteStorage::new(path).await?;
58            Ok(Some(Arc::new(storage)))
59        }
60
61        #[cfg(not(feature = "sqlite"))]
62        StorageConfig::Sqlite { .. } => Err(AgentError::Config(
63            "SQLite storage requires 'sqlite' feature".into(),
64        )),
65
66        #[cfg(feature = "redis-storage")]
67        StorageConfig::Redis {
68            url,
69            prefix,
70            ttl_seconds,
71        } => {
72            let mut storage = RedisStorage::new(url)?;
73            if let Some(p) = prefix {
74                storage = storage.with_prefix(p);
75            }
76            if let Some(ttl) = ttl_seconds {
77                storage = storage.with_ttl(*ttl);
78            }
79            Ok(Some(Arc::new(storage)))
80        }
81
82        #[cfg(not(feature = "redis-storage"))]
83        StorageConfig::Redis { .. } => Err(AgentError::Config(
84            "Redis storage requires 'redis-storage' feature".into(),
85        )),
86    }
87}
88
89#[cfg(not(feature = "sqlite"))]
90mod sqlite_types {
91    use chrono::{DateTime, Utc};
92    use serde::{Deserialize, Serialize};
93    use std::collections::HashMap;
94
95    #[derive(Debug, Clone, Serialize, Deserialize, Default)]
96    pub struct SessionMetadata {
97        #[serde(default)]
98        pub tags: Vec<String>,
99        #[serde(default)]
100        pub user_id: Option<String>,
101        #[serde(default)]
102        pub custom: HashMap<String, serde_json::Value>,
103        #[serde(default)]
104        pub priority: Option<i32>,
105        #[serde(default)]
106        pub expires_at: Option<DateTime<Utc>>,
107    }
108
109    #[derive(Debug, Clone, Serialize, Deserialize)]
110    pub struct SessionInfo {
111        pub session_id: String,
112        pub agent_id: String,
113        pub created_at: DateTime<Utc>,
114        pub updated_at: DateTime<Utc>,
115        pub message_count: usize,
116        pub current_state: Option<String>,
117        #[serde(default)]
118        pub metadata: SessionMetadata,
119    }
120
121    #[derive(Debug, Clone, Default)]
122    pub struct SessionQuery {
123        pub agent_id: Option<String>,
124        pub state: Option<String>,
125        pub tag: Option<String>,
126        pub user_id: Option<String>,
127        pub created_after: Option<DateTime<Utc>>,
128        pub created_before: Option<DateTime<Utc>>,
129        pub updated_after: Option<DateTime<Utc>>,
130        pub limit: Option<u32>,
131        pub offset: Option<u32>,
132        pub order_by: SessionOrderBy,
133    }
134
135    #[derive(Debug, Clone, Default)]
136    pub enum SessionOrderBy {
137        #[default]
138        UpdatedAtDesc,
139        UpdatedAtAsc,
140        CreatedAtDesc,
141        CreatedAtAsc,
142        MessageCountDesc,
143    }
144}
145
146#[cfg(not(feature = "sqlite"))]
147pub use sqlite_types::*;
ai_agents_storage/lib.rs

ai_agents_storage/
lib.rs
