avocado_core/storage/
traits.rs

1//! Core storage backend trait definitions
2
3use async_trait::async_trait;
4use std::sync::Arc;
5
6use crate::storage::vector::VectorSearchProvider;
7use crate::types::{
8    Agent, AgentRelation, AgentRelationSummary, Artifact, CompilerConfig, IngestAction, Message,
9    MessageRole, Result, Session, SessionWithMessages, SessionWorkingSet, Span, Stance, WorkingSet,
10};
11
12/// Configuration for storage backends
13#[derive(Debug, Clone)]
14pub enum StorageConfig {
15    /// SQLite backend (default)
16    Sqlite {
17        /// Path to SQLite database file
18        path: String,
19    },
20    /// PostgreSQL backend with pgvector
21    Postgres {
22        /// PostgreSQL connection string
23        connection_string: String,
24    },
25}
26
27impl StorageConfig {
28    /// Parse configuration from AVOCADO_BACKEND environment variable
29    ///
30    /// # Examples
31    ///
32    /// - `sqlite` or unset -> SQLite with default path
33    /// - `sqlite:/path/to/db.sqlite` -> SQLite with specific path
34    /// - `postgres://user:pass@host/db` -> PostgreSQL
35    pub fn from_env(default_sqlite_path: &str) -> Self {
36        match std::env::var("AVOCADO_BACKEND").ok() {
37            None => StorageConfig::Sqlite {
38                path: default_sqlite_path.to_string(),
39            },
40            Some(ref s) if s == "sqlite" || s.is_empty() => StorageConfig::Sqlite {
41                path: default_sqlite_path.to_string(),
42            },
43            Some(ref s) if s.starts_with("sqlite:") => StorageConfig::Sqlite {
44                path: s.strip_prefix("sqlite:").unwrap().to_string(),
45            },
46            Some(ref s) if s.starts_with("postgres://") || s.starts_with("postgresql://") => {
47                StorageConfig::Postgres {
48                    connection_string: s.clone(),
49                }
50            }
51            Some(s) => {
52                eprintln!(
53                    "[AvocadoDB] Unknown AVOCADO_BACKEND '{}', defaulting to SQLite",
54                    s
55                );
56                StorageConfig::Sqlite {
57                    path: default_sqlite_path.to_string(),
58                }
59            }
60        }
61    }
62}
63
64/// Main storage backend trait
65///
66/// Implementations must be Send + Sync for use in async contexts.
67/// All methods are async to support both sync (SQLite) and async (PostgreSQL) backends.
68#[async_trait]
69pub trait StorageBackend: Send + Sync {
70    // ========== Lifecycle ==========
71
72    /// Get database statistics
73    ///
74    /// # Returns
75    /// Tuple of (artifacts_count, spans_count, total_tokens)
76    async fn get_stats(&self) -> Result<(usize, usize, usize)>;
77
78    /// Clear all data from the database
79    async fn clear(&self) -> Result<()>;
80
81    // ========== Artifacts ==========
82
83    /// Insert an artifact
84    async fn insert_artifact(&self, artifact: &Artifact) -> Result<()>;
85
86    /// Get artifact by ID
87    async fn get_artifact(&self, artifact_id: &str) -> Result<Option<Artifact>>;
88
89    /// Get artifact by path
90    async fn get_artifact_by_path(&self, path: &str) -> Result<Option<Artifact>>;
91
92    /// Delete artifact and associated spans
93    ///
94    /// # Returns
95    /// Number of spans deleted
96    async fn delete_artifact(&self, artifact_id: &str) -> Result<usize>;
97
98    /// Determine what action to take when ingesting a file
99    async fn determine_ingest_action(
100        &self,
101        path: &str,
102        content_hash: &str,
103    ) -> Result<IngestAction>;
104
105    // ========== Spans ==========
106
107    /// Insert multiple spans in a transaction
108    async fn insert_spans(&self, spans: &[Span]) -> Result<()>;
109
110    /// Get all spans (for index building)
111    async fn get_all_spans(&self) -> Result<Vec<Span>>;
112
113    /// Search spans by text (lexical search)
114    async fn search_spans(&self, query: &str, limit: usize) -> Result<Vec<Span>>;
115
116    // ========== Vector Search ==========
117
118    /// Get the vector search provider for this backend
119    ///
120    /// For SQLite, this builds/loads an HNSW index.
121    /// For PostgreSQL, this uses pgvector queries.
122    async fn get_vector_search(&self) -> Result<Arc<dyn VectorSearchProvider>>;
123
124    /// Invalidate cached vector index (called after data changes)
125    async fn invalidate_vector_index(&self);
126
127    // ========== Sessions ==========
128
129    /// Create a new session
130    async fn create_session(
131        &self,
132        user_id: Option<&str>,
133        title: Option<&str>,
134    ) -> Result<Session>;
135
136    /// Get session by ID
137    async fn get_session(&self, session_id: &str) -> Result<Option<Session>>;
138
139    /// List sessions with optional filtering
140    async fn list_sessions(
141        &self,
142        user_id: Option<&str>,
143        limit: Option<usize>,
144    ) -> Result<Vec<Session>>;
145
146    /// Update session metadata
147    async fn update_session(
148        &self,
149        session_id: &str,
150        title: Option<&str>,
151        metadata: Option<&serde_json::Value>,
152    ) -> Result<()>;
153
154    /// Delete session and all associated data
155    async fn delete_session(&self, session_id: &str) -> Result<()>;
156
157    // ========== Messages ==========
158
159    /// Add message to session
160    async fn add_message(
161        &self,
162        session_id: &str,
163        role: MessageRole,
164        content: &str,
165        metadata: Option<&serde_json::Value>,
166    ) -> Result<Message>;
167
168    /// Get messages for session
169    async fn get_messages(
170        &self,
171        session_id: &str,
172        limit: Option<usize>,
173    ) -> Result<Vec<Message>>;
174
175    // ========== Working Sets ==========
176
177    /// Associate working set with session
178    async fn associate_working_set(
179        &self,
180        session_id: &str,
181        message_id: Option<&str>,
182        working_set: &WorkingSet,
183        query: &str,
184        config: &CompilerConfig,
185    ) -> Result<SessionWorkingSet>;
186
187    /// Get full session with messages and working sets
188    async fn get_session_full(&self, session_id: &str) -> Result<Option<SessionWithMessages>>;
189
190    // ========== Agents ==========
191
192    /// Register an agent (insert or update)
193    async fn register_agent(&self, agent: &Agent) -> Result<Agent>;
194
195    /// Get agent by ID
196    async fn get_agent(&self, agent_id: &str) -> Result<Option<Agent>>;
197
198    /// Get agent by name
199    async fn get_agent_by_name(&self, name: &str) -> Result<Option<Agent>>;
200
201    /// List all registered agents
202    async fn list_agents(&self) -> Result<Vec<Agent>>;
203
204    // ========== Agent Relations ==========
205
206    /// Add agent relation (agreement, disagreement, etc.)
207    async fn add_agent_relation(
208        &self,
209        session_id: &str,
210        message_id: &str,
211        from_agent_id: &str,
212        target_message_id: &str,
213        stance: Stance,
214    ) -> Result<AgentRelation>;
215
216    /// Get agent relations for session with resolved names
217    async fn get_agent_relations(&self, session_id: &str) -> Result<AgentRelationSummary>;
218
219    /// Get agents participating in session
220    async fn get_session_agents(&self, session_id: &str) -> Result<Vec<Agent>>;
221}