From 3077bfdff3587aac5ef91a8fb5730c11c8bb5273 Mon Sep 17 00:00:00 2001 From: Gustaf Rydholm Date: Sat, 6 Apr 2024 01:53:41 +0200 Subject: Refactor doc db --- rag/db/document.py | 2 +- rag/rag.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/rag/db/document.py b/rag/db/document.py index c449158..f8c97c2 100644 --- a/rag/db/document.py +++ b/rag/db/document.py @@ -33,7 +33,7 @@ class DocumentDB: document = str.encode("".join([chunk.page_content for chunk in chunks])) return hashlib.sha256(document).hexdigest() - def add_document(self, chunks: List[Document]) -> bool: + def add(self, chunks: List[Document]) -> bool: log.debug("Inserting document hash into documents db...") with self.conn.cursor() as cur: hash = self.__hash(chunks) diff --git a/rag/rag.py b/rag/rag.py index 7242db9..87b44c5 100644 --- a/rag/rag.py +++ b/rag/rag.py @@ -22,7 +22,7 @@ class RAG: def add_pdf(self, filepath: Path): chunks = pdf.parser(filepath) - added = self.document_db.add_document(chunks) + added = self.document_db.add(chunks) if added: log.debug(f"Adding pdf with filepath: {filepath} to vector db") points = self.encoder.encode_document(chunks) -- cgit v1.2.3-70-g09d2