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-4-5",
39            LlmProvider::Ollama => "llama3",
40            LlmProvider::Bedrock => "us.anthropic.claude-haiku-4-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-4-20250514",
50            LlmProvider::Ollama => "llama3",
51            LlmProvider::Bedrock => "us.anthropic.claude-sonnet-4-6",
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}
96
97impl ExtractionConfig {
98    /// Create a config for OpenAI with the given API key.
99    pub fn openai(api_key: impl Into<String>) -> Self {
100        Self {
101            provider: LlmProvider::OpenAI,
102            api_key: Some(api_key.into()),
103            api_url: LlmProvider::OpenAI.default_url().to_string(),
104            model: LlmProvider::OpenAI.default_model().to_string(),
105            ..Self::default()
106        }
107    }
108
109    /// Create a config for Anthropic with the given API key.
110    pub fn anthropic(api_key: impl Into<String>) -> Self {
111        Self {
112            provider: LlmProvider::Anthropic,
113            api_key: Some(api_key.into()),
114            api_url: LlmProvider::Anthropic.default_url().to_string(),
115            model: LlmProvider::Anthropic.default_model().to_string(),
116            ..Self::default()
117        }
118    }
119
120    /// Create a config for a local Ollama instance.
121    pub fn ollama() -> Self {
122        Self {
123            provider: LlmProvider::Ollama,
124            api_key: None,
125            api_url: LlmProvider::Ollama.default_url().to_string(),
126            model: LlmProvider::Ollama.default_model().to_string(),
127            ..Self::default()
128        }
129    }
130
131    /// Create a config for AWS Bedrock (Anthropic Claude on Bedrock) in the
132    /// given region. Credentials are read from the AWS environment variables at
133    /// call time, not stored here. Requires the `bedrock` feature to actually
134    /// dispatch calls.
135    pub fn bedrock(region: impl Into<String>) -> Self {
136        Self {
137            provider: LlmProvider::Bedrock,
138            api_key: None,
139            api_url: LlmProvider::Bedrock.default_url().to_string(),
140            model: LlmProvider::Bedrock.default_model().to_string(),
141            region: Some(region.into()),
142            ..Self::default()
143        }
144    }
145}
146
147impl Default for ExtractionConfig {
148    fn default() -> Self {
149        Self {
150            provider: LlmProvider::OpenAI,
151            api_key: None,
152            api_url: LlmProvider::OpenAI.default_url().to_string(),
153            model: LlmProvider::OpenAI.default_model().to_string(),
154            region: None,
155            max_extractions_per_conversation: 50,
156            quality_threshold: 0.6,
157            deduplication_threshold: 0.85,
158            enable_contradiction_check: true,
159            enable_deduplication: true,
160            extraction_passes: 1,
161        }
162    }
163}
mentedb_extraction/config.rs

mentedb_extraction/
config.rs
