senti/scripts/migrate.py

import json
import os
from backend.database.postgres.db import engine, SessionLocal
from backend.database.postgres.models import Base, TrainingEntry, Inventory
from backend.database.audit_log import AuditLog

def migrate_data():
    print("Creating database tables...")
    Base.metadata.create_all(bind=engine)
    
    db = SessionLocal()
    
    # Migrate fine_tuning_pairs.jsonl to TrainingEntry
    jsonl_path = "data/fine_tuning_pairs.jsonl"
    if os.path.exists(jsonl_path):
        print(f"Migrating {jsonl_path}...")
        added_entries = 0
        with open(jsonl_path, "r", encoding="utf-8") as f:
            for line in f:
                try:
                    data = json.loads(line)
                    if 'user' not in data or 'senti' not in data:
                        continue
                    
                    entry = TrainingEntry(
                        user_id="anonymous",
                        query=data.get('user'),
                        response=data.get('senti'),
                        intent=data.get('intent', 'UNKNOWN'),
                        language=data.get('language', 'en'),
                        country_code=data.get('country_code', 'KE'),
                        flavor=data.get('flavor', 'PERSONAL'),
                        data_type=data.get('data_type', 'general')
                    )
                    db.add(entry)
                    added_entries += 1
                except Exception as e:
                    pass
        
        try:
            db.commit()
            print(f"Successfully migrated {added_entries} training entries to PostgreSQL/SQLite.")
        except Exception as e:
            db.rollback()
            print(f"Error saving training entries: {e}")

    # Migrate inventory.json to Inventory table
    inv_path = "data/inventory.json"
    if os.path.exists(inv_path):
        print(f"Migrating {inv_path}...")
        try:
            with open(inv_path, "r", encoding="utf-8") as f:
                inv_data = json.load(f)
                
            added_inv = 0
            for item, details in inv_data.items():
                inv = Inventory(
                    user_id="user_001", # Assume a default user
                    item_name=item,
                    quantity=details.get("stock", 0),
                    unit=details.get("unit", "units"),
                    reorder_level=details.get("reorder_level", 10)
                )
                db.add(inv)
                added_inv += 1
            
            db.commit()
            print(f"Successfully migrated {added_inv} inventory items.")
        except Exception as e:
            db.rollback()
            print(f"Error migrating inventory: {e}")
            
    # Show total counts
    total_training = db.query(TrainingEntry).count()
    total_inventory = db.query(Inventory).count()
    print(f"Migration Complete.\\nTotal TrainingEntry rows: {total_training}\\nTotal Inventory rows: {total_inventory}")
    db.close()

if __name__ == "__main__":
    migrate_data()