sage_datamart/scripts/init_data.py

#!/usr/bin/env python3
"""
数据集市初始化脚本 — 从 llmusage + llmusage_history 导入全量历史数据到 dm_*

用法: cd /path/to/sage && python /path/to/sage_datamart/scripts/init_data.py
"""
import asyncio
import json
import sys
from datetime import datetime
from sqlor.dbpools import DBPools
from appPublic.jsonConfig import getConfig
from appPublic.uniqueID import getID
from appPublic.dictObject import DictObject

DB_NAME = 'sage'


def parse_usages(usages_str):
    """从 llmusage.usages JSON 中提取 token 数"""
    try:
        u = json.loads(usages_str) if isinstance(usages_str, str) else usages_str
        if isinstance(u, dict):
            return (
                u.get('prompt_tokens', 0) or 0,
                u.get('completion_tokens', 0) or 0,
                u.get('total_tokens', 0) or 0
            )
    except:
        pass
    return 0, 0, 0


async def batch_sync_call_fact(sor, source_table, batch_size=1000):
    """
    批量从 llmusage / llmusage_history 导入 dm_model_call_fact
    自动去重（已存在 luid 跳过）
    """
    offset = 0
    total_inserted = 0
    total_skipped = 0

    while True:
        sql = (
            'SELECT lu.*, l.model as llm_model, l.catelogid, l.providerid, l.ownerid '
            'FROM ' + source_table + ' lu '
            'JOIN llm l ON l.id = lu.llmid '
            'ORDER BY lu.use_time '
            'LIMIT ${limit}$ OFFSET ${offset}$'
        )
        rows = await sor.sqlExe(sql, {'limit': batch_size, 'offset': offset})
        if not rows:
            break

        print(f'  [{source_table}] batch offset={offset} rows={len(rows)}')

        for r in rows:
            r = DictObject(r)
            existing = await sor.sqlExe(
                'SELECT id FROM dm_model_call_fact WHERE luid=${luid}$',
                {'luid': r.id}
            )
            if existing:
                total_skipped += 1
                continue

            prompt_tokens, completion_tokens, _ = parse_usages(r.usages)
            try:
                call_time = datetime.strptime(str(r.use_time)[:19], '%Y-%m-%d %H:%M:%S')
            except Exception:
                total_skipped += 1
                continue

            ns = {
                'id': getID(),
                'luid': r.id,
                'llmid': r.llmid,
                'model': r.llm_model or r.model,
                'catelogid': r.catelogid,
                'userid': r.userid,
                'userorgid': r.userorgid,
                'ownerid': r.ownerid,
                'providerid': r.providerid,
                'call_date': r.use_date or call_time.strftime('%Y-%m-%d'),
                'call_hour': call_time.hour,
                'call_time': r.use_time,
                'ttft_ms': int(r.responsed_seconds * 1000) if r.responsed_seconds else None,
                'ttot_ms': int(r.finish_seconds * 1000) if r.finish_seconds else None,
                'prompt_tokens': prompt_tokens,
                'completion_tokens': completion_tokens,
                'status': r.status or 'SUCCEEDED',
                'amount': r.amount or 0,
                'currency': getattr(r, 'currency', 'CNY') or 'CNY',
            }
            await sor.C('dm_model_call_fact', ns)
            total_inserted += 1

        offset += batch_size

    print(f'  [{source_table}] done: inserted={total_inserted}, skipped={total_skipped}')
    return total_inserted, total_skipped


async def batch_aggregate_history(sor):
    """按 dm_model_call_fact 中所有历史日期重算聚合表"""
    from sage_datamart.etl import aggregate_daily_perf, aggregate_provider_cost

    rows = await sor.sqlExe(
        'SELECT DISTINCT call_date FROM dm_model_call_fact ORDER BY call_date',
        {}
    )
    dates = [r['call_date'].strftime('%Y-%m-%d') if hasattr(r['call_date'], 'strftime') else str(r['call_date']) for r in rows]
    print(f'  aggregating {len(dates)} distinct dates...')

    for d in dates:
        await aggregate_daily_perf(sor, stat_date=d)
        await aggregate_provider_cost(sor, stat_date=d)

    return len(dates)


async def run_init(sor):
    """
    初始化流程:
      1. 建表（如不存在）
      2. 从 llmusage_history 批量导入历史
      3. 从 llmusage 导入当前数据
      4. 重算所有日期的聚合表
    """
    from sage_datamart.etl import ensure_tables

    print('[1/4] Creating tables...')
    await ensure_tables(sor)

    print('[2/4] Importing llmusage_history...')
    ins_hist, skip_hist = await batch_sync_call_fact(sor, 'llmusage_history')

    print('[3/4] Importing llmusage...')
    ins_curr, skip_curr = await batch_sync_call_fact(sor, 'llmusage')

    print('[4/4] Aggregating historical dates...')
    dates = await batch_aggregate_history(sor)

    return {
        'llmusage_history': {'inserted': ins_hist, 'skipped': skip_hist},
        'llmusage': {'inserted': ins_curr, 'skipped': skip_curr},
        'aggregated_dates': dates,
    }


async def main():
    config = getConfig('.')
    db = DBPools(config.databases)

    async with db.sqlorContext(DB_NAME) as sor:
        result = await run_init(sor)

    print('\n=== Initialization Complete ===')
    print(json.dumps(result, indent=2))


if __name__ == '__main__':
    asyncio.run(main())