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::traits::storage::StorageCapability;
17pub use ai_agents_core::{AgentError, AgentSnapshot, AgentStorage, MemorySnapshot, Result};
18pub use snapshot::StateMachineSnapshot;
19pub use storage::FileStorage;
20
21#[cfg(feature = "sqlite")]
22pub use sqlite::{SessionInfo, SessionOrderBy, SessionQuery, SqliteMetadata, SqliteStorage};
23
24#[cfg(feature = "redis-storage")]
25pub use redis::{RedisSessionMeta, RedisStorage};
26
27use serde::{Deserialize, Serialize};
28use std::sync::Arc;
29
30#[derive(Debug, Clone, Serialize, Deserialize, Default)]
31#[serde(tag = "type")]
32pub enum StorageConfig {
33    #[default]
34    #[serde(rename = "none")]
35    None,
36    #[serde(rename = "file")]
37    File { path: String },
38    #[serde(rename = "sqlite")]
39    Sqlite { path: String },
40    #[serde(rename = "redis")]
41    Redis {
42        url: String,
43        #[serde(default)]
44        prefix: Option<String>,
45        #[serde(default)]
46        ttl_seconds: Option<u64>,
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 SqliteMetadata {
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: SqliteMetadata,
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
