mentedb_extraction/
provider.rs

1use crate::config::{ExtractionConfig, LlmProvider};
2use crate::error::ExtractionError;
3
4/// Classify an HTTP error response into a specific ExtractionError variant.
5fn classify_api_error(
6    status: reqwest::StatusCode,
7    body: &str,
8    provider: &str,
9    model: &str,
10) -> ExtractionError {
11    let code = status.as_u16();
12    match code {
13        401 => ExtractionError::AuthError(format!(
14            "{provider} returned 401 Unauthorized. Check your API key (MENTEDB_LLM_API_KEY). \
15             Current provider: {provider}, model: {model}"
16        )),
17        403 => ExtractionError::AuthError(format!(
18            "{provider} returned 403 Forbidden. Your API key may lack permissions for model '{model}'."
19        )),
20        404 => ExtractionError::ModelNotFound(format!(
21            "{provider} returned 404. Model '{model}' may not exist or is not available on your account."
22        )),
23        _ => ExtractionError::ProviderError(format!("{provider} API returned {status}: {body}")),
24    }
25}
26
27/// Trait for LLM providers that can extract memories from conversation text.
28pub trait ExtractionProvider: Send + Sync {
29    /// Send a conversation to the LLM with the given system prompt and return
30    /// the raw response text (expected to be JSON).
31    fn extract(
32        &self,
33        conversation: &str,
34        system_prompt: &str,
35    ) -> impl std::future::Future<Output = Result<String, ExtractionError>> + Send;
36}
37
38/// HTTP-based extraction provider that calls OpenAI, Anthropic, or Ollama APIs.
39pub struct HttpExtractionProvider {
40    client: reqwest::Client,
41    config: ExtractionConfig,
42}
43
44impl HttpExtractionProvider {
45    pub fn new(config: ExtractionConfig) -> Result<Self, ExtractionError> {
46        if config.provider != LlmProvider::Ollama && config.api_key.is_none() {
47            return Err(ExtractionError::ConfigError(
48                "API key is required for this provider".to_string(),
49            ));
50        }
51        let client = reqwest::Client::builder()
52            .timeout(std::time::Duration::from_secs(120))
53            .connect_timeout(std::time::Duration::from_secs(30))
54            .build()
55            .map_err(|e| ExtractionError::ConfigError(format!("HTTP client error: {}", e)))?;
56        Ok(Self { client, config })
57    }
58
59    /// Expand a search query into multiple sub-queries via LLM.
60    ///
61    /// Given a natural language question, identifies the expected answer type
62    /// and extracts 2-3 targeted search queries. The first line of the response
63    /// is the answer type (PLACE, DATE, NUMBER, NAME, PERSON, BRAND, etc.),
64    /// followed by the search queries.
65    ///
66    /// For counting/aggregation/comparison queries, also generates comprehensive
67    /// category synonyms for exhaustive BM25 sweep.
68    pub async fn expand_query(&self, query: &str) -> Result<Vec<String>, ExtractionError> {
69        let system_prompt = "You help search a memory database. Given a question, return a JSON object with:\n\
70            - \"answer_type\": one of PLACE, DATE, TIME, NUMBER, NAME, PERSON, BRAND, ITEM, ACTIVITY, COUNTING, OTHER\n\
71            - \"queries\": array of 2-3 short search queries\n\
72            - For COUNTING only, also include:\n\
73              - \"item_keywords\": comma-separated specific subtypes/instances that would be individually counted\n\
74              - \"broad_keywords\": comma-separated category terms, action verbs, and general synonyms\n\n\
75            Use COUNTING when the question requires COMPLETENESS — counting, listing, aggregating, totaling, \
76            or comparing to find a superlative (most, least, best, worst, first, last, biggest, highest, lowest).\n\n\
77            The distinction matters:\n\
78            - item_keywords: specific things you would COUNT (types of the thing being asked about)\n\
79            - broad_keywords: general terms that help FIND memories but aren't counted themselves\n\n\
80            Examples:\n\
81            Q: \"Where do I take yoga classes?\"\n\
82            {\"answer_type\": \"PLACE\", \"queries\": [\"yoga studio name\", \"yoga class location\"]}\n\n\
83            Q: \"How many doctors did I visit?\"\n\
84            {\"answer_type\": \"COUNTING\", \"queries\": [\"doctor visits appointments\", \"medical specialist visits\"], \
85            \"item_keywords\": \"doctor, Dr., physician, specialist, dermatologist, cardiologist, dentist, surgeon, pediatrician, orthopedist, ophthalmologist\", \
86            \"broad_keywords\": \"medical, clinic, appointment, visit, diagnosed, prescribed, referred, checkup, exam\"}\n\n\
87            Q: \"Which platform did I gain the most followers on?\"\n\
88            {\"answer_type\": \"COUNTING\", \"queries\": [\"social media follower growth\", \"follower count increase\"], \
89            \"item_keywords\": \"TikTok, Instagram, Twitter, YouTube, Facebook, LinkedIn, Snapchat, Reddit, Twitch\", \
90            \"broad_keywords\": \"followers, follower count, gained, growth, platform, social media, increase, jumped, grew\"}";
91        let result = self.call_with_retry(query, system_prompt).await?;
92
93        // Parse JSON response (call_openai forces json_object response format)
94        let mut lines: Vec<String> = Vec::new();
95        let cleaned = result
96            .trim()
97            .trim_start_matches("```json")
98            .trim_end_matches("```")
99            .trim();
100        if let Ok(json) = serde_json::from_str::<serde_json::Value>(cleaned) {
101            if let Some(answer_type) = json.get("answer_type").and_then(|v| v.as_str()) {
102                lines.push(answer_type.to_string());
103            }
104            if let Some(queries) = json.get("queries").and_then(|v| v.as_array()) {
105                for q in queries {
106                    if let Some(s) = q.as_str() {
107                        lines.push(s.to_string());
108                    }
109                }
110            }
111            if let Some(item_kw) = json.get("item_keywords").and_then(|v| v.as_str()) {
112                lines.push(format!("ITEM_KEYWORDS: {}", item_kw));
113            }
114            if let Some(broad_kw) = json.get("broad_keywords").and_then(|v| v.as_str()) {
115                lines.push(format!("BROAD_KEYWORDS: {}", broad_kw));
116            }
117            // Fallback: old single "keywords" field → treat all as item keywords
118            if let Some(keywords) = json.get("keywords").and_then(|v| v.as_str())
119                && json.get("item_keywords").is_none()
120            {
121                lines.push(format!("ITEM_KEYWORDS: {}", keywords));
122            }
123        } else {
124            // Fallback: parse as plain text lines
125            lines = result
126                .lines()
127                .map(|l| l.trim().to_string())
128                .filter(|l| !l.is_empty())
129                .collect();
130        }
131        if std::env::var("MENTEDB_DEBUG").is_ok() {
132            eprintln!("[expand_query] input={:?} parsed={:?}", query, lines);
133        }
134        Ok(lines)
135    }
136
137    async fn call_openai(
138        &self,
139        conversation: &str,
140        system_prompt: &str,
141    ) -> Result<String, ExtractionError> {
142        let body = serde_json::json!({
143            "model": self.config.model,
144            "temperature": 0,
145            "response_format": { "type": "json_object" },
146            "messages": [
147                { "role": "system", "content": system_prompt },
148                { "role": "user", "content": conversation }
149            ]
150        });
151
152        let api_key = self.config.api_key.as_deref().unwrap_or_default();
153
154        let resp = self
155            .client
156            .post(&self.config.api_url)
157            .header("Authorization", format!("Bearer {api_key}"))
158            .header("Content-Type", "application/json")
159            .json(&body)
160            .send()
161            .await?;
162
163        let status = resp.status();
164        let text = resp.text().await?;
165
166        if !status.is_success() {
167            return Err(classify_api_error(
168                status,
169                &text,
170                "OpenAI",
171                &self.config.model,
172            ));
173        }
174
175        let parsed: serde_json::Value = serde_json::from_str(&text)?;
176        parsed["choices"][0]["message"]["content"]
177            .as_str()
178            .map(|s| s.to_string())
179            .ok_or_else(|| {
180                ExtractionError::ParseError("Missing content in OpenAI response".to_string())
181            })
182    }
183
184    /// OpenAI call without forced JSON response format.
185    /// Used for plain text outputs (synthesis, re-ranking, key noun extraction).
186    async fn call_openai_text(
187        &self,
188        conversation: &str,
189        system_prompt: &str,
190    ) -> Result<String, ExtractionError> {
191        let body = serde_json::json!({
192            "model": self.config.model,
193            "temperature": 0,
194            "messages": [
195                { "role": "system", "content": system_prompt },
196                { "role": "user", "content": conversation }
197            ]
198        });
199
200        let api_key = self.config.api_key.as_deref().unwrap_or_default();
201
202        let resp = self
203            .client
204            .post(&self.config.api_url)
205            .header("Authorization", format!("Bearer {api_key}"))
206            .header("Content-Type", "application/json")
207            .json(&body)
208            .send()
209            .await?;
210
211        let status = resp.status();
212        let text = resp.text().await?;
213
214        if !status.is_success() {
215            return Err(classify_api_error(
216                status,
217                &text,
218                "OpenAI",
219                &self.config.model,
220            ));
221        }
222
223        let parsed: serde_json::Value = serde_json::from_str(&text)?;
224        parsed["choices"][0]["message"]["content"]
225            .as_str()
226            .map(|s| s.to_string())
227            .ok_or_else(|| {
228                ExtractionError::ParseError("Missing content in OpenAI response".to_string())
229            })
230    }
231
232    async fn call_anthropic(
233        &self,
234        conversation: &str,
235        system_prompt: &str,
236    ) -> Result<String, ExtractionError> {
237        let body = serde_json::json!({
238            "model": self.config.model,
239            "max_tokens": 4096,
240            "temperature": 0,
241            "system": system_prompt,
242            "messages": [
243                { "role": "user", "content": conversation }
244            ]
245        });
246
247        let api_key = self.config.api_key.as_deref().unwrap_or_default();
248
249        let resp = self
250            .client
251            .post(&self.config.api_url)
252            .header("x-api-key", api_key)
253            .header("anthropic-version", "2023-06-01")
254            .header("Content-Type", "application/json")
255            .json(&body)
256            .send()
257            .await?;
258
259        let status = resp.status();
260        let text = resp.text().await?;
261
262        if !status.is_success() {
263            return Err(classify_api_error(
264                status,
265                &text,
266                "Anthropic",
267                &self.config.model,
268            ));
269        }
270
271        let parsed: serde_json::Value = serde_json::from_str(&text)?;
272
273        // Anthropic may return multiple content blocks; find the first text block
274        let content_text = parsed["content"]
275            .as_array()
276            .and_then(|blocks| {
277                blocks.iter().find_map(|block| {
278                    if block["type"].as_str() == Some("text") {
279                        block["text"].as_str().map(|s| s.to_string())
280                    } else {
281                        None
282                    }
283                })
284            })
285            .or_else(|| {
286                // Fallback: try the old path for backwards compat
287                parsed["content"][0]["text"].as_str().map(|s| s.to_string())
288            });
289
290        match content_text {
291            Some(t) if !t.trim().is_empty() => Ok(t),
292            Some(_) => {
293                tracing::warn!(
294                    model = %self.config.model,
295                    "Anthropic returned empty text content"
296                );
297                Ok("{\"memories\": []}".to_string())
298            }
299            None => {
300                tracing::warn!(
301                    model = %self.config.model,
302                    response_preview = &text[..text.len().min(300)],
303                    "No text block found in Anthropic response"
304                );
305                Ok("{\"memories\": []}".to_string())
306            }
307        }
308    }
309
310    async fn call_ollama(
311        &self,
312        conversation: &str,
313        system_prompt: &str,
314    ) -> Result<String, ExtractionError> {
315        let body = serde_json::json!({
316            "model": self.config.model,
317            "stream": false,
318            "format": "json",
319            "messages": [
320                { "role": "system", "content": system_prompt },
321                { "role": "user", "content": conversation }
322            ]
323        });
324
325        let resp = self
326            .client
327            .post(&self.config.api_url)
328            .header("Content-Type", "application/json")
329            .json(&body)
330            .send()
331            .await?;
332
333        let status = resp.status();
334        let text = resp.text().await?;
335
336        if !status.is_success() {
337            return Err(classify_api_error(
338                status,
339                &text,
340                "Ollama",
341                &self.config.model,
342            ));
343        }
344
345        let parsed: serde_json::Value = serde_json::from_str(&text)?;
346        parsed["message"]["content"]
347            .as_str()
348            .map(|s| s.to_string())
349            .ok_or_else(|| {
350                ExtractionError::ParseError("Missing content in Ollama response".to_string())
351            })
352    }
353
354    /// Execute a request with retry logic for rate limits (HTTP 429).
355    /// Uses exponential backoff: 1s, 2s, 4s.
356    pub async fn call_with_retry(
357        &self,
358        conversation: &str,
359        system_prompt: &str,
360    ) -> Result<String, ExtractionError> {
361        self.call_with_retry_inner(conversation, system_prompt, true)
362            .await
363    }
364
365    /// Like call_with_retry but without forcing JSON response format.
366    /// Use for prompts that expect plain text output (synthesis, re-ranking, etc).
367    pub async fn call_text_with_retry(
368        &self,
369        conversation: &str,
370        system_prompt: &str,
371    ) -> Result<String, ExtractionError> {
372        self.call_with_retry_inner(conversation, system_prompt, false)
373            .await
374    }
375
376    async fn call_with_retry_inner(
377        &self,
378        conversation: &str,
379        system_prompt: &str,
380        force_json: bool,
381    ) -> Result<String, ExtractionError> {
382        let max_attempts = 3;
383        let mut last_err = None;
384
385        for attempt in 0..max_attempts {
386            if attempt > 0 {
387                let delay = std::time::Duration::from_secs(1 << attempt);
388                tracing::warn!(
389                    attempt,
390                    delay_secs = delay.as_secs(),
391                    "retrying after rate limit"
392                );
393                tokio::time::sleep(delay).await;
394            }
395
396            tracing::info!(
397                provider = ?self.config.provider,
398                model = %self.config.model,
399                attempt = attempt + 1,
400                "calling LLM extraction API"
401            );
402
403            let result = match self.config.provider {
404                LlmProvider::OpenAI | LlmProvider::Custom => {
405                    if force_json {
406                        self.call_openai(conversation, system_prompt).await
407                    } else {
408                        self.call_openai_text(conversation, system_prompt).await
409                    }
410                }
411                LlmProvider::Anthropic => self.call_anthropic(conversation, system_prompt).await,
412                LlmProvider::Ollama => self.call_ollama(conversation, system_prompt).await,
413            };
414
415            match result {
416                Ok(text) => {
417                    tracing::info!(response_len = text.len(), "LLM extraction complete");
418                    return Ok(text);
419                }
420                Err(ExtractionError::ProviderError(ref msg))
421                    if msg.contains("429")
422                        || msg.contains("500")
423                        || msg.contains("502")
424                        || msg.contains("503")
425                        || msg.contains("529")
426                        || msg.contains("timeout")
427                        || msg.contains("connection")
428                        || msg.contains("overloaded") =>
429                {
430                    tracing::warn!(attempt = attempt + 1, error = %msg, "retrying transient LLM error");
431                    last_err = Some(result.unwrap_err());
432                    continue;
433                }
434                Err(e) => {
435                    tracing::error!(error = %e, "LLM extraction failed (non-retryable)");
436                    return Err(e);
437                }
438            }
439        }
440
441        match last_err {
442            Some(e) => Err(e),
443            None => Err(ExtractionError::RateLimitExceeded {
444                attempts: max_attempts,
445            }),
446        }
447    }
448}
449
450impl ExtractionProvider for HttpExtractionProvider {
451    async fn extract(
452        &self,
453        conversation: &str,
454        system_prompt: &str,
455    ) -> Result<String, ExtractionError> {
456        self.call_with_retry(conversation, system_prompt).await
457    }
458}
459
460/// Mock extraction provider for testing. Returns a predefined JSON response.
461pub struct MockExtractionProvider {
462    response: String,
463}
464
465impl MockExtractionProvider {
466    /// Create a mock provider that always returns the given JSON string.
467    pub fn new(response: impl Into<String>) -> Self {
468        Self {
469            response: response.into(),
470        }
471    }
472
473    /// Create a mock provider with a realistic extraction response.
474    pub fn with_realistic_response() -> Self {
475        let response = serde_json::json!({
476            "memories": [
477                {
478                    "content": "The team decided to use PostgreSQL 15 as the primary database for the REST API project",
479                    "memory_type": "decision",
480                    "confidence": 0.95,
481                    "entities": ["PostgreSQL", "REST API"],
482                    "tags": ["database", "architecture"],
483                    "reasoning": "Explicitly decided after comparing options"
484                },
485                {
486                    "content": "REST endpoints should follow the /api/v1/ prefix convention",
487                    "memory_type": "decision",
488                    "confidence": 0.9,
489                    "entities": ["REST API"],
490                    "tags": ["api-design", "conventions"],
491                    "reasoning": "Team agreed on URL structure"
492                },
493                {
494                    "content": "User prefers Rust over Go for backend services due to memory safety guarantees",
495                    "memory_type": "preference",
496                    "confidence": 0.85,
497                    "entities": ["Rust", "Go"],
498                    "tags": ["language", "backend"],
499                    "reasoning": "Explicitly stated preference with clear reasoning"
500                },
501                {
502                    "content": "The initial plan to use MongoDB was incorrect; PostgreSQL is the right choice for relational data",
503                    "memory_type": "correction",
504                    "confidence": 0.9,
505                    "entities": ["MongoDB", "PostgreSQL"],
506                    "tags": ["database", "correction"],
507                    "reasoning": "Corrected an earlier wrong assumption"
508                },
509                {
510                    "content": "The project deadline is March 15, 2025",
511                    "memory_type": "fact",
512                    "confidence": 0.8,
513                    "entities": ["REST API project"],
514                    "tags": ["timeline"],
515                    "reasoning": "Confirmed date mentioned in discussion"
516                },
517                {
518                    "content": "Using global mutable state for database connections caused race conditions in testing",
519                    "memory_type": "anti_pattern",
520                    "confidence": 0.85,
521                    "entities": [],
522                    "tags": ["testing", "concurrency"],
523                    "reasoning": "Documented failure pattern to avoid repeating"
524                },
525                {
526                    "content": "Low confidence speculation about maybe using Redis",
527                    "memory_type": "fact",
528                    "confidence": 0.3,
529                    "entities": ["Redis"],
530                    "tags": ["cache"],
531                    "reasoning": "Mentioned but not confirmed"
532                }
533            ]
534        });
535        Self::new(response.to_string())
536    }
537}
538
539impl ExtractionProvider for MockExtractionProvider {
540    async fn extract(
541        &self,
542        _conversation: &str,
543        _system_prompt: &str,
544    ) -> Result<String, ExtractionError> {
545        Ok(self.response.clone())
546    }
547}
mentedb_extraction/provider.rs

mentedb_extraction/
provider.rs
