WIP

notebooks/demo_openai.ipynb

@@ -70,7 +70,7 @@
    "source": [
     "DATA_DIR = Path(\"../tests/test_data\")\n",
     "LETTERS_DIR = DATA_DIR/ \"sample_letters\"\n",
-    "SAMPLE_JOB_AD = DATA_DIR / \"sample_ads/description_5.text\"\n",
+    "SAMPLE_JOB_AD = DATA_DIR / \"sample_ads/description_10.text\"\n",
     "\n",
     "OUTPUT_DIR = Path(\"./output\")\n",
     "OUTPUT_DIR.mkdir(exist_ok=True, parents=True)"
@@ -145,7 +145,7 @@
    "source": [
     "def pprint(text: str):\n",
     "    \"\"\"Pretty print the text as markdown in Jupyter Notebook's output.\"\"\"\n",
-    "    display(Markdown(text.strip()))"
+    "    display(Markdown(text))"
    ],
    "id": "f95d49699943e99b",
    "outputs": [],

tests/test_similarity_search.py

@@ -26,23 +26,22 @@ def test_make_collection():
     # Act
     ranker.make_collection(documents_with_embeddings, collection_name)
 
-    print(ranker.client.get_collection(collection_name))
-
     # Assert
     coll = ranker.client.get_collection(collection_name)
     assert coll is not None
-    print(coll)
-
-#
-# def test_find_similar(ranker, mock_qdrant_client):
-#     results = ranker.find_similar(embedding_to_search,
-#                                   "test_collection",
-#                                   top_k=2, min_similarity=0.1)
-#
-#     mock_qdrant_client.search.assert_called_once_with(
-#         collection_name="test_collection",
-#         query_vector=embedding_to_search,
-#         limit=2,
-#         score_threshold=0.1
-#     )
-#     assert results == search_response
+    assert coll.points_count == 2
+    assert coll.config.params.vectors.size == 3
+    assert coll.config.params.vectors.distance == "Cosine"
+
+
+def test_find_similar():
+    # Arrange
+    ranker = Ranker()
+    collection_name = "test_collection"
+
+    # Act
+    ranker.make_collection(documents_with_embeddings, collection_name)
+    response = ranker.find_similar(embedding_to_search, collection_name, 2)
+
+    # Assert
+    assert response == search_response