27 lines
1.0 KiB
Python
27 lines
1.0 KiB
Python
from ..base import ModelService
|
|
from .config import IMAGE_GENERATION_PRICE, get_config
|
|
from typing import Dict, Any, Tuple, Generator, Union
|
|
import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class NanoBananaService(ModelService):
|
|
def get_api_config(self) -> Tuple[str, str]:
|
|
config = get_config()
|
|
return config['api_url'], config['api_key']
|
|
|
|
def check_balance(self, balance: float) -> Tuple[bool, float, str]:
|
|
if balance < IMAGE_GENERATION_PRICE:
|
|
return False, IMAGE_GENERATION_PRICE, f'余额不足。当前余额: {balance}, 需要: {IMAGE_GENERATION_PRICE}'
|
|
return True, IMAGE_GENERATION_PRICE, ""
|
|
|
|
def calculate_cost(self, usage: Dict[str, Any] = None, stream: bool = False) -> float:
|
|
# 固定按次计费
|
|
return IMAGE_GENERATION_PRICE
|
|
|
|
def prepare_payload(self, data: Dict[str, Any]) -> Dict[str, Any]:
|
|
return {k: v for k, v in data.items() if k not in ['user']}
|
|
|
|
def handle_response(self, response, stream: bool) -> Union[Dict[str, Any], Generator]:
|
|
pass
|