Migrate to service approach for file ingestion

Author: edkaya

genai/rag/ingestion_pipeline.py

@@ -0,0 +1,67 @@
+from typing import List
+from uuid import uuid4
+from time import perf_counter
+
+from langchain_qdrant import QdrantVectorStore
+from langchain_core.documents import Document
+from langchain_community.document_loaders import PyPDFLoader
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+
+from logger import logger
+from metrics import file_ingested_counter, file_ingestion_duration
+
+
+def _load_document(file_path: str) -> List[Document]:
+    """Load a PDF document from the given file path"""
+    loader = PyPDFLoader(file_path)
+    docs = loader.load()
+    return docs
+
+
+def _chunk_documents(docs: List[Document], filename: str) -> List[Document]:
+    """Chunk documents into smaller pieces for
+    better processing and storage."""
+    text_splitter = RecursiveCharacterTextSplitter(
+        chunk_size=1000,
+        chunk_overlap=100
+    )
+    chunked_docs = text_splitter.split_documents(docs)
+    return [
+        Document(
+            page_content=chunk.page_content,
+            metadata={
+                **chunk.metadata,
+                "source": filename
+                }
+        )
+        for chunk in chunked_docs
+    ]
+
+
+def _store_documents(vector_store: QdrantVectorStore, docs: List[Document]):
+    """Store all chunks in vector database"""
+    uuids = [str(uuid4()) for _ in range(len(docs))]
+    vector_store.add_documents(docs, ids=uuids)
+
+
+def ingest(vector_store: QdrantVectorStore, file_path: str, filename: str):
+    """Ingestion method to load, chunk, and store pdf data"""
+    start_time = perf_counter()
+
+    unchunked_docs = _load_document(file_path)
+    logger.info("Documents are loaded for file %s", filename)
+
+    chunked_docs = _chunk_documents(unchunked_docs, filename)
+    logger.info("Documents are chunked for file %s", filename)
+
+    _store_documents(vector_store, chunked_docs)
+    logger.info(
+        "Ingested %s chunks into Qdrant for file %s.",
+        len(chunked_docs),
+        filename
+    )
+
+    duration = perf_counter() - start_time
+    file_ingestion_duration.observe(duration)
+    file_ingested_counter.inc()
+    logger.info("File ingestion duration: %.2f seconds", duration)

genai/service/ingestion_service.py

@@ -1,4 +1,6 @@
-from rag.ingestion_pipeline import IngestionPipeline
+from langchain_qdrant import QdrantVectorStore
+
+from service.ingestion_service import ingest
 from vector_database.qdrant_vdb import QdrantVDB
 from logger import logger
 
@@ -24,7 +26,7 @@ def file_already_uploaded(collection_name: str, filename: str) -> bool:
     return False
 
 
-def get_vector_store(collection_name: str):
+def get_vector_store(collection_name: str) -> QdrantVectorStore:
     """
     Returns the vector store for the specified collection name.
     This function is used to retrieve the vector store instance for a specific
@@ -53,7 +55,5 @@ def ingest_file(collection_name: str, file_path: str, filename: str):
     This function is used to process and store the file content in the vector
     database for later retrieval.
     """
-    pipeline = IngestionPipeline(
-        vector_store=get_vector_store(collection_name)
-        )
-    pipeline.ingest(file_path, filename)
+    vector_store = get_vector_store(collection_name)
+    ingest(vector_store, file_path, filename)

genai/service/qdrant_service.py

@@ -1,6 +1,6 @@
 from pathlib import Path
 from fpdf import FPDF
-from rag.ingestion_pipeline import IngestionPipeline
+from service.ingestion_service import ingest
 from vector_database.qdrant_vdb import QdrantVDB
 
 
@@ -28,10 +28,11 @@ def test_real_ingestion_pipeline(tmp_path):
     generate_sample_pdf(pdf_path)
     filename = pdf_path.name
 
-    # Ingestion
+    # Get vector store
     vector_store = qdrant.create_and_get_vector_storage(collection_name)
