103 lines
3.3 KiB
Python
103 lines
3.3 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Migration script to add WaveSpeed provider fields to UsageSummary table
|
|
Run this script to update the database schema for WaveSpeed usage tracking
|
|
"""
|
|
|
|
import sqlite3
|
|
import os
|
|
import sys
|
|
|
|
def find_database():
|
|
"""Find the database file in common locations"""
|
|
print("🔍 Searching for database files...")
|
|
|
|
# Search in current directory and subdirectories
|
|
for root, dirs, files in os.walk('.'):
|
|
for file in files:
|
|
if file.endswith('.db') or file.endswith('.sqlite'):
|
|
db_path = os.path.join(root, file)
|
|
print(f"📁 Found database: {db_path}")
|
|
return db_path
|
|
|
|
# Check common paths
|
|
search_paths = [
|
|
'./data/alwrity.db',
|
|
'./alwrity.db',
|
|
'./database/alwrity.db',
|
|
'./backend/data/alwrity.db',
|
|
'./backend/alwrity.db'
|
|
]
|
|
|
|
for path in search_paths:
|
|
if os.path.exists(path):
|
|
print(f"📁 Found database at common path: {path}")
|
|
return path
|
|
|
|
print("❌ No database file found in any location")
|
|
return None
|
|
|
|
def run_migration():
|
|
"""Execute the WaveSpeed migration"""
|
|
db_path = find_database()
|
|
|
|
if not db_path:
|
|
print("❌ No database file found")
|
|
print("Please ensure the application has been run at least once to create the database")
|
|
return False
|
|
|
|
print(f"📁 Using database: {db_path}")
|
|
|
|
try:
|
|
conn = sqlite3.connect(db_path)
|
|
cursor = conn.cursor()
|
|
|
|
# Check if columns already exist
|
|
cursor.execute('PRAGMA table_info(usage_summary)')
|
|
columns = [col[1] for col in cursor.fetchall()]
|
|
|
|
wavespeed_cols = [col for col in columns if 'wavespeed' in col]
|
|
|
|
if wavespeed_cols:
|
|
print(f"✅ WaveSpeed columns already exist: {wavespeed_cols}")
|
|
return True
|
|
|
|
print("🔧 Adding WaveSpeed columns to usage_summary table...")
|
|
|
|
# Add the columns
|
|
cursor.execute('ALTER TABLE usage_summary ADD COLUMN wavespeed_calls INTEGER DEFAULT 0')
|
|
cursor.execute('ALTER TABLE usage_summary ADD COLUMN wavespeed_tokens INTEGER DEFAULT 0')
|
|
cursor.execute('ALTER TABLE usage_summary ADD COLUMN wavespeed_cost FLOAT DEFAULT 0.0')
|
|
|
|
conn.commit()
|
|
|
|
# Verify the changes
|
|
cursor.execute('PRAGMA table_info(usage_summary)')
|
|
updated_columns = [col[1] for col in cursor.fetchall()]
|
|
added_wavespeed_cols = [col for col in updated_columns if 'wavespeed' in col]
|
|
|
|
print(f"✅ Successfully added WaveSpeed columns: {added_wavespeed_cols}")
|
|
|
|
return True
|
|
|
|
except sqlite3.Error as e:
|
|
print(f"❌ SQLite error: {e}")
|
|
return False
|
|
except Exception as e:
|
|
print(f"❌ Unexpected error: {e}")
|
|
return False
|
|
finally:
|
|
if 'conn' in locals():
|
|
conn.close()
|
|
|
|
if __name__ == "__main__":
|
|
print("🚀 Running WaveSpeed migration...")
|
|
success = run_migration()
|
|
|
|
if success:
|
|
print("✅ WaveSpeed migration completed successfully!")
|
|
print("The system can now track WaveSpeed LLM usage and costs.")
|
|
else:
|
|
print("❌ Migration failed. Please check the error messages above.")
|
|
sys.exit(1)
|