Small fixes

deepsense-ai · Oct 7, 2024 · d306f84 · d306f84
1 parent 843311f
commit d306f84
Show file tree

Hide file tree

Showing 5 changed files with 22 additions and 10 deletions.
diff --git a/packages/ragbits-core/src/ragbits/core/vector_store/__init__.py b/packages/ragbits-core/src/ragbits/core/vector_store/__init__.py
@@ -2,10 +2,10 @@
 
 from ..utils import get_cls_from_config
 from .base import VectorStore
-# from .chromadb_store import ChromaDBStore
+from .chromadb_store import ChromaDBStore
 from .in_memory import InMemoryVectorStore
 
-__all__ = ["InMemoryVectorStore", "VectorStore"] # , "ChromaDBStore"]
+__all__ = ["InMemoryVectorStore", "VectorStore", "ChromaDBStore"]
 
 module = sys.modules[__name__]
 

diff --git a/packages/ragbits-document-search/src/ragbits/document_search/_main.py b/packages/ragbits-document-search/src/ragbits/document_search/_main.py
@@ -1,16 +1,16 @@
-from ragbits.core.embeddings import get_embeddings
 from typing import Any, Optional, Union
 
 from pydantic import BaseModel, Field
 
+from ragbits.core.embeddings import get_embeddings
 from ragbits.core.embeddings.base import Embeddings
-from ragbits.core.vector_store.base import VectorStore
 from ragbits.core.vector_store import get_vector_store
+from ragbits.core.vector_store.base import VectorStore
 from ragbits.document_search.documents.document import Document, DocumentMeta
 from ragbits.document_search.documents.element import Element
-from ragbits.document_search.retrieval.rephrasers import get_rephraser
 from ragbits.document_search.ingestion.document_processor import DocumentProcessorRouter
 from ragbits.document_search.ingestion.providers.base import BaseProvider
+from ragbits.document_search.retrieval.rephrasers import get_rephraser
 from ragbits.document_search.retrieval.rephrasers.base import QueryRephraser
 from ragbits.document_search.retrieval.rephrasers.noop import NoopQueryRephraser
 from ragbits.document_search.retrieval.rerankers import get_reranker

diff --git a/...ages/ragbits-document-search/src/ragbits/document_search/retrieval/rephrasers/__init__.py b/...ages/ragbits-document-search/src/ragbits/document_search/retrieval/rephrasers/__init__.py
@@ -2,6 +2,7 @@
 from typing import Optional
 
 from ragbits.core.utils import get_cls_from_config
+
 from .base import QueryRephraser
 from .noop import NoopQueryRephraser
 

diff --git a/packages/ragbits-document-search/src/ragbits/document_search/retrieval/rerankers/__init__.py b/packages/ragbits-document-search/src/ragbits/document_search/retrieval/rerankers/__init__.py
@@ -2,6 +2,7 @@
 from typing import Optional
 
 from ragbits.core.utils import get_cls_from_config
+
 from .base import Reranker
 from .noop import NoopReranker
 

diff --git a/packages/ragbits-document-search/tests/unit/test_document_search.py b/packages/ragbits-document-search/tests/unit/test_document_search.py
@@ -1,10 +1,10 @@
-from ragbits.core.embeddings.noop import NoopEmbeddings
 from pathlib import Path
 from typing import Union
 from unittest.mock import AsyncMock
 
 import pytest
 
+from ragbits.core.embeddings.noop import NoopEmbeddings
 from ragbits.core.vector_store.in_memory import InMemoryVectorStore
 from ragbits.document_search import DocumentSearch
 from ragbits.document_search._main import SearchConfig
@@ -14,9 +14,7 @@
 
 CONFIG = {
     "embedder": {"type": "NoopEmbeddings"},
-    "vector_store": {
-        "type": "packages.ragbits-core.src.ragbits.core.vector_store.in_memory:InMemoryVectorStore"
-    },
+    "vector_store": {"type": "packages.ragbits-core.src.ragbits.core.vector_store.in_memory:InMemoryVectorStore"},
     "reranker": {"type": "NoopReranker"},
 }
 
@@ -34,7 +32,7 @@ async def test_document_search():
 
     assert isinstance(first_result, TextElement)
     assert first_result.content == "Name of Peppa's brother is George"
-    
+
 
 async def test_document_search_from_config():
     document_search = DocumentSearch.from_config(CONFIG)
@@ -49,6 +47,7 @@ async def test_document_search_from_config():
     assert isinstance(first_result, TextElement)
     assert first_result.content == "Name of Peppa's brother is George"
 
+
 @pytest.mark.parametrize(
     "document",
     [
@@ -62,6 +61,17 @@ async def test_document_search_ingest_document(document: Union[DocumentMeta, Doc
     embeddings_mock = AsyncMock()
     embeddings_mock.embed_text.return_value = [[0.1, 0.1]]
 
+    document_search = DocumentSearch(embedder=embeddings_mock, vector_store=InMemoryVectorStore())
+
+    await document_search.ingest_document(document, document_processor=DummyProvider())
+
+    results = await document_search.search("Peppa's brother")
+
+    first_result = results[0]
+
+    assert isinstance(first_result, TextElement)
+    assert first_result.content == "Name of Peppa's brother is George"
+
 
 async def test_document_search_insert_elements():
     embeddings_mock = AsyncMock()