-    pipeline = IngestionPipeline(vector_store=vector_store)
-    pipeline.ingest(str(pdf_path), filename)
+
+    # Ingestion
+    ingest(vector_store, str(pdf_path), filename)
 
     found = qdrant.collection_contains_file(
         qdrant.client,

genai/tests/integration/ingestion_test.py

@@ -1,30 +1,22 @@
 import io
-from unittest.mock import patch, MagicMock
+from unittest.mock import patch
 from fastapi.testclient import TestClient
 from main import app
 from service.auth_service import UserInfo, get_current_user
 
 client = TestClient(app)
 
 
-@patch(
-    "service.qdrant_service.qdrant.client.collection_exists",
-    return_value=False
-    )
-@patch("service.qdrant_service.qdrant.create_and_get_vector_storage")
-@patch("service.qdrant_service.IngestionPipeline")
+@patch("routes.routes.ingest_file")
+@patch("routes.routes.file_already_uploaded", return_value=False)
 def test_upload_file_success(
-        mock_pipeline_class,
-        _mock_vector_store,
-        _mock_exists
+        mock_file_uploaded,
+        mock_ingest_file
 ):
     # Mock the authentication by overriding the dependency
     mock_user = UserInfo(user_id=123, username="test_user")
     app.dependency_overrides[get_current_user] = lambda: mock_user
 
-    mock_pipeline = MagicMock()
-    mock_pipeline_class.return_value = mock_pipeline
-
     file_content = b"%PDF-1.4 dummy content"
     file = io.BytesIO(file_content)
     file.name = "test.pdf"
@@ -37,8 +29,8 @@ def test_upload_file_success(
     assert response.status_code == 200
     assert response.json() == {"message": "File processed successfully."}
 
-    mock_pipeline_class.assert_called_once()
-    mock_pipeline.ingest.assert_called_once()
+    mock_ingest_file.assert_called_once()
+    mock_file_uploaded.assert_called_once()
 
     # Clean up the dependency override
     app.dependency_overrides.clear()

genai/tests/integration/upload_test.py

@@ -1,38 +1,40 @@
 from unittest.mock import MagicMock, patch
 from uuid import UUID
-from rag.ingestion_pipeline import IngestionPipeline
+from service.ingestion_service import (
+    ingest,
+    _load_document,
+    _chunk_documents,
+    _store_documents
+)
 from langchain_core.documents import Document
 
 
 def test_load_document_returns_documents():
-    with patch("rag.ingestion_pipeline.PyPDFLoader") as mock_loader:
+    with patch("service.ingestion_service.PyPDFLoader") as mock_loader:
         mock_loader.return_value.load.return_value = [
             Document(page_content="Test")
         ]
-        pipeline = IngestionPipeline(vector_store=MagicMock())
-        docs = pipeline.load_document("fake_path.pdf")
+        docs = _load_document("fake_path.pdf")
         assert isinstance(docs, list)
         assert isinstance(docs[0], Document)
         assert docs[0].page_content == "Test"
 
 
 def test_chunk_documents_returns_chunks():
-    pipeline = IngestionPipeline(vector_store=MagicMock())
     dummy_doc = Document(
         page_content="This is a long text. " * 100,
         metadata={}
     )
-    chunks = pipeline.chunk_documents([dummy_doc], filename="sample.pdf")
+    chunks = _chunk_documents([dummy_doc], filename="sample.pdf")
     assert isinstance(chunks, list)
     assert all(isinstance(doc, Document) for doc in chunks)
     assert all(doc.metadata["source"] == "sample.pdf" for doc in chunks)
 
 
 def test_store_documents_calls_add_documents_with_uuids():
     mock_vector_store = MagicMock()
-    pipeline = IngestionPipeline(vector_store=mock_vector_store)
     docs = [Document(page_content="Chunk", metadata={}) for _ in range(3)]
-    pipeline.store_documents(docs)
+    _store_documents(mock_vector_store, docs)
     args, kwargs = mock_vector_store.add_documents.call_args
     passed_docs = args[0]
     passed_ids = kwargs["ids"]
@@ -42,19 +44,18 @@ def test_store_documents_calls_add_documents_with_uuids():
 
 
 def test_ingest_calls_all_steps():
-    pipeline = IngestionPipeline(vector_store=MagicMock())
-
-    with patch.object(pipeline, "load_document") as mock_load, \
-         patch.object(pipeline, "chunk_documents") as mock_chunk, \
-         patch.object(pipeline, "store_documents") as mock_store, \
-         patch("rag.ingestion_pipeline.logger") as mock_logger, \
-         patch("rag.ingestion_pipeline.file_ingestion_duration.observe"), \
-         patch("rag.ingestion_pipeline.file_ingested_counter.inc"):
+    mock_vector_store = MagicMock()
+    with patch("service.ingestion_service._load_document") as mock_load, \
+         patch("service.ingestion_service._chunk_documents") as mock_chunk, \
+         patch("service.ingestion_service._store_documents") as mock_store, \
+         patch("service.ingestion_service.logger") as mock_logger, \
+         patch("service.ingestion_service.file_ingestion_duration.observe"), \
+         patch("service.ingestion_service.file_ingested_counter.inc"):
 
         mock_load.return_value = [Document(page_content="Doc")]
         mock_chunk.return_value = [Document(page_content="Chunk")]
 
-        pipeline.ingest("test.pdf", "testfile.pdf")
+        ingest(mock_vector_store, "test.pdf", "testfile.pdf")
 
         mock_load.assert_called_once_with("test.pdf")
         mock_chunk.assert_called_once()