diff --git a/rig-neo4j/src/lib.rs b/rig-neo4j/src/lib.rs index a7143c3..3868643 100644 --- a/rig-neo4j/src/lib.rs +++ b/rig-neo4j/src/lib.rs @@ -339,71 +339,3 @@ impl Neo4jClient { Ok(()) } } - -#[allow(dead_code)] -#[cfg(test)] -mod tests { - use super::*; - use neo4rs::ConfigBuilder; - use rig::{ - providers::openai::{Client, TEXT_EMBEDDING_ADA_002}, - vector_store::VectorStoreIndex, - }; - use serde::Deserialize; - use std::env; - - const NEO4J_URI: &str = "neo4j+s://demo.neo4jlabs.com:7687"; - const NEO4J_DB: &str = "recommendations"; - const NEO4J_USERNAME: &str = "recommendations"; - const NEO4J_PASSWORD: &str = "recommendations"; - - #[derive(Debug, Deserialize)] - struct Movie { - title: String, - plot: String, - } - - #[tokio::test] - async fn test_connect() { - let result = Neo4jClient::from_config( - ConfigBuilder::default() - .uri(NEO4J_URI) - .db(NEO4J_DB) - .user(NEO4J_USERNAME) - .password(NEO4J_PASSWORD) - .build() - .unwrap(), - ) - .await; - assert!(result.is_ok()); - } - - #[tokio::test] - async fn test_vector_search_no_display() { - let results = vector_search().await.unwrap(); - assert!(results.len() > 0); - } - - async fn vector_search() -> Result, VectorStoreError> { - let openai_api_key = env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY not set"); - let openai_client = Client::new(&openai_api_key); - let model = openai_client.embedding_model(TEXT_EMBEDDING_ADA_002); - - let client = Neo4jClient::from_config( - ConfigBuilder::default() - .uri(NEO4J_URI) - .db(NEO4J_DB) - .user(NEO4J_USERNAME) - .password(NEO4J_PASSWORD) - .build() - .unwrap(), - ) - .await - .unwrap(); - - let index = client - .get_index(model, "moviePlotsEmbedding", SearchParams::default()) - .await?; - Ok(index.top_n::("Batman", 3).await?) - } -}