From 7480582f2e006e134a687e53021aed5ee83492db Mon Sep 17 00:00:00 2001
From: anthonymonforte <32207192+anthonymonforte@users.noreply.github.com>
Date: Wed, 21 Aug 2024 14:17:40 -0700
Subject: [PATCH] Update database_module.py

Replaced some hard coded strings with constants
---
 src/modules/database_module.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/modules/database_module.py b/src/modules/database_module.py
index 0a0120d..6951750 100644
--- 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: