mentedb-extraction 0.5.0

LLM-powered memory extraction engine for MenteDB
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119

use serde::{Deserialize, Serialize};

/// Which LLM provider to use for extraction.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum LlmProvider {
    OpenAI,
    Anthropic,
    Ollama,
    Custom,
}

impl LlmProvider {
    /// Default API URL for this provider.
    pub fn default_url(&self) -> &str {
        match self {
            LlmProvider::OpenAI => "https://api.openai.com/v1/chat/completions",
            LlmProvider::Anthropic => "https://api.anthropic.com/v1/messages",
            LlmProvider::Ollama => "http://localhost:11434/api/chat",
            LlmProvider::Custom => "http://localhost:8080/v1/chat/completions",
        }
    }

    /// Default extraction model for this provider (cheap, high-volume).
    pub fn default_model(&self) -> &str {
        match self {
            LlmProvider::OpenAI => "gpt-4o-mini",
            LlmProvider::Anthropic => "claude-haiku-4-5",
            LlmProvider::Ollama => "llama3",
            LlmProvider::Custom => "default",
        }
    }

    /// Default reader model for this provider (smart, low-volume).
    pub fn default_reader_model(&self) -> &str {
        match self {
            LlmProvider::OpenAI => "gpt-4o",
            LlmProvider::Anthropic => "claude-sonnet-4-20250514",
            LlmProvider::Ollama => "llama3",
            LlmProvider::Custom => "default",
        }
    }
}

/// Configuration for the extraction pipeline.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExtractionConfig {
    /// Which LLM provider to use.
    pub provider: LlmProvider,
    /// API key for the provider (not needed for Ollama).
    pub api_key: Option<String>,
    /// API endpoint URL (defaults based on provider).
    pub api_url: String,
    /// Model name to use.
    pub model: String,
    /// Maximum number of memories to extract from a single conversation.
    pub max_extractions_per_conversation: usize,
    /// Minimum confidence score for a memory to be accepted (0.0 to 1.0).
    pub quality_threshold: f32,
    /// Embedding similarity above which a memory is considered a duplicate (0.0 to 1.0).
    pub deduplication_threshold: f32,
    /// Whether to check new memories against existing ones for contradictions.
    pub enable_contradiction_check: bool,
    /// Whether to check new memories against existing ones for duplicates.
    pub enable_deduplication: bool,
    /// Number of extraction passes (1 = single pass, 2 = first pass + verification).
    pub extraction_passes: usize,
}

impl ExtractionConfig {
    /// Create a config for OpenAI with the given API key.
    pub fn openai(api_key: impl Into<String>) -> Self {
        Self {
            provider: LlmProvider::OpenAI,
            api_key: Some(api_key.into()),
            api_url: LlmProvider::OpenAI.default_url().to_string(),
            model: LlmProvider::OpenAI.default_model().to_string(),
            ..Self::default()
        }
    }

    /// Create a config for Anthropic with the given API key.
    pub fn anthropic(api_key: impl Into<String>) -> Self {
        Self {
            provider: LlmProvider::Anthropic,
            api_key: Some(api_key.into()),
            api_url: LlmProvider::Anthropic.default_url().to_string(),
            model: LlmProvider::Anthropic.default_model().to_string(),
            ..Self::default()
        }
    }

    /// Create a config for a local Ollama instance.
    pub fn ollama() -> Self {
        Self {
            provider: LlmProvider::Ollama,
            api_key: None,
            api_url: LlmProvider::Ollama.default_url().to_string(),
            model: LlmProvider::Ollama.default_model().to_string(),
            ..Self::default()
        }
    }
}

impl Default for ExtractionConfig {
    fn default() -> Self {
        Self {
            provider: LlmProvider::OpenAI,
            api_key: None,
            api_url: LlmProvider::OpenAI.default_url().to_string(),
            model: LlmProvider::OpenAI.default_model().to_string(),
            max_extractions_per_conversation: 50,
            quality_threshold: 0.6,
            deduplication_threshold: 0.85,
            enable_contradiction_check: true,
            enable_deduplication: true,
            extraction_passes: 1,
        }
    }
}