vv_agent/memory/
provider.rs

1use std::future::Future;
2use std::pin::Pin;
3
4use crate::events::RunEvent;
5use crate::types::Metadata;
6
7pub type MemoryFuture<T> = Pin<Box<dyn Future<Output = Result<T, MemoryError>> + Send>>;
8
9pub trait MemoryProvider: Send + Sync {
10    fn search(&self, request: MemorySearchRequest) -> MemoryFuture<Vec<MemorySearchResult>>;
11    fn save(&self, request: MemorySaveRequest) -> MemoryFuture<MemorySaveResult>;
12
13    fn before_compact(&self, _event: &RunEvent) -> MemoryFuture<MemoryProviderResult> {
14        Box::pin(async { Ok(MemoryProviderResult::default()) })
15    }
16
17    fn after_compact(&self, _event: &RunEvent) -> MemoryFuture<()> {
18        Box::pin(async { Ok(()) })
19    }
20}
21
22#[derive(Debug, Clone, Default, PartialEq)]
23pub struct MemorySearchRequest {
24    pub query: String,
25    pub metadata: Metadata,
26}
27
28#[derive(Debug, Clone, Default, PartialEq)]
29pub struct MemorySearchResult {
30    pub content: String,
31    pub score: Option<f64>,
32    pub metadata: Metadata,
33}
34
35#[derive(Debug, Clone, Default, PartialEq)]
36pub struct MemorySaveRequest {
37    pub content: String,
38    pub metadata: Metadata,
39}
40
41#[derive(Debug, Clone, Default, PartialEq)]
42pub struct MemorySaveResult {
43    pub id: Option<String>,
44    pub metadata: Metadata,
45}
46
47#[derive(Debug, Clone, Default, PartialEq)]
48pub struct MemoryProviderResult {
49    pub metadata: Metadata,
50}
51
52#[derive(Debug, Clone, PartialEq, Eq)]
53pub struct MemoryError {
54    message: String,
55}
56
57impl MemoryError {
58    pub fn new(message: impl Into<String>) -> Self {
59        Self {
60            message: message.into(),
61        }
62    }
63}
64
65impl std::fmt::Display for MemoryError {
66    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
67        formatter.write_str(&self.message)
68    }
69}
70
71impl std::error::Error for MemoryError {}
72
73pub(crate) fn block_on_memory_future<T: Send>(future: MemoryFuture<T>) -> Result<T, MemoryError> {
74    if let Ok(handle) = tokio::runtime::Handle::try_current() {
75        if handle.runtime_flavor() == tokio::runtime::RuntimeFlavor::MultiThread {
76            tokio::task::block_in_place(|| handle.block_on(future))
77        } else {
78            handle.block_on(future)
79        }
80    } else {
81        tokio::runtime::Builder::new_current_thread()
82            .enable_all()
83            .build()
84            .map_err(|error| MemoryError::new(error.to_string()))?
85            .block_on(future)
86    }
87}
vv_agent/memory/provider.rs

vv_agent/memory/
provider.rs
