Phase 2.2: Update AIEngine for token-based billing + GlobalModuleSettings

This commit is contained in:
IGNY8 VPS (Salman)
2025-12-23 07:03:35 +00:00
parent 12c82e78f9
commit 01a42b1549
3 changed files with 114 additions and 6 deletions

View File

@@ -376,18 +376,34 @@ class AIEngine:
# Map function name to operation type
operation_type = self._get_operation_type(function_name)
# Calculate actual amount based on results
# Extract token usage from AI response (standardize key names)
tokens_input = raw_response.get('input_tokens', 0)
tokens_output = raw_response.get('output_tokens', 0)
total_tokens = tokens_input + tokens_output
# Get model_config for token-based billing
model_config = None
if tokens_input > 0 or tokens_output > 0:
# Get model from response or use account default
model_config = CreditService.get_model_for_operation(
account=self.account,
operation_type=operation_type,
task_model_override=None # TODO: Support task-level model override
)
# Calculate actual amount based on results (for non-token operations)
actual_amount = self._get_actual_amount(function_name, save_result, parsed, data)
# Deduct credits using the new convenience method
# Deduct credits using token-based calculation if tokens available
CreditService.deduct_credits_for_operation(
account=self.account,
operation_type=operation_type,
amount=actual_amount,
amount=actual_amount, # Fallback for non-token operations
tokens_input=tokens_input,
tokens_output=tokens_output,
model_config=model_config,
cost_usd=raw_response.get('cost'),
model_used=raw_response.get('model', ''),
tokens_input=raw_response.get('tokens_input', 0),
tokens_output=raw_response.get('tokens_output', 0),
related_object_type=self._get_related_object_type(function_name),
related_object_id=save_result.get('id') or save_result.get('cluster_id') or save_result.get('task_id'),
metadata={
@@ -399,7 +415,7 @@ class AIEngine:
}
)
logger.info(f"[AIEngine] Credits deducted: {operation_type}, amount: {actual_amount}")
logger.info(f"[AIEngine] Credits deducted: {operation_type}, tokens: {total_tokens} ({tokens_input} in, {tokens_output} out), model: {model_config.model_name if model_config else 'legacy'}")
except InsufficientCreditsError as e:
# This shouldn't happen since we checked before, but log it
logger.error(f"[AIEngine] Insufficient credits during deduction: {e}")