use anyhow::Result;
#[cfg(feature = "surrealdb-storage")]
use post_cortex::storage::traits::StorageBackendType;
use post_cortex_memory::{ConversationMemorySystem, SystemConfig};
use std::sync::Arc;
use tempfile::{TempDir, tempdir};
use uuid::Uuid;
pub struct TestFixture {
pub system: Arc<ConversationMemorySystem>,
pub session_id: Uuid,
pub _temp_dir: TempDir,
}
impl TestFixture {
pub async fn new() -> Result<Self> {
let temp_dir = tempdir()?;
let mut config = SystemConfig::default();
config.data_directory = temp_dir.path().to_str().unwrap().to_string();
config.enable_embeddings = true;
config.auto_vectorize_on_update = true;
#[cfg(feature = "surrealdb-storage")]
{
config.storage_backend = StorageBackendType::SurrealDB;
config.surrealdb_endpoint = Some("mem://".to_string());
config.surrealdb_namespace = Some("test".to_string());
config.surrealdb_database = Some("test".to_string());
}
let system = Arc::new(
ConversationMemorySystem::new(config)
.await
.map_err(|e| anyhow::anyhow!(e))?,
);
let session_id = system
.create_session(
Some("test-session".to_string()),
Some("Test session".to_string()),
)
.await
.map_err(|e| anyhow::anyhow!(e))?;
Ok(Self {
system,
session_id,
_temp_dir: temp_dir,
})
}
pub async fn with_content(content: &[&str]) -> Result<Self> {
let fixture = Self::new().await?;
for text in content {
fixture
.system
.add_incremental_update(fixture.session_id, text.to_string(), None)
.await
.map_err(|e| anyhow::anyhow!(e))?;
}
tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
Ok(fixture)
}
pub async fn search_with_bias(
&self,
query: &str,
bias: f32,
) -> Result<Vec<post_cortex_memory::content_vectorizer::SemanticSearchResult>> {
let engine = self
.system
.ensure_semantic_engine_initialized()
.await
.map_err(|e| anyhow::anyhow!(e))?;
let results = engine
.semantic_search_session(self.session_id, query, None, None, Some(bias))
.await
.map_err(|e| anyhow::anyhow!(e))?;
Ok(results)
}
}