From 9705e26cfffef2a27d12a80decad62d787affe08 Mon Sep 17 00:00:00 2001
From: Kyle Kelley <rgbkrk@gmail.com>
Date: Thu, 4 Apr 2024 16:48:34 -0700
Subject: [PATCH] create a not-exactly-a-benchmark

---
 crates/semantic_index/src/embedding.rs | 41 ++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/crates/semantic_index/src/embedding.rs b/crates/semantic_index/src/embedding.rs
index b3f0b2b0d1..06fccb3784 100644
--- a/crates/semantic_index/src/embedding.rs
+++ b/crates/semantic_index/src/embedding.rs
@@ -164,6 +164,47 @@ mod test {
         }
     }
 
+    #[gpui::test]
+    async fn test_ollama_embedding_not_exactly_a_benchmark(executor: BackgroundExecutor) {
+        executor.allow_parking();
+
+        let client = Arc::new(HttpClientWithUrl::new("http://localhost:11434/"));
+        let provider =
+            OllamaEmbeddingProvider::new(client.into(), EmbeddingModel::OllamaNomicEmbedText);
+
+        let t_nomic = std::time::Instant::now();
+        for i in 0..100 {
+            let embedding = provider
+                .get_embedding(format!("Hello, world! {}", i))
+                .await
+                .unwrap();
+
+            match embedding {
+                Embedding::OllamaNomicEmbedText(e) => assert_eq!(e.len(), EMBEDDING_SIZE_TINY),
+                _ => panic!("Invalid embedding size"),
+            }
+        }
+        dbg!(t_nomic.elapsed());
+
+        let client = Arc::new(HttpClientWithUrl::new("http://localhost:11434/"));
+        let provider =
+            OllamaEmbeddingProvider::new(client.into(), EmbeddingModel::OllamaMxbaiEmbedLarge);
+
+        let t_mxbai = std::time::Instant::now();
+        for i in 0..100 {
+            let embedding = provider
+                .get_embedding(format!("Hello, world! {}", i))
+                .await
+                .unwrap();
+
+            match embedding {
+                Embedding::OllamaMxbaiEmbedLarge(e) => assert_eq!(e.len(), EMBEDDING_SIZE_XSMALL),
+                _ => panic!("Invalid embedding size"),
+            }
+        }
+        dbg!(t_mxbai.elapsed());
+    }
+
     #[gpui::test]
     fn test_normalize_embedding() {
         // Create an vector of size EMBEDDING_SIZE_TINY with all values set to 1.0