From dc3c1cf58956c63d8ffff4a85e62152dd428505c Mon Sep 17 00:00:00 2001
From: HavenDV <havendv@gmail.com>
Date: Tue, 9 Apr 2024 05:06:06 +0400
Subject: [PATCH] test: Fixed.

---
 .../OpenSearch/test/OpenSearchTests.cs        | 38 ++++++++-----------
 1 file changed, 15 insertions(+), 23 deletions(-)

diff --git a/src/Databases/OpenSearch/test/OpenSearchTests.cs b/src/Databases/OpenSearch/test/OpenSearchTests.cs
index 0d9af470..46cf6fd2 100644
--- a/src/Databases/OpenSearch/test/OpenSearchTests.cs
+++ b/src/Databases/OpenSearch/test/OpenSearchTests.cs
@@ -22,7 +22,7 @@ public class OpenSearchTests
 
     #region Query Images
 
-    public Task setup_image_tests()
+    public void setup_image_tests()
     {
         _indexName = "images-index";
         var username = Environment.GetEnvironmentVariable("OPENSEARCH_USERNAME");
@@ -47,15 +47,13 @@ public Task setup_image_tests()
                 Dimensions = _options.Dimensions
             }
         };
-        _vectorDatabase = new OpenSearchVectorStore(_embeddingModel, _options);
-
-        return Task.CompletedTask;
+        _vectorDatabase = new OpenSearchVectorStore(_options);
     }
 
     [Test]
     public async Task index_test_images()
     {
-        await setup_image_tests();
+        setup_image_tests();
 
         string[] extensions = { ".bmp",".gif", ".jpg", ".jpeg", ".png", ".tiff" };
         var files = Directory.EnumerateFiles(@"[images directory]", "*.*", SearchOption.AllDirectories)
@@ -87,13 +85,13 @@ public async Task index_test_images()
             documents.Add(document);
         }
 
-        var pages = await _vectorStore!.AddImagesAsync(documents);
+        var pages = await _vectorDatabase!.AddDocumentsAsync(_embeddingModel!, documents);
     }
 
     [Test]
     public async Task can_query_image_against_images()
     {
-        await setup_image_tests();
+        setup_image_tests();
       
         var path = Path.Combine(Path.GetTempPath(), "test_image.jpg");
         var imageData = await File.ReadAllBytesAsync(path);
@@ -104,22 +102,21 @@ public async Task can_query_image_against_images()
             Strings = new List<string>(),
             Images = new List<Data> { Data.FromBytes(binaryData.ToArray()) }
         };
-        var embedding = await _embeddings!.CreateEmbeddingsAsync(embeddingRequest)
+        var embedding = await _embeddingModel!.CreateEmbeddingsAsync(embeddingRequest)
             .ConfigureAwait(false);
 
         var floats = embedding.ToSingleArray();
-        IEnumerable<Document> similaritySearchByVectorAsync = await (_vectorStore?.SimilaritySearchByVectorAsync(floats)!).ConfigureAwait(false);
+        var similaritySearchByVectorAsync = await _vectorDatabase!.SearchAsync(floats).ConfigureAwait(false);
 
-        Console.WriteLine("Count: " + similaritySearchByVectorAsync.Count());
+        Console.WriteLine("Count: " + similaritySearchByVectorAsync.Items.Count);
     }
 
     [Test]
     public async Task can_query_text_against_images()
     {
-        await setup_image_tests();
+        setup_image_tests();
 
         var llm = new Claude3SonnetModel(_provider!);
-        var index = new VectorStoreIndexWrapper(_vectorStore!);
 
         var promptText =
             @"Use the following pieces of context to answer the question at the end. If the answer is not in context then just say that you don't know, don't try to make up an answer. Keep the answer as short as possible.
@@ -131,7 +128,7 @@ public async Task can_query_text_against_images()
 
         var chain =
             Set("tell me about the orange shirt", outputKey: "question")                     // set the question
-            | RetrieveDocuments(index, inputKey: "question", outputKey: "documents", amount: 10) // take 5 most similar documents
+            | RetrieveDocuments(_vectorDatabase!, _embeddingModel!, inputKey: "question", outputKey: "documents", amount: 10) // take 5 most similar documents
             | StuffDocuments(inputKey: "documents", outputKey: "context")                       // combine documents together and put them into context
             | Template(promptText)                                                              // replace context and question in the prompt with their values
             | LLM(llm);                                                                       // send the result to the language model
@@ -144,7 +141,7 @@ public async Task can_query_text_against_images()
 
     #region Query Simple Documents
 
-    public Task setup_document_tests()
+    public void setup_document_tests()
     {
         _indexName = "test-index";
         var username = Environment.GetEnvironmentVariable("OPENSEARCH_USERNAME");
@@ -172,13 +169,10 @@ public Task setup_document_tests()
         _vectorDatabase = new OpenSearchVectorStore(_options);
     }
 
-        return Task.CompletedTask;
-    }
-
     [Test]
     public async Task index_test_documents()
     {
-        await setup_document_tests();
+        setup_document_tests();
 
         var documents = new[]
         {
@@ -196,10 +190,9 @@ public async Task index_test_documents()
     [Test]
     public async Task can_query_test_documents()
     {
-        await setup_document_tests();
+        setup_document_tests();
 
         var llm = new Claude3SonnetModel(_provider!);
-        var index = new VectorStoreIndexWrapper(_vectorStore!);
 
         const string question = "what color is the car?";
 
@@ -229,7 +222,7 @@ public async Task can_query_test_documents()
     [Test]
     public async Task index_harry_potter_book()
     {
-        await setup_document_tests();
+        setup_document_tests();
 
         var pdfSource = new PdfPigPdfSource("x:\\Harry-Potter-Book-1.pdf");
         var documents = await pdfSource.LoadAsync();
@@ -241,10 +234,9 @@ public async Task index_harry_potter_book()
     [Test]
     public async Task can_query_harry_potter_book()
     {
-        await setup_document_tests();
+        setup_document_tests();
 
         var llm = new Claude3SonnetModel(_provider!);
-        var index = new VectorStoreIndexWrapper(_vectorStore!);
 
         var promptText =
             @"Use the following pieces of context to answer the question at the end. If the answer is not in context then just say that you don't know, don't try to make up an answer. Keep the answer as short as possible.