mentedb_extraction/
config.rs

1use serde::{Deserialize, Serialize};
2
3/// Which LLM provider to use for extraction.
4#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
5pub enum LlmProvider {
6    OpenAI,
7    Anthropic,
8    Ollama,
9    /// AWS Bedrock (Anthropic Claude on Bedrock) signed with SigV4. Uses AWS
10    /// credentials from the environment, not an API key. The endpoint is built
11    /// per-request from `ExtractionConfig::region` plus the model id, so
12    /// `default_url` is empty for this provider. Requires the `bedrock` feature.
13    Bedrock,
14    Custom,
15}
16
17impl LlmProvider {
18    /// Default API URL for this provider.
19    ///
20    /// Bedrock returns an empty string: its endpoint is region-specific and is
21    /// built per-request from `ExtractionConfig::region` and the model id
22    /// (`https://bedrock-runtime.{region}.amazonaws.com/model/{model}/invoke`).
23    pub fn default_url(&self) -> &str {
24        match self {
25            LlmProvider::OpenAI => "https://api.openai.com/v1/chat/completions",
26            LlmProvider::Anthropic => "https://api.anthropic.com/v1/messages",
27            LlmProvider::Ollama => "http://localhost:11434/api/chat",
28            // Region-specific; built per-request in call_bedrock.
29            LlmProvider::Bedrock => "",
30            LlmProvider::Custom => "http://localhost:8080/v1/chat/completions",
31        }
32    }
33
34    /// Default extraction model for this provider (cheap, high-volume).
35    pub fn default_model(&self) -> &str {
36        match self {
37            LlmProvider::OpenAI => "gpt-4o-mini",
38            LlmProvider::Anthropic => "claude-haiku-5-5",
39            LlmProvider::Ollama => "llama3",
40            LlmProvider::Bedrock => "us.anthropic.claude-haiku-5-5",
41            LlmProvider::Custom => "default",
42        }
43    }
44
45    /// Default reader model for this provider (smart, low-volume).
46    pub fn default_reader_model(&self) -> &str {
47        match self {
48            LlmProvider::OpenAI => "gpt-4o",
49            LlmProvider::Anthropic => "claude-sonnet-5-5",
50            LlmProvider::Ollama => "llama3",
51            LlmProvider::Bedrock => "us.anthropic.claude-sonnet-5-5",
52            LlmProvider::Custom => "default",
53        }
54    }
55}
56
57/// Resolve the AWS region for Bedrock from the environment, defaulting to
58/// `us-east-1`. Checks `MENTEDB_LLM_REGION` first (MenteDB-specific override),
59/// then the standard `AWS_REGION`.
60pub fn default_bedrock_region() -> String {
61    std::env::var("MENTEDB_LLM_REGION")
62        .or_else(|_| std::env::var("AWS_REGION"))
63        .unwrap_or_else(|_| "us-east-1".to_string())
64}
65
66/// Configuration for the extraction pipeline.
67#[derive(Debug, Clone, Serialize, Deserialize)]
68pub struct ExtractionConfig {
69    /// Which LLM provider to use.
70    pub provider: LlmProvider,
71    /// API key for the provider (not needed for Ollama).
72    pub api_key: Option<String>,
73    /// API endpoint URL (defaults based on provider).
74    ///
75    /// Ignored for the Bedrock provider, whose endpoint is derived from
76    /// `region` and `model` per-request.
77    pub api_url: String,
78    /// Model name to use.
79    pub model: String,
80    /// AWS region for the Bedrock provider (e.g. "us-east-1"). Ignored by all
81    /// other providers. Defaults from `MENTEDB_LLM_REGION`/`AWS_REGION`.
82    pub region: Option<String>,
83    /// Maximum number of memories to extract from a single conversation.
84    pub max_extractions_per_conversation: usize,
85    /// Minimum confidence score for a memory to be accepted (0.0 to 1.0).
86    pub quality_threshold: f32,
87    /// Embedding similarity above which a memory is considered a duplicate (0.0 to 1.0).
88    pub deduplication_threshold: f32,
89    /// Whether to check new memories against existing ones for contradictions.
90    pub enable_contradiction_check: bool,
91    /// Whether to check new memories against existing ones for duplicates.
92    pub enable_deduplication: bool,
93    /// Number of extraction passes (1 = single pass, 2 = first pass + verification).
94    pub extraction_passes: usize,
95    /// Output token ceiling for each LLM call. Claude 5.5 models think before
96    /// answering by default and the thinking counts against this budget, so it
97    /// must leave room for both. Only tokens actually produced are billed.
98    #[serde(default = "default_max_output_tokens")]
99    pub max_output_tokens: u32,
100    /// Reasoning effort for Claude models, sent as `output_config.effort`
101    /// ("low", "medium", "high"). `None` keeps the model default. Lower effort
102    /// thinks less, which is faster and cheaper.
103    #[serde(default)]
104    pub effort: Option<String>,
105    /// Model for judge and rerank calls (reading and comparing memories),
106    /// when it should differ from `model`, which stays the extraction and
107    /// enrichment model. `None` uses `model` for both.
108    #[serde(default)]
109    pub reader_model: Option<String>,
110    /// Whole request timeout for one LLM call, in seconds.
111    #[serde(default = "default_request_timeout_secs")]
112    pub request_timeout_secs: u64,
113    /// Connection timeout for one LLM call, in seconds.
114    #[serde(default = "default_connect_timeout_secs")]
115    pub connect_timeout_secs: u64,
116}
117
118/// Default whole request timeout for an LLM call.
119pub fn default_request_timeout_secs() -> u64 {
120    120
121}
122
123/// Default connection timeout for an LLM call.
124pub fn default_connect_timeout_secs() -> u64 {
125    30
126}
127
128/// Default output token ceiling, sized so thinking plus a full extraction fits.
129pub fn default_max_output_tokens() -> u32 {
130    16_000
131}
132
133impl ExtractionConfig {
134    /// Apply `MENTEDB_LLM_MAX_OUTPUT_TOKENS` and `MENTEDB_LLM_EFFORT` from the
135    /// environment, keeping the current values when they are unset or invalid.
136    pub fn with_output_env(mut self) -> Self {
137        if let Some(n) = std::env::var("MENTEDB_LLM_MAX_OUTPUT_TOKENS")
138            .ok()
139            .and_then(|v| v.parse::<u32>().ok())
140            .filter(|n| *n > 0)
141        {
142            self.max_output_tokens = n;
143        }
144        if let Ok(v) = std::env::var("MENTEDB_LLM_EFFORT") {
145            let v = v.trim().to_lowercase();
146            self.effort = (!v.is_empty()).then_some(v);
147        }
148        if let Some(n) = std::env::var("MENTEDB_LLM_TIMEOUT_SECS")
149            .ok()
150            .and_then(|v| v.parse::<u64>().ok())
151            .filter(|n| *n > 0)
152        {
153            self.request_timeout_secs = n;
154        }
155        self
156    }
157
158    /// The config for judge and rerank calls: this one with `model` set to
159    /// `reader_model` when that is given.
160    pub fn reader_config(&self) -> Self {
161        let mut c = self.clone();
162        if let Some(m) = self.reader_model.as_ref().filter(|m| !m.is_empty()) {
163            c.model = m.clone();
164        }
165        c
166    }
167
168    /// Create a config for OpenAI with the given API key.
169    pub fn openai(api_key: impl Into<String>) -> Self {
170        Self {
171            provider: LlmProvider::OpenAI,
172            api_key: Some(api_key.into()),
173            api_url: LlmProvider::OpenAI.default_url().to_string(),
174            model: LlmProvider::OpenAI.default_model().to_string(),
175            ..Self::default()
176        }
177    }
178
179    /// Create a config for Anthropic with the given API key.
180    pub fn anthropic(api_key: impl Into<String>) -> Self {
181        Self {
182            provider: LlmProvider::Anthropic,
183            api_key: Some(api_key.into()),
184            api_url: LlmProvider::Anthropic.default_url().to_string(),
185            model: LlmProvider::Anthropic.default_model().to_string(),
186            ..Self::default()
187        }
188    }
189
190    /// Create a config for a local Ollama instance.
191    pub fn ollama() -> Self {
192        Self {
193            provider: LlmProvider::Ollama,
194            api_key: None,
195            api_url: LlmProvider::Ollama.default_url().to_string(),
196            model: LlmProvider::Ollama.default_model().to_string(),
197            ..Self::default()
198        }
199    }
200
201    /// Create a config for AWS Bedrock (Anthropic Claude on Bedrock) in the
202    /// given region. Credentials are read from the AWS environment variables at
203    /// call time, not stored here. Requires the `bedrock` feature to actually
204    /// dispatch calls.
205    pub fn bedrock(region: impl Into<String>) -> Self {
206        Self {
207            provider: LlmProvider::Bedrock,
208            api_key: None,
209            api_url: LlmProvider::Bedrock.default_url().to_string(),
210            model: LlmProvider::Bedrock.default_model().to_string(),
211            region: Some(region.into()),
212            ..Self::default()
213        }
214    }
215}
216
217impl Default for ExtractionConfig {
218    fn default() -> Self {
219        Self {
220            provider: LlmProvider::OpenAI,
221            api_key: None,
222            api_url: LlmProvider::OpenAI.default_url().to_string(),
223            model: LlmProvider::OpenAI.default_model().to_string(),
224            region: None,
225            max_extractions_per_conversation: 50,
226            quality_threshold: 0.6,
227            deduplication_threshold: 0.85,
228            enable_contradiction_check: true,
229            enable_deduplication: true,
230            extraction_passes: 1,
231            max_output_tokens: default_max_output_tokens(),
232            effort: None,
233            reader_model: None,
234            request_timeout_secs: default_request_timeout_secs(),
235            connect_timeout_secs: default_connect_timeout_secs(),
236        }
237    }
238}
mentedb_extraction/config.rs

mentedb_extraction/
config.rs
