1use std::sync::Arc;
2
3use anyhow::{Result, anyhow};
4use axum::http::HeaderValue;
5use locus_core_rs::application::services::{
6 CalibrationService, ContextQueryService, MonthlyRollupService, MoodCatalogService,
7 RekeyScopeService, StoreContextService,
8};
9use locus_core_rs::application::validation::TreeSitterValidator;
10use locus_core_rs::domain::contracts::{
11 EmbeddingProvider, NodeStore, NodeStoreInitializer, NodeValidator, SemanticIndexStore,
12 SemanticIndexStoreInitializer,
13};
14use locus_core_rs::parsing::SttpNodeParser;
15use locus_core_rs::storage::{
16 InMemoryNodeStore, InMemorySemanticIndexStore, SurrealDbEndpointsSettings, SurrealDbNodeStore,
17 SurrealDbRuntimeOptions, SurrealDbSemanticIndexStore, SurrealDbSettings,
18};
19#[cfg(feature = "local-embedding")]
20use locus_sdk::infrastructure::embeddings::LocalEmbeddingProvider;
21use locus_sdk::infrastructure::embeddings::OllamaEmbeddingProvider;
22use locus_sqlite_adapter::{SqliteDatabase, SqliteNodeStore, SqliteSemanticIndexStore};
23use locus_surreal_adapter::RuntimeSurrealDbClient;
24use tracing::{error, info};
25
26use crate::app_state::AppState;
27use crate::gateway_args::{EmbeddingsProviderKind, GatewayArgs, GatewayBackend};
28use crate::http_models::CorsAllowedOrigins;
29use crate::providers::{AvecScorer, OllamaAvecScorer};
30
31pub(crate) async fn build_state(args: &GatewayArgs) -> Result<AppState> {
32 build_state_with_backend(&args.backend, Some(args)).await
33}
34
35#[cfg(test)]
36pub(crate) async fn build_in_memory_state() -> Result<AppState> {
37 build_in_memory_state_with_args(None).await
38}
39
40pub(crate) fn parse_cors_allowed_origins(value: &str) -> Result<CorsAllowedOrigins> {
41 let trimmed = value.trim();
42 if trimmed.is_empty() {
43 return Err(anyhow!(
44 "CORS allowed origins cannot be empty when CORS is enabled"
45 ));
46 }
47
48 if trimmed == "*" {
49 return Ok(CorsAllowedOrigins::Any);
50 }
51
52 let mut origins = Vec::new();
53 for origin in trimmed
54 .split(',')
55 .map(str::trim)
56 .filter(|part| !part.is_empty())
57 {
58 let header = HeaderValue::from_str(origin)
59 .map_err(|_| anyhow!("Invalid CORS origin value: {origin}"))?;
60 origins.push(header);
61 }
62
63 if origins.is_empty() {
64 return Err(anyhow!(
65 "CORS allowed origins must include at least one origin or '*'"
66 ));
67 }
68
69 Ok(CorsAllowedOrigins::Explicit(origins))
70}
71
72pub(crate) async fn shutdown_signal() {
73 if let Err(err) = tokio::signal::ctrl_c().await {
74 error!(error = %err, "Failed waiting for ctrl_c signal");
75 }
76}
77
78async fn build_state_with_backend(
79 backend: &GatewayBackend,
80 options: Option<&GatewayArgs>,
81) -> Result<AppState> {
82 match backend {
83 GatewayBackend::InMemory => build_in_memory_state_with_args(options).await,
84 GatewayBackend::Sqlite => {
85 let options = options.ok_or_else(|| {
86 anyhow!("SQLite backend selected, but no gateway runtime options were provided.")
87 })?;
88 build_sqlite_state(options).await
89 }
90 GatewayBackend::Surreal => {
91 let options = options.ok_or_else(|| {
92 anyhow!("Surreal backend selected, but no gateway runtime options were provided.")
93 })?;
94 build_surreal_state(options).await
95 }
96 }
97}
98
99async fn build_in_memory_state_with_args(args: Option<&GatewayArgs>) -> Result<AppState> {
100 let store = Arc::new(InMemoryNodeStore::new());
101 let semantic_index = Arc::new(InMemorySemanticIndexStore::new());
102
103 let initializer: Arc<dyn NodeStoreInitializer> = store.clone();
104 initializer.initialize_async().await?;
105
106 let semantic_initializer: Arc<dyn SemanticIndexStoreInitializer> = semantic_index.clone();
107 semantic_initializer.initialize_async().await?;
108
109 let store_trait: Arc<dyn NodeStore> = store;
110 let semantic_trait: Arc<dyn SemanticIndexStore> = semantic_index;
111 let validator: Arc<dyn NodeValidator> = Arc::new(TreeSitterValidator);
112 let embedding_provider = build_embedding_provider(args)?;
113 let avec_scorer = build_avec_scorer(args);
114
115 Ok(build_services(
116 store_trait,
117 semantic_trait,
118 validator,
119 embedding_provider,
120 avec_scorer,
121 ))
122}
123
124fn build_services(
125 store_trait: Arc<dyn NodeStore>,
126 semantic_index: Arc<dyn SemanticIndexStore>,
127 validator: Arc<dyn NodeValidator>,
128 embedding_provider: Option<Arc<dyn EmbeddingProvider>>,
129 avec_scorer: Option<Arc<dyn AvecScorer>>,
130) -> AppState {
131 let parser = SttpNodeParser::new();
132 let store_context = match embedding_provider.as_ref() {
133 Some(provider) => Arc::new(
134 StoreContextService::with_embedding_provider(
135 store_trait.clone(),
136 validator.clone(),
137 provider.clone(),
138 parser,
139 )
140 .with_semantic_index(semantic_index.clone()),
141 ),
142 None => Arc::new(
143 StoreContextService::new(
144 store_trait.clone(),
145 validator.clone(),
146 SttpNodeParser::new(),
147 )
148 .with_semantic_index(semantic_index.clone()),
149 ),
150 };
151
152 let mut monthly_rollup = MonthlyRollupService::new(store_trait.clone(), validator.clone())
153 .with_semantic_index(semantic_index.clone());
154 if let Some(provider) = embedding_provider.as_ref() {
155 monthly_rollup = monthly_rollup.with_embedding_provider(provider.clone());
156 }
157
158 AppState {
159 node_store: store_trait.clone(),
160 semantic_index,
161 embedding_provider: embedding_provider.clone(),
162 avec_scorer,
163 calibration: Arc::new(CalibrationService::new(store_trait.clone())),
164 context_query: Arc::new(ContextQueryService::new(store_trait.clone())),
165 mood_catalog: Arc::new(MoodCatalogService::new()),
166 store_context,
167 monthly_rollup: Arc::new(monthly_rollup),
168 rekey_scope: Arc::new(RekeyScopeService::new(store_trait)),
169 }
170}
171
172async fn build_surreal_state(args: &GatewayArgs) -> Result<AppState> {
173 let mut settings = SurrealDbSettings::default();
174 settings.endpoints = SurrealDbEndpointsSettings {
175 embedded: args
176 .surreal_embedded_endpoint
177 .clone()
178 .or(settings.endpoints.embedded),
179 remote: args
180 .surreal_remote_endpoint
181 .clone()
182 .or(settings.endpoints.remote),
183 };
184 settings.namespace = args.surreal_namespace.clone();
185 settings.database = args.surreal_database.clone();
186 settings.user = Some(args.surreal_user.clone());
187 settings.password = Some(args.surreal_password.clone());
188
189 let mut runtime_args = Vec::new();
190 if args.remote {
191 runtime_args.push("--remote".to_string());
192 }
193
194 let runtime = SurrealDbRuntimeOptions::from_args(
195 &runtime_args,
196 &settings,
197 Some(args.root_dir_name.as_str()),
198 )?;
199
200 info!(
201 backend = "surreal",
202 root_dir = runtime.root_dir,
203 mode = if runtime.use_remote {
204 "remote"
205 } else {
206 "embedded"
207 },
208 endpoint = runtime.endpoint,
209 namespace = runtime.namespace,
210 database = runtime.database,
211 "Surreal backend requested"
212 );
213
214 let client = Arc::new(
215 RuntimeSurrealDbClient::connect(
216 &runtime,
217 settings.user.as_deref(),
218 settings.password.as_deref(),
219 )
220 .await?,
221 );
222
223 let semantic_index = Arc::new(SurrealDbSemanticIndexStore::new(client.clone()));
224 let store = Arc::new(SurrealDbNodeStore::new(client));
225
226 let initializer: Arc<dyn NodeStoreInitializer> = store.clone();
227 initializer.initialize_async().await?;
228
229 let semantic_initializer: Arc<dyn SemanticIndexStoreInitializer> = semantic_index.clone();
230 semantic_initializer.initialize_async().await?;
231
232 let store_trait: Arc<dyn NodeStore> = store;
233 let semantic_trait: Arc<dyn SemanticIndexStore> = semantic_index;
234 let validator: Arc<dyn NodeValidator> = Arc::new(TreeSitterValidator);
235 let embedding_provider = build_embedding_provider(Some(args))?;
236 let avec_scorer = build_avec_scorer(Some(args));
237
238 Ok(build_services(
239 store_trait,
240 semantic_trait,
241 validator,
242 embedding_provider,
243 avec_scorer,
244 ))
245}
246
247async fn build_sqlite_state(args: &GatewayArgs) -> Result<AppState> {
248 info!(backend = "sqlite", path = %args.sqlite_path, "SQLite backend requested");
249 let db = Arc::new(SqliteDatabase::open(&args.sqlite_path)?);
250 let semantic_index = Arc::new(SqliteSemanticIndexStore::new(db.clone()));
251 let store = Arc::new(SqliteNodeStore::new(db));
252
253 let initializer: Arc<dyn NodeStoreInitializer> = store.clone();
254 initializer.initialize_async().await?;
255
256 let semantic_initializer: Arc<dyn SemanticIndexStoreInitializer> = semantic_index.clone();
257 semantic_initializer.initialize_async().await?;
258
259 let store_trait: Arc<dyn NodeStore> = store;
260 let semantic_trait: Arc<dyn SemanticIndexStore> = semantic_index;
261 let validator: Arc<dyn NodeValidator> = Arc::new(TreeSitterValidator);
262 let embedding_provider = build_embedding_provider(Some(args))?;
263 let avec_scorer = build_avec_scorer(Some(args));
264
265 Ok(build_services(
266 store_trait,
267 semantic_trait,
268 validator,
269 embedding_provider,
270 avec_scorer,
271 ))
272}
273
274fn build_avec_scorer(args: Option<&GatewayArgs>) -> Option<Arc<dyn AvecScorer>> {
275 let args = args?;
276 if !args.avec_scoring_enabled {
277 return None;
278 }
279
280 Some(Arc::new(OllamaAvecScorer::new(
281 args.avec_scoring_endpoint.clone(),
282 args.avec_scoring_model.clone(),
283 )))
284}
285
286fn build_embedding_provider(
287 args: Option<&GatewayArgs>,
288) -> Result<Option<Arc<dyn EmbeddingProvider>>> {
289 let Some(args) = args else {
290 return Ok(None);
291 };
292
293 if !args.embeddings_enabled {
294 return Ok(None);
295 }
296
297 let provider: Arc<dyn EmbeddingProvider> = match args.embeddings_provider {
298 EmbeddingsProviderKind::Ollama => Arc::new(OllamaEmbeddingProvider::new(
299 args.embeddings_endpoint.clone(),
300 args.embeddings_model.clone(),
301 )),
302 #[cfg(feature = "local-embedding")]
303 EmbeddingsProviderKind::Local => Arc::new(LocalEmbeddingProvider::new(
304 args.embeddings_model.clone(),
305 args.embeddings_repo.clone(),
306 )?),
307 };
308
309 Ok(Some(provider))
310}