517 lines
17 KiB
Python
517 lines
17 KiB
Python
|
"""
|
||
|
|
Fanfic Writer v2.0 - Price Table Manager
|
||
|
|
Multi-platform pricing management with version control
|
||
|
|
"""
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Dict, Any, Optional, List, Tuple
|
||
|
|
from datetime import datetime
|
||
|
|
|
||
|
|
from .atomic_io import atomic_write_json, atomic_append_jsonl
|
||
|
|
from .utils import get_timestamp_iso
|
||
|
|
|
||
|
|
|
||
|
|
# Default pricing table (RMB per 1M tokens)
|
||
|
|
DEFAULT_PRICE_TABLE = {
|
||
|
|
"version": "1.0.0",
|
||
|
|
"updated_at": "2026-02-16T00:00:00+08:00",
|
||
|
|
"source": "default",
|
||
|
|
"usd_cny_rate": 6.90,
|
||
|
|
"currency": "CNY",
|
||
|
|
"models": [
|
||
|
|
{
|
||
|
|
"key": "moonshot:kimi-k2.5:standard:<=128k:off:none",
|
||
|
|
"provider": "moonshot",
|
||
|
|
"model_id": "kimi-k2.5",
|
||
|
|
"model_name": "Kimi K2.5",
|
||
|
|
"tier": "standard",
|
||
|
|
"context_bucket": "<=128k",
|
||
|
|
"thinking_mode": "off",
|
||
|
|
"cache_mode": "none",
|
||
|
|
"currency": "CNY",
|
||
|
|
"input_rate": 4.14,
|
||
|
|
"output_rate": 20.70,
|
||
|
|
"cached_input_rate": 0.69,
|
||
|
|
"cache_write_rate": 0.0
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"key": "nvidia:moonshotai/kimi-k2.5:standard:<=128k:off:none",
|
||
|
|
"provider": "nvidia",
|
||
|
|
"model_id": "moonshotai/kimi-k2.5",
|
||
|
|
"model_name": "Kimi 2.5 (NVIDIA)",
|
||
|
|
"tier": "standard",
|
||
|
|
"context_bucket": "<=128k",
|
||
|
|
"thinking_mode": "off",
|
||
|
|
"cache_mode": "none",
|
||
|
|
"currency": "USD",
|
||
|
|
"input_rate": 0.60, # USD per 1M
|
||
|
|
"output_rate": 3.00,
|
||
|
|
"cached_input_rate": 0.0,
|
||
|
|
"cache_write_rate": 0.0
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"key": "zhipu:glm-5:standard:<=32k:off:none",
|
||
|
|
"provider": "zhipu",
|
||
|
|
"model_id": "glm-5",
|
||
|
|
"model_name": "GLM-5",
|
||
|
|
"tier": "standard",
|
||
|
|
"context_bucket": "<=32k",
|
||
|
|
"thinking_mode": "off",
|
||
|
|
"cache_mode": "none",
|
||
|
|
"currency": "CNY",
|
||
|
|
"input_rate": 4.00,
|
||
|
|
"output_rate": 18.00,
|
||
|
|
"cached_input_rate": 0.0,
|
||
|
|
"cache_write_rate": 0.0
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"key": "zhipu:glm-5:standard:>=32k:off:none",
|
||
|
|
"provider": "zhipu",
|
||
|
|
"model_id": "glm-5",
|
||
|
|
"model_name": "GLM-5 (Long)",
|
||
|
|
"tier": "standard",
|
||
|
|
"context_bucket": ">=32k",
|
||
|
|
"thinking_mode": "off",
|
||
|
|
"cache_mode": "none",
|
||
|
|
"currency": "CNY",
|
||
|
|
"input_rate": 6.00,
|
||
|
|
"output_rate": 22.00,
|
||
|
|
"cached_input_rate": 0.0,
|
||
|
|
"cache_write_rate": 0.0
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"key": "google:gemini-3-flash-preview:standard:<=128k:off:none",
|
||
|
|
"provider": "google",
|
||
|
|
"model_id": "gemini-3-flash-preview",
|
||
|
|
"model_name": "Gemini 3 Flash",
|
||
|
|
"tier": "standard",
|
||
|
|
"context_bucket": "<=128k",
|
||
|
|
"thinking_mode": "off",
|
||
|
|
"cache_mode": "none",
|
||
|
|
"currency": "USD",
|
||
|
|
"input_rate": 0.35,
|
||
|
|
"output_rate": 0.70,
|
||
|
|
"cached_input_rate": 0.0,
|
||
|
|
"cache_write_rate": 0.0
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"key": "openai:gpt-5.2:standard:<=128k:off:none",
|
||
|
|
"provider": "openai",
|
||
|
|
"model_id": "gpt-5.2",
|
||
|
|
"model_name": "GPT-5.2",
|
||
|
|
"tier": "standard",
|
||
|
|
"context_bucket": "<=128k",
|
||
|
|
"thinking_mode": "off",
|
||
|
|
"cache_mode": "none",
|
||
|
|
"currency": "USD",
|
||
|
|
"input_rate": 1.75,
|
||
|
|
"output_rate": 14.00,
|
||
|
|
"cached_input_rate": 0.0,
|
||
|
|
"cache_write_rate": 0.0
|
||
|
|
}
|
||
|
|
]
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
class PriceTableManager:
|
||
|
|
"""
|
||
|
|
Manages pricing table with version control
|
||
|
|
Supports multi-platform model selection and cost calculation
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, run_dir: Path):
|
||
|
|
self.run_dir = Path(run_dir)
|
||
|
|
self.config_path = self.run_dir / "0-config" / "price-table.json"
|
||
|
|
self.cost_report_path = self.run_dir / "logs" / "cost-report.jsonl"
|
||
|
|
self._price_table: Optional[Dict[str, Any]] = None
|
||
|
|
|
||
|
|
def initialize(self, usd_cny_rate: float = 6.90) -> bool:
|
||
|
|
"""Initialize default price table"""
|
||
|
|
table = DEFAULT_PRICE_TABLE.copy()
|
||
|
|
table['usd_cny_rate'] = usd_cny_rate
|
||
|
|
table['updated_at'] = get_timestamp_iso()
|
||
|
|
|
||
|
|
return atomic_write_json(self.config_path, table)
|
||
|
|
|
||
|
|
def load(self) -> Dict[str, Any]:
|
||
|
|
"""Load price table"""
|
||
|
|
if self._price_table is not None:
|
||
|
|
return self._price_table
|
||
|
|
|
||
|
|
if not self.config_path.exists():
|
||
|
|
self.initialize()
|
||
|
|
|
||
|
|
with open(self.config_path, 'r', encoding='utf-8') as f:
|
||
|
|
self._price_table = json.load(f)
|
||
|
|
|
||
|
|
return self._price_table
|
||
|
|
|
||
|
|
def update_price_table(
|
||
|
|
self,
|
||
|
|
new_table: Dict[str, Any],
|
||
|
|
keep_old: bool = True
|
||
|
|
) -> bool:
|
||
|
|
"""
|
||
|
|
Update price table with hot-swap support
|
||
|
|
|
||
|
|
Args:
|
||
|
|
new_table: New price table
|
||
|
|
keep_old: If True, backup old version
|
||
|
|
"""
|
||
|
|
current = self.load()
|
||
|
|
|
||
|
|
# Backup old version
|
||
|
|
if keep_old:
|
||
|
|
backup_name = f"price-table-v{current['version']}.json"
|
||
|
|
backup_path = self.config_path.parent / backup_name
|
||
|
|
atomic_write_json(backup_path, current)
|
||
|
|
|
||
|
|
# Update with new version
|
||
|
|
new_table['previous_version'] = current['version']
|
||
|
|
new_table['updated_at'] = get_timestamp_iso()
|
||
|
|
|
||
|
|
# Log the update
|
||
|
|
self._log_price_update(current['version'], new_table['version'])
|
||
|
|
|
||
|
|
self._price_table = new_table
|
||
|
|
return atomic_write_json(self.config_path, new_table)
|
||
|
|
|
||
|
|
def _log_price_update(self, old_version: str, new_version: str):
|
||
|
|
"""Log price table update to cost-report"""
|
||
|
|
record = {
|
||
|
|
'timestamp': get_timestamp_iso(),
|
||
|
|
'event': 'price_table_update',
|
||
|
|
'old_version': old_version,
|
||
|
|
'new_version': new_version,
|
||
|
|
'event_id': f"CP-UPDATE-{get_timestamp_iso()}"
|
||
|
|
}
|
||
|
|
atomic_append_jsonl(self.cost_report_path, record)
|
||
|
|
|
||
|
|
def find_price_item(
|
||
|
|
self,
|
||
|
|
provider: str,
|
||
|
|
model_id: str,
|
||
|
|
tier: str = "standard",
|
||
|
|
context_bucket: str = "<=128k",
|
||
|
|
thinking_mode: str = "off",
|
||
|
|
cache_mode: str = "none"
|
||
|
|
) -> Optional[Dict[str, Any]]:
|
||
|
|
"""
|
||
|
|
Find pricing item with fallback matching
|
||
|
|
|
||
|
|
Matching order (strict to loose):
|
||
|
|
1. Exact match
|
||
|
|
2. cache_mode=none
|
||
|
|
3. thinking_mode=off
|
||
|
|
4. closest context_bucket
|
||
|
|
"""
|
||
|
|
table = self.load()
|
||
|
|
models = table.get('models', [])
|
||
|
|
|
||
|
|
# Build key components
|
||
|
|
exact_key = f"{provider}:{model_id}:{tier}:{context_bucket}:{thinking_mode}:{cache_mode}"
|
||
|
|
|
||
|
|
# Try exact match
|
||
|
|
for model in models:
|
||
|
|
if model.get('key') == exact_key:
|
||
|
|
return model
|
||
|
|
|
||
|
|
# Try with cache_mode=none
|
||
|
|
if cache_mode != "none":
|
||
|
|
for model in models:
|
||
|
|
if model.get('provider') == provider and \
|
||
|
|
model.get('model_id') == model_id and \
|
||
|
|
model.get('tier') == tier and \
|
||
|
|
model.get('context_bucket') == context_bucket and \
|
||
|
|
model.get('thinking_mode') == thinking_mode and \
|
||
|
|
model.get('cache_mode') == "none":
|
||
|
|
return model
|
||
|
|
|
||
|
|
# Try with thinking_mode=off
|
||
|
|
if thinking_mode != "off":
|
||
|
|
for model in models:
|
||
|
|
if model.get('provider') == provider and \
|
||
|
|
model.get('model_id') == model_id and \
|
||
|
|
model.get('tier') == tier and \
|
||
|
|
model.get('context_bucket') == context_bucket and \
|
||
|
|
model.get('thinking_mode') == "off":
|
||
|
|
return model
|
||
|
|
|
||
|
|
# Try closest context_bucket (use larger one to avoid underestimating)
|
||
|
|
context_buckets = ["<=32k", "<=128k", ">128k"]
|
||
|
|
current_idx = context_buckets.index(context_bucket) if context_bucket in context_buckets else 1
|
||
|
|
|
||
|
|
for idx in range(current_idx, len(context_buckets)):
|
||
|
|
bucket = context_buckets[idx]
|
||
|
|
for model in models:
|
||
|
|
if model.get('provider') == provider and \
|
||
|
|
model.get('model_id') == model_id and \
|
||
|
|
model.get('tier') == tier and \
|
||
|
|
model.get('context_bucket') == bucket:
|
||
|
|
return model
|
||
|
|
|
||
|
|
# No match found - blocking error
|
||
|
|
raise RuntimeError(
|
||
|
|
f"No pricing match for {provider}:{model_id}:{tier}:{context_bucket}. "
|
||
|
|
"Please update price-table.json"
|
||
|
|
)
|
||
|
|
|
||
|
|
def calculate_cost(
|
||
|
|
self,
|
||
|
|
provider: str,
|
||
|
|
model_id: str,
|
||
|
|
prompt_tokens: int,
|
||
|
|
completion_tokens: int,
|
||
|
|
cached_prompt_tokens: int = 0,
|
||
|
|
**kwargs
|
||
|
|
) -> Dict[str, Any]:
|
||
|
|
"""
|
||
|
|
Calculate cost for a model call
|
||
|
|
|
||
|
|
Returns cost in both original currency and RMB
|
||
|
|
"""
|
||
|
|
table = self.load()
|
||
|
|
usd_cny_rate = table.get('usd_cny_rate', 6.90)
|
||
|
|
|
||
|
|
# Find price item
|
||
|
|
price_item = self.find_price_item(provider, model_id, **kwargs)
|
||
|
|
|
||
|
|
currency = price_item.get('currency', 'USD')
|
||
|
|
input_rate = price_item.get('input_rate', 0)
|
||
|
|
output_rate = price_item.get('output_rate', 0)
|
||
|
|
cached_rate = price_item.get('cached_input_rate', 0)
|
||
|
|
|
||
|
|
# Calculate cost
|
||
|
|
prompt_cost = (prompt_tokens / 1_000_000) * input_rate
|
||
|
|
cached_cost = (cached_prompt_tokens / 1_000_000) * cached_rate
|
||
|
|
completion_cost = (completion_tokens / 1_000_000) * output_rate
|
||
|
|
|
||
|
|
total_cost = prompt_cost + cached_cost + completion_cost
|
||
|
|
|
||
|
|
# Convert to RMB if needed
|
||
|
|
if currency == 'USD':
|
||
|
|
cost_rmb = total_cost * usd_cny_rate
|
||
|
|
else:
|
||
|
|
cost_rmb = total_cost
|
||
|
|
|
||
|
|
return {
|
||
|
|
'currency': currency,
|
||
|
|
'cost_original': total_cost,
|
||
|
|
'cost_rmb': cost_rmb,
|
||
|
|
'usd_cny_rate': usd_cny_rate,
|
||
|
|
'price_table_version': table['version'],
|
||
|
|
'breakdown': {
|
||
|
|
'prompt_tokens': prompt_tokens,
|
||
|
|
'prompt_cost': prompt_cost,
|
||
|
|
'cached_tokens': cached_prompt_tokens,
|
||
|
|
'cached_cost': cached_cost,
|
||
|
|
'completion_tokens': completion_tokens,
|
||
|
|
'completion_cost': completion_cost
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
def log_cost(
|
||
|
|
self,
|
||
|
|
event_id: str,
|
||
|
|
phase: str,
|
||
|
|
chapter: Optional[int],
|
||
|
|
event: str,
|
||
|
|
provider: str,
|
||
|
|
model_id: str,
|
||
|
|
prompt_tokens: int,
|
||
|
|
completion_tokens: int,
|
||
|
|
cached_tokens: int = 0,
|
||
|
|
**kwargs
|
||
|
|
) -> Dict[str, Any]:
|
||
|
|
"""
|
||
|
|
Log cost to cost-report.jsonl
|
||
|
|
|
||
|
|
Returns the logged record
|
||
|
|
"""
|
||
|
|
cost_result = self.calculate_cost(
|
||
|
|
provider, model_id,
|
||
|
|
prompt_tokens, completion_tokens, cached_tokens,
|
||
|
|
**kwargs
|
||
|
|
)
|
||
|
|
|
||
|
|
record = {
|
||
|
|
'timestamp': get_timestamp_iso(),
|
||
|
|
'event_id': event_id,
|
||
|
|
'run_id': self.run_dir.name,
|
||
|
|
'phase': phase,
|
||
|
|
'chapter': chapter,
|
||
|
|
'event': event,
|
||
|
|
'provider': provider,
|
||
|
|
'model_id': model_id,
|
||
|
|
'prompt_tokens': prompt_tokens,
|
||
|
|
'completion_tokens': completion_tokens,
|
||
|
|
'cached_prompt_tokens': cached_tokens,
|
||
|
|
'total_tokens': prompt_tokens + completion_tokens,
|
||
|
|
'currency': cost_result['currency'],
|
||
|
|
'cost_original': round(cost_result['cost_original'], 6),
|
||
|
|
'cost_rmb': round(cost_result['cost_rmb'], 6),
|
||
|
|
'usd_cny_rate': cost_result['usd_cny_rate'],
|
||
|
|
'price_table_version': cost_result['price_table_version'],
|
||
|
|
'pricing_source': 'price-table.json'
|
||
|
|
}
|
||
|
|
|
||
|
|
atomic_append_jsonl(self.cost_report_path, record)
|
||
|
|
|
||
|
|
return record
|
||
|
|
|
||
|
|
def get_total_cost(self) -> Dict[str, float]:
|
||
|
|
"""Get total cost for this run"""
|
||
|
|
if not self.cost_report_path.exists():
|
||
|
|
return {'total_rmb': 0.0, 'total_usd': 0.0, 'record_count': 0}
|
||
|
|
|
||
|
|
total_rmb = 0.0
|
||
|
|
total_usd = 0.0
|
||
|
|
record_count = 0
|
||
|
|
|
||
|
|
with open(self.cost_report_path, 'r', encoding='utf-8') as f:
|
||
|
|
for line in f:
|
||
|
|
try:
|
||
|
|
record = json.loads(line.strip())
|
||
|
|
if record.get('currency') == 'USD':
|
||
|
|
total_usd += record.get('cost_original', 0)
|
||
|
|
total_rmb += record.get('cost_rmb', 0)
|
||
|
|
record_count += 1
|
||
|
|
except:
|
||
|
|
pass
|
||
|
|
|
||
|
|
return {
|
||
|
|
'total_rmb': round(total_rmb, 4),
|
||
|
|
'total_usd': round(total_usd, 4),
|
||
|
|
'record_count': record_count
|
||
|
|
}
|
||
|
|
|
||
|
|
def get_model_alias_mapping(self, alias: str) -> Optional[str]:
|
||
|
|
"""Map old aliases to current model_ids"""
|
||
|
|
alias_map = {
|
||
|
|
'qwen3.5': 'qwen3-max',
|
||
|
|
'qwen-3.5': 'qwen3-max',
|
||
|
|
'kimi': 'kimi-k2.5',
|
||
|
|
'gpt4': 'gpt-5.2',
|
||
|
|
'claude': 'claude-sonnet-4.5'
|
||
|
|
}
|
||
|
|
return alias_map.get(alias.lower())
|
||
|
|
|
||
|
|
|
||
|
|
# ============================================================================
|
||
|
|
# Cost Budget Manager
|
||
|
|
# ============================================================================
|
||
|
|
|
||
|
|
class CostBudgetManager:
|
||
|
|
"""
|
||
|
|
Manages cost budget and alerts
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, run_dir: Path, price_manager: PriceTableManager):
|
||
|
|
self.run_dir = Path(run_dir)
|
||
|
|
self.price_manager = price_manager
|
||
|
|
self.budget_file = self.run_dir / "0-config" / "cost-budget.json"
|
||
|
|
|
||
|
|
def set_budget(self, max_rmb: float, warning_threshold: float = 0.8) -> bool:
|
||
|
|
"""Set cost budget"""
|
||
|
|
budget = {
|
||
|
|
'max_rmb': max_rmb,
|
||
|
|
'warning_threshold': warning_threshold,
|
||
|
|
'created_at': get_timestamp_iso()
|
||
|
|
}
|
||
|
|
return atomic_write_json(self.budget_file, budget)
|
||
|
|
|
||
|
|
def check_budget(self) -> Dict[str, Any]:
|
||
|
|
"""Check current spending against budget"""
|
||
|
|
if not self.budget_file.exists():
|
||
|
|
return {'has_budget': False}
|
||
|
|
|
||
|
|
with open(self.budget_file, 'r', encoding='utf-8') as f:
|
||
|
|
budget = json.load(f)
|
||
|
|
|
||
|
|
total = self.price_manager.get_total_cost()
|
||
|
|
spent = total['total_rmb']
|
||
|
|
max_rmb = budget.get('max_rmb', float('inf'))
|
||
|
|
warning_threshold = budget.get('warning_threshold', 0.8)
|
||
|
|
|
||
|
|
ratio = spent / max_rmb if max_rmb > 0 else 0
|
||
|
|
|
||
|
|
return {
|
||
|
|
'has_budget': True,
|
||
|
|
'spent_rmb': spent,
|
||
|
|
'max_rmb': max_rmb,
|
||
|
|
'remaining_rmb': max_rmb - spent,
|
||
|
|
'ratio': ratio,
|
||
|
|
'status': 'exceeded' if ratio >= 1.0 else 'warning' if ratio >= warning_threshold else 'ok'
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
# ============================================================================
|
||
|
|
# Module Test
|
||
|
|
# ============================================================================
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
import tempfile
|
||
|
|
|
||
|
|
print("=== Price Table Manager Test ===\n")
|
||
|
|
|
||
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
||
|
|
run_dir = Path(tmpdir) / "run"
|
||
|
|
run_dir.mkdir()
|
||
|
|
config_dir = run_dir / "0-config"
|
||
|
|
config_dir.mkdir()
|
||
|
|
logs_dir = run_dir / "logs"
|
||
|
|
logs_dir.mkdir()
|
||
|
|
|
||
|
|
# Test initialization
|
||
|
|
mgr = PriceTableManager(run_dir)
|
||
|
|
mgr.initialize(usd_cny_rate=7.0)
|
||
|
|
print("[Test] Price table initialized: PASS")
|
||
|
|
|
||
|
|
# Test load
|
||
|
|
table = mgr.load()
|
||
|
|
print(f"[Test] Loaded version: {table['version']}")
|
||
|
|
print(f"[Test] USD/CNY rate: {table['usd_cny_rate']}")
|
||
|
|
|
||
|
|
# Test find price
|
||
|
|
item = mgr.find_price_item('moonshot', 'kimi-k2.5')
|
||
|
|
print(f"[Test] Found price item: {item['model_name']}")
|
||
|
|
print(f" Input: {item['input_rate']} CNY/1M")
|
||
|
|
print(f" Output: {item['output_rate']} CNY/1M")
|
||
|
|
|
||
|
|
# Test cost calculation
|
||
|
|
cost = mgr.calculate_cost(
|
||
|
|
'moonshot', 'kimi-k2.5',
|
||
|
|
prompt_tokens=1000,
|
||
|
|
completion_tokens=2000
|
||
|
|
)
|
||
|
|
print(f"\n[Test] Cost calculation:")
|
||
|
|
print(f" Currency: {cost['currency']}")
|
||
|
|
print(f" Original: {cost['cost_original']:.6f}")
|
||
|
|
print(f" RMB: {cost['cost_rmb']:.6f}")
|
||
|
|
|
||
|
|
# Test cost logging
|
||
|
|
record = mgr.log_cost(
|
||
|
|
event_id='test-001',
|
||
|
|
phase='6.3',
|
||
|
|
chapter=1,
|
||
|
|
event='draft_generate',
|
||
|
|
provider='moonshot',
|
||
|
|
model_id='kimi-k2.5',
|
||
|
|
prompt_tokens=1000,
|
||
|
|
completion_tokens=2000
|
||
|
|
)
|
||
|
|
print(f"\n[Test] Cost logged: {record['event_id']}")
|
||
|
|
|
||
|
|
# Test total cost
|
||
|
|
total = mgr.get_total_cost()
|
||
|
|
print(f"\n[Test] Total cost: {total['total_rmb']:.4f} RMB")
|
||
|
|
|
||
|
|
# Test budget
|
||
|
|
budget_mgr = CostBudgetManager(run_dir, mgr)
|
||
|
|
budget_mgr.set_budget(max_rmb=100.0)
|
||
|
|
status = budget_mgr.check_budget()
|
||
|
|
print(f"\n[Test] Budget status: {status['status']}")
|
||
|
|
|
||
|
|
print("\n=== All tests completed ===")
|