From 0702e36b8e1f6dc20a2532f3c20f0e65c4e2838c Mon Sep 17 00:00:00 2001 From: Victor <194116445+dodekapod@users.noreply.github.com> Date: Thu, 10 Jul 2025 10:27:46 +0200 Subject: [PATCH 1/2] Use PostgreSQL instead of SQLite --- .hydra_config/config.yaml | 5 +++++ Dockerfile | 3 ++- docker-compose.yaml | 12 ++++++++++-- extern/infinity.yaml | 3 ++- openrag/components/indexer/indexer.py | 2 +- openrag/components/indexer/vectordb/utils.py | 6 ++++++ openrag/components/indexer/vectordb/vectordb.py | 12 +++++++++--- pyproject.toml | 2 ++ 8 files changed, 37 insertions(+), 8 deletions(-) diff --git a/.hydra_config/config.yaml b/.hydra_config/config.yaml index 73a3da9f9..58c2fb4bb 100644 --- a/.hydra_config/config.yaml +++ b/.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} + reranker: enable: ${oc.decode:${oc.env:RERANKER_ENABLED, true}} model_name: ${oc.env:RERANKER_MODEL, Alibaba-NLP/gte-multilingual-reranker-base} diff --git a/Dockerfile b/Dockerfile index 6e7a1a601..f6dd94f60 100644 --- a/Dockerfile +++ b/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 \ No newline at end of file +ENTRYPOINT ../entrypoint.sh diff --git a/docker-compose.yaml b/docker-compose.yaml index d8bb9cb89..ab9de00c1 100644 --- a/docker-compose.yaml +++ b/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: @@ -67,6 +66,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 @@ -90,4 +97,5 @@ services: image: openrag-vllm-openai-cpu deploy: {} profiles: - - 'cpu' \ No newline at end of file + - 'cpu' + diff --git a/extern/infinity.yaml b/extern/infinity.yaml index 252aaed83..8fbd3e10a 100644 --- a/extern/infinity.yaml +++ b/extern/infinity.yaml @@ -37,4 +37,5 @@ services: --model-id ${RERANKER_MODEL:-Alibaba-NLP/gte-multilingual-reranker-base} --port ${RERANKER_PORT:-7997} profiles: - - 'cpu' \ No newline at end of file + - 'cpu' + diff --git a/openrag/components/indexer/indexer.py b/openrag/components/indexer/indexer.py index c4c86860d..cc2182e60 100644 --- a/openrag/components/indexer/indexer.py +++ b/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 diff --git a/openrag/components/indexer/vectordb/utils.py b/openrag/components/indexer/vectordb/utils.py index 7cfb48301..fc49315a3 100644 --- a/openrag/components/indexer/vectordb/utils.py +++ b/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 diff --git a/openrag/components/indexer/vectordb/vectordb.py b/openrag/components/indexer/vectordb/vectordb.py index 0c49110df..044017f3d 100644 --- a/openrag/components/indexer/vectordb/vectordb.py +++ b/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, ) @@ -1025,6 +1028,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) diff --git a/pyproject.toml b/pyproject.toml index af7130396..effe5c9dd 100644 --- a/pyproject.toml +++ b/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] From 81f102b96708d5a18842f9e3c08499cd9c500b64 Mon Sep 17 00:00:00 2001 From: antncy Date: Fri, 11 Jul 2025 16:28:32 +0000 Subject: [PATCH 2/2] default password when the var isn't supplied --- .hydra_config/config.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.hydra_config/config.yaml b/.hydra_config/config.yaml index fcbcdcb57..33ef6b243 100644 --- a/.hydra_config/config.yaml +++ b/.hydra_config/config.yaml @@ -43,7 +43,7 @@ vectordb: rdb: port: ${oc.env:POSTGRES_PORT, 5432} user: ${oc.env:POSTGRES_USER, root} - password: ${oc.env:POSTGRES_PASSWORD} + password: ${oc.env:POSTGRES_PASSWORD, root_password} reranker: enable: ${oc.decode:${oc.env:RERANKER_ENABLED, true}}