ID-based RAG FastAPI

About this fork

This is a fork of danny-avila/rag_api, tailored for legal-office document workflows. It keeps full upstream compatibility (same API, same LibreChat integration) while changing how several document types are loaded and adding new ones. Key differences from upstream:

• Remote OCR via Mistral — PDFs and standalone images (.png, .jpg, .jpeg, .gif, .bmp, .tif/.tiff, .webp) are OCR'd through the Mistral OCR API instead of a bundled local OCR engine. Set MISTRAL_API_KEY to enable. Because OCR is remote, the local OCR/computer-vision stack (rapidocr-onnxruntime, opencv, onnxruntime) has been removed, which also makes the image considerably smaller.
• Richer DOCX text extraction — .docx files routed to /text are converted with pandoc, preserving tracked changes, comments, and (optionally) headers/footers. Configurable via DOCX_TEXT_USE_PANDOC, DOCX_TEXT_TRACK_CHANGES, and DOCX_TEXT_INCLUDE_HEADERS_FOOTERS.
• Email support — .eml and .msg (Outlook) files are parsed, with sender/recipient/subject/date headers optionally prepended (EMAIL_INCLUDE_HEADERS). Bcc is never exposed.
• .rtf support added.
• Legacy .doc is rejected with a clear error asking the user to convert to .docx (the old binary format could not be loaded reliably).

Prebuilt image: georgx22/rag_api_lite on Docker Hub — the lite build (no local embedding stack; uses remote embeddings + remote Mistral OCR).

Overview

This project integrates Langchain with FastAPI in an Asynchronous, Scalable manner, providing a framework for document indexing and retrieval, using PostgreSQL/pgvector.

Files are organized into embeddings by file_id. The primary use case is for integration with LibreChat, but this simple API can be used for any ID-based use case.

The main reason to use the ID approach is to work with embeddings on a file-level. This makes for targeted queries when combined with file metadata stored in a database, such as is done by LibreChat.

The API will evolve over time to employ different querying/re-ranking methods, embedding models, and vector stores.

Features

• Document Management: Methods for adding, retrieving, and deleting documents.
• Vector Store: Utilizes Langchain's vector store for efficient document retrieval.
• Asynchronous Support: Offers async operations for enhanced performance.

Setup

Getting Started

• Configure .env file based on section below
• Setup pgvector database:
  • Run an existing PSQL/PGVector setup, or,
  • Docker: docker compose up (also starts RAG API)
    • or, use docker just for DB: docker compose -f ./db-compose.yaml up
• Run API:
  • Docker: docker compose up (also starts PSQL/pgvector)
    • or, use docker just for RAG API: docker compose -f ./api-compose.yaml up
  • Local:
    • Make sure to setup DB_HOST to the correct database hostname
    • Run the following commands (preferably in a virtual environment)

pip install -r requirements.txt
uvicorn main:app

Clean Install (Local Development)

To do a clean reinstall of all dependencies (e.g., after updating requirements.txt):

# Remove existing virtual environment and recreate it
rm -rf venv
python3 -m venv venv
source venv/bin/activate
pip install -r requirements.txt

For the lite version (without sentence_transformers/huggingface):

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

For Docker, rebuild without cache:

docker compose build --no-cache

Environment Variables

The following environment variables are required to run the application:

• RAG_OPENAI_API_KEY: The API key for OpenAI API Embeddings (if using default settings).

  • Note: OPENAI_API_KEY will work but RAG_OPENAI_API_KEY will override it in order to not conflict with LibreChat setting.

• RAG_OPENAI_BASEURL: (Optional) The base URL for your OpenAI API Embeddings

• RAG_OPENAI_PROXY: (Optional) Proxy for OpenAI API Embeddings

  • Note: When using with LibreChat, you can also set HTTP_PROXY and HTTPS_PROXY environment variables in the docker-compose.override.yml file (see Proxy Configuration section below)

• VECTOR_DB_TYPE: (Optional) select vector database type, default to pgvector.

