From a1603d4c6c29f414304fc379074eb81b5b00c5d0 Mon Sep 17 00:00:00 2001
From: Gustaf Rydholm <gustaf.rydholm@gmail.com>
Date: Sat, 6 Apr 2024 00:18:57 +0200
Subject: Add logging in dbs

---
 rag/db/documents.py  |  8 ++++++--
 rag/db/embeddings.py | 47 ----------------------------------------------
 rag/db/vectors.py    | 53 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 59 insertions(+), 49 deletions(-)
 delete mode 100644 rag/db/embeddings.py
 create mode 100644 rag/db/vectors.py

diff --git a/rag/db/documents.py b/rag/db/documents.py
index 7d088da..6f83b1f 100644
--- a/rag/db/documents.py
+++ b/rag/db/documents.py
@@ -4,6 +4,7 @@ from typing import List
 
 import psycopg
 from langchain_core.documents.base import Document
+from loguru import logger as log
 
 TABLES = """
 CREATE TABLE IF NOT EXISTS document (
@@ -16,21 +17,24 @@ class Documents:
         self.conn = psycopg.connect(
             f"dbname={os.environ['RAG_DB_NAME']} user={os.environ['RAG_DB_USER']}"
         )
-        self.__create_content_table()
+        self.__configure()
 
     def close(self):
         self.conn.close()
 
-    def __create_content_table(self):
+    def __configure(self):
+        log.debug("Creating documents table if it does not exist...")
         with self.conn.cursor() as cur:
             cur.execute(TABLES)
             self.conn.commit()
 
     def __hash(self, chunks: List[Document]) -> str:
+        log.debug("Generating sha256 hash for pdf document")
         document = str.encode("".join([chunk.page_content for chunk in chunks]))
         return hashlib.sha256(document).hexdigest()
 
     def add_document(self, chunks: List[Document]) -> bool:
+        log.debug("Inserting document hash into documents db...")
         with self.conn.cursor() as cur:
             hash = self.__hash(chunks)
             cur.execute(
diff --git a/rag/db/embeddings.py b/rag/db/embeddings.py
deleted file mode 100644
index 4f61dcc..0000000
--- a/rag/db/embeddings.py
+++ /dev/null
@@ -1,47 +0,0 @@
-import os
-from dataclasses import dataclass
-from typing import Dict, List
-from uuid import uuid4
-
-from qdrant_client import QdrantClient
-from qdrant_client.http.models import StrictFloat
-from qdrant_client.models import Distance, VectorParams, PointStruct
-
-
-@dataclass
-class Point:
-    id: str
-    vector: List[StrictFloat]
-    payload: Dict[str, str]
-
-
-class Embeddings:
-    def __init__(self):
-        self.dim = int(os.environ["EMBEDDING_DIM"])
-        self.collection_name = os.environ["QDRANT_COLLECTION_NAME"]
-        self.client = QdrantClient(url=os.environ["QDRANT_URL"])
-        self.client.delete_collection(
-            collection_name=self.collection_name,
-        )
-        self.client.create_collection(
-            collection_name=self.collection_name,
-            vectors_config=VectorParams(size=self.dim, distance=Distance.COSINE),
-        )
-
-    def add(self, points: List[Point]):
-        print(len(points))
-        self.client.upload_points(
-            collection_name=self.collection_name,
-            points=[
-                PointStruct(id=point.id, vector=point.vector, payload=point.payload)
-                for point in points
-            ],
-            parallel=4,
-            max_retries=3,
-        )
-
-    def search(self, query: List[float], limit: int = 4):
-        hits = self.client.search(
-            collection_name=self.collection_name, query_vector=query, limit=limit
-        )
-        return hits
diff --git a/rag/db/vectors.py b/rag/db/vectors.py
new file mode 100644
index 0000000..9e8becb
--- /dev/null
+++ b/rag/db/vectors.py
@@ -0,0 +1,53 @@
+import os
+from dataclasses import dataclass
+from typing import Dict, List
+
+from qdrant_client import QdrantClient
+from qdrant_client.http.models import StrictFloat
+from qdrant_client.models import Distance, ScoredPoint, VectorParams, PointStruct
+from loguru import logger as log
+
+
+@dataclass
+class Point:
+    id: str
+    vector: List[StrictFloat]
+    payload: Dict[str, str]
+
+
+class Vectors:
+    def __init__(self):
+        self.dim = int(os.environ["EMBEDDING_DIM"])
+        self.collection_name = os.environ["QDRANT_COLLECTION_NAME"]
+        self.client = QdrantClient(url=os.environ["QDRANT_URL"])
+        self.__configure()
+
+    def __configure(self):
+        collections = list(map(lambda col: col.name, self.client.get_collections()))
+        if self.collection_name not in collections:
+            log.debug(f"Configuring collection {self.collection_name}...")
+            self.client.create_collection(
+                collection_name=self.collection_name,
+                vectors_config=VectorParams(size=self.dim, distance=Distance.COSINE),
+            )
+        else:
+            log.debug(f"Collection {self.collection_name} already exists...")
+
+    def add(self, points: List[Point]):
+        log.debug(f"Inserting {len(points)} vectors into the vector db...")
+        self.client.upload_points(
+            collection_name=self.collection_name,
+            points=[
+                PointStruct(id=point.id, vector=point.vector, payload=point.payload)
+                for point in points
+            ],
+            parallel=4,
+            max_retries=3,
+        )
+
+    def search(self, query: List[float], limit: int = 4) -> List[ScoredPoint]:
+        log.debug("Searching for vectors...")
+        hits = self.client.search(
+            collection_name=self.collection_name, query_vector=query, limit=limit
+        )
+        return hits
-- 
cgit v1.2.3-70-g09d2