scripts/migrate/migrate_chroma_db.py

import uuid
from datetime import datetime

import chromadb
from ktem.index.models import Index
from sqlalchemy import (
    JSON,
    Column,
    DateTime,
    Integer,
    String,
    UniqueConstraint,
    create_engine,
    select,
)
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.ext.mutable import MutableDict
from sqlalchemy.orm import Session
from tzlocal import get_localzone


def _init_resource(private: bool = True, id: int = 1):
    """Init schemas. Hard-code"""
    Base = declarative_base()

    if private:
        Source = type(
            "Source",
            (Base,),
            {
                "__tablename__": f"index__{id}__source",
                "__table_args__": (
                    UniqueConstraint("name", "user", name="_name_user_uc"),
                ),
                "id": Column(
                    String,
                    primary_key=True,
                    default=lambda: str(uuid.uuid4()),
                    unique=True,
                ),
                "name": Column(String),
                "path": Column(String),
                "size": Column(Integer, default=0),
                "date_created": Column(
                    DateTime(timezone=True), default=datetime.now(get_localzone())
                ),
                "user": Column(Integer, default=1),
                "note": Column(
                    MutableDict.as_mutable(JSON),  # type: ignore
                    default={},
                ),
            },
        )
    else:
        Source = type(
            "Source",
            (Base,),
            {
                "__tablename__": f"index__{id}__source",
                "id": Column(
                    String,
                    primary_key=True,
                    default=lambda: str(uuid.uuid4()),
                    unique=True,
                ),
                "name": Column(String, unique=True),
                "path": Column(String),
                "size": Column(Integer, default=0),
                "date_created": Column(
                    DateTime(timezone=True), default=datetime.now(get_localzone())
                ),
                "user": Column(Integer, default=1),
                "note": Column(
                    MutableDict.as_mutable(JSON),  # type: ignore
                    default={},
                ),
            },
        )
    Index = type(
        "IndexTable",
        (Base,),
        {
            "__tablename__": f"index__{id}__index",
            "id": Column(Integer, primary_key=True, autoincrement=True),
            "source_id": Column(String),
            "target_id": Column(String),
            "relation_type": Column(String),
            "user": Column(Integer, default=1),
        },
    )

    return {"Source": Source, "Index": Index}


def get_chromadb_collection(
    db_dir: str = "../ktem_app_data/user_data/vectorstore",
    collection_name: str = "index_1",
):
    """Extract collection from chromadb"""
    client = chromadb.PersistentClient(path=db_dir)
    collection = client.get_or_create_collection(collection_name)

    return collection


def update_metadata(metadata, file_id):
    """Update file_id"""
    metadata["file_id"] = file_id
    return metadata


def migrate_chroma_db(
    chroma_db_dir: str, sqlite_path: str, is_private: bool = True, int_index: int = 1
):
    chroma_collection_name = f"index_{int_index}"

    """Update chromadb with metadata.file_id"""
    engine = create_engine(sqlite_path)
    resource = _init_resource(private=is_private, id=int_index)
    print("Load sqlalchemy engine successfully!")

    chroma_db_collection = get_chromadb_collection(
        db_dir=chroma_db_dir, collection_name=chroma_collection_name
    )
    print(
        f"Load chromadb collection: {chroma_collection_name}, "
        f"path: {chroma_db_dir} successfully!"
    )

    # Load docs id of user
    with Session(engine) as session:
        stmt = select(resource["Source"])
        results = session.execute(stmt)
        doc_ids = [r[0].id for r in results.all()]
    print(f"Retrieve n-docs: {len(doc_ids)}")
    print(doc_ids)

    for doc_id in doc_ids:
        print("-")
        # Find corresponding vector ids
        with Session(engine) as session:
            stmt = select(resource["Index"]).where(
                resource["Index"].relation_type == "vector",
                resource["Index"].source_id.in_([doc_id]),
            )
            results = session.execute(stmt)
            vs_ids = [r[0].target_id for r in results.all()]

        print(f"Got {len(vs_ids)} vs_ids for doc {doc_id}")

        # Update file_id
        if len(vs_ids) > 0:
            batch = chroma_db_collection.get(ids=vs_ids, include=["metadatas"])
            batch.update(
                ids=batch["ids"],
                metadatas=[
                    update_metadata(metadata, doc_id) for metadata in batch["metadatas"]
                ],
            )

        # Assert file_id. Skip
        print(f"doc-{doc_id} got updated")


def main(chroma_db_dir: str, sqlite_path: str):
    engine = create_engine(sqlite_path)

    with Session(engine) as session:
        stmt = select(Index)

        results = session.execute(stmt)
        file_indices = [r[0] for r in results.all()]

        for file_index in file_indices:
            _id = file_index.id
            _is_private = file_index.config["private"]

            print(f"Migrating for Index id: {_id}, is_private: {_is_private}")

            migrate_chroma_db(
                chroma_db_dir=chroma_db_dir,
                sqlite_path=sqlite_path,
                is_private=_is_private,
                int_index=_id,
            )


if __name__ == "__main__":
    chrome_db_dir: str = "./vectorstore/kan_db"
    sqlite_path: str = "sqlite:///../ktem_app_data/user_data/sql.db"

    main(chrome_db_dir, sqlite_path)