Addressed review comments:

- added Postgres test_util,
- completed unit tests for persist and retrieval

Signed-off-by: shamb0 <r.raajey@gmail.com>

Cargo.toml

@@ -72,6 +72,8 @@ temp-dir = "0.1.13"
 wiremock = "0.6.0"
 test-case = "3.3.1"
 insta = { version = "1.39.0", features = ["yaml"] }
+tempfile = "3.10.1"
+portpicker = "0.1.1"
 
 [workspace.lints.rust]
 unsafe_code = "forbid"

examples/Cargo.toml

@@ -31,6 +31,7 @@ qdrant-client = { workspace = true }
 fluvio = { workspace = true }
 temp-dir = { workspace = true }
 sqlx = { workspace = true }
+swiftide-test-utils = { path = "../swiftide-test-utils" }
 
 [[example]]
 doc-scrape-examples = true

examples/index_md_into_pgvector.rs

@@ -1,5 +1,5 @@
 /**
-* This example demonstrates how to use the Pgvector integration with Swiftide
+* This example demonstrates how to use the Pgvector integration with SwiftIDE
 */
 use std::path::PathBuf;
 use swiftide::{
@@ -11,9 +11,40 @@ use swiftide::{
         },
         EmbeddedField,
     },
-    integrations::{self, pgvector::PgVector},
+    integrations::{self, fastembed::FastEmbed, pgvector::PgVector},
+    query::{self, answers, query_transformers, response_transformers},
+    traits::SimplePrompt,
 };
 
+async fn ask_query(
+    llm_client: impl SimplePrompt + Clone + 'static,
+    embed: FastEmbed,
+    vector_store: PgVector,
+    question: String,
+) -> Result<String, Box<dyn std::error::Error>> {
+    // By default the search strategy is SimilaritySingleEmbedding
+    // which takes the latest query, embeds it, and does a similarity search
+    //
+    // Pgvector will return an error if multiple embeddings are set
+    //
+    // The pipeline generates subquestions to increase semantic coverage, embeds these in a single
+    // embedding, retrieves the default top_k documents, summarizes them and uses that as context
+    // for the final answer.
+    let pipeline = query::Pipeline::default()
+        .then_transform_query(query_transformers::GenerateSubquestions::from_client(
+            llm_client.clone(),
+        ))
+        .then_transform_query(query_transformers::Embed::from_client(embed))
+        .then_retrieve(vector_store.clone())
+        .then_transform_response(response_transformers::Summary::from_client(
+            llm_client.clone(),
+        ))
+        .then_answer(answers::Simple::from_client(llm_client.clone()));
+
+    let result = pipeline.query(question).await?;
+    Ok(result.answer().into())
+}
+
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
     tracing_subscriber::fmt::init();
@@ -23,15 +54,15 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let manifest_dir = std::env::var("CARGO_MANIFEST_DIR").expect("CARGO_MANIFEST_DIR not set");
 
     // Create a PathBuf to test dataset from the manifest directory
-    let test_dataset_path = PathBuf::from(manifest_dir).join("test_dataset");
+    let test_dataset_path = PathBuf::from(manifest_dir).join("../README.md");
+
     tracing::info!("Test Dataset path: {:?}", test_dataset_path);
 
-    let pgv_db_url = std::env::var("DATABASE_URL")
-        .as_deref()
-        .unwrap_or("postgresql://myuser:mypassword@localhost:5432/mydatabase")
-        .to_owned();
+    let (_pgv_db_container, pgv_db_url, _temp_dir) = swiftide_test_utils::start_postgres().await;
+
+    tracing::info!("pgv_db_url :: {:#?}", pgv_db_url);
 
-    let ollama_client = integrations::ollama::Ollama::default()
+    let llm_client = integrations::ollama::Ollama::default()
         .with_default_prompt_model("llama3.2:latest")
         .to_owned();
 
@@ -41,7 +72,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     // Configure Pgvector with a default vector size, a single embedding
     // and in addition to embedding the text metadata, also store it in a field
     let pgv_storage = PgVector::builder()
-        .try_from_url(pgv_db_url, Some(10))
+        .try_connect_to_pool(pgv_db_url, Some(10))
         .await
         .expect("Failed to connect to postgres server")
         .vector_size(384)
@@ -52,24 +83,44 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         .unwrap();
 
     // Drop the existing test table before running the test
-    tracing::info!("Dropping existing test table if it exists");
+    tracing::info!("Dropping existing test table & index if it exists");
     let drop_table_sql = "DROP TABLE IF EXISTS swiftide_pgvector_test";
+    let drop_index_sql = "DROP INDEX IF EXISTS swiftide_pgvector_test_embedding_idx";
 
-    if let Some(pool) = pgv_storage.get_pool() {
-        sqlx::query(drop_table_sql).execute(pool).await?;
+    if let Ok(pool) = pgv_storage.get_pool() {
+        sqlx::query(drop_table_sql).execute(&pool).await?;
+        sqlx::query(drop_index_sql).execute(&pool).await?;
     } else {
         return Err("Failed to get database connection pool".into());
     }
 
     tracing::info!("Starting indexing pipeline");
     indexing::Pipeline::from_loader(FileLoader::new(test_dataset_path).with_extensions(&["md"]))
         .then_chunk(ChunkMarkdown::from_chunk_range(10..2048))
-        .then(MetadataQAText::new(ollama_client.clone()))
-        .then_in_batch(Embed::new(fastembed).with_batch_size(100))
+        .then(MetadataQAText::new(llm_client.clone()))
+        .then_in_batch(Embed::new(fastembed.clone()).with_batch_size(100))
         .then_store_with(pgv_storage.clone())
         .run()
         .await?;
 
-    tracing::info!("Indexing test completed successfully");
+    for (i, question) in [
+            "What is SwiftIDE? Provide a clear, comprehensive summary in under 50 words.",
+            "How can I use SwiftIDE to connect with the Ethereum blockchain? Please provide a concise, comprehensive summary in less than 50 words.",
+        ]
+        .iter()
+        .enumerate()
+        {
+            let result = ask_query(
+                llm_client.clone(),
+                fastembed.clone(),
+                pgv_storage.clone(),
+                question.to_string(),
+            ).await?;
+            tracing::info!("*** Answer Q{} ***", i + 1);
+            tracing::info!("{}", result);
+            tracing::info!("===X===");
+        }
+
+    tracing::info!("PgVector Indexing & retrieval test completed successfully");
     Ok(())
 }

swiftide-indexing/src/loaders/file_loader.rs

@@ -109,7 +109,7 @@ impl Loader for FileLoader {
             .filter(|entry| entry.file_type().is_some_and(|ft| ft.is_file()))
             .filter(move |entry| self.file_has_extension(entry.path()))
             .map(|entry| {
-                tracing::debug!("Reading file: {:?}", entry);
+                tracing::info!("Reading file: {:?}", entry);
                 let content =
                     std::fs::read_to_string(entry.path()).context("Failed to read file")?;
                 let original_size = content.len();