# Database Setup Guide

This guide explains how to set up and use the PostgreSQL database with pgvector for semantic table routing in Quber.

## Overview

The database infrastructure enables:
- **Document-scoped semantic search** across extracted tables
- **Vector similarity** for routing metric queries to relevant tables
- **Persistent storage** for table analyses with metadata
- **Future embedding generation** for semantic search capabilities

## Prerequisites

- Docker and Docker Compose (for running PostgreSQL)
- Python 3.13+ with uv package manager
- Database dependencies installed (see Installation)

## Installation

### 1. Install Python Dependencies

```bash
# Install all dependencies including database packages
uv sync

# Or with pip
pip install -e .
```

This installs:
- `sqlalchemy>=2.0.0` - Database ORM
- `psycopg[binary]>=3.2.0` - PostgreSQL driver
- `pgvector>=0.3.6` - Vector extension support
- `sentence-transformers>=3.0.0` - Local embedding models
- `openai>=1.0.0` - OpenAI embeddings (optional)

### 2. Configure Environment Variables

The following environment variables should already be in your `.env` file:

```bash
# PostgreSQL Database Configuration
POSTGRES_USER="quber"
POSTGRES_PASSWORD="quber_dev"
POSTGRES_DB="quber_rag"
POSTGRES_HOST="localhost"
POSTGRES_PORT="5432"

# Embedding Configuration (optional)
EMBEDDING_PROVIDER="local"  # Options: "local" or "openai"
EMBEDDING_DEVICE="cuda"     # Options: "cuda" or "cpu" (for local model)
# OPENAI_API_KEY="sk-..."   # Required only if using EMBEDDING_PROVIDER="openai"
```

**Note**: Change these values for production deployments!

### 3. Start PostgreSQL with Docker Compose

```bash
# Start the database (runs in background)
docker compose up -d

# Check status
docker compose ps

# View logs
docker compose logs -f pgvector-db
```

The database will:
- Run PostgreSQL with pgvector extension
- Persist data in a Docker volume named `pgdata`
- Auto-initialize schema from `sql/init.sql` on first start
- Expose port 5432 on localhost

### 4. Initialize Database Schema

If starting fresh or the auto-init didn't work:

```bash
# Initialize schema (safe - won't drop existing data)
quber db init

# Or force recreate (WARNING: deletes all data)
quber db init --drop
```

## Database Schema

### Documents Table

Stores metadata about processed PDF documents:

```sql
CREATE TABLE documents (
  id SERIAL PRIMARY KEY,
  filename VARCHAR(255) NOT NULL UNIQUE,
  extraction_date TIMESTAMP NOT NULL,
  total_pages INT,
  total_tables INT,
  executive_summary TEXT,
  model_provider VARCHAR(50),
  model VARCHAR(100),
  created_at TIMESTAMP DEFAULT NOW()
);
```

### Extracted Tables Table

Stores individual tables with vector embeddings:

```sql
CREATE TABLE extracted_tables (
  id SERIAL PRIMARY KEY,
  document_id INT REFERENCES documents(id) ON DELETE CASCADE,
  table_id INT NOT NULL,
  page_number INT NOT NULL,
  procedural_title TEXT,
  llm_title TEXT,
  llm_description TEXT,
  table_markdown TEXT,
  headers JSONB,
  table_metadata JSONB,
  -- Embeddings (1024 dims for bge-large-en-v1.5)
  title_embedding VECTOR(1024),
  description_embedding VECTOR(1024),
  created_at TIMESTAMP DEFAULT NOW(),
  UNIQUE(document_id, table_id)
);
```

**Indexes**:
- `idx_tables_document` - Fast document lookups
- `idx_tables_desc_embedding` - IVFFlat index for description similarity
- `idx_tables_title_embedding` - IVFFlat index for title similarity
- `idx_documents_filename` - Fast filename lookups

## Importing Data

### Import Existing JSON Analyses

#### Using CLI Commands

```bash
# Import a single file (without embeddings)
quber db import-json output/TMUS_Q225_991_analysis.json

# Import with embedding generation
quber db import-json output/TMUS_Q225_991_analysis.json --generate-embeddings

# Import all files from a directory
quber db import-dir output/

# Import directory with embeddings
quber db import-dir output/ --generate-embeddings

# Import with custom pattern
quber db import-dir output/ --pattern "*.json"

# View database statistics
quber db stats
```

#### Using Standalone Script

```bash
# Import single file
python scripts/import_json_to_db.py output/TMUS_Q225_991_analysis.json

# Import directory
python scripts/import_json_to_db.py output/ --pattern "*_analysis.json"

# Initialize DB and import in one command
python scripts/import_json_to_db.py output/ --init-db --stats

# Show help
python scripts/import_json_to_db.py --help
```

### Data Format

The import utility expects JSON files with this structure:

```json
{
  "document": "filename.pdf",
  "extraction_date": "2025-10-05T21:10:49.397791",
  "total_pages": 9,
  "total_tables": 7,
  "model_provider": "anthropic",
  "model": "claude-3-5-sonnet-latest",
  "tables": [
    {
      "table_id": 0,
      "page": 2,
      "procedural_title": "...",
      "llm_title": "...",
      "llm_description": "...",
      "table_markdown": "| ... |",
      "headers": [...],
      "metadata": {...}
    }
  ]
}
```