• POSTGRES_USE_UNIX_SOCKET: (Optional) Set to "True" when connecting to the PostgreSQL database server with Unix Socket.

• POSTGRES_DB: (Optional) The name of the PostgreSQL database, used when VECTOR_DB_TYPE=pgvector.

• POSTGRES_USER: (Optional) The username for connecting to the PostgreSQL database.

• POSTGRES_PASSWORD: (Optional) The password for connecting to the PostgreSQL database.

• DB_HOST: (Optional) The hostname or IP address of the PostgreSQL database server.

• DB_PORT: (Optional) The port number of the PostgreSQL database server.

• PGVECTOR_CREATE_EXTENSION: (Optional) Set to "False" to skip the CREATE EXTENSION IF NOT EXISTS vector call on startup. Default is "True". Use this when the vector extension is already installed on a managed Postgres (e.g. RDS, Azure Database for PostgreSQL) and the application user is not a superuser.

• PG_POOL_PRE_PING: (Optional) Set to "False" to disable SQLAlchemy's pre-ping check. Default is "True". When enabled, the connection pool issues a lightweight SELECT 1 before handing out a pooled connection, so stale connections dropped by a remote server or middlebox idle timeout are transparently replaced instead of surfacing as query errors. Recommended for any deployment that connects to a remote PostgreSQL instance (managed Postgres, connections that traverse a load balancer, etc.).

• PG_POOL_RECYCLE: (Optional) Maximum age in seconds of a pooled connection before it is recycled. Default is "-1" (disabled). Set to a positive value when the server enforces a hard idle or max-lifetime limit (e.g. "1800" for a 30-minute cap).

• POSTGRES_SCHEMA: (Optional) Prepend this schema to the Postgres search_path so langchain's pgvector tables live in (and are read from) it. Unset by default (uses the user's default schema, typically public). Useful when sharing a database with other services — create the schema out-of-band first (CREATE SCHEMA IF NOT EXISTS <name>; GRANT USAGE, CREATE ON SCHEMA <name> TO <app_user>;); the RAG API will not create it for you and fails fast at startup if the schema is missing. public is always appended to the resulting search path so the vector data type stays resolvable when the extension was installed there (the common case). Multiple schemas may be supplied as a comma-separated list (e.g. myapp,extensions) when the vector extension lives in a non-public schema.

• RAG_HOST: (Optional) The hostname or IP address where the API server will run. Defaults to "0.0.0.0"

• RAG_PORT: (Optional) The port number where the API server will run. Defaults to port 8000.

• JWT_SECRET: (Optional) The secret key used for verifying JWT tokens for requests.

  • The secret is only used for verification. This basic approach assumes a signed JWT from elsewhere.
  • Omit to run API without requiring authentication

• 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_DISTANCE_THRESHOLD: (Optional, VECTOR_DB_TYPE=pgvector only) Drop results whose vector distance is greater than this value, after the top-k search. Unset by default (no filtering). Lower distance = more similar, so e.g. 0.5 keeps only hits with distance ≤ 0.5 and discards weaker matches. Useful for reducing downstream LLM token cost when the top-k call returns loosely-related chunks. Appropriate values depend on the embedding model and distance strategy — inspect your actual scores before choosing one. Ignored (with a startup warning) under VECTOR_DB_TYPE=atlas-mongo, because Atlas returns a similarity score (higher = better) with inverted semantics.

• RAG_UPLOAD_DIR: (Optional) The directory where uploaded files are stored. Default value is "./uploads/".

• DOCX_TEXT_USE_PANDOC: (Optional) Boolean. When "True" (default), .docx files sent to the /text endpoint are converted via pandoc so tracked changes and comments are preserved (output is Markdown). The embedding path always uses Docx2txtLoader. Pandoc only reads OOXML .docx; legacy binary .doc files are rejected with an error asking you to convert to .docx. Requires the pandoc binary (already installed in the provided Docker images).

