Coverage for src/kimi_code_usage/mcp.py: 86%
122 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-14 14:57 +0800
1import os
2from fastmcp import FastMCP
3from kimi_code_usage.config import ConfigResolver
4from kimi_code_usage.providers import ProviderUsage, dispatch_all
6# Initialize FastMCP for Kimi Code Usage
7mcp = FastMCP("Kimi Code Usage")
9L = {
10 "error_key": "No API keys configured. Please set environment variables (e.g. KIMI_API_KEY) or configure ~/.kimi-usage/config.json",
11 "no_data": "No usage data found.",
12}
14def _format_tokens(n: float) -> str:
15 if n >= 1_000_000:
16 return f"{n/1_000_000:.1f}M"
17 if n >= 1_000:
18 return f"{n/1_000:.1f}K"
19 return f"{n:.0f}"
22def _model_short_name(model: str) -> str:
23 if "/" in model:
24 return model.rsplit("/", 1)[-1]
25 return model
28def _format_activity_lines(item: ProviderUsage, provider_title: str, metric: str = "requests") -> list[str]:
29 lines = []
30 if item.activity_totals:
31 t = item.activity_totals
32 parts = [f"{t.requests:,.0f} requests"]
33 tok = f"In: {_format_tokens(t.prompt_tokens)} / Out: {_format_tokens(t.completion_tokens)}"
34 if t.reasoning_tokens:
35 tok += f" (+ {_format_tokens(t.reasoning_tokens)} reason)"
36 parts.append(tok)
37 parts.append(f"Spend: ${t.spend:.2f}")
38 lines.append(f"{provider_title} - Activity: {' | '.join(parts)}")
40 if item.daily_activity:
41 metric_label = metric.capitalize()
42 lines.append(f"{provider_title} - Daily {metric_label}:")
43 if metric == "requests":
44 max_val = max(sum(m.requests for m in d.models) for d in item.daily_activity) or 1.0
45 for day in item.daily_activity[-14:]:
46 bar_len = int(sum(m.requests for m in day.models) / max_val * 20) if max_val > 0 else 0
47 date = day.date[5:10] if len(day.date) >= 10 else day.date
48 lines.append(f" {date} {'█' * bar_len}{'·' * (20 - bar_len)} {sum(m.requests for m in day.models):,.0f}")
49 elif metric == "tokens":
50 max_val = max(sum(m.prompt_tokens + m.completion_tokens + m.reasoning_tokens for m in d.models) for d in item.daily_activity) or 1.0
51 for day in item.daily_activity[-14:]:
52 val = sum(m.prompt_tokens + m.completion_tokens + m.reasoning_tokens for m in day.models)
53 bar_len = int(val / max_val * 20) if max_val > 0 else 0
54 date = day.date[5:10] if len(day.date) >= 10 else day.date
55 lines.append(f" {date} {'█' * bar_len}{'·' * (20 - bar_len)} {_format_tokens(val)}")
56 else:
57 max_val = max(d.total for d in item.daily_activity) or 1.0
58 for day in item.daily_activity[-14:]:
59 bar_len = int(day.total / max_val * 20) if max_val > 0 else 0
60 date = day.date[5:10] if len(day.date) >= 10 else day.date
61 lines.append(f" {date} {'█' * bar_len}{'·' * (20 - bar_len)} ${day.total:.2f}")
63 if item.top_models:
64 lines.append(f"{provider_title} - Top Models:")
65 for m in item.top_models[:7]:
66 name = _model_short_name(m.model)
67 if metric == "requests":
68 val_str = f"{m.requests:,.0f} req"
69 elif metric == "tokens":
70 val_str = _format_tokens(m.prompt_tokens + m.completion_tokens + m.reasoning_tokens)
71 else:
72 val_str = f"${m.spend:.2f}"
73 lines.append(f" {name}: {val_str}")
75 return lines
78@mcp.tool()
79async def get_usage(provider: str | None = None) -> str:
80 """
81 Get API usage/quota from configured LLM providers.
82 获取已配置的 LLM 服务商的 API 使用量与配额限制。
84 Args:
85 provider: Optional filter - "kimi", "openai", "anthropic", "openrouter", or comma-separated list.
86 """
87 resolver = ConfigResolver()
88 config = resolver.resolve()
90 allowed = None
91 if provider:
92 allowed = [p.strip().lower() for p in provider.split(",") if p.strip()]
93 for p in list(config.providers.keys()):
94 if p not in allowed:
95 config.providers[p].api_key = None
97 if not config.enabled_providers:
98 return L["error_key"]
100 try:
101 results, errors = await dispatch_all(config)
103 if allowed:
104 results = {k: v for k, v in results.items() if k in allowed}
105 errors = {k: v for k, v in errors.items() if k in allowed}
107 if not results and not errors:
108 return L["no_data"]
110 lines = []
111 order = config.provider_order
113 for p in order:
114 p_items = results.get(p)
115 if not p_items:
116 continue
117 for item in p_items:
118 provider_title = "Kimi" if p == "kimi" else p.capitalize()
120 # Structured OpenRouter activity data
121 if item.activity_totals or item.daily_activity or item.top_models:
122 lines.extend(_format_activity_lines(item, provider_title))
123 continue
125 if item.unit == "text":
126 if item.text_value:
127 line = f"{provider_title} - {item.label}: {item.text_value}"
128 else:
129 line = f"{provider_title} - {item.label}"
130 else:
131 line = f"{provider_title} - {item.label}: "
133 if item.limit is not None and item.limit > 0:
134 pct_rem = max(0.0, 100.0 - (item.used / item.limit * 100.0))
135 if item.unit == "$":
136 line += f"${item.used:.2f}/${item.limit:.2f} used ({pct_rem:.0f}% remaining)"
137 else:
138 line += f"{item.used:,.0f}/{item.limit:,.0f} used ({pct_rem:.0f}% remaining)"
139 else:
140 if item.unit == "$":
141 line += f"${item.used:.2f} used"
142 elif item.unit == "text":
143 # For text unit, label already printed above
144 pass
145 else:
146 line += f"{item.used:,.0f} used ({item.unit})"
148 if item.countdown and item.reset_at:
149 line += f" | Reset in {item.countdown} (at {item.reset_at})"
150 lines.append(line)
152 for p in order:
153 err = errors.get(p)
154 if err:
155 provider_title = "Kimi" if p == "kimi" else p.capitalize()
156 lines.append(f"{provider_title} - Error: {err}")
158 return "\n".join(lines)
160 except Exception as e:
161 return f"Error fetching usage: {str(e)}"
163@mcp.tool()
164async def get_kimi_usage() -> str:
165 """
166 Get the current Kimi Coding Plan API usage and quota limits.
167 获取当前 Kimi Coding Plan API 的使用量和配额限制。
168 """
169 return await get_usage("kimi")
171def run_mcp():
172 mcp.run()
174if __name__ == "__main__":
175 run_mcp()