colbert integration test on cassandra (#396)

* colbert integration test on cassandra

* use nested event_loop

* fix file name typo

* add comments

Author: zzzming

libs/colbert/pyproject.toml

@@ -15,6 +15,7 @@ colbert-ai = "0.2.19"
 pyarrow = "14.0.1"
 torch = "2.2.1"
 cassio = "~0.1.7"
+nest-asyncio = "^1.6.0"
 
 [tool.poetry.group.test.dependencies]
 ragstack-ai-tests-utils = { path = "../tests-utils", develop = true }

libs/colbert/ragstack_colbert/__init__.py

@@ -15,7 +15,7 @@
 """
 
 from .cassandra_vector_store import CassandraVectorStore
-from .cobert_retriever import ColbertRetriever
+from .colbert_retriever import ColbertRetriever
 from .colbert_embedding_model import ColbertEmbeddingModel
 from .constant import DEFAULT_COLBERT_DIM, DEFAULT_COLBERT_MODEL
 from .objects import ChunkData, EmbeddedChunk, RetrievedChunk

libs/colbert/ragstack_colbert/cobert_retriever.py renamed to libs/colbert/ragstack_colbert/colbert_retriever.py

@@ -11,6 +11,7 @@
 import asyncio
 import logging
 import math
+import nest_asyncio
 from typing import Any, Dict, List, Optional, Set
 
 import torch
@@ -336,5 +337,8 @@ def retrieve(
             The actual retrieval process involves encoding the query, performing an ANN search to find relevant
             embeddings, scoring these embeddings for similarity, and retrieving the corresponding text chunks.
         """
+        # nest_asyncio does not a new event loop to be created
+        # in the case there is already an event loop such as colab, it's required
+        nest_asyncio.apply()
         loop = asyncio.get_event_loop()
         return loop.run_until_complete(self.aretrieve(query=query, k=k, query_maxlen=query_maxlen))

libs/colbert/tests/integration_tests/test_colbert_embedding_retrieval.py

@@ -25,7 +25,7 @@ def astra_db():
 
 
 #@pytest.mark.parametrize("vector_store", ["cassandra", "astra_db"])
-@pytest.mark.parametrize("vector_store", ["astra_db"])
+@pytest.mark.parametrize("vector_store", ["cassandra"])
 def test_embedding_cassandra_retriever(request, vector_store: str):
     vector_store = request.getfixturevalue(vector_store)
     narrative = """

libs/colbert/tests/unit_tests/test_colbert_retriever.py

@@ -1,6 +1,6 @@
 import torch
 
-from ragstack_colbert.cobert_retriever import max_similarity_torch
+from ragstack_colbert.colbert_retriever import max_similarity_torch
 from ragstack_colbert.colbert_embedding_model import calculate_query_maxlen