211 lines
6.8 KiB
Python
211 lines
6.8 KiB
Python
"""Thin async billing client for the third-party proxy.
|
|
|
|
Calls the same reserve/finalize HTTP endpoints as BillingMiddleware,
|
|
but with semantics appropriate for third-party task calls:
|
|
- estimatedTokens = 0 (not applicable)
|
|
- finalAmount = actual provider monetary charge (thirdPartyConsumeMoney)
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from datetime import datetime, timedelta
|
|
from typing import Any
|
|
|
|
import httpx
|
|
|
|
from deerflow.config.app_config import get_app_config
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_SUCCESS_STATUS_CODES = {200, 1000}
|
|
|
|
|
|
async def reserve(
|
|
*,
|
|
thread_id: str | None,
|
|
call_id: str,
|
|
provider: str,
|
|
operation: str,
|
|
frozen_amount: float,
|
|
frozen_type: int | None,
|
|
frozen_token: int = 0,
|
|
request_payload: dict[str, Any] | None = None,
|
|
) -> str | None:
|
|
"""Reserve billing before forwarding a submit call.
|
|
|
|
Returns the frozen_id string on success, or None if billing is disabled
|
|
or the reserve call fails (non-blocking — proxy continues in that case).
|
|
"""
|
|
cfg = get_app_config().billing
|
|
if not cfg.enabled or not cfg.reserve_url:
|
|
logger.info(
|
|
"[ThirdPartyProxy][Billing] reserve skipped: enabled=%s reserve_url=%s call_id=%s",
|
|
cfg.enabled,
|
|
cfg.reserve_url,
|
|
call_id,
|
|
)
|
|
return None
|
|
|
|
resolved_frozen_type = frozen_type if frozen_type is not None else cfg.frozen_type
|
|
expire_at = datetime.now() + timedelta(seconds=cfg.default_expire_seconds)
|
|
payload: dict[str, Any] = {
|
|
"sessionId": thread_id,
|
|
"callId": call_id,
|
|
"modelName": _extract_model_name(request_payload) or provider,
|
|
"question": f"skill invokes {operation.split('/')[-1]}",
|
|
"frozenType": resolved_frozen_type,
|
|
"expireAt": expire_at.strftime("%Y-%m-%d %H:%M:%S"),
|
|
}
|
|
|
|
if resolved_frozen_type == 1:
|
|
payload["estimatedInputTokens"] = int(frozen_token)
|
|
payload["estimatedOutputTokens"] = int(frozen_token)
|
|
else:
|
|
payload["frozenAmount"] = frozen_amount
|
|
payload["estimatedInputTokens"] = 0
|
|
payload["estimatedOutputTokens"] = 0
|
|
|
|
logger.info(
|
|
"[ThirdPartyProxy][Billing] reserve request: url=%s call_id=%s provider=%s thread_id=%s",
|
|
cfg.reserve_url,
|
|
call_id,
|
|
provider,
|
|
thread_id,
|
|
)
|
|
logger.debug("[ThirdPartyProxy][Billing] reserve payload: %s", payload)
|
|
try:
|
|
async with httpx.AsyncClient(timeout=cfg.timeout_seconds) as client:
|
|
resp = await client.post(cfg.reserve_url, headers=cfg.headers, json=payload)
|
|
resp.raise_for_status()
|
|
data: dict = resp.json()
|
|
except Exception as exc:
|
|
logger.warning("[ThirdPartyProxy][Billing] reserve HTTP error: %s", exc)
|
|
return None
|
|
|
|
logger.info(
|
|
"[ThirdPartyProxy][Billing] reserve response: call_id=%s status_code=%s",
|
|
call_id,
|
|
resp.status_code,
|
|
)
|
|
logger.debug("[ThirdPartyProxy][Billing] reserve response body: %s", data)
|
|
|
|
if not _is_success(data):
|
|
logger.warning(
|
|
"[ThirdPartyProxy][Billing] reserve rejected: call_id=%s status=%s payload=%s",
|
|
call_id,
|
|
data.get("status") or data.get("code"),
|
|
data,
|
|
)
|
|
return None
|
|
|
|
frozen_id = (data.get("data") or {}).get("frozenId")
|
|
if not isinstance(frozen_id, str) or not frozen_id:
|
|
logger.warning(
|
|
"[ThirdPartyProxy][Billing] reserve response missing frozenId: call_id=%s payload=%s",
|
|
call_id,
|
|
data,
|
|
)
|
|
return None
|
|
|
|
logger.info("[ThirdPartyProxy][Billing] reserve ok: call_id=%s frozen_id=%s", call_id, frozen_id)
|
|
logger.debug(
|
|
"[ThirdPartyProxy][Billing] reserve success details: provider=%s operation=%s expire_at=%s",
|
|
provider,
|
|
operation,
|
|
payload["expireAt"],
|
|
)
|
|
return frozen_id
|
|
|
|
|
|
async def finalize(
|
|
*,
|
|
frozen_id: str,
|
|
final_amount: float,
|
|
finalize_reason: str,
|
|
usage_input_tokens: int = 0,
|
|
usage_output_tokens: int = 0,
|
|
) -> bool:
|
|
"""Finalize billing after a third-party call reaches a terminal state.
|
|
|
|
final_amount is the actual provider charge (e.g. thirdPartyConsumeMoney from RunningHub).
|
|
Pass 0 for failed/cancelled calls.
|
|
Returns True on success.
|
|
"""
|
|
cfg = get_app_config().billing
|
|
if not cfg.enabled or not cfg.finalize_url:
|
|
# Billing not configured — treat as success so the caller marks the record finalized
|
|
logger.info(
|
|
"[ThirdPartyProxy][Billing] finalize skipped: enabled=%s finalize_url=%s frozen_id=%s",
|
|
cfg.enabled,
|
|
cfg.finalize_url,
|
|
frozen_id,
|
|
)
|
|
return True
|
|
|
|
payload = {
|
|
"frozenId": frozen_id,
|
|
"finalAmount": final_amount,
|
|
"usageInputTokens": usage_input_tokens,
|
|
"usageOutputTokens": usage_output_tokens,
|
|
"usageTotalTokens": usage_input_tokens + usage_output_tokens,
|
|
"finalizeReason": finalize_reason,
|
|
}
|
|
|
|
logger.info(
|
|
"[ThirdPartyProxy][Billing] finalize request: frozen_id=%s amount=%s reason=%s url=%s",
|
|
frozen_id,
|
|
final_amount,
|
|
finalize_reason,
|
|
cfg.finalize_url,
|
|
)
|
|
logger.debug("[ThirdPartyProxy][Billing] finalize payload: %s", payload)
|
|
try:
|
|
async with httpx.AsyncClient(timeout=cfg.timeout_seconds) as client:
|
|
resp = await client.post(cfg.finalize_url, headers=cfg.headers, json=payload)
|
|
resp.raise_for_status()
|
|
data: dict = resp.json()
|
|
except Exception as exc:
|
|
logger.warning("[ThirdPartyProxy][Billing] finalize HTTP error: frozen_id=%s err=%s", frozen_id, exc)
|
|
return False
|
|
|
|
logger.info(
|
|
"[ThirdPartyProxy][Billing] finalize response: frozen_id=%s status_code=%s",
|
|
frozen_id,
|
|
resp.status_code,
|
|
)
|
|
logger.debug("[ThirdPartyProxy][Billing] finalize response body: %s", data)
|
|
|
|
if not _is_success(data):
|
|
logger.warning(
|
|
"[ThirdPartyProxy][Billing] finalize rejected: frozen_id=%s status=%s payload=%s",
|
|
frozen_id,
|
|
data.get("status") or data.get("code"),
|
|
data,
|
|
)
|
|
return False
|
|
|
|
logger.info("[ThirdPartyProxy][Billing] finalize ok: frozen_id=%s", frozen_id)
|
|
logger.debug(
|
|
"[ThirdPartyProxy][Billing] finalize success details: amount=%s reason=%s",
|
|
final_amount,
|
|
finalize_reason,
|
|
)
|
|
return True
|
|
|
|
|
|
def _is_success(data: dict) -> bool:
|
|
status = data.get("status") or data.get("code")
|
|
if isinstance(status, int) and status in _SUCCESS_STATUS_CODES:
|
|
return True
|
|
return data.get("success") is True
|
|
|
|
|
|
def _extract_model_name(request_payload: dict[str, Any] | None) -> str | None:
|
|
if not isinstance(request_payload, dict):
|
|
return None
|
|
model = request_payload.get("model")
|
|
if isinstance(model, str) and model:
|
|
return model
|
|
return None
|