Skip to main content

locus_mcp/
composition.rs

1use std::sync::Arc;
2
3use anyhow::Result;
4use locus_core_rs::domain::contracts::EmbeddingProvider;
5use locus_core_rs::{ParseProfile, SurrealDbSettings};
6#[cfg(feature = "local-embedding")]
7use locus_sdk::infrastructure::embeddings::LocalEmbeddingProvider;
8use locus_sdk::infrastructure::embeddings::OllamaEmbeddingProvider;
9use tracing::info;
10
11pub(crate) use locus_surreal_adapter::RuntimeSurrealDbClient;
12
13#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14pub(crate) enum McpStorage {
15    InMemory,
16    Sqlite,
17    Surreal,
18}
19
20pub(crate) fn resolve_storage(args: &[String]) -> McpStorage {
21    let forced_memory = env_flag("LOCUS_MCP_IN_MEMORY")
22        || args
23            .iter()
24            .any(|arg| arg.eq_ignore_ascii_case("--in-memory"));
25    if forced_memory {
26        return McpStorage::InMemory;
27    }
28
29    match env_or_arg("LOCUS_MCP_STORAGE", args, "--storage")
30        .map(|value| value.to_ascii_lowercase())
31        .as_deref()
32    {
33        Some("sqlite") => McpStorage::Sqlite,
34        Some("inmemory") | Some("in-memory") | Some("memory") => McpStorage::InMemory,
35        _ => McpStorage::Surreal,
36    }
37}
38
39pub(crate) fn sqlite_path(args: &[String]) -> String {
40    env_or_arg("LOCUS_MCP_SQLITE_PATH", args, "--sqlite-path")
41        .unwrap_or_else(|| "data/locus-mcp.sqlite".to_string())
42}
43
44#[derive(Debug, Clone)]
45enum EmbeddingsProviderKind {
46    Ollama,
47    #[cfg(feature = "local-embedding")]
48    Local,
49}
50
51impl EmbeddingsProviderKind {
52    fn parse(value: &str) -> Option<Self> {
53        match value.trim().to_ascii_lowercase().as_str() {
54            "ollama" => Some(Self::Ollama),
55            #[cfg(feature = "local-embedding")]
56            "local" | "local-embedding" | "candle" => Some(Self::Local),
57            _ => None,
58        }
59    }
60}
61
62pub(crate) fn init_logging() {
63    let _ = tracing_subscriber::fmt()
64        .with_writer(std::io::stderr)
65        .with_env_filter(
66            tracing_subscriber::EnvFilter::try_from_default_env()
67                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
68        )
69        .try_init();
70}
71
72pub(crate) fn load_surreal_settings(args: &[String]) -> Result<SurrealDbSettings> {
73    let mut settings = SurrealDbSettings::default();
74    settings.endpoints.embedded = Some("surrealkv://data/locus-mcp".to_string());
75    settings.database = "locus_mcp".to_string();
76
77    if let Some(value) = env_or_arg(
78        "LOCUS_MCP_SURREAL_REMOTE_ENDPOINT",
79        args,
80        "--remote-endpoint",
81    ) {
82        settings.endpoints.remote = Some(value);
83    }
84    if let Some(value) = env_or_arg(
85        "LOCUS_MCP_SURREAL_EMBEDDED_ENDPOINT",
86        args,
87        "--embedded-endpoint",
88    ) {
89        settings.endpoints.embedded = Some(value);
90    }
91    if let Some(value) = env_or_arg("LOCUS_MCP_SURREAL_ENDPOINT", args, "--endpoint") {
92        settings.endpoints.remote = Some(value.clone());
93        settings.endpoints.embedded = Some(value);
94    }
95    if let Some(value) = env_or_arg("LOCUS_MCP_SURREAL_NAMESPACE", args, "--namespace") {
96        settings.namespace = value;
97    }
98    if let Some(value) = env_or_arg("LOCUS_MCP_SURREAL_DATABASE", args, "--database") {
99        settings.database = value;
100    }
101    if let Some(value) = env_or_arg("LOCUS_MCP_SURREAL_USERNAME", args, "--username") {
102        settings.user = Some(value);
103    }
104    if let Some(value) = env_or_arg("LOCUS_MCP_SURREAL_PASSWORD", args, "--password") {
105        settings.password = Some(value);
106    }
107
108    Ok(settings)
109}
110
111pub(crate) fn runtime_args(args: &[String]) -> Vec<String> {
112    let mut runtime_args = args.to_vec();
113    if env_flag("LOCUS_MCP_REMOTE") && !runtime_args.iter().any(|value| value == "--remote") {
114        runtime_args.push("--remote".to_string());
115    }
116    runtime_args
117}
118
119pub(crate) fn build_embedding_provider(
120    args: &[String],
121) -> Result<Option<Arc<dyn EmbeddingProvider>>> {
122    let embeddings_enabled = env_flag("LOCUS_MCP_EMBEDDINGS_ENABLED")
123        || args
124            .iter()
125            .any(|arg| arg.eq_ignore_ascii_case("--embeddings-enabled"));
126
127    if !embeddings_enabled {
128        return Ok(None);
129    }
130
131    let provider_kind_raw = env_or_arg(
132        "LOCUS_MCP_EMBEDDINGS_PROVIDER",
133        args,
134        "--embeddings-provider",
135    )
136    .unwrap_or_else(|| "ollama".to_string());
137    let provider_kind = EmbeddingsProviderKind::parse(&provider_kind_raw).ok_or_else(|| {
138        anyhow::anyhow!(
139            "unsupported embeddings provider '{}'; expected 'ollama'{}",
140            provider_kind_raw,
141            if cfg!(feature = "local-embedding") {
142                " or 'local'"
143            } else {
144                ""
145            }
146        )
147    })?;
148
149    let endpoint = env_or_arg(
150        "LOCUS_MCP_EMBEDDINGS_ENDPOINT",
151        args,
152        "--embeddings-endpoint",
153    )
154    .unwrap_or_else(|| "http://127.0.0.1:11434/api/embeddings".to_string());
155    let model = env_or_arg("LOCUS_MCP_EMBEDDINGS_MODEL", args, "--embeddings-model")
156        .unwrap_or_else(|| "sttp-encoder".to_string());
157    #[cfg(feature = "local-embedding")]
158    let repo = env_or_arg("LOCUS_MCP_EMBEDDINGS_REPO", args, "--embeddings-repo")
159        .unwrap_or_else(|| "sentence-transformers/all-MiniLM-L6-v2".to_string());
160
161    let provider: Arc<dyn EmbeddingProvider> = match provider_kind {
162        EmbeddingsProviderKind::Ollama => {
163            info!(
164                provider = "ollama",
165                endpoint = %endpoint,
166                model = %model,
167                "auto-embedding enabled for store_context"
168            );
169            Arc::new(OllamaEmbeddingProvider::new(endpoint, model))
170        }
171        #[cfg(feature = "local-embedding")]
172        EmbeddingsProviderKind::Local => {
173            info!(
174                provider = "local",
175                model = %model,
176                repo = %repo,
177                "auto-embedding enabled for store_context"
178            );
179            Arc::new(LocalEmbeddingProvider::new(model, repo)?)
180        }
181    };
182
183    Ok(Some(provider))
184}
185
186pub(crate) fn resolve_parser_profile(args: &[String]) -> Result<ParseProfile> {
187    let raw = env_or_arg("LOCUS_MCP_PARSE_PROFILE", args, "--parse-profile")
188        .unwrap_or_else(|| "strict_typed_ir".to_string());
189
190    parse_profile(raw.as_str()).ok_or_else(|| {
191        anyhow::anyhow!(
192            "unsupported parse profile '{}'; expected one of: strict_typed_ir, strict, tolerant",
193            raw
194        )
195    })
196}
197
198fn parse_profile(value: &str) -> Option<ParseProfile> {
199    match value.trim().to_ascii_lowercase().as_str() {
200        "strict_typed_ir" | "strict-typed-ir" | "stricttypedir" | "typed_ir" | "typed-ir" => {
201            Some(ParseProfile::StrictTypedIr)
202        }
203        "strict" => Some(ParseProfile::Strict),
204        "tolerant" | "default" => Some(ParseProfile::Tolerant),
205        _ => None,
206    }
207}
208
209fn env_or_arg(env_key: &str, args: &[String], arg_name: &str) -> Option<String> {
210    if let Ok(value) = std::env::var(env_key) {
211        let trimmed = value.trim();
212        if !trimmed.is_empty() {
213            return Some(trimmed.to_string());
214        }
215    }
216
217    arg_value(args, arg_name)
218}
219
220fn arg_value(args: &[String], key: &str) -> Option<String> {
221    args.windows(2)
222        .find(|window| window[0].eq_ignore_ascii_case(key))
223        .map(|window| window[1].clone())
224}
225
226fn env_flag(key: &str) -> bool {
227    std::env::var(key)
228        .map(|value| {
229            let normalized = value.trim().to_ascii_lowercase();
230            normalized == "1" || normalized == "true" || normalized == "yes"
231        })
232        .unwrap_or(false)
233}