langchain-rust 6.0.0

LangChain for Rust, the easiest way to write LLM-based programs in Rust
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

use std::sync::Arc;

use async_trait::async_trait;
use tokio::sync::Mutex;

use crate::schemas::{
    memory::{BaseMemory, MemoryError},
    messages::Message,
};

pub struct SimpleMemory {
    messages: Vec<Message>,
}

impl SimpleMemory {
    pub fn new() -> Self {
        Self {
            messages: Vec::new(),
        }
    }
}

impl Default for SimpleMemory {
    fn default() -> Self {
        Self::new()
    }
}

impl Into<Arc<dyn BaseMemory>> for SimpleMemory {
    fn into(self) -> Arc<dyn BaseMemory> {
        Arc::new(self)
    }
}

impl Into<Arc<Mutex<dyn BaseMemory>>> for SimpleMemory {
    fn into(self) -> Arc<Mutex<dyn BaseMemory>> {
        Arc::new(Mutex::new(self))
    }
}

#[async_trait]
impl BaseMemory for SimpleMemory {
    async fn messages(&self) -> Result<Vec<Message>, MemoryError> {
        Ok(self.messages.clone())
    }

    async fn add_message(&mut self, message: Message) -> Result<(), MemoryError> {
        self.messages.push(message);
        Ok(())
    }

    async fn clear(&mut self) -> Result<(), MemoryError> {
        self.messages.clear();
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::schemas::MessageType;

    #[tokio::test]
    async fn stores_all_messages_in_order() {
        let mut memory = SimpleMemory::new();
        memory.add_user_message(&"question").await.unwrap();
        memory.add_ai_message(&"answer").await.unwrap();
        memory
            .add_message(Message::new_system_message("sys"))
            .await
            .unwrap();

        let messages = memory.messages().await.unwrap();
        let types: Vec<_> = messages.iter().map(|m| m.message_type.clone()).collect();
        assert_eq!(
            types,
            vec![
                MessageType::HumanMessage,
                MessageType::AIMessage,
                MessageType::SystemMessage
            ]
        );
        assert_eq!(
            memory.to_string().await.unwrap(),
            "human: question\nai: answer\nsystem: sys"
        );
    }

    #[tokio::test]
    async fn clear_empties_memory() {
        let mut memory = SimpleMemory::new();
        memory.add_user_message(&"x").await.unwrap();
        memory.clear().await.unwrap();
        assert!(memory.messages().await.unwrap().is_empty());
    }

    #[tokio::test]
    async fn converts_into_shared_mutex() {
        let memory: Arc<Mutex<dyn BaseMemory>> = SimpleMemory::new().into();
        memory
            .lock()
            .await
            .add_user_message(&"shared")
            .await
            .unwrap();
        assert_eq!(memory.lock().await.messages().await.unwrap().len(), 1);
    }
}