**Note**: Use `--generate-embeddings` flag to populate embedding fields during import. Without this flag, embedding fields will be NULL and can be generated later using `quber db generate-embeddings`.

## Database Operations

### Connection Management

```python
from quber.db import get_engine, get_session, Document, ExtractedTable

# Get database engine
engine = get_engine()

# Use context manager for sessions (recommended)
with get_session() as session:
    documents = session.query(Document).all()
    for doc in documents:
        print(f"{doc.filename}: {len(doc.tables)} tables")
```

### Querying Data

```python
# Get a specific document
with get_session() as session:
    doc = session.query(Document).filter_by(
        filename="TMUS_Q225_991.pdf"
    ).first()

    # Access related tables
    for table in doc.tables:
        print(f"Table {table.table_id} on page {table.page_number}")
        print(f"Title: {table.llm_title}")
```

### Embedding Generation

Generate embeddings for existing tables in the database:

```bash
# Generate embeddings for all tables (uses local model by default)
quber db generate-embeddings

# Use OpenAI instead of local model
quber db generate-embeddings --provider openai

# Force CPU instead of GPU for local model
quber db generate-embeddings --device cpu

# Adjust batch size for processing
quber db generate-embeddings --batch-size 64

# Enable verbose logging
quber db generate-embeddings --verbose
```

**Embedding Models**:

- **Local (default)**: `BAAI/bge-large-en-v1.5`
  - 1024-dimensional vectors
  - Runs on GPU (CUDA) if available, falls back to CPU
  - No API key required
  - First run downloads ~1.3GB model

- **OpenAI**: `text-embedding-3-small`
  - 1536-dimensional vectors
  - Requires `OPENAI_API_KEY` environment variable
  - API costs apply

**Performance Notes**:
- Local model on RTX 3090: ~50 tables in 15 seconds (initial model load)
- Embeddings combine table title + description: `"{title}\n\n{description}"`
- Batch processing reuses model instance for efficiency

### Document-Scoped Semantic Search

Once embeddings are generated, you can perform vector similarity searches:

```python
from pgvector.sqlalchemy import Vector
from sqlalchemy import select

# Find similar tables within a specific document
query_embedding = [0.1, 0.2, ...]  # 1024-dim vector

with get_session() as session:
    # Document-scoped search
    similar_tables = session.query(ExtractedTable).filter(
        ExtractedTable.document_id == doc_id
    ).order_by(
        ExtractedTable.description_embedding.cosine_distance(query_embedding)
    ).limit(5).all()
```

## Docker Compose Management

### Start/Stop Database

```bash
# Start database
docker compose up -d

# Stop database (preserves data)
docker compose down

# Stop and remove volumes (WARNING: data loss)
docker compose down -v
```

### Database Maintenance

```bash
# View database logs
docker compose logs -f pgvector-db

# Access PostgreSQL shell
docker compose exec pgvector-db psql -U quber -d quber_rag

# Backup database
docker compose exec pgvector-db pg_dump -U quber quber_rag > backup.sql

# Restore database
cat backup.sql | docker compose exec -T pgvector-db psql -U quber quber_rag
```

### Health Check

The Docker Compose configuration includes a health check:

```yaml
healthcheck:
  test: ["CMD-SHELL", "pg_isready -U quber"]
  interval: 10s
  timeout: 5s
  retries: 5
```

Check health status:

```bash
docker compose ps
# Look for "healthy" in the STATUS column
```

## Troubleshooting

### Connection Errors

If you can't connect to the database:

1. **Check if container is running**:
   ```bash
   docker compose ps
   ```

2. **Check logs for errors**:
   ```bash
   docker compose logs pgvector-db
   ```

3. **Verify environment variables**:
   ```bash
   cat .env | grep POSTGRES
   ```

4. **Test connection manually**:
   ```bash
   docker compose exec pgvector-db psql -U quber -d quber_rag -c "SELECT 1;"
   ```

### Import Failures

If imports fail:

1. **Check database is running** (see above)
2. **Verify JSON file structure** matches expected format
3. **Check for duplicate documents** - files with same filename are skipped
4. **Enable verbose logging**:
   ```bash
   quber db import-json file.json --verbose
   ```

### Reset Database

To completely reset the database:

```bash
# Method 1: Drop and recreate tables
quber db init --drop

# Method 2: Destroy and recreate container
docker compose down -v
docker compose up -d
quber db init
```

## GPU Requirements

For optimal performance with local embeddings:

**CUDA Support**:
- PyTorch with CUDA 12.8+
- NVIDIA GPU with compute capability 6.0+ (Pascal or newer)
- Tested on RTX 3090 with 24GB VRAM

**Check CUDA Availability**:
```python
import torch
print(f"CUDA available: {torch.cuda.is_available()}")
print(f"GPU: {torch.cuda.get_device_name(0)}")
```

If CUDA is unavailable, the local model automatically falls back to CPU (slower but functional).

## Next Steps

After setting up the database infrastructure:

1. ✅ **Generate embeddings** - Use `quber db generate-embeddings` for existing data
2. **Implement semantic search** - Build query interface for document-scoped table retrieval
3. **Add caching layer** - Optimize repeated queries with Redis or similar
4. **Production deployment** - Migrate to managed PostgreSQL with proper security

## References

- [pgvector Documentation](https://github.com/pgvector/pgvector)
- [SQLAlchemy Documentation](https://docs.sqlalchemy.org/)
- [PostgreSQL Documentation](https://www.postgresql.org/docs/)