DOCSEARCHbyAlgolia
Submit Docs

ID-based RAG FastAPI

repository·main·Indexed 21 days ago

https://github.com/danny-avila/rag_api

An asynchronous, scalable FastAPI framework for Retrieval-Augmented Generation (RAG) using Langchain and PostgreSQL/pgvector. It organizes document embeddings by file_id for targeted queries and is optimized for integration with LibreChat. The framework supports multiple embedding providers (OpenAI, Bedrock, Azure, HuggingFace, Google GenAI, VertexAI, Ollama) and can optionally use Atlas MongoDB as a vector database.

Tokens
3.4K
Snippets
12
Records
17
Agent score
73%

What's inside rag_api

  1. Overview of ID-based RAG FastAPI

    main

    ID-based RAG FastAPI is a framework for document indexing and retrieval that integrates Langchain with FastAPI. It is designed to be asynchronous and scalable, utilizing PostgreSQL with the pgvector extension as the vector store.

    Key architectural concept: Files are organized into embeddings by file_id. This allows for targeted queries by combining embeddings with file metadata stored in a database. While optimized for integration with LibreChat, the API is suitable for any use case requiring ID-based document retrieval.

  2. Core features of ID-based RAG FastAPI

    main

    The API provides the following core capabilities:

    • Document Management: APIs for adding, retrieving, and deleting documents.
    • Vector Store: Integration with Langchain's vector store for efficient retrieval.
    • Asynchronous Support: Built on FastAPI to provide high-performance asynchronous operations.

  3. Configure Embedding Batch Processing

    main

    For large files, enable batching to reduce memory consumption. This is critical for memory-constrained environments like Kubernetes.

    Configuration Variables

    • EMBEDDING_BATCH_SIZE: Number of document chunks per batch. Set to 0 to disable batching.
    • EMBEDDING_MAX_QUEUE_SIZE: Max batches to buffer in memory.
    • PARALLEL_EXECUTION: Max async embedding/insertion consumers per file.

    Recommended Settings

    • For text-embedding-3-small: Set EMBEDDING_BATCH_SIZE=750.
    • For low memory (< 2GB RAM): Set EMBEDDING_BATCH_SIZE=100-250.
    • For high throughput: Set EMBEDDING_BATCH_SIZE=1000-2000 and EMBEDDING_MAX_QUEUE_SIZE=5.

    Behavior

    • When EMBEDDING_BATCH_SIZE > 0, documents are processed in batches.
    • Memory usage is roughly: EMBEDDING_BATCH_SIZE * (EMBEDDING_MAX_QUEUE_SIZE + PARALLEL_EXECUTION).
    • On failure, remaining batch work is stopped and successfully inserted documents are rolled back.

  4. Setup and Run the RAG API

    main

    To get started with the RAG API, follow these steps:

    1. Configure Environment

    Create a .env file and populate it with the required environment variables (see Environment Variables).

    2. Setup Database (pgvector)

    You can use an existing PostgreSQL/pgvector setup or use Docker:

    • Docker (Full Stack): docker compose up (starts both DB and RAG API).
    • Docker (DB only): docker compose -f ./db-compose.yaml up.

    3. Run the API

    • Docker: docker compose -f ./api-compose.yaml up (if using a separate DB).
    • Local Development: Ensure DB_HOST is set correctly, then run:
    npip install -r requirements.txt
uvicorn main:app

  5. Clean Install for Local Development

    main

    To perform a clean reinstall of dependencies (e.g., after updating requirements.txt), recreate your virtual environment:

    Standard Install:

    rm -rf venv
python3 -m venv venv
source venv/bin/activate
pip install -r requirements.txt

    Lite Install (without sentence_transformers/huggingface):

    rm -rf venv
