diff options
| -rw-r--r-- | rag/db/document.py (renamed from rag/db/documents.py) | 2 | ||||
| -rw-r--r-- | rag/db/vector.py (renamed from rag/db/vectors.py) | 10 |
2 files changed, 7 insertions, 5 deletions
diff --git a/rag/db/documents.py b/rag/db/document.py index 6f83b1f..8e4d208 100644 --- a/rag/db/documents.py +++ b/rag/db/document.py @@ -12,7 +12,7 @@ CREATE TABLE IF NOT EXISTS document ( """ -class Documents: +class DocumentDB: def __init__(self) -> None: self.conn = psycopg.connect( f"dbname={os.environ['RAG_DB_NAME']} user={os.environ['RAG_DB_USER']}" diff --git a/rag/db/vectors.py b/rag/db/vector.py index 9e8becb..4aa62cc 100644 --- a/rag/db/vectors.py +++ b/rag/db/vector.py @@ -2,10 +2,10 @@ import os from dataclasses import dataclass from typing import Dict, List +from loguru import logger as log 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 +from qdrant_client.models import Distance, PointStruct, ScoredPoint, VectorParams @dataclass @@ -15,7 +15,7 @@ class Point: payload: Dict[str, str] -class Vectors: +class VectorDB: def __init__(self): self.dim = int(os.environ["EMBEDDING_DIM"]) self.collection_name = os.environ["QDRANT_COLLECTION_NAME"] @@ -23,7 +23,9 @@ class Vectors: self.__configure() def __configure(self): - collections = list(map(lambda col: col.name, self.client.get_collections())) + collections = list( + map(lambda col: col.name, self.client.get_collections().collections) + ) if self.collection_name not in collections: log.debug(f"Configuring collection {self.collection_name}...") self.client.create_collection( |