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

vv_agent/memory/
provider.rs
