Update database_module.py

Replaced some hard coded strings with constants
anthonymonforte · Aug 21, 2024 · 7480582 · 7480582
1 parent 656bb12
commit 7480582
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/src/modules/database_module.py b/src/modules/database_module.py
@@ -12,6 +12,8 @@
 #import chromadb
 #from chromadb.utils.data_loaders import ImageLoader
 
+from src.modules.constants import DictionaryKeys
+
 class EmbeddingDb:
 
     def __init__(self, db_path: str) -> None:
@@ -60,8 +62,8 @@ def generate_chunk_ids(self, chunks: list[Document]) -> list[Document]:
         current_chunk_index = 0
 
         for chunk in chunks:
-            source = chunk.metadata.get("source")
-            page = chunk.metadata.get("page")
+            source = chunk.metadata.get(DictionaryKeys.Source)
+            page = chunk.metadata.get(DictionaryKeys.Page)
             current_page_id = f"{source}:{page}"
 
             if current_page_id == last_page_id: