sage_datamart/scripts/init_data.py

#!/usr/bin/env python3
"""
数据集市初始化脚本 — 从 llmusage + llmusage_history 导入全量历史数据到 dm_*
建表由 models/*.json + json2ddl + build.sh 部署时完成，本脚本只负责数据读写。

用法: cd /path/to/sage && python /path/to/sage_datamart/scripts/init_data.py
"""
import asyncio
import json
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 Exception:
        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.providerid, '
            '(SELECT llmcatelogid FROM llm_api_map '
            'WHERE llmid = lu.llmid AND isdefaultcatelog = \'1\' LIMIT 1) as catelogid '
            '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('  [' + source_table + '] batch offset=' + str(offset) + ' rows=' + str(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

            currency = getattr(r, 'amount_currency', None) or 'CNY'
            amount = getattr(r, 'amount', 0) or 0
            userorgid = getattr(r, 'userorgid', None)
            distributor_orgid = getattr(r, 'distributor_orgid', None)

            ns = {
                'id': getID(),
                'luid': r.id,
                'llmid': r.llmid,
                'model': r.llm_model or '',
                'catelogid': r.catelogid or '',
                'userid': r.userid or '',
                'userorgid': userorgid or '',
                'ownerid': r.ownerid or '',
                'providerid': r.providerid or '',
                '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': amount,
                'currency': currency,
                'amount_cny': amount,
                'distributor_orgid': distributor_orgid or '',
                'sale_userid': '',
            }
            await sor.C('dm_model_call_fact', ns)
            total_inserted += 1

        offset += batch_size

    print('  [' + source_table + '] done: inserted=' + str(total_inserted) + ', skipped=' + str(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('  aggregating ' + str(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. 从 llmusage_history 批量导入历史
      2. 从 llmusage 导入当前数据
      3. 重算所有日期的聚合表
    """
    print('[1/3] Importing llmusage_history...')
    ins_hist, skip_hist = await batch_sync_call_fact(sor, 'llmusage_history')

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

    print('[3/3] 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())