refactor: pass callable instead of object

server/api/v1/chat.py

@@ -105,7 +105,7 @@ async def upload_files(
         chunk_generator = store_chunks(
             redis,
             chat_id,
-            embedder,
+            embedder.encode_normalise,
             extract_documents_from_pdfs([{'data': BytesIO(await file.read()), 'name': file.filename} for file in data]),
             chunk_document,
             text_splitter,

server/databases/redis/features/store_chunks.py

@@ -1,10 +1,5 @@
 from asyncio import gather
-from typing import (
-    AsyncGenerator,
-    Callable,
-    Iterable,
-    Protocol,
-)
+from typing import AsyncIterator, Callable, Iterable
 
 from numpy import float32
 from numpy.typing import NDArray
@@ -15,30 +10,14 @@
 from server.features.extraction.models import Document
 
 
-class Embedder(Protocol):
-    """
-    Summary
-    -------
-    a generic protocol for embedding text
-    """
-
-    def encode_normalise(self, sentences: str | list[str]) -> NDArray[float32]:
-        """
-        Summary
-        -------
-        encode a sentence or list of sentences into a normalised embedding
-        """
-        ...
-
-
 async def store_chunks(
     redis: RedisAsync,
     chat_id: str,
-    embedder: Embedder,
+    embedder: Callable[[str], NDArray[float32]],
     documents: Iterable[Document | None],
     chunker: Callable[[Document, TextSplitter], Iterable[Chunk]],
     text_splitter: TextSplitter,
-) -> AsyncGenerator[tuple[str, str] | tuple[None, None], None]:
+) -> AsyncIterator[tuple[str, str] | tuple[None, None]]:
     """
     Summary
     -------
@@ -70,7 +49,7 @@ async def store_chunks(
                     chunk.source_id,
                     chunk.id,
                     mapping={
-                        'vector': embedder.encode_normalise(chunk.content).tobytes(),
+                        'vector': embedder(chunk.content).tobytes(),
                         'content': chunk.content,
                         'chat_id': chat_id,
                     },