Evolve Project Aether: Implement Postgres persistence, RQ background jobs for ingestion, and containerization

Author: gabaoun

Dockerfile

@@ -0,0 +1,22 @@
+FROM python:3.11-slim
+
+WORKDIR /app
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    build-essential \
+    libpq-dev \
+    && rm -rf /var/lib/apt/lists/*
+
+# Install Python dependencies
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy application code
+COPY . .
+
+# Expose API port
+EXPOSE 8000
+
+# Default command (will be overridden by docker-compose)
+CMD ["uvicorn", "src.api.app:app", "--host", "0.0.0.0", "--port", "8000"]

docker-compose.yml

@@ -1,32 +1,48 @@
-version: '3.8'
-
 services:
-  qdrant:
-    image: qdrant/qdrant:latest
-    container_name: qdrant
+  api:
+    build: .
     ports:
-      - "6333:6333"
-      - "6334:6334"
+      - "8000:8000"
+    env_file:
+      - .env
+    environment:
+      - REDIS_HOST=redis
+      - DATABASE_URL=postgresql://user:password@postgres:5432/aether
+    depends_on:
+      - redis
+      - postgres
+    volumes:
+      - ./data:/app/data
+
+  worker:
+    build: .
+    command: python -m src.worker
+    env_file:
+      - .env
+    environment:
+      - REDIS_HOST=redis
+      - DATABASE_URL=postgresql://user:password@postgres:5432/aether
+    depends_on:
+      - redis
+      - postgres
     volumes:
-      - qdrant_data:/qdrant/storage
-    networks:
-      - aether_network
+      - ./data:/app/data
 
   redis:
     image: redis:7-alpine
-    container_name: redis_cache
     ports:
       - "6379:6379"
-    volumes:
-      - redis_data:/data
-    command: redis-server --save 60 1 --loglevel warning
-    networks:
-      - aether_network
 
-networks:
-  aether_network:
-    driver: bridge
+  postgres:
+    image: postgres:15-alpine
+    environment:
+      POSTGRES_DB: aether
+      POSTGRES_USER: user
+      POSTGRES_PASSWORD: password
+    ports:
+      - "5432:5432"
+    volumes:
+      - postgres_data:/var/lib/postgresql/data
 
 volumes:
-  qdrant_data:
-  redis_data:
+  postgres_data:

requirements.txt

@@ -8,6 +8,10 @@ llama-index-vector-stores-qdrant
 qdrant-client
 chromadb>=0.6.0
 redis
+sqlalchemy
+psycopg2-binary
+rq
+alembic
 pydantic>=2.0
 pydantic-settings
 python-dotenv

src/api/app.py

@@ -1,12 +1,15 @@
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, Depends
 from pydantic import BaseModel
 from src.services.chroma import ChromaService
-from src.pipeline.ingestion import IngestionWorkflow
 from src.pipeline.retrieval import RetrievalWorkflow
 from src.config.settings import settings
 from src.utils.logger import logger
+from src.db.session import SessionLocal, get_db
+from src.models.db import IngestionJob
+from src.worker import queue
 import os
 from contextlib import asynccontextmanager
+from sqlalchemy.orm import Session
 
 # Global variables for chroma service and workflow
 chroma_service = None
@@ -17,19 +20,13 @@ async def lifespan(app: FastAPI):
     global chroma_service, retrieval_wf
 
     try:
-        ingestion_wf = IngestionWorkflow()
-        
-        if os.path.exists(settings.data_dir) and os.listdir(settings.data_dir):
-            logger.info(f"Initializing ingestion from {settings.data_dir}...")
-            # ingestion_wf.run returns chroma_service (StopEvent result)
-            chroma_service = await ingestion_wf.run(input_dir=settings.data_dir)
-            retrieval_wf = RetrievalWorkflow(chroma_service=chroma_service)
-            logger.info("API Startup: Ingestion complete and Chroma Cloud index ready.")
-        else:
-            logger.warning(f"Data directory '{settings.data_dir}' is empty or missing. API in degraded mode.")
+        # Initialize search infrastructure on startup
+        chroma_service = ChromaService()
+        retrieval_wf = RetrievalWorkflow(chroma_service=chroma_service)
+        logger.info("API Startup: Chroma Cloud retrieval ready.")
     except Exception as e:
         logger.error(f"Startup failed: {e}")
-        logger.warning("API starting in degraded mode due to infrastructure or ingestion error.")
+        logger.warning("API starting in degraded mode due to infrastructure error.")
 
     yield
 
@@ -42,9 +39,48 @@ class QueryResponse(BaseModel):
     answer: str
     from_cache: bool
 
+class IngestResponse(BaseModel):
+    job_id: str
+
+class JobStatusResponse(BaseModel):
+    id: str
+    status: str
+
 @app.get("/health")
 async def health():
-    return {"status": "ok", "index_ready": chroma_service is not None}
+    return {"status": "ok", "retrieval_ready": retrieval_wf is not None}
+
+@app.post("/ingest", response_model=IngestResponse, status_code=202)
+async def ingest_docs(db: Session = Depends(get_db)):
+    """
+    Trigger document ingestion as a background job.
+    """
+    try:
+        # Create a job record in Postgres
+        job = IngestionJob(status="PENDING")
+        db.add(job)
+        db.commit()
+        db.refresh(job)
+        
+        # Enqueue the background task
+        queue.enqueue("src.jobs.ingestion.process_ingestion", str(job.id))
+        
+        logger.info(f"Ingestion job {job.id} enqueued.")
+        return IngestResponse(job_id=str(job.id))
+    except Exception as e:
+        logger.error(f"Failed to enqueue ingestion job: {e}")
+        raise HTTPException(status_code=500, detail="Failed to trigger ingestion.")
+
+@app.get("/jobs/{job_id}", response_model=JobStatusResponse)
+async def get_job_status(job_id: str, db: Session = Depends(get_db)):
+    """
+    Check the status of an ingestion job.
+    """
+    job = db.query(IngestionJob).filter(IngestionJob.id == job_id).first()
+    if not job:
+        raise HTTPException(status_code=404, detail="Job not found.")
+    
+    return JobStatusResponse(id=str(job.id), status=job.status)
 
 @app.post("/query", response_model=QueryResponse)
 async def query_docs(request: QueryRequest):

src/config/settings.py

@@ -20,6 +20,8 @@ class Settings(BaseSettings):
     redis_port: int = 6379
     semantic_cache_threshold: float = 0.85
 
+    database_url: str = "postgresql://user:password@postgres:5432/aether"
+    
     phoenix_collector_endpoint: str = "http://localhost:6006"
     log_level: str = "INFO"
 

src/core/splitter.py

@@ -0,0 +1,13 @@
+from sqlalchemy import create_engine
+from sqlalchemy.orm import sessionmaker
+from src.config.settings import settings
+
+engine = create_engine(settings.database_url)
+SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
+
+def get_db():
+    db = SessionLocal()
+    try:
+        yield db
+    finally:
+        db.close()

src/db/session.py

@@ -0,0 +1,44 @@
+import asyncio
+from src.db.session import SessionLocal
+from src.models.db import IngestionJob, Document
+from src.pipeline.ingestion import IngestionWorkflow
+from src.utils.logger import logger
+from src.config.settings import settings
+
+def process_ingestion(job_id: str):
+    """
+    Background job to process document ingestion.
+    """
+    db = SessionLocal()
+    job = db.query(IngestionJob).filter(IngestionJob.id == job_id).first()
+    
+    if not job:
+        logger.error(f"Job {job_id} not found in database.")
+        db.close()
+        return
+
+    job.status = "PROCESSING"
+    db.commit()
+
+    try:
+        logger.info(f"Starting ingestion workflow for job {job_id}")
+        workflow = IngestionWorkflow()
+        
+        # Run the workflow
+        # Note: workflow.run is async, so we need to run it in an event loop
+        loop = asyncio.get_event_loop()
+        if loop.is_running():
+            # This shouldn't happen in a standard RQ worker but just in case
+            future = asyncio.ensure_future(workflow.run(input_dir=settings.data_dir))
+            loop.run_until_complete(future)
+        else:
+            asyncio.run(workflow.run(input_dir=settings.data_dir))
+
+        job.status = "COMPLETED"
+        logger.info(f"Ingestion job {job_id} completed successfully.")
+    except Exception as e:
+        logger.error(f"Ingestion job {job_id} failed: {e}")
+        job.status = "FAILED"
+    finally:
+        db.commit()
+        db.close()

src/jobs/ingestion.py

@@ -0,0 +1,24 @@
+import uuid
+from datetime import datetime
+from sqlalchemy import Column, String, DateTime, ForeignKey
+from sqlalchemy.dialects.postgresql import UUID
+from sqlalchemy.orm import declarative_base
+
+Base = declarative_base()
+
+class Document(Base):
+    __tablename__ = "documents"
+
+    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
+    filename = Column(String, nullable=True)
+    status = Column(String, default="PENDING") # PENDING, COMPLETED, FAILED
+    created_at = Column(DateTime, default=datetime.utcnow)
+
+class IngestionJob(Base):
+    __tablename__ = "ingestion_jobs"
+
+    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
+    document_id = Column(UUID(as_uuid=True), ForeignKey("documents.id"), nullable=True)
+    status = Column(String, default="PENDING") # PENDING, PROCESSING, COMPLETED, FAILED
+    trace_id = Column(String, nullable=True)
+    created_at = Column(DateTime, default=datetime.utcnow)

src/models/db.py

@@ -1,6 +1,5 @@
 from typing import List
 from llama_index.core import Document
-from llama_index.embeddings.huggingface import HuggingFaceEmbedding
 from llama_index.core.workflow import (
     Workflow,
     Event,
@@ -9,7 +8,7 @@
     step,
 )
 from tenacity import retry, stop_after_attempt, wait_exponential
-from src.core.splitter import SemanticDoubleMergingSplitter
+from llama_index.core.node_parser import SentenceSplitter
 from src.config.settings import settings
 from src.utils.token_counter import TokenCounter
 from src.core.pii import PIIMasker
@@ -33,10 +32,9 @@ class IngestionWorkflow(Workflow):
     def __init__(self, **kwargs: dict) -> None:
         super().__init__(**kwargs)
         self.chroma_service = ChromaService()
-        self.embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")
-        self.node_parser = SemanticDoubleMergingSplitter(
-            embed_model=self.embed_model,
-            min_chunk_size=500
+        self.node_parser = SentenceSplitter(
+            chunk_size=1024,
+            chunk_overlap=20
         )
         self.token_counter = TokenCounter()
         self.pii_masker = PIIMasker()