yumoqing 6351ee9b56 feat: Redis atomic balance reservation — prevent concurrent overspend
- balance.py: reserve_balance/finalize_balance/refund_balance with Lua
- utils.py: get_model_max_cost from llmusage history, update_model_max_cost
- init.py: register reserve_balance, finalize_balance, refund_balance on env
- chat/completions/index.dspy: reserve before inference, refund on exception
- DB migration: llm.max_cost DECIMAL(10,4) for historical max charge
2026-07-31 14:57:17 +08:00

60 lines
1.8 KiB
Plaintext

async def gen():
env = request._run_ns.copy()
f = partial(inference_generator, request, params_kw=params_kw)
try:
if params_kw.stream:
async for l in f():
yield f'data: {l}\n'
yield 'data: [DONE]\n\n'
else:
async for l in f():
yield l
except Exception as e:
luid = params_kw.get('_luid')
if luid:
await env.refund_balance(luid)
raise
debug_params('params_kw', params_kw)
catelogid = params_kw.catelogid or 't2t'
if params_kw.off_peak:
off_peak = params_kw.off_peak
if off_peak in [True, "Y" "y", 1, "1"]:
off_peak = True
else:
off_peak = False
params_kw.off_peak = off_peak
userid = await get_user()
userorgid = await get_userorgid()
if userid is None:
debug(f'need login')
return openai_403()
if not params_kw.prompt and not params_kw.messages:
debug(f'missing prompt and messages, model={params_kw.model}')
d = return_error('Missing need data(prompt or messages)')
return json_response(d, status=400)
env = request._run_ns
llmid = await env.get_llmid_cached(env, params_kw.model or 'qwen3-max', catelogid)
if not llmid:
debug(f'{params_kw.model=} not found for catelogid={catelogid}')
return openai_400()
params_kw.llmid = llmid
params_kw.llmcatelogid = catelogid
debug(f'{params_kw.llmid=}')
luid = getID()
reserved = await env.reserve_balance(params_kw.llmid, userorgid, luid)
if not reserved.get('ok'):
debug(f'{userid=} balance not enough: {reserved}')
return openai_429()
params_kw._luid = luid
params_kw._reserved = reserved
f = await checkCustomerBalance(params_kw.llmid, userid, userorgid)
if not f and reserved.get('no_redis'):
debug(f'{userid=} balance not enough (DB fallback)')
return openai_429()
# debug(f'{tools=}, {request._run_ns.tools=}')
return await env.stream_response(request, gen, content_type='application/json')