app/db/database.py

"""Database setup for SQLite and ChromaDB"""
from sqlalchemy import create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
import chromadb
from chromadb.config import Settings as ChromaSettings
from app.core.config import get_settings

settings = get_settings()

# ==================== SQLite Setup ====================

# Create SQLite engine
engine = create_engine(
    settings.database_url,
    connect_args={"check_same_thread": False},  # Needed for SQLite
    echo=False  # Set to True for SQL query logging
)

# Create SessionLocal class
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

# Create Base class for models
Base = declarative_base()


def get_db():
    """Dependency to get SQLite database session"""
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


def init_db():
    """Initialize SQLite database tables"""
    # Import models to register them with Base
    from app.models import document  # noqa: F401
    Base.metadata.create_all(bind=engine)


# ==================== ChromaDB Setup ====================

# Initialize ChromaDB client (singleton pattern)
_chroma_client = None


def get_chroma_client():
    """Get ChromaDB client instance (singleton)"""
    global _chroma_client
    
    if _chroma_client is None:
        _chroma_client = chromadb.PersistentClient(
            path=settings.chroma_db_path,
            settings=ChromaSettings(
                anonymized_telemetry=False,
                allow_reset=True
            )
        )
    
    return _chroma_client


def get_chroma_collection():
    """Get or create ChromaDB collection"""
    client = get_chroma_client()
    
    # Get or create collection
    collection = client.get_or_create_collection(
        name=settings.chroma_collection_name,
        metadata={"description": "Document embeddings for RAG"}
    )
    
    return collection


def init_chroma():
    """Initialize ChromaDB collection"""
    collection = get_chroma_collection()
    print(f"ChromaDB collection '{settings.chroma_collection_name}' initialized")
    print(f"Collection count: {collection.count()}")