Use PostgreSQL instead of SQLite

.hydra_config/config.yaml

@@ -40,6 +40,11 @@ vectordb:
   hybrid_mode: true
   enable: true
 
+rdb:
+  port: ${oc.env:POSTGRES_PORT, 5432}
+  user: ${oc.env:POSTGRES_USER, root}
+  password: ${oc.env:POSTGRES_PASSWORD, root_password}
+
 reranker:
   enable: ${oc.decode:${oc.env:RERANKER_ENABLED, true}}
   model_name: ${oc.env:RERANKER_MODEL, Alibaba-NLP/gte-multilingual-reranker-base}

Dockerfile

@@ -10,6 +10,7 @@ RUN apt-get update && apt-get install -y \
     gcc \
     cmake \
     make \
+    libpq-dev python3-dev \
     && rm -rf /var/lib/apt/lists/*
 
 # install ffmpeg
@@ -44,4 +45,4 @@ COPY prompts/ /app/prompts/
 COPY .hydra_config/ /app/.hydra_config/
 ENV PYTHONPATH=/app/openrag/
 ENV APP_iPORT=${APP_iPORT:-8080}
-ENTRYPOINT ../entrypoint.sh
+ENTRYPOINT ../entrypoint.sh

docker-compose.yaml

@@ -13,7 +13,6 @@ x-openrag: &openrag_template
     - ${DATA_VOLUME:-./data}:/app/data
     - ${MODEL_WEIGHTS_VOLUME:-~/.cache/huggingface}:/app/model_weights # Model weights for RAG
     - ./openrag:/app/openrag # For dev mode
-    - ${DB_VOLUME:-./db}:/app/db
     - /$SHARED_ENV:/ray_mount/.env # Shared environment variables
     - /ray_mount/logs:/app/logs
   ports:
@@ -66,6 +65,14 @@ services:
     profiles:
       - 'cpu'
 
+  rdb:
+    image: postgres:15
+    environment:
+      - POSTGRES_PASSWORD=${POSTGRES_PASSWORD:-root_password}
+      - POSTGRES_USER=${POSTGRES_USER:-root}
+    volumes:
+      - ${DB_VOLUME:-./db}:/var/lib/postgresql/data
+
   vllm-gpu:
     <<: *vllm_template
     image: vllm/vllm-openai:latest
@@ -89,4 +96,5 @@ services:
     image: openrag-vllm-openai-cpu
     deploy: {}
     profiles:
-      - 'cpu'
+      - 'cpu'
+

extern/infinity.yaml

@@ -37,4 +37,5 @@ services:
       --model-id ${RERANKER_MODEL:-Alibaba-NLP/gte-multilingual-reranker-base}
       --port ${RERANKER_PORT:-7997}
     profiles:
-      - 'cpu'
+      - 'cpu'
+

openrag/components/indexer/indexer.py

@@ -18,7 +18,7 @@
 save_uploaded_files = os.environ.get("SAVE_UPLOADED_FILES", "true").lower() == "true"
 
 
-@ray.remote(max_restarts=-1, max_concurrency=1000, concurrency_groups={"insertion": 1})
+@ray.remote(max_restarts=-1, max_concurrency=1000, concurrency_groups={"insertion": 8})
 class Indexer:
     def __init__(self):
         from utils.logger import get_logger

openrag/components/indexer/vectordb/utils.py

@@ -12,6 +12,10 @@
     UniqueConstraint,
     create_engine,
 )
+from sqlalchemy_utils import (
+    database_exists,
+    create_database,
+)
 from sqlalchemy.orm import (
     declarative_base,
     relationship,
@@ -91,6 +95,8 @@ def __repr__(self):
 class PartitionFileManager:
     def __init__(self, database_url: str, logger=logger):
         self.engine = create_engine(database_url)
+        if not database_exists(database_url):
+            create_database(database_url)
         Base.metadata.create_all(self.engine)
         self.Session = sessionmaker(bind=self.engine)
         self.logger = logger

openrag/components/indexer/vectordb/vectordb.py

@@ -170,7 +170,10 @@ def __init__(
         self.vector_store = None
         self.partition_file_manager: PartitionFileManager = None
         self.default_partition = "_default"
-        self.db_dir = kwargs.get("db_dir", None)
+        self.rdb_host = kwargs.get("rdb_host", None)
+        self.rdb_port = kwargs.get("rdb_port", None)
+        self.rdb_user = kwargs.get("rdb_user", None)
+        self.rdb_password = kwargs.get("rdb_password", None)
 
         # Set the initial collection name (if provided)
         if collection_name:
@@ -203,7 +206,7 @@ def collection_name(self, name: str):
         )
 
         self.partition_file_manager = PartitionFileManager(
-            database_url=f"sqlite:///{self.db_dir}/partitions_for_collection_{name}.db",
+            database_url=f"postgresql://{self.rdb_user}:{self.rdb_password}@{self.rdb_host}:{self.rdb_port}/partitions_for_collection_{name}",
             logger=self.logger,
         )
 
@@ -1028,6 +1031,9 @@ def create_vdb(config, logger, embeddings) -> ABCVectorDB:
 
         dbconfig["embeddings"] = embeddings
         dbconfig["logger"] = logger
-        dbconfig["db_dir"] = config.paths.db_dir
+        dbconfig["rdb_host"] = "rdb"
+        dbconfig["rdb_port"] = config["rdb"].port
+        dbconfig["rdb_user"] = config["rdb"].user
+        dbconfig["rdb_password"] = config["rdb"].password
 
         return vdb_cls(**dbconfig)

pyproject.toml

@@ -40,6 +40,8 @@ dependencies = [
     "umap-learn>=0.5.7",
     "eml-parser>=2.0.0",
     "psutil>=7.0.0",
+    "psycopg2",
+    "sqlalchemy_utils",
 ]
 
 [dependency-groups]