summaryrefslogtreecommitdiff
path: root/rag
diff options
context:
space:
mode:
Diffstat (limited to 'rag')
-rw-r--r--rag/retriever/document.py13
1 files changed, 8 insertions, 5 deletions
diff --git a/rag/retriever/document.py b/rag/retriever/document.py
index 67ff936..8a50f01 100644
--- a/rag/retriever/document.py
+++ b/rag/retriever/document.py
@@ -7,14 +7,17 @@ from loguru import logger as log
TABLES = """
CREATE TABLE IF NOT EXISTS document (
- hash text PRIMARY KEY)
+ hash VARCHAR(64) PRIMARY KEY)
"""
class DocumentDB:
def __init__(self) -> None:
self.conn = psycopg.connect(
- f"dbname={os.environ['DOCUMENT_DB_NAME']} user={os.environ['DOCUMENT_DB_USER']}"
+ (
+ f"dbname={os.environ['DOCUMENT_DB_NAME']} "
+ f"user={os.environ['DOCUMENT_DB_USER']}"
+ )
)
self.__configure()
@@ -36,7 +39,7 @@ class DocumentDB:
hash = self.__hash(blob)
cur.execute(
"""
- SELECT * FROM document
+ SELECT * FROM document
WHERE
hash = %s
""",
@@ -47,8 +50,8 @@ class DocumentDB:
log.debug("Inserting document hash into documents db...")
cur.execute(
"""
- INSERT INTO document
- (hash) VALUES
+ INSERT INTO document
+ (hash) VALUES
(%s)
""",
(hash,),