renamed models

backend/core/quivr_core/api/models/databases/supabase/files.py → backend/core/quivr_core/api/models/file_repository.py

@@ -1,27 +1,31 @@
-from quivr_core.api.models.databases.repository import Repository
+from uuid import UUID
 
+from quivr_core.api.models.settings import get_supabase_client
 
-class File(Repository):
-    def __init__(self, supabase_client):
-        self.db = supabase_client
 
-    def set_file_vectors_ids(self, file_sha1):
+class FileRepository:
+    def __init__(self):
+        self.db = get_supabase_client()
+
+    def set_file_vectors_ids(self, file_sha1: bytes):
         response = (
             self.db.table("vectors")
             .select("id")
-            .filter("file_sha1", "eq", file_sha1)
+            .filter("file_sha1", "eq", str(file_sha1))
             .execute()
         )
         return response.data
 
-    def get_brain_vectors_by_brain_id_and_file_sha1(self, brain_id, file_sha1):
+    def get_brain_vectors_by_brain_id_and_file_sha1(
+        self, brain_id: UUID, file_sha1: bytes
+    ):
         self.set_file_vectors_ids(file_sha1)
         # Check if file exists in that brain
         response = (
             self.db.table("brains_vectors")
             .select("brain_id, vector_id")
             .filter("brain_id", "eq", str(brain_id))
-            .filter("file_sha1", "eq", file_sha1)
+            .filter("file_sha1", "eq", str(file_sha1))
             .execute()
         )
 

backend/core/quivr_core/api/models/files.py

@@ -6,8 +6,8 @@
 from pydantic import BaseModel
 
 from quivr_core.api.logger import get_logger
-from quivr_core.api.models.databases.supabase.supabase import SupabaseDB
-from quivr_core.api.models.settings import get_supabase_db
+from quivr_core.api.models.file_repository import FileRepository
+from quivr_core.api.models.vectors import Vector
 from quivr_core.api.modules.brain.service.brain_vector_service import BrainVectorService
 from quivr_core.api.packages.files.file import compute_sha1_from_content
 
@@ -28,12 +28,10 @@ class File(BaseModel):
 
     def __init__(self, **data):
         super().__init__(**data)
+        self.vectors_repository = Vector()
+        self.file_repository = FileRepository()
         data["file_sha1"] = compute_sha1_from_content(data["bytes_content"])
 
-    @property
-    def supabase_db(self) -> SupabaseDB:
-        return get_supabase_db()
-
     def compute_documents(self, loader_class):
         """
         Compute the documents from the file
@@ -56,7 +54,7 @@ def set_file_vectors_ids(self):
         Set the vectors_ids property with the ids of the vectors
         that are associated with the file in the vectors table
         """
-        self.vectors_ids = self.supabase_db.get_vectors_by_file_sha1(
+        self.vectors_ids = self.vectors_repository.get_vectors_by_file_sha1(
             self.file_sha1
         ).data
 
@@ -79,7 +77,7 @@ def file_already_exists_in_brain(self, brain_id):
         Args:
             brain_id (str): Brain id
         """
-        response = self.supabase_db.get_brain_vectors_by_brain_id_and_file_sha1(
+        response = self.file_repository.get_brain_vectors_by_brain_id_and_file_sha1(
             brain_id,
             self.file_sha1,  # type: ignore
         )