🚰 feat: Stream Document Embeddings to Database in Batches (#214)

* Process documents using async pipeline with concurrent embedding generation and DB insertion.
This provides better memory efficiency by immediately inserting embeddings as they're generated.

* - mad maxsize of embedding-queue configurable with EMBEDIING_MAX_QUEUE_SIZE env var

* implemented fixes for copilot code review

* implemented fixes for copilot code review
- white space changes and remove unneeded import

* implemented fixes for copilot code review
- added documentation in function header for executor argument

* feat: improve streaming document embeddings with bug fixes and tests

This PR implements several critical fixes and improvements for the batch
embedding feature that streams document embeddings to reduce memory consumption.

## Critical Fixes
- Fix blocking async in _process_documents_batched_sync by wrapping sync
  calls in run_in_executor()
- Unify rollback behavior: add `if all_ids:` guard to async path to only
  rollback when documents were actually inserted
- Fix producer error handling: move `put(None)` to finally block to always
  signal completion
- Add task_done() call for None signal in consumer
- Fix consumer error handling: put exception object in results_queue instead
  of empty list, enabling proper error propagation

## Code Quality Improvements
- Remove dead `# temp_vector_store` comment
- Remove ineffective `del` statements (don't help with memory in Python)
- Fix all f-string logging to use %-style formatting
- Update docstrings to accurately describe function behavior
- Add type hints (AsyncPgVector, PgVector, ThreadPoolExecutor)
- Rename functions for clarity: embedding_producer → batch_producer,
  database_consumer → embedding_consumer
- Add early return for empty document lists

## Tests (31 new tests)
- tests/test_batch_processing.py: 23 unit tests covering async pipeline,
  sync batched processing, edge cases, and producer-consumer pattern
- tests/test_batch_processing_integration.py: 8 integration tests including
  memory optimization verification with tracemalloc

## Documentation
- Add EMBEDDING_BATCH_SIZE and EMBEDDING_MAX_QUEUE_SIZE to README env vars
- Add "Embedding Batch Processing" section with configuration guide
- Add "Running Tests" section with commands and test categories
- Improve inline comments in config.py explaining batch processing trade-offs
- Add pytest.ini for test configuration

Closes #213

* feat: add utility function for batch calculation and update tests

This commit introduces a new utility function `calculate_num_batches` to streamline batch size calculations across the document processing routes. The function handles various edge cases, including zero items and batch sizes of zero. Additionally, existing tests have been updated to utilize this new function, ensuring accurate batch count calculations and improving test coverage for edge cases.

* refactor: update vector store mocking in tests to use AsyncPgVector

This commit modifies the test setup in `test_main.py` to replace the mocking of vector store methods with the `AsyncPgVector` class. The changes include updating the signatures of dummy functions to accept `self` and ensuring that all relevant methods are patched correctly. This enhances the accuracy of the tests by aligning them with the asynchronous implementation of the vector store.

* test: clear cache for embedding function in vector store tests

This commit updates the test setup in `test_main.py` to clear the cache of the `get_cached_query_embedding` function before running tests. This ensures that the mock embedding function is used consistently, improving the reliability of the tests. Additionally, the comment for overriding the embedding function has been clarified to indicate that it uses a dummy implementation that does not call OpenAI.

* test: enhance vector store tests with dummy embedding function

This commit updates the test setup in `test_main.py` to clear the cache of the `get_cached_query_embedding` function and replaces it with a dummy implementation that returns predefined embeddings. This change improves the reliability of the tests by ensuring consistent behavior during test execution.

---------

Co-authored-by: Marc Amick <[email protected]>
Co-authored-by: Danny Avila <[email protected]>

Author: 3 people

README.md

@@ -59,6 +59,8 @@ The following environment variables are required to run the application:
 - `COLLECTION_NAME`: (Optional) The name of the collection in the vector store. Default value is "testcollection".
 - `CHUNK_SIZE`: (Optional) The size of the chunks for text processing. Default value is "1500".
 - `CHUNK_OVERLAP`: (Optional) The overlap between chunks during text processing. Default value is "100".
+- `EMBEDDING_BATCH_SIZE`: (Optional) Number of document chunks to process per batch. Set to `0` (default) to disable batching. Recommended value is `750` for `text-embedding-3-small`.
+- `EMBEDDING_MAX_QUEUE_SIZE`: (Optional) Maximum number of batches to buffer in memory during async processing. Default value is "3".
 - `RAG_UPLOAD_DIR`: (Optional) The directory where uploaded files are stored. Default value is "./uploads/".
 - `PDF_EXTRACT_IMAGES`: (Optional) A boolean value indicating whether to extract images from PDF files. Default value is "False".
 - `DEBUG_RAG_API`: (Optional) Set to "True" to show more verbose logging output in the server console, and to enable postgresql database routes
@@ -95,6 +97,41 @@ The following environment variables are required to run the application:
 
 Make sure to set these environment variables before running the application. You can set them in a `.env` file or as system environment variables.
 
+### Embedding Batch Processing
+
+For large files, you can enable batched embedding processing to reduce memory consumption. This is particularly useful in memory-constrained environments like Kubernetes pods with memory limits.
+
+#### Configuration
+
+| Variable | Default | Description |
+|----------|---------|-------------|
+| `EMBEDDING_BATCH_SIZE` | `0` | Number of document chunks to process per batch. `0` disables batching (original behavior). |
+| `EMBEDDING_MAX_QUEUE_SIZE` | `3` | Maximum number of batches to buffer in memory during async processing. |
+
+#### Recommended Settings
+
+For `text-embedding-3-small` model:
+- `EMBEDDING_BATCH_SIZE=750` - Good balance of throughput and memory
+
+For memory-constrained environments (< 2GB RAM):
+- `EMBEDDING_BATCH_SIZE=100-250`
+
+For high-throughput environments:
+- `EMBEDDING_BATCH_SIZE=1000-2000`
+- `EMBEDDING_MAX_QUEUE_SIZE=5`
+
+#### Behavior
+
+When `EMBEDDING_BATCH_SIZE > 0`:
+- Documents are processed in batches of the specified size
+- Each batch is embedded and inserted before the next batch starts
+- On failure, successfully inserted documents are rolled back
+- Memory usage is bounded by `EMBEDDING_BATCH_SIZE * EMBEDDING_MAX_QUEUE_SIZE`
+
+When `EMBEDDING_BATCH_SIZE = 0` (default):
+- All documents are processed at once (original behavior)
+- Better for small files or memory-rich environments
+
 ### Use Atlas MongoDB as Vector Database
 
 Instead of using the default pgvector, we could use [Atlas MongoDB](https://www.mongodb.com/products/platform/atlas-vector-search) as the vector database. To do so, set the following environment variables
@@ -169,6 +206,81 @@ Notes:
 
 ### Dev notes:
 
+#### Running Tests
+
+##### Prerequisites
+
+Install test dependencies:
+
+```bash
+pip install -r test_requirements.txt
+```
+
+##### Running All Tests
+
+```bash
+# Run all tests
+pytest
+
+# Run with verbose output
+pytest -v
+
+# Run with coverage (if pytest-cov is installed)
+pytest --cov=app
+```
+
+##### Running Specific Test Files
+
+```bash
+# Run batch processing unit tests
+pytest tests/test_batch_processing.py -v
+
+# Run batch processing integration tests (memory optimization tests)
+pytest tests/test_batch_processing_integration.py -v
+
+# Run main API tests
+pytest tests/test_main.py -v
+```
+
+##### Running Tests by Category
+
+```bash
+# Run only integration tests (marked with @pytest.mark.integration)
+pytest -m integration -v
+
+# Skip integration tests
+pytest -m "not integration" -v
+
+# Run only async tests
+pytest -k "async" -v
+```
+
+##### Test Categories
+
+| Test File | Description |
+|-----------|-------------|
+| `test_batch_processing.py` | Unit tests for batch processing functions |
+| `test_batch_processing_integration.py` | Memory optimization and integration tests |
+| `test_main.py` | API endpoint tests |
+| `test_config.py` | Configuration tests |
+| `test_middleware.py` | Middleware tests |
+| `test_models.py` | Model tests |
+
+##### Memory Optimization Tests
+
+The `test_batch_processing_integration.py` file includes tests that verify the memory optimization behavior:
+
+- **`test_memory_bounded_by_batch_size`**: Verifies that the number of documents in memory at any time is bounded by `EMBEDDING_BATCH_SIZE`
+- **`test_memory_tracking_with_tracemalloc`**: Uses Python's `tracemalloc` to monitor memory usage during batch processing
+- **`test_sync_memory_bounded_by_batch_size`**: Same verification for the synchronous code path
+
+Run memory tests specifically:
+
+```bash
+pytest tests/test_batch_processing_integration.py::TestMemoryOptimization -v
+pytest tests/test_batch_processing_integration.py::TestSyncBatchedMemory -v
+```
+
 #### Installing pre-commit formatter
 
 Run the following commands to install pre-commit formatter, which uses [black](https://github.com/psf/black) code formatter:

app/config.py

@@ -70,6 +70,22 @@ def get_env_variable(
 CHUNK_SIZE = int(get_env_variable("CHUNK_SIZE", "1500"))
 CHUNK_OVERLAP = int(get_env_variable("CHUNK_OVERLAP", "100"))
 
+# Batch processing configuration for memory-constrained environments.
+# When EMBEDDING_BATCH_SIZE > 0, documents are processed in batches to reduce
+# peak memory usage. This is useful for Kubernetes pods with memory limits.
+#
+# Trade-offs:
+# - Smaller batch size = lower memory, more DB round trips
+# - Larger batch size = higher memory, fewer DB round trips
+# - 0 = disable batching, process all at once (original behavior)
+#
+# Recommended: 750 for text-embedding-3-small (good balance of speed and memory)
+EMBEDDING_BATCH_SIZE = int(get_env_variable("EMBEDDING_BATCH_SIZE", "0"))
+
+# Maximum number of batches to buffer in memory during async processing.
+# Higher values allow more parallelism but use more memory.
+EMBEDDING_MAX_QUEUE_SIZE = int(get_env_variable("EMBEDDING_MAX_QUEUE_SIZE", "3"))
+
 env_value = get_env_variable("PDF_EXTRACT_IMAGES", "False").lower()
 PDF_EXTRACT_IMAGES = True if env_value == "true" else False