use cortex_mem_config::Config;
use cortex_mem_core::{init::initialize_memory_system, memory::MemoryManager};
use cortex_mem_rig::tool::MemoryToolConfig;
use cortex_mem_rig::{
GetMemoryArgs, ListMemoriesArgs, MemoryToolOutput, QueryMemoryArgs, StoreMemoryArgs,
create_memory_tools,
};
use rig::tool::Tool;
use std::sync::Arc;
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
let config = Config::load("config.toml")?;
println!("Loaded configuration");
let (vector_store, llm_client) = initialize_memory_system(&config).await?;
println!("Initialized vector store and LLM client");
let memory_manager = Arc::new(MemoryManager::new(
vector_store,
llm_client,
config.memory.clone(),
));
println!("Created memory manager");
let memory_tools = create_memory_tools(
memory_manager,
&config,
Some(MemoryToolConfig {
default_user_id: Some("example_user".to_string()),
default_agent_id: Some("example_agent".to_string()),
..Default::default()
}),
);
println!("Created memory tools");
println!("\n=== Store Memory Example ===");
let store_args = StoreMemoryArgs {
content: "The user prefers to work in the morning and is interested in Rust programming."
.to_string(),
user_id: Some("example_user".to_string()),
agent_id: Some("example_agent".to_string()),
memory_type: Some("personal".to_string()),
topics: Some(vec!["preferences".to_string(), "programming".to_string()]),
};
match memory_tools.store_memory().call(store_args).await {
Ok(MemoryToolOutput {
success,
message,
data,
}) => {
println!(
"Store memory result: success={}, message={}",
success, message
);
if let Some(data) = data {
println!("Data: {}", serde_json::to_string_pretty(&data)?);
}
}
Err(e) => println!("Error: {}", e),
}
println!("\n=== Query Memory Example ===");
let query_args = QueryMemoryArgs {
query: "What are the user's programming preferences?".to_string(),
k: Some(5),
memory_type: None,
min_salience: Some(0.5),
topics: Some(vec!["programming".to_string(), "preferences".to_string()]),
user_id: Some("example_user".to_string()),
agent_id: Some("example_agent".to_string()),
};
match memory_tools.query_memory().call(query_args).await {
Ok(MemoryToolOutput {
success,
message,
data,
}) => {
println!(
"Query memory result: success={}, message={}",
success, message
);
if let Some(data) = data {
println!("Data: {}", serde_json::to_string_pretty(&data)?);
}
}
Err(e) => println!("Error: {}", e),
}
println!("\n=== List Memories Example ===");
let list_args = ListMemoriesArgs {
limit: Some(10),
memory_type: Some("personal".to_string()),
user_id: Some("example_user".to_string()),
agent_id: Some("example_agent".to_string()),
};
match memory_tools.list_memories().call(list_args).await {
Ok(MemoryToolOutput {
success,
message,
data,
}) => {
println!(
"List memories result: success={}, message={}",
success, message
);
if let Some(data) = data {
println!("Data: {}", serde_json::to_string_pretty(&data)?);
}
}
Err(e) => println!("Error: {}", e),
}
println!("\n=== Get Memory Example ===");
let get_args = GetMemoryArgs {
memory_id: "example_memory_id".to_string(),
};
match memory_tools.get_memory().call(get_args).await {
Ok(MemoryToolOutput {
success,
message,
data,
}) => {
println!(
"Get memory result: success={}, message={}",
success, message
);
if let Some(data) = data {
println!("Data: {}", serde_json::to_string_pretty(&data)?);
}
}
Err(e) => println!("Error: {}", e),
}
println!("\nAll examples completed successfully!");
Ok(())
}