ai_agents_storage/
lib.rs

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

ai_agents_storage/
lib.rs
