diff --git a/llmage/accounting.py b/llmage/accounting.py index 6ab3042..fd50ec0 100644 --- a/llmage/accounting.py +++ b/llmage/accounting.py @@ -195,6 +195,13 @@ async def llm_accounting(llmusage): 'accounting_status': 'accounted' } await sor.U('llmusage', ns) + # Finalize Redis balance reservation (non-critical) + try: + from .balance import finalize_balance + await finalize_balance(env, llmusage.id, trans_amount, + llmid=llmusage.llmid) + except Exception as e: + exception(f'finalize_balance failed (non-critical): {e}') async def get_accounting_llmusages(luid=None): env = ServerEnv() diff --git a/llmage/asyncinference.py b/llmage/asyncinference.py index 34b06d0..621a72c 100644 --- a/llmage/asyncinference.py +++ b/llmage/asyncinference.py @@ -13,6 +13,7 @@ from appPublic.base64_to_file import base64_to_file, getFilenameFromBase64 from ahserver.serverenv import get_serverenv, ServerEnv from ahserver.filestorage import FileStorage from .accounting import llm_accounting, llm_charging +from .balance import refund_balance from .utils import * # Global set to keep references to background tasks @@ -67,7 +68,7 @@ async def async_uapi_request(request, llm, uapi = env.UpAppApi(request) userid = await env.uapi_data.get_calluserid(llm.upappid, orgid=llm.ownerid) b = None - luid = getID() + luid = params_kw.get('_luid') or getID() try: start_timestamp = time.time() if llm.callbackurl: @@ -77,6 +78,12 @@ async def async_uapi_request(request, llm, try: b = await uapi.call(llm.upappid, llm.apiname, userid, params=params_kw) except Exception as e: + # Refund balance reservation on submission failure + try: + if luid: + await refund_balance(ServerEnv(), luid) + except Exception: + pass estr = erase_apikey(e) ed = {"error": f"ERROR:{estr}", "status": "FAILED"} exception(f'{ed}') @@ -123,6 +130,12 @@ async def async_uapi_request(request, llm, task.add_done_callback(_background_tasks.discard) except Exception as e: + # Refund balance reservation on outer failure + try: + if luid: + await refund_balance(ServerEnv(), luid) + except Exception: + pass ed = {"error": f"ERROR:{e}", "status": "FAILED"} s = json.dumps(ed, ensure_ascii=False) s = ''.join(s.split('\n')) @@ -229,6 +242,12 @@ async def query_task_status(request, luid, onetime=False): ns['usages'] = json.dumps(new_output['usage']) await append_new_llmoutput(llmusage.ioinfo, new_output) await modify_llmusage(ns) + if llmusage.status == 'FAILED': + # Async task failed — refund the balance reservation + try: + await refund_balance(ServerEnv(), luid) + except Exception: + pass if llmusage.status in ['UNKNOWN', 'FAILED', 'SUCCEEDED']: critical(f'finished .. {llmusage.status=}') return diff --git a/llmage/balance.py b/llmage/balance.py index b706e24..7ba1be2 100644 --- a/llmage/balance.py +++ b/llmage/balance.py @@ -2,14 +2,54 @@ Redis atomic balance reservation for llmage. Prevents concurrent overspend: pre-deduct before inference, settle after accounting. +Uses a module-level redis.asyncio singleton so the atomic reserve works in +EVERY process (web workers + backend_accounting) without needing env.redis +to be injected. This fixes the previous fatal issue where env.redis was never +set, so reserve always silently fell back to the non-atomic DB check. + Key patterns: balance:{userorgid} — current pre-deducted balance (int, cents) - reserve:{luid} — {userorgid}|{llmid}|{max_cost} (TTL 600s) + reserve:{luid} — {userorgid}|{llmid}|{max_cost} (TTL) model:max_cost:{llmid} — historical max actual customer charge (int, cents) + +Reserve TTL: 600s for fast (stream/sync) calls, 3600s for async tasks that +may run for minutes before finalize/refund. """ import asyncio from appPublic.log import debug, exception +# ── Module-level async Redis singleton ─────────────────────── +_redis = None +_redis_lock = asyncio.Lock() + + +def _redis_url(): + try: + from appPublic.jsonConfig import getConfig + config = getConfig() + url = getattr(config.website, 'session_redis', None) + if url: + u = getattr(url, 'url', None) + if u: + return u + except Exception: + pass + return "redis://127.0.0.1:6379" + + +async def _get_redis(): + """Return a shared redis.asyncio client (lazy singleton).""" + global _redis + if _redis is not None: + return _redis + async with _redis_lock: + if _redis is None: + import redis.asyncio as aioredis + _redis = await aioredis.from_url( + _redis_url(), decode_responses=True) + return _redis + + # ── Lua scripts ────────────────────────────────────────────── RESERVE_LUA = """ @@ -19,6 +59,7 @@ local reserve_key = KEYS[3] local max_cost = tonumber(ARGV[1]) local db_balance = tonumber(ARGV[2]) -- 0 means no DB fallback local reserve_val = ARGV[3] +local ttl = tonumber(ARGV[4]) -- Get or init max_cost: use stored if higher, else set from arg local stored_max = redis.call('GET', cost_key) @@ -48,7 +89,7 @@ if balance < max_cost then return {0, 'insufficient balance'} end redis.call('DECRBY', bal_key, max_cost) -redis.call('SETEX', reserve_key, 600, reserve_val) +redis.call('SETEX', reserve_key, ttl, reserve_val) return {1, max_cost} """ @@ -123,24 +164,69 @@ def _from_cents(c): return round(int(c) / 100, 2) -async def _redis_eval(env, script, nkeys, *keys_and_args): - """Eval Lua script via env's Redis connection.""" - if not hasattr(env, 'redis') or env.redis is None: - raise RuntimeError('Redis unavailable') - return await asyncio.to_thread(env.redis.eval, script, nkeys, *keys_and_args) +async def _update_max_cost(llmid, actual_cost): + """Bump model:max_cost if actual exceeds stored value; persist to DB. - -async def reserve_balance(env, llmid, userorgid, luid): - """Atomically check balance and deduct max_cost via Redis Lua.""" + Keeps the pre-deduct baseline growing even when the original reserve was + skipped (no_history cold start) — otherwise reserve would never activate. + """ try: + redis = await _get_redis() + cost_key = f'model:max_cost:{llmid}' + cost_cents = _cents(actual_cost) + stored = await redis.get(cost_key) + if stored and int(stored) >= cost_cents: + return + await redis.set(cost_key, cost_cents) + from .utils import update_model_max_cost + await update_model_max_cost(llmid, actual_cost) + except Exception as e: + debug(f'_update_max_cost failed: {e}') + + +async def reserve_balance(env, llmid, userorgid, luid, ttl=600, userid=None): + """Atomically check balance and deduct max_cost via Redis Lua. + + Policy (centralized so every entry behaves the same): + - Self-owned org (llm.ownerid == userorgid): skip reserve. + - tpac user (external balance system): skip reserve. + - No pricing (ppid empty): skip reserve (availability handled elsewhere). + + Returns: + {'ok': True, 'max_cost': X} — reserved X + {'ok': True, 'max_cost': 0, 'skip': ...} — reserve skipped by policy + {'ok': True, 'max_cost': 0, 'no_history': True} — no max_cost data yet + {'ok': False, 'reason': ...} — insufficient balance + {'ok': True, 'max_cost': 0, 'no_redis': True} — Redis down, DB fallback + """ + try: + # ── Policy checks (need llm info) ── + try: + from .utils import get_llmage_llm, get_user_tpac + llm = await get_llmage_llm(llmid) + if llm and llm.ownerid == userorgid: + return {'ok': True, 'max_cost': 0, 'skip': 'self_org'} + if not llm or not llm.ppid: + return {'ok': True, 'max_cost': 0, 'skip': 'no_ppid'} + # tpac user: balance lives in external system, skip redis reserve + if userid: + try: + tpac = await get_user_tpac(userid) + if tpac: + return {'ok': True, 'max_cost': 0, 'skip': 'tpac'} + except Exception as e: + debug(f'reserve_balance: tpac check failed: {e}') + except Exception as e: + debug(f'reserve_balance: policy check failed: {e}') + + redis = await _get_redis() bal_key = f'balance:{userorgid}' cost_key = f'model:max_cost:{llmid}' reserve_key = f'reserve:{luid}' - - # Load max_cost from DB if not cached - redis = env.redis + + # Load max_cost from Redis if cached, else from DB history max_cost = 0 - stored = await asyncio.to_thread(redis.get, cost_key) + stored = await redis.get(cost_key) if stored: max_cost = int(stored) else: @@ -151,13 +237,13 @@ async def reserve_balance(env, llmid, userorgid, luid): max_cost = _cents(mc) except Exception as e: debug(f'reserve_balance: get_model_max_cost failed: {e}') - + if max_cost <= 0: return {'ok': True, 'max_cost': 0, 'no_history': True} - - # Load balance from DB if not cached + + # Load balance from DB if not cached in Redis db_balance = 0 - stored_bal = await asyncio.to_thread(redis.get, bal_key) + stored_bal = await redis.get(bal_key) if not stored_bal: try: from accounting.getaccount import getCustomerBalance @@ -168,70 +254,62 @@ async def reserve_balance(env, llmid, userorgid, luid): db_balance = _cents(float(bal)) except Exception as e: debug(f'reserve_balance: getCustomerBalance failed: {e}') - + reserve_val = f'{userorgid}|{llmid}|{max_cost}' - result = await _redis_eval(env, RESERVE_LUA, 3, + result = await redis.eval(RESERVE_LUA, 3, bal_key, cost_key, reserve_key, - str(max_cost), str(db_balance), reserve_val) - + str(max_cost), str(db_balance), reserve_val, str(ttl)) + if result[0] == 0: return {'ok': False, 'reason': result[1]} return {'ok': True, 'max_cost': _from_cents(result[1])} - - except RuntimeError: - debug('reserve_balance: Redis not configured, skip') - return {'ok': True, 'max_cost': 0, 'no_redis': True} + except Exception as e: # Redis down or unreachable — fall back to DB check, don't block debug(f'reserve_balance: Redis error, falling back to DB: {e}') return {'ok': True, 'max_cost': 0, 'no_redis': True} -async def finalize_balance(env, luid, actual_cost): - """After accounting: adjust balance, update max_cost.""" +async def finalize_balance(env, luid, actual_cost, llmid=None): + """After accounting: adjust balance, update max_cost. + + llmid: needed for cold-start max_cost updates when no reserve existed + (reserve skipped due to no_history) — keeps the pre-deduct baseline + growing so future calls can reserve. + """ try: - cost_key = 'model:max_cost:dummy' + redis = await _get_redis() reserve_key = f'reserve:{luid}' - - # We need the llmid to get the correct cost_key, but we won't know it - # until parsing the reserve. So use a placeholder and let Lua handle it. - # Actually FINALIZE_LUA only uses cost_key for updating, so pass a dummy. - # The real cost_key needs the llmid which comes from reserve parsing. - - redis = env.redis - reserve_val = await asyncio.to_thread(redis.get, reserve_key) + + reserve_val = await redis.get(reserve_key) if not reserve_val: debug(f'finalize_balance: no reserve for luid={luid}') + if llmid: + await _update_max_cost(llmid, actual_cost) return - - # Parse to get llmid + + # Parse to get llmid (reserve_val: userorgid|llmid|max_cost) parts = reserve_val.split('|') if len(parts) < 3: return llmid = parts[1] cost_key = f'model:max_cost:{llmid}' - - result = await _redis_eval(env, FINALIZE_LUA, 2, + + result = await redis.eval(FINALIZE_LUA, 2, cost_key, reserve_key, str(_cents(actual_cost))) - + if result[0] == 0: debug(f'finalize_balance: {result[1]}') return - + max_c = _from_cents(result[2]) diff = _from_cents(result[4]) debug(f'finalize_balance: luid={luid} max={max_c} actual={actual_cost} diff={diff}') - + # Persist max_cost to DB if actual exceeded previous max if actual_cost > max_c: - try: - from .utils import update_model_max_cost - await update_model_max_cost(llmid, actual_cost) - except Exception as e: - debug(f'finalize_balance: DB persist failed: {e}') - - except RuntimeError: - debug('finalize_balance: Redis unavailable, skip') + await _update_max_cost(llmid, actual_cost) + except Exception as e: exception(f'finalize_balance error: {e}') @@ -239,13 +317,39 @@ async def finalize_balance(env, luid, actual_cost): async def refund_balance(env, luid): """Full refund on API failure.""" try: + redis = await _get_redis() reserve_key = f'reserve:{luid}' - result = await _redis_eval(env, REFUND_LUA, 1, reserve_key) + result = await redis.eval(REFUND_LUA, 1, reserve_key) if result[0] == 1: debug(f'refund_balance: luid={luid} refund={_from_cents(result[1])}') else: debug(f'refund_balance: luid={luid} skipped ({result[1]})') - except RuntimeError: - pass except Exception as e: exception(f'refund_balance error: {e}') + + +async def extend_reserve(env, luid, ttl): + """Extend TTL of an existing reserve. + + Used when an entry pre-reserved with the short (600s) TTL but the + request is dispatched to async mode, where the task may run longer. + """ + try: + redis = await _get_redis() + await redis.expire(f'reserve:{luid}', ttl) + except Exception as e: + debug(f'extend_reserve failed: {e}') + + +async def invalidate_balance_cache(userorgid): + """Delete the cached Redis balance so the next reserve reloads from DB. + + Must be called after recharge / recharge reversal, otherwise the + pre-deduct baseline stays stale and valid requests get rejected. + """ + try: + redis = await _get_redis() + await redis.delete(f'balance:{userorgid}') + debug(f'invalidate_balance_cache: balance:{userorgid} deleted') + except Exception as e: + debug(f'invalidate_balance_cache failed: {e}') diff --git a/llmage/init.py b/llmage/init.py index c8f9891..e2ce390 100644 --- a/llmage/init.py +++ b/llmage/init.py @@ -173,8 +173,8 @@ def load_llmage(): env._reserve_balance = reserve_balance env._finalize_balance = finalize_balance env._refund_balance = refund_balance - env.reserve_balance = lambda llmid, userorgid, luid: reserve_balance(env, llmid, userorgid, luid) - env.finalize_balance = lambda luid, actual_cost: finalize_balance(env, luid, actual_cost) + env.reserve_balance = lambda llmid, userorgid, luid, ttl=600, userid=None: reserve_balance(env, llmid, userorgid, luid, ttl=ttl, userid=userid) + env.finalize_balance = lambda luid, actual_cost, llmid=None: finalize_balance(env, luid, actual_cost, llmid=llmid) env.refund_balance = lambda luid: refund_balance(env, luid) env.backup_accounted_llmusage = backup_accounted_llmusage env.read_ioinfo_content = read_ioinfo_content diff --git a/llmage/llmclient.py b/llmage/llmclient.py index 60b3b49..7e71a2f 100644 --- a/llmage/llmclient.py +++ b/llmage/llmclient.py @@ -14,7 +14,7 @@ from ahserver.filestorage import FileStorage from .asyncinference import async_uapi_request from .syncinference import sync_uapi_request from .accounting import llm_accounting, llm_charging -from .balance import refund_balance +from .balance import refund_balance, reserve_balance, extend_reserve from .utils import * async def uapi_request(request, llm, callerid, callerorgid, params_kw=None): @@ -27,7 +27,7 @@ async def uapi_request(request, llm, callerid, callerorgid, params_kw=None): userid = await env.uapi_data.get_calluserid(llm.upappid, orgid=llm.ownerid) outlines = [] txt = '' - luid = getID() + luid = params_kw.get('_luid') or getID() llmusage = None try: start_timestamp = time.time() @@ -161,6 +161,29 @@ async def _inference_generator(request, callerid, callerorgid, yield errmsg return params_kw.model = llm.model + # ── Unified balance reserve ──────────────────────────────────────── + # Covers every entry that did not pre-reserve. v1/chat/completions + # reserves at dspy level and passes _luid; everyone else reserves here. + # Reusing _luid downstream also fixes the old mismatch where the + # reserve key could never equal the llmusage id used by finalize. + if params_kw.get('_luid'): + if llm.stream == 'async': + # Entry reserved with the short TTL; async tasks run longer + await extend_reserve(env, params_kw._luid, 3600) + else: + _luid = getID() + _ttl = 3600 if llm.stream == 'async' else 600 + reserved = await reserve_balance(env, llm.id, callerorgid, _luid, + ttl=_ttl, userid=callerid) + if not reserved.get('ok'): + debug(f'balance reserve rejected: {reserved}') + errmsg = json.dumps({'status': 'FAILED', + 'error': f'余额不足(balance reserve rejected): {reserved.get("reason")}'}, + ensure_ascii=False) + '\n' + yield errmsg + return + params_kw._luid = _luid + params_kw._reserved = reserved if llm.stream == 'async': if llm.callbackurl: cb_url = env.entire_url(llm.callbackurl) diff --git a/llmage/product_interface.py b/llmage/product_interface.py index c06208d..416ac1e 100644 --- a/llmage/product_interface.py +++ b/llmage/product_interface.py @@ -22,6 +22,7 @@ from .utils import ( write_llmusage, ) from .accounting import llm_charging +from .balance import reserve_balance, refund_balance async def _resolve_llm(resource_ref_id): @@ -219,24 +220,46 @@ async def execute_product_service(resource_ref_id, user_id, user_org_id, request if not params_kw.get('transno'): params_kw.transno = luid + # ── Balance reservation (product path) ───────────────────────── + # Same atomic pre-deduct as the direct dspy path; policy checks + # (self-org / tpac / no-ppid) live inside reserve_balance. + ttl = 3600 if full_llm.stream == 'async' else 600 + reserved = await reserve_balance(env, full_llm.id, user_org_id, luid, + ttl=ttl, userid=user_id) + if not reserved.get('ok'): + return {'success': False, + 'message': f'余额不足(balance reserve rejected): {reserved.get("reason")}', + 'status': 'FAILED', 'task_id': luid} + params_kw._luid = luid + try: if full_llm.stream == 'async': from .asyncinference import async_uapi_request_product result = await async_uapi_request_product( full_llm, userid, user_id, user_org_id, params_kw, luid) - return result - elif not full_llm.stream: from .syncinference import sync_uapi_request_product result = await sync_uapi_request_product( full_llm, userid, user_id, user_org_id, params_kw, luid) - return result else: result = await _collect_stream(full_llm, userid, user_id, user_org_id, params_kw, luid) - return result + + # Sub-callers swallow exceptions and return success=False — + # refund here too (refund is idempotent via GETDEL) + if not result.get('success'): + try: + await refund_balance(env, luid) + except Exception: + pass + return result except Exception as e: + # Refund balance reservation on failure + try: + await refund_balance(env, luid) + except Exception: + pass exception(f'execute_product_service error: {e}') return {'success': False, 'message': str(e), 'status': 'FAILED', 'task_id': luid} @@ -349,6 +372,16 @@ async def execute_product_service_stream(resource_ref_id, user_id, user_org_id, if not params_kw.get('transno'): params_kw.transno = luid + # ── Balance reservation (product stream path) ────────────────── + reserved = await reserve_balance(env, full_llm.id, user_org_id, luid, + ttl=600, userid=user_id) + if not reserved.get('ok'): + yield {'chunk': None, 'usage_data': None, 'done': True, + 'error': f'余额不足(balance reserve rejected): {reserved.get("reason")}', + 'task_id': luid, 'status': 'FAILED'} + return + params_kw._luid = luid + outlines = [] txt = '' usage = None @@ -416,6 +449,11 @@ async def execute_product_service_stream(resource_ref_id, user_id, user_org_id, } except Exception as e: + # Refund balance reservation on failure + try: + await refund_balance(env, luid) + except Exception: + pass exception(f'stream error: {e}') yield {'chunk': None, 'usage_data': None, 'done': True, 'error': str(e), 'task_id': luid, 'status': 'FAILED'} diff --git a/llmage/syncinference.py b/llmage/syncinference.py index d6c13c3..977a1b8 100644 --- a/llmage/syncinference.py +++ b/llmage/syncinference.py @@ -14,6 +14,7 @@ from appPublic.base64_to_file import base64_to_file, getFilenameFromBase64 from ahserver.serverenv import get_serverenv, ServerEnv from ahserver.filestorage import FileStorage from .accounting import llm_accounting, llm_charging +from .balance import refund_balance from .utils import * async def sync_uapi_request(request, llm, callerid, callerorgid, params_kw=None): @@ -28,7 +29,7 @@ async def sync_uapi_request(request, llm, callerid, callerorgid, params_kw=None) outlines = [] b = None d = None - luid = getID() + luid = params_kw.get('_luid') or getID() try: start_timestamp = time.time() responsed_seconds = None @@ -69,6 +70,12 @@ async def sync_uapi_request(request, llm, callerid, callerorgid, params_kw=None) yield b # await write_llmusage(llmusage) except Exception as e: + # Refund balance reservation on failure + try: + if luid: + await refund_balance(ServerEnv(), luid) + except Exception: + pass exception(f'{e=},{format_exc()}, {b=}') estr = erase_apikey(e) ed = {"error": f"ERROR:{estr}", "status": "FAILED" ,"llmusageid": luid} diff --git a/wwwroot/v1/chat/completions/index.dspy b/wwwroot/v1/chat/completions/index.dspy index f15918b..c40ac0d 100644 --- a/wwwroot/v1/chat/completions/index.dspy +++ b/wwwroot/v1/chat/completions/index.dspy @@ -49,7 +49,7 @@ params_kw.llmcatelogid = catelogid debug(f'{params_kw.llmid=}') luid = getID() -reserved = await env.reserve_balance(params_kw.llmid, userorgid, luid) +reserved = await env.reserve_balance(params_kw.llmid, userorgid, luid, userid=userid) if not reserved.get('ok'): debug(f'{userid=} balance not enough: {reserved}') return openai_429()