mem0-rust 0.2.0

Rust implementation of mem0 - Universal memory layer for AI Agents
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

//! Ollama LLM provider.

use async_trait::async_trait;
use ollama_rs::{generation::completion::request::GenerationRequest, Ollama};

use super::traits::{GenerateOptions, LLM};
use crate::config::OllamaLLMConfig;
use crate::errors::LLMError;
use crate::models::{Message, Role};

/// Ollama LLM provider for local models
pub struct OllamaLLM {
    client: Ollama,
    model: String,
    default_temperature: f32,
}

impl OllamaLLM {
    /// Create a new Ollama LLM
    pub fn new(config: OllamaLLMConfig) -> Self {
        let url = url::Url::parse(&config.base_url).unwrap_or_else(|_| {
            url::Url::parse("http://localhost:11434").unwrap()
        });

        let host = url.host_str().unwrap_or("localhost").to_string();
        let port = url.port().unwrap_or(11434);

        let client = Ollama::new(host, port);

        Self {
            client,
            model: config.model,
            default_temperature: config.temperature,
        }
    }

    /// Format messages into a prompt string
    fn format_messages(messages: &[Message]) -> String {
        let mut prompt = String::new();

        for msg in messages {
            match msg.role {
                Role::System => {
                    prompt.push_str(&format!("System: {}\n\n", msg.content));
                }
                Role::User => {
                    prompt.push_str(&format!("User: {}\n\n", msg.content));
                }
                Role::Assistant => {
                    prompt.push_str(&format!("Assistant: {}\n\n", msg.content));
                }
            }
        }

        prompt.push_str("Assistant: ");
        prompt
    }
}

#[async_trait]
impl LLM for OllamaLLM {
    async fn generate(
        &self,
        messages: &[Message],
        options: GenerateOptions,
    ) -> Result<String, LLMError> {
        let prompt = Self::format_messages(messages);

        let mut request = GenerationRequest::new(self.model.clone(), prompt);

        // Set options
        let temperature = options.temperature.unwrap_or(self.default_temperature);
        request = request.options(
            ollama_rs::generation::options::GenerationOptions::default()
                .temperature(temperature as f64),
        );

        if options.json_mode {
            request = request.format(ollama_rs::generation::completion::request::FormatType::Json);
        }

        let response = self
            .client
            .generate(request)
            .await
            .map_err(|e| LLMError::Api(e.to_string()))?;

        Ok(response.response)
    }

    fn model_name(&self) -> &str {
        &self.model
    }
}