diff --git a/.hydra_config/config.yaml b/.hydra_config/config.yaml index 51538047c..33ef6b243 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, root_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 4e5af31e9..4684a559c 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: @@ -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' \ 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 078d11da3..1866eef31 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, ) @@ -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) 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]