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}
106
107/// Default output token ceiling, sized so thinking plus a full extraction fits.
108pub fn default_max_output_tokens() -> u32 {
109    16_000
110}
111
112impl ExtractionConfig {
113    /// Apply `MENTEDB_LLM_MAX_OUTPUT_TOKENS` and `MENTEDB_LLM_EFFORT` from the
114    /// environment, keeping the current values when they are unset or invalid.
115    pub fn with_output_env(mut self) -> Self {
116        if let Some(n) = std::env::var("MENTEDB_LLM_MAX_OUTPUT_TOKENS")
117            .ok()
118            .and_then(|v| v.parse::<u32>().ok())
119            .filter(|n| *n > 0)
120        {
121            self.max_output_tokens = n;
122        }
123        if let Ok(v) = std::env::var("MENTEDB_LLM_EFFORT") {
124            let v = v.trim().to_lowercase();
125            self.effort = (!v.is_empty()).then_some(v);
126        }
127        self
128    }
129
130    /// Create a config for OpenAI with the given API key.
131    pub fn openai(api_key: impl Into<String>) -> Self {
132        Self {
133            provider: LlmProvider::OpenAI,
134            api_key: Some(api_key.into()),
135            api_url: LlmProvider::OpenAI.default_url().to_string(),
136            model: LlmProvider::OpenAI.default_model().to_string(),
137            ..Self::default()
138        }
139    }
140
141    /// Create a config for Anthropic with the given API key.
142    pub fn anthropic(api_key: impl Into<String>) -> Self {
143        Self {
144            provider: LlmProvider::Anthropic,
145            api_key: Some(api_key.into()),
146            api_url: LlmProvider::Anthropic.default_url().to_string(),
147            model: LlmProvider::Anthropic.default_model().to_string(),
148            ..Self::default()
149        }
150    }
151
152    /// Create a config for a local Ollama instance.
153    pub fn ollama() -> Self {
154        Self {
155            provider: LlmProvider::Ollama,
156            api_key: None,
157            api_url: LlmProvider::Ollama.default_url().to_string(),
158            model: LlmProvider::Ollama.default_model().to_string(),
159            ..Self::default()
160        }
161    }
162
163    /// Create a config for AWS Bedrock (Anthropic Claude on Bedrock) in the
164    /// given region. Credentials are read from the AWS environment variables at
165    /// call time, not stored here. Requires the `bedrock` feature to actually
166    /// dispatch calls.
167    pub fn bedrock(region: impl Into<String>) -> Self {
168        Self {
169            provider: LlmProvider::Bedrock,
170            api_key: None,
171            api_url: LlmProvider::Bedrock.default_url().to_string(),
172            model: LlmProvider::Bedrock.default_model().to_string(),
173            region: Some(region.into()),
174            ..Self::default()
175        }
176    }
177}
178
179impl Default for ExtractionConfig {
180    fn default() -> Self {
181        Self {
182            provider: LlmProvider::OpenAI,
183            api_key: None,
184            api_url: LlmProvider::OpenAI.default_url().to_string(),
185            model: LlmProvider::OpenAI.default_model().to_string(),
186            region: None,
187            max_extractions_per_conversation: 50,
188            quality_threshold: 0.6,
189            deduplication_threshold: 0.85,
190            enable_contradiction_check: true,
191            enable_deduplication: true,
192            extraction_passes: 1,
193            max_output_tokens: default_max_output_tokens(),
194            effort: None,
195        }
196    }
197}
mentedb_extraction/config.rs

mentedb_extraction/
config.rs
