Skip to main content

distri_types/configuration/
config.rs

1use crate::a2a::{AgentCapabilities, AgentProvider, SecurityScheme};
2use schemars::JsonSchema;
3use serde::{Deserialize, Serialize};
4use std::collections::HashMap;
5use utoipa::ToSchema;
6
7#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, Default, ToSchema)]
8pub struct StoreConfig {
9    /// Metadata store (agent_config, tool_auth) - persistent across sessions
10    #[serde(default)]
11    pub metadata: MetadataStoreConfig,
12
13    /// Memory store configuration (for vector search) - persistent cross-session memory
14    #[serde(default)]
15    pub memory: Option<MemoryStoreConfig>,
16
17    /// Session stores (threads, tasks, scratchpad, session) - ephemeral by default
18    #[serde(default)]
19    pub session: SessionStoreConfig,
20}
21#[derive(
22    Debug, Clone, Serialize, Deserialize, JsonSchema, Default, PartialEq, Eq, Hash, ToSchema,
23)]
24#[serde(tag = "type", content = "config", rename_all = "lowercase")]
25pub enum StoreType {
26    #[default]
27    Sqlite,
28    Postgres,
29    Custom {
30        name: String,
31    },
32}
33
34impl StoreType {
35    pub fn label(&self) -> &str {
36        match self {
37            StoreType::Sqlite => "sqlite",
38            StoreType::Postgres => "postgres",
39            StoreType::Custom { name } => name.as_str(),
40        }
41    }
42}
43
44#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, ToSchema)]
45pub struct MetadataStoreConfig {
46    #[serde(default)]
47    pub store_type: StoreType,
48    #[serde(default)]
49    pub db_config: Option<DbConnectionConfig>,
50}
51
52impl Default for MetadataStoreConfig {
53    fn default() -> Self {
54        Self {
55            store_type: StoreType::Sqlite,
56            db_config: Some(DbConnectionConfig::default()),
57        }
58    }
59}
60
61#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, ToSchema)]
62pub struct MemoryStoreConfig {
63    #[serde(default)]
64    pub store_type: StoreType,
65    #[serde(default)]
66    pub db_config: Option<DbConnectionConfig>,
67    #[serde(default = "default_embedding_dimension")]
68    pub embedding_dimension: usize,
69    #[serde(default = "default_similarity_threshold")]
70    pub similarity_threshold: f32,
71    #[serde(default = "default_max_results")]
72    pub max_results: usize,
73    pub openai_api_key: Option<String>,
74}
75
76impl Default for MemoryStoreConfig {
77    fn default() -> Self {
78        Self {
79            store_type: StoreType::Sqlite,
80            db_config: Some(DbConnectionConfig::default()),
81            embedding_dimension: default_embedding_dimension(),
82            similarity_threshold: default_similarity_threshold(),
83            max_results: default_max_results(),
84            openai_api_key: None,
85        }
86    }
87}
88
89#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, ToSchema)]
90pub struct SessionStoreConfig {
91    /// If true, creates a new ephemeral in-memory SQLite connection for each thread execution
92    /// When ephemeral, store_type and db_config are ignored (always uses in-memory SQLite)
93    #[serde(default = "default_ephemeral")]
94    pub ephemeral: bool,
95    /// Store type (only used when ephemeral=false)
96    #[serde(default)]
97    pub store_type: StoreType,
98    /// Database config (only used when ephemeral=false)
99    #[serde(default)]
100    pub db_config: Option<DbConnectionConfig>,
101}
102
103fn default_ephemeral() -> bool {
104    true
105}
106
107impl Default for SessionStoreConfig {
108    fn default() -> Self {
109        Self {
110            ephemeral: true,
111            store_type: StoreType::Sqlite,
112            db_config: Some(DbConnectionConfig::default()),
113        }
114    }
115}
116
117#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, ToSchema)]
118pub struct PostgresConfig {
119    pub database_url: String,
120    #[serde(default = "default_postgres_max_connections")]
121    pub max_connections: u32,
122    #[serde(default = "default_postgres_min_connections")]
123    pub min_connections: u32,
124    #[serde(default = "default_postgres_connection_timeout")]
125    pub connection_timeout: u64,
126    #[serde(default = "default_postgres_idle_timeout")]
127    pub idle_timeout: u64,
128}
129
130fn default_postgres_max_connections() -> u32 {
131    3
132}
133
134fn default_postgres_min_connections() -> u32 {
135    1
136}
137
138fn default_postgres_connection_timeout() -> u64 {
139    30
140}
141
142fn default_postgres_idle_timeout() -> u64 {
143    600
144}
145
146#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, ToSchema)]
147pub struct DbConnectionConfig {
148    #[serde(default = "default_database_url")]
149    pub database_url: String,
150    #[serde(default = "default_connections")]
151    pub max_connections: u32,
152}
153
154impl Default for DbConnectionConfig {
155    fn default() -> Self {
156        DbConnectionConfig {
157            database_url: default_database_url(),
158            max_connections: default_connections(),
159        }
160    }
161}
162
163fn default_database_url() -> String {
164    // "file::memory:".to_string()
165    ".distri/distri.db".to_string()
166}
167
168fn default_connections() -> u32 {
169    3
170}
171
172#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, ToSchema)]
173#[serde(tag = "type", content = "config", rename_all = "lowercase")]
174pub enum ObjectStorageConfig {
175    /// Local filesystem storage
176    #[serde(rename = "filesystem")]
177    FileSystem {
178        /// Base directory for storing objects
179        base_path: String,
180    },
181
182    /// Google Cloud Storage
183    #[serde(rename = "gcs")]
184    GoogleCloudStorage {
185        bucket: String,
186        project_id: String,
187        /// Path to service account key file
188        service_account_key: Option<String>,
189        /// Base64 encoded service account key
190        service_account_key_base64: Option<String>,
191    },
192
193    /// AWS S3 (future implementation)
194    #[serde(rename = "s3")]
195    S3 {
196        bucket: String,
197        region: String,
198        endpoint: Option<String>,
199        access_key_id: String,
200        secret_access_key: String,
201        path_style: Option<bool>,
202    },
203}
204
205impl Default for ObjectStorageConfig {
206    fn default() -> Self {
207        Self::FileSystem {
208            base_path: "./data/artifacts".to_string(),
209        }
210    }
211}
212
213fn default_embedding_dimension() -> usize {
214    1536
215}
216
217fn default_similarity_threshold() -> f32 {
218    0.7
219}
220
221fn default_max_results() -> usize {
222    10
223}
224
225#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, ToSchema)]
226pub struct ServerConfig {
227    #[serde(default = "default_server_url")]
228    pub base_url: String,
229    #[serde(default)]
230    pub port: Option<u16>,
231    #[serde(default)]
232    pub host: Option<String>,
233    #[serde(default = "default_agent_provider")]
234    #[schema(value_type = Object)]
235    pub agent_provider: AgentProvider,
236    #[serde(default)]
237    pub default_input_modes: Vec<String>,
238    #[serde(default)]
239    pub default_output_modes: Vec<String>,
240    #[serde(default)]
241    #[schema(value_type = Object)]
242    pub security_schemes: HashMap<String, SecurityScheme>,
243    #[serde(default)]
244    pub security: Vec<HashMap<String, Vec<String>>>,
245    #[serde(default = "default_capabilities")]
246    #[schema(value_type = Object)]
247    pub capabilities: AgentCapabilities,
248    #[serde(default = "default_preferred_transport")]
249    pub preferred_transport: Option<String>,
250    #[serde(default = "default_documentation_url")]
251    pub documentation_url: Option<String>,
252}
253
254fn default_capabilities() -> AgentCapabilities {
255    AgentCapabilities {
256        streaming: true,
257        push_notifications: false,
258        state_transition_history: true,
259        extensions: vec![],
260    }
261}
262
263impl Default for ServerConfig {
264    fn default() -> Self {
265        Self {
266            base_url: default_server_url(),
267            agent_provider: default_agent_provider(),
268            default_input_modes: vec![],
269            default_output_modes: vec![],
270            security_schemes: HashMap::new(),
271            security: vec![],
272            capabilities: default_capabilities(),
273            port: None,
274            host: None,
275            preferred_transport: default_preferred_transport(),
276            documentation_url: default_documentation_url(),
277        }
278    }
279}
280
281fn default_agent_provider() -> AgentProvider {
282    AgentProvider {
283        organization: "Distri".to_string(),
284        url: "https://distri.ai".to_string(),
285    }
286}
287
288fn default_server_url() -> String {
289    "http://localhost:8081/v1".to_string()
290}
291
292fn default_documentation_url() -> Option<String> {
293    Some("https://github.com/distrihub/distri".to_string())
294}
295
296fn default_preferred_transport() -> Option<String> {
297    Some("JSONRPC".to_string())
298}
299
300#[derive(Debug, Serialize, Deserialize, Clone, JsonSchema, ToSchema)]
301pub struct ExternalMcpServer {
302    pub name: String,
303    #[serde(default, flatten)]
304    #[schema(value_type = Object)]
305    pub config: crate::McpServerMetadata,
306}