python3 -m venv venv
source venv/bin/activate
pip install -r requirements.lite.txt

    Docker Rebuild:

    docker compose build --no-cache

  6. Configure Proxy Settings for LibreChat

    main

    If you are running the RAG API with LibreChat and need to route requests through a proxy, set the HTTP_PROXY and HTTPS_PROXY environment variables in your docker-compose.override.yml file:

    rag_api:
    environment:
        - HTTP_PROXY=<your-proxy>
        - HTTPS_PROXY=<your-proxy>

  7. Run Tests for RAG API

    main

    The project uses pytest for testing.

    Prerequisites

    Install test dependencies:

    pip install -r test_requirements.txt

    Running Tests

    • All tests: pytest
    • Verbose output: pytest -v
    • Specific files:
      • Batch processing unit tests: pytest tests/test_batch_processing.py -v
      • Batch processing integration tests: pytest tests/test_batch_processing_integration.py -v
      • Main API tests: pytest tests/test_main.py -v
    • By Category:
      • Integration tests: pytest -m integration -v
      • Async tests: pytest -k "async"

    Test Categories

    Test FileDescription
    test_batch_processing.pyUnit tests for batch processing
    test_batch_processing_integration.pyMemory optimization and integration tests
    test_main.pyAPI endpoint tests
    test_config.pyConfiguration tests
    test_middleware.pyMiddleware tests
    test_models.pyModel tests
    pip install -r test_requirements.txt
pytest -v

  8. Use Atlas MongoDB as Vector Database

    main

    To use Atlas MongoDB instead of pgvector, configure the following environment variables:

    VECTOR_DB_TYPE=atlas-mongo
ATLAS_MONGO_DB_URI=<mongodb+srv://...>
COLLECTION_NAME=<vector collection>
ATLAS_SEARCH_INDEX=<vector search index>

    1. Create the Vector Search Index

    In Atlas, create a vector search index for your collection using this JSON configuration:

    {
  "fields": [
    {
      "numDimensions": 1536,
      "path": "embedding",
      "similarity": "cosine",
      "type": "vector"
    },
    {
      "path": "file_id",
      "type": "filter"
    }
  ]
}

    2. Create a file_id Index

    To keep lookups fast, create a standard MongoDB index on file_id via Atlas UI, Compass, or mongosh:

    db.getCollection("<COLLECTION_NAME>").createIndex({ file_id: 1 })

    Replace <COLLECTION_NAME> with your actual collection name.

    VECTOR_DB_TYPE=atlas-mongo
ATLAS_MONGO_DB_URI=mongodb+srv://user:pass@cluster.mongodb.net/
COLLECTION_NAME=my_rag_collection
ATLAS_SEARCH_INDEX=my_vector_index

  9. Understand the RAG API Lifespan and Resource Management

    main

    The API manages several critical resources through a FastAPI lifespan context manager:

    1. Thread Pool: On startup, a bounded ThreadPoolExecutor is created (named rag-worker) and attached to app.state.thread_pool. This pool is used for background tasks. On shutdown, the pool is drained (shutdown(wait=True)).
    2. Vector Database Connections:
      • If VECTOR_DB_TYPE is set to VectorDBType.PGVECTOR, the application initializes the PSQLDatabase connection pool and ensures vector indexes exist during startup. It closes the pool during shutdown.
      • The application calls close_vector_store_connections(vector_store) during shutdown to clean up other vector store connections (like MongoDB or SQLAlchemy engines).
    3. Application State: Global configuration values like CHUNK_SIZE, CHUNK_OVERLAP, and PDF_EXTRACT_IMAGES are attached to app.state for access within route handlers.

  10. Run the RAG API with Uvicorn

    main

    The application is designed to be executed as a standard FastAPI application using uvicorn. It uses environment variables RAG_HOST and RAG_PORT to determine the binding address and port. When running in debug mode, additional routes (such as pgvector_routes) are included in the application.

    if __name__ == "__main__":
    uvicorn.run(app, host=RAG_HOST, port=RAG_PORT, log_config=None)

  11. Configure the FastAPI service via Docker Compose

    main

    The fastapi service runs the core RAG API. It depends on a db service and uses several environment variables for configuration. You can customize the embedding batch size using the EMBEDDING_BATCH_SIZE environment variable, which defaults to 500 if not specified.

    fastapi:
    build: .
    environment:
      - DB_HOST=db
      - DB_PORT=5432
      - EMBEDDING_BATCH_SIZE=${EMBEDDING_BATCH_SIZE:-500}
    ports:
      - "8000:8000"
    volumes:
      - ./uploads:/app/uploads
    depends_on:
      - db
    env_file:
      - .env

  12. Configure RAG API via Environment Variables

    main

    The application's behavior and resource management are controlled by several environment variables. Key variables include:

    • RAG_HOST: The host address for the Uvicorn server.
    • RAG_PORT: The port for the Uvicorn server.
    • RAG_THREAD_POOL_SIZE: Determines the number of workers in the ThreadPoolExecutor. It defaults to the system's CPU count but is capped at a maximum of 8 workers. If not set, it uses os.cpu_count().