WIP

tests/test_connectors.py

@@ -1,53 +1,53 @@
-from unittest.mock import MagicMock
-
-import numpy as np
-import pytest
-
-from easy_letters import OpenAIConnector, LanguageModels
-
-# Sample data for testing
-documents = ["Document 1", "Document 2"]
-embedding_response = {
-    "data": [
-        {"embedding": [0.1, 0.2, 0.3]},
-        {"embedding": [0.4, 0.5, 0.6]}
-    ]
-}
-chat_response = {
-    "choices": [
-        {"message": {"content": "This is a test response"}}
-    ]
-}
-
-
-@pytest.fixture
-def mock_openai_client(mocker):
-    mock_client = MagicMock()
-    mock_client.embeddings.create.return_value = embedding_response
-    mock_client.chat.completions.create.return_value = chat_response
-    mocker.patch("openai.Client", return_value=mock_client)
-    return mock_client
-
-
-@pytest.fixture
-def openai_connector(mock_openai_client):
-    return OpenAIConnector(api_key="fake-api-key")
-
-
-def test_embed(openai_connector):
-    model = LanguageModels.OPENAI_GPT35TURBO
-    embeddings = openai_connector.embed(documents, model)
-
-    assert isinstance(embeddings, list)
-    assert all(isinstance(e, np.ndarray) for e in embeddings)
-    assert np.array_equal(embeddings[0], np.array([0.1, 0.2, 0.3]))
-    assert np.array_equal(embeddings[1], np.array([0.4, 0.5, 0.6]))
-
-
-def test_chat(openai_connector):
-    prompt = "Hey, how are you?"
-    model = LanguageModels.OPENAI_GPT35TURBO
-    response = openai_connector.chat(prompt, model)
-
-    assert isinstance(response, str)
-    assert response == "This is a test response"
+# from unittest.mock import MagicMock
+#
+# import numpy as np
+# import pytest
+#
+# from easy_letters import OpenAIConnector, LanguageModels
+#
+# # Sample data for testing
+# documents = ["Document 1", "Document 2"]
+# embedding_response = {
+#     "data": [
+#         {"embedding": [0.1, 0.2, 0.3]},
+#         {"embedding": [0.4, 0.5, 0.6]}
+#     ]
+# }
+# chat_response = {
+#     "choices": [
+#         {"message": {"content": "This is a test response"}}
+#     ]
+# }
+#
+#
+# @pytest.fixture
+# def mock_openai_client(mocker):
+#     mock_client = MagicMock()
+#     mock_client.embeddings.create.return_value = embedding_response
+#     mock_client.chat.completions.create.return_value = chat_response
+#     mocker.patch("openai.Client", return_value=mock_client)
+#     return mock_client
+#
+#
+# @pytest.fixture
+# def openai_connector(mock_openai_client):
+#     return OpenAIConnector(api_key="fake-api-key")
+#
+#
+# def test_embed(openai_connector):
+#     model = LanguageModels.OPENAI_GPT35TURBO
+#     embeddings = openai_connector.embed(documents, model)
+#
+#     assert isinstance(embeddings, list)
+#     assert all(isinstance(e, np.ndarray) for e in embeddings)
+#     assert np.array_equal(embeddings[0], np.array([0.1, 0.2, 0.3]))
+#     assert np.array_equal(embeddings[1], np.array([0.4, 0.5, 0.6]))
+#
+#
+# def test_chat(openai_connector):
+#     prompt = "Hey, how are you?"
+#     model = LanguageModels.OPENAI_GPT35TURBO
+#     response = openai_connector.chat(prompt, model)
+#
+#     assert isinstance(response, str)
+#     assert response == "This is a test response"

tests/test_similarity_search.py

@@ -1,62 +1,48 @@
-from unittest.mock import MagicMock
-
 import numpy as np
-import pytest
-from qdrant_client.http.models import PointStruct, VectorParams, Distance
 
 from easy_letters import Ranker
 
-# Sample data for testing
+# Sample documents and their embeddings for testing
 documents_with_embeddings = {
     'text': ["Document 1", "Document 2"],
     'embedding': [np.array([0.1, 0.2, 0.3]), np.array([0.4, 0.5, 0.6])]
 }
+
+# Sample embedding to search for similar documents
 embedding_to_search = np.array([0.1, 0.2, 0.3])
+
+# The expected response (score is Cosine similarity)
 search_response = [
     {"id": 1, "score": 0.95, "payload": {"text": "Document 1"}},
     {"id": 2, "score": 0.85, "payload": {"text": "Document 2"}}
 ]
 
 
-@pytest.fixture
-def mock_qdrant_client(mocker):
-    mock_client = MagicMock()
-    mock_client.search.return_value = search_response
-    mocker.patch("qdrant_client.QdrantClient", return_value=mock_client)
-    return mock_client
-
-
-@pytest.fixture
-def ranker(mock_qdrant_client):
-    return Ranker()
-
-
-def test_make_collection(ranker, mock_qdrant_client):
+def test_make_collection():
+    # Arrange
+    ranker = Ranker()
     collection_name = "test_collection"
+
+    # Act
     ranker.make_collection(documents_with_embeddings, collection_name)
 
-    mock_qdrant_client.create_collection.assert_called_once_with(
-        collection_name=collection_name,
-        vectors_config=VectorParams(size=3, distance=Distance.COSINE)
-    )
-    points = [
-        PointStruct(id=0, vector=documents_with_embeddings['embedding'][0],
-                    payload={'text': "Document 1"}),
-        PointStruct(id=1, vector=documents_with_embeddings['embedding'][1],
-                    payload={'text': "Document 2"})
-    ]
-    mock_qdrant_client.upsert.assert_called_once_with(collection_name, points)
-
-
-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
+    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