• DOCX_TEXT_TRACK_CHANGES: (Optional) pandoc --track-changes mode for the .docx /text path: all (default; keep insertions/deletions, record each edit's author/date, and emit comments), accept (final text), or reject (original text).

• DOCX_TEXT_INCLUDE_HEADERS_FOOTERS: (Optional) Boolean. When "True" (default), .docx headers/footers (matter numbers, "PRIVILEGED & CONFIDENTIAL", "DRAFT", etc.) — which pandoc drops — are extracted via python-docx and prepended to the /text output.

• EMAIL_INCLUDE_HEADERS: (Optional) Boolean. When "True" (default), the From/To/Cc/Subject/Date headers are prepended to extracted email text for .eml and .msg files. Set "False" to extract the body only.

• Standalone image uploads (.png, .jpg/.jpeg, .gif, .bmp, .tif/.tiff, .webp) — including multi-page TIFF — are run through Mistral OCR, yielding one result per page. Requires MISTRAL_API_KEY (same key used for PDF OCR). This is separate from PDF handling; images embedded inside PDFs are not extracted.

• IMAGE_OCR_MAX_PAGES: (Optional) Max number of frames/pages of a single image upload (e.g. animated GIF or large multi-page TIFF) sent to OCR; frames beyond the cap are skipped with a warning. Default value is 100.

• DEBUG_RAG_API: (Optional) Set to "True" to show more verbose logging output in the server console, and to enable postgresql database routes

• DEBUG_PGVECTOR_QUERIES: (Optional) Set to "True" to enable detailed PostgreSQL query logging for pgvector operations. Useful for debugging performance issues with vector database queries.

• CONSOLE_JSON: (Optional) Set to "True" to log as json for Cloud Logging aggregations

• EMBEDDINGS_PROVIDER: (Optional) either "openai", "bedrock", "azure", "huggingface", "huggingfacetei", "google_genai", "vertexai", or "ollama", where "huggingface" uses sentence_transformers; defaults to "openai"

• EMBEDDINGS_MODEL: (Optional) Set a valid embeddings model to use from the configured provider.

  • Defaults
  • openai: "text-embedding-3-small"
  • azure: "text-embedding-3-small" (will be used as your Azure Deployment)
  • huggingface: "sentence-transformers/all-MiniLM-L6-v2"
  • huggingfacetei: "http://huggingfacetei:3000". Hugging Face TEI uses model defined on TEI service launch.
  • vertexai: "gemini-embedding-001"
  • ollama: "nomic-embed-text"
  • bedrock: "amazon.titan-embed-text-v1"
  • google_genai: "gemini-embedding-001"

• EMBEDDINGS_CHUNK_SIZE: (Optional) The chunk size used by the OpenAI and Azure embeddings clients to limit the number of inputs per request. Default value is 200.

• EMBEDDINGS_DIMENSIONS: (Optional) Output vector size to request from the embedding model. Only honored by the openai and azure providers, and only supported by text-embedding-3-* models. Leave unset to use the model's native dimensionality (1536 for text-embedding-3-small, 3072 for text-embedding-3-large). Setting a smaller value (e.g. 512, 1024) trades some retrieval quality for lower storage cost and faster similarity search. Note: do not change this on an existing collection — all vectors in a pgvector column must share the same dimensionality.

• RAG_AZURE_OPENAI_API_VERSION: (Optional) Default is 2023-05-15. The version of the Azure OpenAI API.

• RAG_AZURE_OPENAI_API_KEY: (Optional) The API key for Azure OpenAI service.

  • Note: AZURE_OPENAI_API_KEY will work but RAG_AZURE_OPENAI_API_KEY will override it in order to not conflict with LibreChat setting.

• RAG_AZURE_OPENAI_ENDPOINT: (Optional) The endpoint URL for Azure OpenAI service, including the resource.

  • Example: https://YOUR_RESOURCE_NAME.openai.azure.com.
  • Note: AZURE_OPENAI_ENDPOINT will work but RAG_AZURE_OPENAI_ENDPOINT will override it in order to not conflict with LibreChat setting.

• HF_TOKEN: (Optional) if needed for huggingface option.

• OLLAMA_BASE_URL: (Optional) defaults to http://ollama:11434.

• ATLAS_SEARCH_INDEX: (Optional) the name of the vector search index if using Atlas MongoDB, defaults to vector_index

• MONGO_VECTOR_COLLECTION: Deprecated for MongoDB, please use ATLAS_SEARCH_INDEX and COLLECTION_NAME

• AWS_DEFAULT_REGION: (Optional) defaults to us-east-1

• AWS_ACCESS_KEY_ID: (Optional) needed for bedrock embeddings

• AWS_SECRET_ACCESS_KEY: (Optional) needed for bedrock embeddings

• GOOGLE_API_KEY, GOOGLE_KEY, RAG_GOOGLE_API_KEY: (Optional) Google API key for Google GenAI embeddings. Priority order: RAG_GOOGLE_API_KEY > GOOGLE_KEY > GOOGLE_API_KEY

• AWS_SESSION_TOKEN: (Optional) may be needed for bedrock embeddings

• GOOGLE_APPLICATION_CREDENTIALS: (Optional) needed for Google VertexAI embeddings. This should be a path to a service account credential file in JSON format.

• GOOGLE_CLOUD_PROJECT: (Optional) Google Cloud project ID, needed for VertexAI embeddings.

• GOOGLE_CLOUD_LOCATION: (Optional) Google Cloud region for VertexAI embeddings. Defaults to us-central1.

• RAG_CHECK_EMBEDDING_CTX_LENGTH (Optional) Default is true, disabling this will send raw input to the embedder, use this for custom embedding models.

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 as the vector database. To do so, set 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>

The ATLAS_MONGO_DB_URI could be the same or different from what is used by LibreChat. Even if it is the same, the $COLLECTION_NAME collection needs to be a completely new one, separate from all collections used by LibreChat. In addition, create a vector search index for collection above (remember to assign $ATLAS_SEARCH_INDEX) with the following json:

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

Follow one of the four documented methods to create the vector index.

Create a file_id Index (recommended)

We recommend creating a standard MongoDB index on file_id to keep lookups fast. After creating the collection, run the following once (via Atlas UI, Compass, or mongosh):

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

Replace <COLLECTION_NAME> with the same collection used by the RAG API. This ensures lookups remain fast even as the number of embedded documents grows.

Proxy Configuration

When using the RAG API with LibreChat and you need to configure proxy settings, you can set the HTTP_PROXY and HTTPS_PROXY environment variables in the docker-compose.override.yml file (from the LibreChat repository):

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

This configuration will ensure that all HTTP/HTTPS requests from the RAG API container are routed through your specified proxy server.

Cloud Installation Settings:

AWS:

Make sure your RDS Postgres instance adheres to this requirement:

The pgvector extension version 0.5.0 is available on database instances in Amazon RDS running PostgreSQL 15.4-R2 and higher, 14.9-R2 and higher, 13.12-R2 and higher, and 12.16-R2 and higher in all applicable AWS Regions, including the AWS GovCloud (US) Regions.

In order to setup RDS Postgres with RAG API, you can follow these steps:

• Create a RDS Instance/Cluster using the provided AWS Documentation.

• Login to the RDS Cluster using the Endpoint connection string from the RDS Console or from your IaC Solution output.

• The login is via the Master User.

• Create a dedicated database for rag_api: create database rag_api;.

• Create a dedicated user\role for that database: create role rag;

• Switch to the database you just created: \c rag_api

• Enable the Vector extension: create extension vector;

• Use the documentation provided above to set up the connection string to the RDS Postgres Instance\Cluster.

Notes:

• Even though you're logging with a Master user, it doesn't have all the super user privileges, that's why we cannot use the command: create role x with superuser;
• If you do not enable the extension, rag_api service will throw an error that it cannot create the extension due to the note above.

Dev notes:

Running Tests

Prerequisites

Install test dependencies:

pip install -r test_requirements.txt

Running All Tests

# 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

# 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

# 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:

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 code formatter:

pip install pre-commit
pre-commit install