Use selenium for testing

Author: abhikdps

.gitignore

@@ -7,4 +7,6 @@ __pycache__/
 .env.*
 .idea/
 .vscode/
-*.db
+*.db
+.mypy_cache/
+.cache/

config.yaml

@@ -0,0 +1,18 @@
+version: v1
+ingest_threads: 8
+collections:
+  - name: "Source Collection"
+    id: "source_collection"
+    mode: "overwrite"
+    chunk_size: 500
+    chunk_overlap: 250
+    embedding_model: "all-MiniLM-L6-v2"
+    metadata:
+      key: "value"
+    sources:
+      - type: "source"
+        url_fragment: "/departments"
+        recursive: true
+        attachments: true
+        metadata:
+          key: "value"

config/config.yaml

@@ -7,16 +7,20 @@
 load_dotenv()
 
 # PATHS
-DIRECTORY_PATH = pathlib.Path(os.path.dirname(__file__)).parent
+DIRECTORY_PATH = pathlib.Path.cwd()
 KNOWLEDGE_REPOSITORY_PATH = DIRECTORY_PATH / "knowledge"
 SOURCE_RESPOSITORY_PATH = KNOWLEDGE_REPOSITORY_PATH / "source"
 
 # INGEST
-DEVICE = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")
+DEVICE = (
+    "cuda"
+    if torch.cuda.is_available()
+    else ("mps" if torch.backends.mps.is_available() else "cpu")
+)
 
 # PGVECTOR
-PGVECTOR_USER = os.environ.get("PGVECTOR_USER", None)
-PGVECTOR_PASS = os.environ.get("PGVECTOR_PASS", None)
-PGVECTOR_DATABASE_NAME = os.environ.get("PGVECTOR_DATABASE_NAME", None)
+PGVECTOR_USER = os.environ.get("PGVECTOR_USER")
+PGVECTOR_PASS = os.environ.get("PGVECTOR_PASS")
+PGVECTOR_DATABASE_NAME = os.environ.get("PGVECTOR_DATABASE_NAME")
 PGVECTOR_HOST = os.environ.get("PGVECTOR_URI", "localhost")
 PGVECTOR_PORT = int(os.environ.get("PGVECTOR_PORT", 5432))

scripts/constants.py

@@ -23,17 +23,14 @@
 logger = logging.getLogger(__name__)
 
 
-def get_embedder(embedding_model_name: str):
-    """Define embedder to convert text into vectors."""
-    model_kwargs = {"device": DEVICE}
-    embedder = HuggingFaceEmbeddings(
+def get_embedder(embedding_model_name: str) -> HuggingFaceEmbeddings:
+    """Initialize an embedder to convert text into vectors."""
+    return HuggingFaceEmbeddings(
         model_name=embedding_model_name,
-        model_kwargs=model_kwargs,
+        model_kwargs={"device": DEVICE},
         show_progress=True,
     )
 
-    return embedder
-
 
 def ingest(
     meta_lookup: dict[pathlib.Path, dict],
@@ -44,7 +41,6 @@ def ingest(
     embedding_model_name: str = "sentence-transformers/all-MiniLM-L6-v2",
     mode: str = "overwrite",
     collection_metadata: dict = {},
-    logs_folder_id: str = None,
 ):
     """Load documents into a vectorstore."""
     # Get documents
@@ -59,7 +55,9 @@ def ingest(
         file_name = source.stem
         document.metadata["_source"] = document.metadata["source"]
         document.metadata["source"] = file_name
-        chunks = split_document(document, extension, chunk_size=chunk_size, chunk_overlap=chunk_overlap)
+        chunks = split_document(
+            document, extension, chunk_size=chunk_size, chunk_overlap=chunk_overlap
+        )
         # Attach metadata to each chunk
         for chunk in chunks:
             path_metadata = meta_lookup.get(source, {})
@@ -101,11 +99,15 @@ def ingest(
         logger.info(f"Collection {collection_name} created")
 
     # Load the documents
-    logger.info(f"Loading {len(all_documents)} embeddings to {PGVECTOR_HOST} - {PGVECTOR_DATABASE_NAME} - {collection_name}")
+    logger.info(
+        f"Loading {len(all_documents)} embeddings to {PGVECTOR_HOST} - {PGVECTOR_DATABASE_NAME} - {collection_name}"
+    )
     db.add_documents(documents=all_documents)
     logger.info(f"Successfully loaded {len(all_documents)} embeddings")
 
-    directory_source_url_chunks = [list(origin_url) + [chunks] for origin_url, chunks in origin_urls.items()]
+    directory_source_url_chunks = [
+        list(origin_url) + [chunks] for origin_url, chunks in origin_urls.items()
+    ]
     df = pd.DataFrame(directory_source_url_chunks, columns=["origin", "url", "chunks"])
     filename = f"{PGVECTOR_HOST} - {collection_name} - {datetime.now()}.csv"
     outpath = DIRECTORY_PATH / "logs" / filename