pub mod manager;
pub mod types;
pub mod vector_db;
pub use types::{
AccessLevel, AgentContext, ContextError, ContextId, ContextPersistence, ContextQuery,
FilePersistenceConfig, HierarchicalMemory, Knowledge, KnowledgeBase, KnowledgeId,
KnowledgeItem, KnowledgeSource, KnowledgeType, MemoryItem, MemoryType, QueryType,
RetentionPolicy, SessionId, StorageStats, VectorBatchItem, VectorBatchOperation,
VectorContentType, VectorId, VectorMetadata, VectorOperationType, VectorSearchResult,
};
pub use manager::{ContextManager, ContextManagerConfig, FilePersistence, StandardContextManager};
pub use vector_db::{
EmbeddingService, MockEmbeddingService, QdrantClientWrapper, QdrantConfig, QdrantDistance,
TfIdfEmbeddingService, VectorDatabase, VectorDatabaseStats,
};
#[cfg(test)]
mod tests {
use super::*;
use crate::types::AgentId;
use std::time::SystemTime;
#[tokio::test]
async fn test_context_manager_creation() {
let config = ContextManagerConfig::default();
let agent_id = AgentId::new();
let manager = StandardContextManager::new(config, &agent_id.to_string())
.await
.unwrap();
assert!(manager.initialize().await.is_ok());
}
#[tokio::test]
async fn test_session_creation() {
let config = ContextManagerConfig::default();
let agent_id = AgentId::new();
let manager = StandardContextManager::new(config, &agent_id.to_string())
.await
.unwrap();
manager.initialize().await.unwrap();
let session_id = manager.create_session(agent_id).await.unwrap();
let context = manager
.retrieve_context(agent_id, Some(session_id))
.await
.unwrap();
assert!(context.is_some());
assert_eq!(context.unwrap().session_id, session_id);
}
#[tokio::test]
async fn test_memory_operations() {
let config = ContextManagerConfig::default();
let agent_id = AgentId::new();
let manager = StandardContextManager::new(config, &agent_id.to_string())
.await
.unwrap();
manager.initialize().await.unwrap();
let _session_id = manager.create_session(agent_id).await.unwrap();
let memory_updates = vec![types::MemoryUpdate {
target: types::MemoryTarget::Working("test_key".to_string()),
operation: types::UpdateOperation::Add,
data: serde_json::Value::String("test_value".to_string()),
}];
assert!(manager
.update_memory(agent_id, memory_updates)
.await
.is_ok());
}
#[tokio::test]
async fn test_knowledge_operations() {
let config = ContextManagerConfig::default();
let agent_id = AgentId::new();
let manager = StandardContextManager::new(config, &agent_id.to_string())
.await
.unwrap();
manager.initialize().await.unwrap();
let _session_id = manager.create_session(agent_id).await.unwrap();
let fact = types::KnowledgeFact {
id: KnowledgeId::new(),
subject: "test".to_string(),
predicate: "is".to_string(),
object: "example".to_string(),
confidence: 0.9,
source: types::KnowledgeSource::UserProvided,
created_at: SystemTime::now(),
verified: true,
};
let _knowledge_id = manager
.add_knowledge(agent_id, Knowledge::Fact(fact))
.await
.unwrap();
let results = manager
.search_knowledge(agent_id, "test", 10)
.await
.unwrap();
assert!(!results.is_empty());
}
#[tokio::test]
async fn test_context_query() {
let config = ContextManagerConfig::default();
let agent_id = AgentId::new();
let manager = StandardContextManager::new(config, &agent_id.to_string())
.await
.unwrap();
manager.initialize().await.unwrap();
let _session_id = manager.create_session(agent_id).await.unwrap();
let query = ContextQuery {
query_type: QueryType::Semantic,
search_terms: vec!["test".to_string()],
max_results: 10,
time_range: None,
memory_types: vec![],
relevance_threshold: 0.7,
include_embeddings: false,
};
let results = manager.query_context(agent_id, query).await.unwrap();
assert!(results.len() <= 10);
}
#[tokio::test]
async fn test_context_stats() {
let config = ContextManagerConfig::default();
let agent_id = AgentId::new();
let manager = StandardContextManager::new(config, &agent_id.to_string())
.await
.unwrap();
manager.initialize().await.unwrap();
let _session_id = manager.create_session(agent_id).await.unwrap();
let stats = manager.get_context_stats(agent_id).await.unwrap();
assert_eq!(stats.total_memory_items, 0);
assert_eq!(stats.total_knowledge_items, 0);
assert_eq!(stats.total_conversations, 0);
}
#[tokio::test]
async fn test_secrets_integration() {
let config = ContextManagerConfig::default();
let agent_id = AgentId::new();
let manager = StandardContextManager::new(config, &agent_id.to_string())
.await
.unwrap();
let _secrets = manager.secrets();
}
}