From 8e1f4ba09b08ecb5f4590afa1106f1df8372437d Mon Sep 17 00:00:00 2001 From: xueZhixin Date: Mon, 16 Mar 2026 17:18:19 +0800 Subject: [PATCH 1/3] added loading mode for assigning keys --- src/aish/i18n/en-US.yaml | 42 +++ src/aish/i18n/zh-CN.yaml | 42 +++ src/aish/shell.py | 20 ++ src/aish/wizard/provider_helpers.py | 43 ++- src/aish/wizard/providers.py | 9 +- src/aish/wizard/setup_wizard.py | 519 ++++++++++++++++++++++++++-- src/aish/wizard/verification.py | 3 +- tests/test_cli.py | 218 ++++++++++++ 8 files changed, 846 insertions(+), 50 deletions(-) diff --git a/src/aish/i18n/en-US.yaml b/src/aish/i18n/en-US.yaml index 3d50591..92d27bc 100644 --- a/src/aish/i18n/en-US.yaml +++ b/src/aish/i18n/en-US.yaml @@ -27,11 +27,43 @@ cli: no_spaces: "API key must not contain spaces." setup: + entry_title: "Initial setup" + entry_header: "Choose how you want to configure your API access." + step_free_key: "Free key setup" step_provider: "Step 1/3: Choose provider" step_provider_endpoint: "Step 1.5/3: Choose endpoint" step_zai_endpoint: "Step 1.5/3: Choose Z.AI endpoint" step_key: "Step 2/3: Enter API key" step_model: "Step 3/3: Choose model" + free_key_header: "Get a free API key using this device fingerprint." + free_key_detecting_location: "Detecting service region..." + free_key_location_detected: "Detected region: {location}" + free_key_location_cn: "Regional" + free_key_location_overseas: "Global" + free_key_registering: "Registering free API key..." + free_key_success: "Free API key acquired" + free_key_missing_key: "Registration succeeded but no API key was returned" + free_key_already_registered: "This device has already registered a free key" + free_key_fingerprint_failed: "Unable to collect device information for fingerprint" + free_key_failed_with_reason: "Free API key registration failed: {reason}" + free_key_provider_label: "Free API" + free_key_privacy_title: "Privacy Notice" + free_key_privacy_notice: | + To provide the free API Key service, we need to collect the following device information: + + • Device fingerprint (a hash generated from hardware info like MAC address, disk serial number, motherboard serial number, etc.) + • Your service region (for optimal server assignment) + + Purpose: + • Device fingerprint is used to prevent abuse, ensuring each device can only obtain one free key + • Service region is used to assign you the nearest server node + + We promise: + • No personal identity information will be collected + • Device fingerprint is an irreversible hash, original hardware info cannot be recovered + • Data is used solely for this service and will not be shared with third parties + free_key_quota_exhausted: "Free quota exhausted" + free_key_quota_exhausted_hint: "Your free API Key quota has been exhausted. Please configure your own API Key to continue." provider_header: "Choose a provider (type to filter)" provider_filter_prompt: "Filter: " provider_filter_hint: "Type to filter, use ↑/↓ to select, Enter to confirm, Esc to cancel." @@ -121,9 +153,16 @@ cli: action_retry_model: "Choose another model" action_retry_api_base: "Re-enter API Base" action_retry_api_key: "Re-enter API Key" + action_retry_free_key: "Retry free key registration" + action_fallback_manual: "Switch to manual setup" + action_use_existing_key: "Use existing key" + action_use_free_key: "Use free API key" + action_manual_setup: "Manual provider setup" action_change_provider: "Change provider" action_continue: "Continue anyway (not recommended)" action_exit: "Exit setup" + action_agree: "Agree and continue" + action_disagree: "Disagree" retry_api_base_prompt: "Enter API Base (type back to return)" retry_api_base_prompt_with_current: "Enter API Base (current: {current}, type back to return)" retry_api_base_required: "API Base is required" @@ -264,6 +303,9 @@ shell: llm_error_details_title: "Debug Details" tool_args_json_invalid_hint: "❌ Request failed, please retry." + hint: + run_setup: "Type /setup to reconfigure your API Key" + prompt: confirm_execute: "?Execute this command? [y/n]:" ai_hint: "Use ';' to ask AI" diff --git a/src/aish/i18n/zh-CN.yaml b/src/aish/i18n/zh-CN.yaml index a19955d..05c9cd9 100644 --- a/src/aish/i18n/zh-CN.yaml +++ b/src/aish/i18n/zh-CN.yaml @@ -27,11 +27,43 @@ cli: no_spaces: "API Key 不能包含空格。" setup: + entry_title: "初始化配置" + entry_header: "请选择 API 接入方式。" + step_free_key: "免费 Key 配置" step_provider: "步骤 1/3：选择 Provider" step_provider_endpoint: "步骤 1.5/3：选择端点" step_zai_endpoint: "步骤 1.5/3：选择 Z.AI 端点" step_key: "步骤 2/3：填写 API Key" step_model: "步骤 3/3：选择模型" + free_key_header: "通过本机设备指纹获取免费 API Key。" + free_key_detecting_location: "正在检测服务区域..." + free_key_location_detected: "检测到区域：{location}" + free_key_location_cn: "国内节点" + free_key_location_overseas: "国际节点" + free_key_registering: "正在注册免费 API Key..." + free_key_success: "已获取免费 API Key" + free_key_missing_key: "注册成功，但服务端未返回 API Key" + free_key_already_registered: "该设备已注册过免费Key" + free_key_fingerprint_failed: "无法采集设备信息，指纹生成失败" + free_key_failed_with_reason: "免费 API Key 注册失败：{reason}" + free_key_provider_label: "免费 API" + free_key_privacy_title: "隐私声明" + free_key_privacy_notice: | + 为了提供免费 API Key 服务，我们需要收集以下设备信息： + + • 设备指纹（基于 MAC 地址、磁盘序列号、主板序列号等硬件信息生成的哈希值） + • 服务区域（用于分配最近的服务节点） + + 用途说明： + • 设备指纹用于防止滥用，确保每个设备只能获取一次免费 Key + • 服务区域用于为您分配最近的服务节点 + + 我们承诺： + • 不会收集任何个人身份信息 + • 设备指纹为不可逆的哈希值，无法还原原始硬件信息 + • 数据仅用于本服务，不会与第三方共享 + free_key_quota_exhausted: "免费额度已用完" + free_key_quota_exhausted_hint: "您的免费 API Key 额度已耗尽，请配置您自己的 API Key 继续使用。" provider_header: "选择 Provider" provider_filter_prompt: "过滤：" provider_filter_hint: "输入关键词实时过滤，↑/↓ 选择，回车确认，Esc 取消" @@ -121,9 +153,16 @@ cli: action_retry_model: "重新选择模型" action_retry_api_base: "重新输入 API Base" action_retry_api_key: "重新输入 API Key" + action_retry_free_key: "重试免费 Key 注册" + action_fallback_manual: "切换到手动配置" + action_use_existing_key: "使用已有的 Key" + action_use_free_key: "使用免费 API Key" + action_manual_setup: "手动配置 Provider" action_change_provider: "更换 Provider" action_continue: "继续保存（不推荐）" action_exit: "退出配置" + action_agree: "同意并继续" + action_disagree: "不同意" retry_api_base_prompt: "请输入 API Base（输入 back 返回）" retry_api_base_prompt_with_current: "请输入 API Base（当前：{current}，输入 back 返回）" retry_api_base_required: "API Base 不能为空" @@ -264,6 +303,9 @@ shell: llm_error_details_title: "调试详情" tool_args_json_invalid_hint: "❌ 请求失败，请重试。" + hint: + run_setup: "输入 /setup 重新配置 API Key" + prompt: confirm_execute: "?是否执行此命令? [y/n]: " ai_hint: "';' Ask AI ..." diff --git a/src/aish/shell.py b/src/aish/shell.py index 4e19f02..74274fc 100644 --- a/src/aish/shell.py +++ b/src/aish/shell.py @@ -3094,6 +3094,13 @@ def handle_error_event(self, event: LLMEvent): error_type = event.data.get("error_type", "general") error_details = event.data.get("error_details") + # Check for quota/rate limit errors + error_lower = str(error_message).lower() + quota_exhausted = any( + keyword in error_lower + for keyword in ["rate limit", "quota", "insufficient", "429", "exceeded"] + ) + if error_type == "streaming_error": self.console.print(f"❌ Streaming Error: {error_message}", style="red") elif error_type == "litellm_error": @@ -3112,6 +3119,19 @@ def handle_error_event(self, event: LLMEvent): border_style="dim", ) ) + # Show quota exhausted hint + if quota_exhausted: + self.console.print( + Panel( + t("cli.setup.free_key_quota_exhausted_hint"), + title=t("cli.setup.free_key_quota_exhausted"), + border_style="yellow", + ) + ) + self.console.print( + f"💡 {t('shell.hint.run_setup')}", + style="dim" + ) else: self.console.print(f"❌ Error: {error_message}", style="red") diff --git a/src/aish/wizard/provider_helpers.py b/src/aish/wizard/provider_helpers.py index 4cd47a5..56cc04b 100644 --- a/src/aish/wizard/provider_helpers.py +++ b/src/aish/wizard/provider_helpers.py @@ -4,19 +4,36 @@ from dataclasses import dataclass -from .constants import (_AI_GATEWAY_DEFAULT_MODEL, _AI_GATEWAY_MODELS, - _HUGGINGFACE_DEFAULT_MODEL, _HUGGINGFACE_MODELS, - _KILOCODE_DEFAULT_MODEL, _KILOCODE_MODELS, - _MINIMAX_DEFAULT_MODELS, _MINIMAX_ENDPOINTS, - _MINIMAX_MODELS, _MISTRAL_DEFAULT_MODEL, - _MISTRAL_MODELS, _MOONSHOT_DEFAULT_MODELS, - _MOONSHOT_ENDPOINTS, _MOONSHOT_MODELS, - _OLLAMA_DEFAULT_MODEL, _OLLAMA_MODELS, _QIANFAN_MODELS, - _QWEN_DEFAULT_MODEL, _QWEN_MODELS, - _TOGETHER_DEFAULT_MODEL, _TOGETHER_MODELS, - _VLLM_DEFAULT_MODEL, _VLLM_MODELS, _XAI_DEFAULT_MODEL, - _XAI_MODELS, _ZAI_DEFAULT_MODELS, _ZAI_ENDPOINTS, - _ZAI_MODELS) +from .constants import ( + _AI_GATEWAY_DEFAULT_MODEL, + _AI_GATEWAY_MODELS, + _HUGGINGFACE_DEFAULT_MODEL, + _HUGGINGFACE_MODELS, + _KILOCODE_DEFAULT_MODEL, + _KILOCODE_MODELS, + _MINIMAX_DEFAULT_MODELS, + _MINIMAX_ENDPOINTS, + _MINIMAX_MODELS, + _MISTRAL_DEFAULT_MODEL, + _MISTRAL_MODELS, + _MOONSHOT_DEFAULT_MODELS, + _MOONSHOT_ENDPOINTS, + _MOONSHOT_MODELS, + _OLLAMA_DEFAULT_MODEL, + _OLLAMA_MODELS, + _QIANFAN_MODELS, + _QWEN_DEFAULT_MODEL, + _QWEN_MODELS, + _TOGETHER_DEFAULT_MODEL, + _TOGETHER_MODELS, + _VLLM_DEFAULT_MODEL, + _VLLM_MODELS, + _XAI_DEFAULT_MODEL, + _XAI_MODELS, + _ZAI_DEFAULT_MODELS, + _ZAI_ENDPOINTS, + _ZAI_MODELS, +) @dataclass diff --git a/src/aish/wizard/providers.py b/src/aish/wizard/providers.py index c4a9198..1bfb8c3 100644 --- a/src/aish/wizard/providers.py +++ b/src/aish/wizard/providers.py @@ -6,8 +6,13 @@ from ..i18n import t from ..litellm_loader import load_litellm -from .constants import (_PROVIDER_ALIASES, _PROVIDER_BASES, _PROVIDER_ENV_KEYS, - _PROVIDER_LABELS, _PROVIDER_PRIORITY) +from .constants import ( + _PROVIDER_ALIASES, + _PROVIDER_BASES, + _PROVIDER_ENV_KEYS, + _PROVIDER_LABELS, + _PROVIDER_PRIORITY, +) from .helpers import _is_valid_url, _looks_like_api_base, _matches_filter_query from .types import ProviderOption diff --git a/src/aish/wizard/setup_wizard.py b/src/aish/wizard/setup_wizard.py index 4291e7c..1fcde66 100644 --- a/src/aish/wizard/setup_wizard.py +++ b/src/aish/wizard/setup_wizard.py @@ -6,8 +6,11 @@ import json import os import shutil +import subprocess import sys -from typing import Optional +from dataclasses import dataclass +from pathlib import Path +from typing import Optional, Tuple from urllib.error import HTTPError, URLError from urllib.request import Request, urlopen @@ -15,32 +18,282 @@ from rich import box from rich.console import Console from rich.panel import Panel -from rich.progress import (Progress, SpinnerColumn, TextColumn, - TimeElapsedColumn) +from rich.progress import Progress, SpinnerColumn, TextColumn, TimeElapsedColumn from rich.table import Table from ..config import Config, ConfigModel from ..i18n import t from ..litellm_loader import load_litellm, preload_litellm from ..providers.openai_codex import is_openai_codex_model -from .constants import (_HUGGINGFACE_DEFAULT_MODEL, - _KILOCODE_DEFAULT_MODEL, _MISTRAL_DEFAULT_MODEL, - _OLLAMA_DEFAULT_MODEL, _QIANFAN_MODELS, - _QWEN_DEFAULT_MODEL, _STATIC_FILTER_SKIP_PROVIDERS, - _TOGETHER_DEFAULT_MODEL, _VLLM_DEFAULT_MODEL, - _XAI_DEFAULT_MODEL) -from .helpers import (_ask_value, _display_width, _is_blank, _is_valid_url, - _mask_secret, _matches_filter_query, - _prompt_secret_with_mask, _sanitize_filter_input) -from .provider_helpers import (ProviderEndpointInfo, get_provider_endpoints, - get_provider_models, has_multi_endpoints) -from .providers import (_filter_provider_options, _get_provider_options, - _maybe_resolve_api_base, _provider_note, - _with_api_base) +from .constants import ( + _HUGGINGFACE_DEFAULT_MODEL, + _KILOCODE_DEFAULT_MODEL, + _MISTRAL_DEFAULT_MODEL, + _OLLAMA_DEFAULT_MODEL, + _QIANFAN_MODELS, + _QWEN_DEFAULT_MODEL, + _STATIC_FILTER_SKIP_PROVIDERS, + _TOGETHER_DEFAULT_MODEL, + _VLLM_DEFAULT_MODEL, + _XAI_DEFAULT_MODEL, +) + +# ============================================================================ +# Free API Key Module - supports both Python package and standalone binary +# ============================================================================ + +# Sentinel for fallback to manual setup +FALLBACK_MANUAL_SETUP = object() + + +@dataclass +class RegisterResult: + """Result of free key registration attempt.""" + success: bool = False + api_key: Optional[str] = None + api_base: Optional[str] = None + model: Optional[str] = None + error_message: Optional[str] = None + already_registered: bool = False + + @classmethod + def from_dict(cls, data: dict) -> "RegisterResult": + return cls( + success=data.get("success", False), + api_key=data.get("api_key"), + api_base=data.get("api_base"), + model=data.get("model"), + error_message=data.get("error_message"), + already_registered=data.get("already_registered", False), + ) + + +def _find_freekey_binary() -> Optional[str]: + """Find the standalone aish_freekey binary. + + Searches in: + 1. PATH environment + 2. ~/.local/bin/ + 3. /usr/local/bin/ + + Returns: + Path to binary if found, None otherwise + """ + binary_name = "aish_freekey_bin" + if sys.platform == "win32": + binary_name += ".exe" + + # 1. Check PATH + binary = shutil.which(binary_name) + if binary: + return binary + + # 2. Check common installation locations + common_paths = [ + Path.home() / ".local" / "bin" / binary_name, + Path("/usr/local/bin") / binary_name, + ] + + for path in common_paths: + if path.exists() and os.access(path, os.X_OK): + return str(path) + + return None + + +def _run_binary(binary_path: str, command: str, *args: str) -> str: + """Run the binary and return stdout.""" + try: + result = subprocess.run( + [binary_path, command, *args], + capture_output=True, + text=True, + timeout=60, + ) + return result.stdout.strip() + except Exception: + return "" + + +def _run_binary_json(binary_path: str, command: str, *args: str) -> dict: + """Run the binary and return JSON result.""" + output = _run_binary(binary_path, command, *args) + if not output: + return {} + try: + return json.loads(output) + except json.JSONDecodeError: + return {} + + +# Module-level: detect if free key functionality is available +_FREEKEY_BINARY_PATH: Optional[str] = None +_HAS_FREEKEY_PYTHON_PACKAGE = False + +# Try Python package first +try: + from aish_freekey import ( + FALLBACK_MANUAL_SETUP as _pkg_fallback, + RegisterResult as _pkg_register_result, + detect_geo_location as _pkg_detect_geo_location, + extract_free_key_info as _pkg_extract_free_key_info, + generate_device_fingerprint as _pkg_generate_device_fingerprint, + get_default_config_for_location as _pkg_get_default_config_for_location, + register_free_key_with_retry as _pkg_register_free_key_with_retry, + request_free_api_key as _pkg_request_free_api_key, + ) + _HAS_FREEKEY_PYTHON_PACKAGE = True +except ImportError: + pass + +# If no Python package, try standalone binary +if not _HAS_FREEKEY_PYTHON_PACKAGE: + _FREEKEY_BINARY_PATH = _find_freekey_binary() + +# Module is available if either Python package or binary is found +HAS_FREE_KEY_MODULE = _HAS_FREEKEY_PYTHON_PACKAGE or (_FREEKEY_BINARY_PATH is not None) + + +# Provide unified interface functions +def generate_device_fingerprint() -> str: + """Generate a SHA256 fingerprint from device hardware information.""" + if _HAS_FREEKEY_PYTHON_PACKAGE: + return _pkg_generate_device_fingerprint() + if _FREEKEY_BINARY_PATH: + return _run_binary(_FREEKEY_BINARY_PATH, "fp") + return "" + + +def detect_geo_location() -> str: + """Detect if user is in China (cn) or overseas.""" + if _HAS_FREEKEY_PYTHON_PACKAGE: + return _pkg_detect_geo_location() + if _FREEKEY_BINARY_PATH: + result = _run_binary_json(_FREEKEY_BINARY_PATH, "loc") + return result.get("location", "cn") + return "cn" + + +def extract_free_key_info( + payload: dict, +) -> Tuple[Optional[str], Optional[str], Optional[str]]: + """Extract API key, API base, and model from the registration response.""" + if _HAS_FREEKEY_PYTHON_PACKAGE: + return _pkg_extract_free_key_info(payload) + + # Inline implementation for binary mode + if not isinstance(payload, dict): + return (None, None, None) + + api_key = payload.get("apikey") or payload.get("api_key") + if not isinstance(api_key, str) or not api_key.strip(): + return (None, None, None) + api_key = api_key.strip() + + api_base = payload.get("api_base") or payload.get("api_base_url") + if isinstance(api_base, str): + api_base = api_base.strip() or None + else: + api_base = None + + model = payload.get("model") + if isinstance(model, str): + model = model.strip() or None + else: + model = None + + return (api_key, api_base, model) + + +def get_default_config_for_location(location: str) -> Tuple[Optional[str], Optional[str]]: + """Get default api_base and model for a location. + + Returns (None, None) when using binary mode - the binary/server will provide + the appropriate values. + """ + if _HAS_FREEKEY_PYTHON_PACKAGE: + return _pkg_get_default_config_for_location(location) + + # Binary mode: server will provide api_base and model + return (None, None) + + +def register_free_key_with_retry( + location: Optional[str] = None, +) -> RegisterResult | object: + """Register a free API key.""" + if _HAS_FREEKEY_PYTHON_PACKAGE: + result = _pkg_register_free_key_with_retry(location) + # Convert package RegisterResult to our RegisterResult + if hasattr(result, 'success'): + return RegisterResult( + success=result.success, + api_key=getattr(result, 'api_key', None), + api_base=getattr(result, 'api_base', None), + model=getattr(result, 'model', None), + error_message=getattr(result, 'error_message', None), + already_registered=getattr(result, 'already_registered', False), + ) + return result + + if _FREEKEY_BINARY_PATH: + result = _run_binary_json(_FREEKEY_BINARY_PATH, "reg") + if not result: + return RegisterResult( + success=False, + error_message="Failed to communicate with registration service", + ) + return RegisterResult.from_dict(result) + + return FALLBACK_MANUAL_SETUP + + +def request_free_api_key( + fingerprint: str, + quota: int = 2000000, + location: str = "cn", +) -> dict: + """Request a free API key from the registration server. + + Note: This function is deprecated. Use register_free_key_with_retry instead. + """ + if _HAS_FREEKEY_PYTHON_PACKAGE: + return _pkg_request_free_api_key(fingerprint, quota, location) + return {"status": "error", "message": "Use register_free_key_with_retry instead"} + + +# End of Free API Key Module +from .helpers import ( + _ask_value, + _display_width, + _is_blank, + _is_valid_url, + _mask_secret, + _matches_filter_query, + _prompt_secret_with_mask, + _sanitize_filter_input, +) +from .provider_helpers import ( + ProviderEndpointInfo, + get_provider_endpoints, + get_provider_models, + has_multi_endpoints, +) +from .providers import ( + _filter_provider_options, + _get_provider_options, + _maybe_resolve_api_base, + _provider_note, + _with_api_base, +) from .types import ProviderOption, ToolSupportResult -from .verification import (_check_tool_support, _quick_static_check, - _status_text, build_failure_reason, - run_verification) +from .verification import ( + _check_tool_support, + _quick_static_check, + _status_text, + build_failure_reason, + run_verification, +) console = Console() @@ -89,8 +342,7 @@ def _select_provider_realtime( from prompt_toolkit.formatted_text import ANSI from prompt_toolkit.key_binding import KeyBindings from prompt_toolkit.layout import HSplit, Layout, VSplit, Window - from prompt_toolkit.layout.controls import (BufferControl, - FormattedTextControl) + from prompt_toolkit.layout.controls import BufferControl, FormattedTextControl from prompt_toolkit.styles import Style except Exception: return _REALTIME_UNAVAILABLE @@ -359,8 +611,7 @@ def _select_endpoint_realtime( from prompt_toolkit.formatted_text import ANSI from prompt_toolkit.key_binding import KeyBindings from prompt_toolkit.layout import HSplit, Layout, VSplit, Window - from prompt_toolkit.layout.controls import (BufferControl, - FormattedTextControl) + from prompt_toolkit.layout.controls import BufferControl, FormattedTextControl from prompt_toolkit.styles import Style except Exception: return _REALTIME_UNAVAILABLE @@ -585,8 +836,7 @@ def _prompt_url_with_inline_validation( from prompt_toolkit.buffer import Buffer from prompt_toolkit.key_binding import KeyBindings from prompt_toolkit.layout import HSplit, Layout, VSplit, Window - from prompt_toolkit.layout.controls import (BufferControl, - FormattedTextControl) + from prompt_toolkit.layout.controls import BufferControl, FormattedTextControl from prompt_toolkit.styles import Style error_text = "" @@ -946,8 +1196,7 @@ def _select_model_realtime( from prompt_toolkit.formatted_text import ANSI from prompt_toolkit.key_binding import KeyBindings from prompt_toolkit.layout import HSplit, Layout, VSplit, Window - from prompt_toolkit.layout.controls import (BufferControl, - FormattedTextControl) + from prompt_toolkit.layout.controls import BufferControl, FormattedTextControl from prompt_toolkit.styles import Style except Exception: return _REALTIME_UNAVAILABLE @@ -1225,6 +1474,123 @@ def _prompt_model(provider: ProviderOption, api_key: Optional[str]) -> Optional[ return normalized +def _prompt_setup_entry_mode() -> str: + """Prompt user to choose between free key setup or manual setup. + + Returns 'free_key', 'manual', or 'exit'. + When aish_freekey module is not available, directly returns 'manual'. + """ + if not HAS_FREE_KEY_MODULE: + return "manual" + + console.print( + Panel( + t("cli.setup.entry_header"), + title=t("cli.setup.entry_title"), + border_style="blue", + ) + ) + return _prompt_setup_action( + [ + ("free_key", t("cli.setup.action_use_free_key")), + ("manual", t("cli.setup.action_manual_setup")), + ("exit", t("cli.setup.action_exit")), + ] + ) + + +def _handle_free_key_registration() -> tuple[str, str, str] | object | None: + """Handle free key registration flow with UI. + + Returns: + tuple[str, str, str]: (api_key, api_base, model) on success + FALLBACK_MANUAL_SETUP: user chose to fallback to manual setup + None: user cancelled + """ + if not HAS_FREE_KEY_MODULE: + return FALLBACK_MANUAL_SETUP + + while True: + console.print( + Panel( + t("cli.setup.free_key_header"), + title=t("cli.setup.step_free_key"), + border_style="blue", + ) + ) + + # Show privacy notice and get user consent + console.print( + Panel( + t("cli.setup.free_key_privacy_notice"), + title=t("cli.setup.free_key_privacy_title"), + border_style="yellow", + ) + ) + + consent = _prompt_setup_action( + [ + ("agree", t("cli.setup.action_agree")), + ("disagree", t("cli.setup.action_disagree")), + ] + ) + + if consent == "disagree": + return None # Return to entry selection + + # Detect geo location + console.print(t("cli.setup.free_key_detecting_location"), style="dim") + location = detect_geo_location() + if location == "cn": + location_display = t("cli.setup.free_key_location_cn") + elif location == "overseas": + location_display = t("cli.setup.free_key_location_overseas") + else: + location_display = location + console.print( + t("cli.setup.free_key_location_detected", location=location_display), + style="dim", + ) + + console.print(t("cli.setup.free_key_registering"), style="dim") + + # Call the private module's registration function + result = register_free_key_with_retry(location=location) + + # Handle fallback sentinel + if result is FALLBACK_MANUAL_SETUP: + return FALLBACK_MANUAL_SETUP + + # Handle RegisterResult + if result.success: + console.print(t("cli.setup.free_key_success"), style="green") + if result.already_registered: + console.print( + t("cli.setup.free_key_already_registered"), + style="yellow", + ) + return (result.api_key, result.api_base, result.model) + + # Handle failure + failure_reason = result.error_message or t("cli.setup.verify_failed_unknown") + console.print( + t("cli.setup.free_key_failed_with_reason", reason=failure_reason), + style="red", + ) + action = _prompt_setup_action( + [ + ("retry", t("cli.setup.action_retry_free_key")), + ("manual", t("cli.setup.action_fallback_manual")), + ("exit", t("cli.setup.action_exit")), + ] + ) + if action == "retry": + continue + if action == "manual": + return FALLBACK_MANUAL_SETUP + return None + + def _prompt_setup_action(options: list[tuple[str, str]]) -> str: """Prompt user to select an action using arrow keys.""" return _select_action_realtime(options) @@ -1355,16 +1721,103 @@ def _interactive_setup(config: Config) -> Optional[ConfigModel]: preload_litellm() while True: - provider = _select_provider() - if provider is None: + provider: Optional[ProviderOption] = None + api_key: Optional[str] = None + + # Prompt user to choose setup mode + setup_mode = _prompt_setup_entry_mode() + if setup_mode == "exit": console.print(t("cli.setup.cancelled"), style="yellow") return None - api_key = _prompt_api_key(provider.env_key) - if api_key is None: + if setup_mode == "free_key": + free_result = _handle_free_key_registration() + if free_result is None: + console.print(t("cli.setup.cancelled"), style="yellow") + return None + if free_result is FALLBACK_MANUAL_SETUP: + setup_mode = "manual" + else: + api_key, api_base, model = free_result + provider = ProviderOption( + key="custom", + label=t("cli.setup.free_key_provider_label"), + api_base=api_base, + env_key=None, + requires_api_base=False, + ) + provider = _maybe_resolve_api_base(provider, api_key=api_key) + + # For free key, use the server-provided model directly + # Skip the model selection step and proceed to verification + if model: + console.print( + t("cli.setup.model_saved_as", model=model), + style="dim", + ) + # Proceed directly to verification with the server-provided model + provider = _maybe_resolve_api_base( + provider, + api_key=api_key, + model_hint=model, + ) + + connectivity, tool_support = run_verification( + model=model, + api_base=provider.api_base, + api_key=api_key, + ) + + if connectivity.ok and tool_support.supports: + console.print(t("cli.setup.verify_simple_success"), style="green") + _persist_setup_config( + config=config, + api_base=provider.api_base, + api_key=api_key, + model=model, + ) + return config.model_config + else: + failure_reason = build_failure_reason( + connectivity, tool_support + ) + console.print( + t( + "cli.setup.verify_simple_failed_with_reason", + reason=failure_reason, + ), + style="red", + ) + action = _prompt_setup_action( + [ + ("retry", t("cli.setup.action_retry_free_key")), + ("manual", t("cli.setup.action_fallback_manual")), + ("exit", t("cli.setup.action_exit")), + ] + ) + if action == "retry": + continue + if action == "manual": + setup_mode = "manual" + else: + console.print(t("cli.setup.cancelled"), style="yellow") + return None + + if setup_mode == "manual": + provider = _select_provider() + if provider is None: + console.print(t("cli.setup.cancelled"), style="yellow") + return None + + api_key = _prompt_api_key(provider.env_key) + if api_key is None: + console.print(t("cli.setup.cancelled"), style="yellow") + return None + provider = _maybe_resolve_api_base(provider, api_key=api_key) + + if provider is None or api_key is None: console.print(t("cli.setup.cancelled"), style="yellow") return None - provider = _maybe_resolve_api_base(provider, api_key=api_key) while True: model = _prompt_model(provider, api_key) diff --git a/src/aish/wizard/verification.py b/src/aish/wizard/verification.py index 3952600..c30da74 100644 --- a/src/aish/wizard/verification.py +++ b/src/aish/wizard/verification.py @@ -14,8 +14,7 @@ import anyio from rich.console import Console -from rich.progress import (Progress, ProgressColumn, SpinnerColumn, Task, - TextColumn) +from rich.progress import Progress, ProgressColumn, SpinnerColumn, Task, TextColumn from rich.text import Text from ..i18n import t diff --git a/tests/test_cli.py b/tests/test_cli.py index 8272a84..1dcee1c 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -4,6 +4,7 @@ from pathlib import Path from unittest.mock import Mock, patch +import pytest from typer.testing import CliRunner from aish.cli import app, run @@ -18,6 +19,15 @@ class _FakeAuthState: auth_path: Path +def _has_free_key_module() -> bool: + """Check if free key functionality is available (binary or Python package).""" + try: + from aish.wizard.setup_wizard import HAS_FREE_KEY_MODULE + return HAS_FREE_KEY_MODULE + except ImportError: + return False + + class TestCLI: """Tests for CLI commands""" @@ -385,3 +395,211 @@ def test_models_auth_login_dispatches_through_provider_contract( login_with_browser.assert_called_once() assert mock_config.config_model.model == "fake-provider/model-x" assert mock_config.config_model.codex_auth_path == "/tmp/fake-auth.json" + + +@pytest.mark.skipif( + not _has_free_key_module(), + reason="Free key module not available - these tests require the binary or Python package", +) +class TestSetupWizardFreeKeyHelpers: + """Tests for free API key registration helper functions. + + Note: These tests require the aish_freekey binary or Python package. + They will be skipped if neither is available. + """ + + def test_extract_free_key_info_from_data_payload(self): + """Test extracting API key and base from a successful response.""" + from aish.wizard.setup_wizard import extract_free_key_info + + payload = { + "status": "success", + "apikey": " test-key ", + "api_base": " https://example.com/v1 ", + } + + api_key, api_base, model = extract_free_key_info(payload) + + assert api_key == "test-key" + assert api_base == "https://example.com/v1" + + def test_extract_free_key_info_from_fixed_payload(self): + """Test extracting API key when only apikey is present.""" + from aish.wizard.setup_wizard import extract_free_key_info + + payload = { + "status": "success", + "apikey": "k-123", + } + + api_key, api_base, model = extract_free_key_info(payload) + + assert api_key == "k-123" + assert api_base is None + + def test_extract_free_key_info_from_non_fixed_payload(self): + """Test that api_key field (different from apikey) is also accepted.""" + from aish.wizard.setup_wizard import extract_free_key_info + + payload = { + "status": "success", + "api_key": "legacy-field", + } + + api_key, api_base, model = extract_free_key_info(payload) + + # The implementation supports both 'apikey' and 'api_key' + assert api_key == "legacy-field" + assert api_base is None + + def test_extract_free_key_info_empty_apikey(self): + """Test that empty apikey returns None.""" + from aish.wizard.setup_wizard import extract_free_key_info + + payload = { + "status": "success", + "apikey": " ", + } + + api_key, api_base, model = extract_free_key_info(payload) + + assert api_key is None + assert api_base is None + + def test_request_free_api_key_returns_stub(self): + """Test request_free_api_key returns stub message (Go binary handles actual requests).""" + import aish.wizard.setup_wizard as setup_module + + # Force binary mode + setup_module._HAS_FREEKEY_PYTHON_PACKAGE = False + setup_module._FREEKEY_BINARY_PATH = "/fake/path" + + try: + result = setup_module.request_free_api_key("fingerprint") + assert result["status"] == "error" + assert "Use register_free_key_with_retry" in result["message"] + finally: + # Restore - try to import package again + try: + from aish_freekey import request_free_api_key as _pkg_func + setup_module._HAS_FREEKEY_PYTHON_PACKAGE = True + except ImportError: + pass + + def test_register_free_key_with_retry_success(self, monkeypatch): + """Test successful free key registration in binary mode.""" + import aish.wizard.setup_wizard as setup_module + + # Force binary mode + setup_module._HAS_FREEKEY_PYTHON_PACKAGE = False + setup_module._FREEKEY_BINARY_PATH = "/fake/path/aish_freekey_bin" + + # Mock the binary JSON response + def mock_run_binary_json(binary_path, cmd, *args): + return { + "success": True, + "api_key": "free-key", + "api_base": "https://free.example.com/v1", + "model": "test-model", + } + + monkeypatch.setattr(setup_module, "_run_binary_json", mock_run_binary_json) + + try: + result = setup_module.register_free_key_with_retry() + assert result.success is True + assert result.api_key == "free-key" + assert result.api_base == "https://free.example.com/v1" + finally: + # Restore + try: + from aish_freekey import register_free_key_with_retry as _pkg_func + setup_module._HAS_FREEKEY_PYTHON_PACKAGE = True + except ImportError: + pass + + def test_register_free_key_with_retry_default_api_base(self, monkeypatch): + """Test free key registration uses default API base when not returned.""" + import aish.wizard.setup_wizard as setup_module + + # Force binary mode + setup_module._HAS_FREEKEY_PYTHON_PACKAGE = False + setup_module._FREEKEY_BINARY_PATH = "/fake/path/aish_freekey_bin" + + # Mock the binary JSON response (no api_base returned) + def mock_run_binary_json(binary_path, cmd, *args): + return { + "success": True, + "api_key": "free-key", + "api_base": "", + "model": "", + } + + monkeypatch.setattr(setup_module, "_run_binary_json", mock_run_binary_json) + + try: + result = setup_module.register_free_key_with_retry(location="cn") + assert result.success is True + assert result.api_key == "free-key" + finally: + # Restore + try: + from aish_freekey import register_free_key_with_retry as _pkg_func + setup_module._HAS_FREEKEY_PYTHON_PACKAGE = True + except ImportError: + pass + + def test_register_free_key_with_retry_failure(self, monkeypatch): + """Test registration failure in binary mode.""" + import aish.wizard.setup_wizard as setup_module + + # Force binary mode + setup_module._HAS_FREEKEY_PYTHON_PACKAGE = False + setup_module._FREEKEY_BINARY_PATH = "/fake/path/aish_freekey_bin" + + # Mock the binary JSON response for failure + def mock_run_binary_json(binary_path, cmd, *args): + return { + "success": False, + "error_message": "Registration failed", + } + + monkeypatch.setattr(setup_module, "_run_binary_json", mock_run_binary_json) + + try: + result = setup_module.register_free_key_with_retry() + assert result.success is False + assert result.error_message == "Registration failed" + finally: + # Restore + try: + from aish_freekey import register_free_key_with_retry as _pkg_func + setup_module._HAS_FREEKEY_PYTHON_PACKAGE = True + except ImportError: + pass + + def test_register_free_key_with_retry_empty_response(self, monkeypatch): + """Test registration with empty response in binary mode.""" + import aish.wizard.setup_wizard as setup_module + + # Force binary mode + setup_module._HAS_FREEKEY_PYTHON_PACKAGE = False + setup_module._FREEKEY_BINARY_PATH = "/fake/path/aish_freekey_bin" + + # Mock empty response from binary + def mock_run_binary_json(binary_path, cmd, *args): + return {} + + monkeypatch.setattr(setup_module, "_run_binary_json", mock_run_binary_json) + + try: + result = setup_module.register_free_key_with_retry() + assert result.success is False + assert "Failed to communicate" in result.error_message + finally: + # Restore + try: + from aish_freekey import register_free_key_with_retry as _pkg_func + setup_module._HAS_FREEKEY_PYTHON_PACKAGE = True + except ImportError: + pass From 53006b288c3bb44769c1b7f47b6d203c11c2ffa0 Mon Sep 17 00:00:00 2001 From: xueZhixin Date: Mon, 16 Mar 2026 17:41:57 +0800 Subject: [PATCH 2/3] delet _pkg_fallback and _pkg_register_result --- src/aish/wizard/setup_wizard.py | 13 ++++++------- tests/test_cli.py | 10 +++++----- 2 files changed, 11 insertions(+), 12 deletions(-) diff --git a/src/aish/wizard/setup_wizard.py b/src/aish/wizard/setup_wizard.py index 1fcde66..0a2c2cd 100644 --- a/src/aish/wizard/setup_wizard.py +++ b/src/aish/wizard/setup_wizard.py @@ -133,8 +133,6 @@ def _run_binary_json(binary_path: str, command: str, *args: str) -> dict: # Try Python package first try: from aish_freekey import ( - FALLBACK_MANUAL_SETUP as _pkg_fallback, - RegisterResult as _pkg_register_result, detect_geo_location as _pkg_detect_geo_location, extract_free_key_info as _pkg_extract_free_key_info, generate_device_fingerprint as _pkg_generate_device_fingerprint, @@ -263,7 +261,7 @@ def request_free_api_key( # End of Free API Key Module -from .helpers import ( +from .helpers import ( # noqa: E402 _ask_value, _display_width, _is_blank, @@ -273,21 +271,21 @@ def request_free_api_key( _prompt_secret_with_mask, _sanitize_filter_input, ) -from .provider_helpers import ( +from .provider_helpers import ( # noqa: E402 ProviderEndpointInfo, get_provider_endpoints, get_provider_models, has_multi_endpoints, ) -from .providers import ( +from .providers import ( # noqa: E402 _filter_provider_options, _get_provider_options, _maybe_resolve_api_base, _provider_note, _with_api_base, ) -from .types import ProviderOption, ToolSupportResult -from .verification import ( +from .types import ProviderOption, ToolSupportResult # noqa: E402 +from .verification import ( # noqa: E402 _check_tool_support, _quick_static_check, _status_text, @@ -1738,6 +1736,7 @@ def _interactive_setup(config: Config) -> Optional[ConfigModel]: if free_result is FALLBACK_MANUAL_SETUP: setup_mode = "manual" else: + assert isinstance(free_result, tuple) api_key, api_base, model = free_result provider = ProviderOption( key="custom", diff --git a/tests/test_cli.py b/tests/test_cli.py index 1dcee1c..19d0a45 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -481,7 +481,7 @@ def test_request_free_api_key_returns_stub(self): finally: # Restore - try to import package again try: - from aish_freekey import request_free_api_key as _pkg_func + from aish_freekey import request_free_api_key as _pkg_func # noqa: F401 setup_module._HAS_FREEKEY_PYTHON_PACKAGE = True except ImportError: pass @@ -513,7 +513,7 @@ def mock_run_binary_json(binary_path, cmd, *args): finally: # Restore try: - from aish_freekey import register_free_key_with_retry as _pkg_func + from aish_freekey import register_free_key_with_retry as _pkg_func # noqa: F401 setup_module._HAS_FREEKEY_PYTHON_PACKAGE = True except ImportError: pass @@ -544,7 +544,7 @@ def mock_run_binary_json(binary_path, cmd, *args): finally: # Restore try: - from aish_freekey import register_free_key_with_retry as _pkg_func + from aish_freekey import register_free_key_with_retry as _pkg_func # noqa: F401 setup_module._HAS_FREEKEY_PYTHON_PACKAGE = True except ImportError: pass @@ -573,7 +573,7 @@ def mock_run_binary_json(binary_path, cmd, *args): finally: # Restore try: - from aish_freekey import register_free_key_with_retry as _pkg_func + from aish_freekey import register_free_key_with_retry as _pkg_func # noqa: F401 setup_module._HAS_FREEKEY_PYTHON_PACKAGE = True except ImportError: pass @@ -599,7 +599,7 @@ def mock_run_binary_json(binary_path, cmd, *args): finally: # Restore try: - from aish_freekey import register_free_key_with_retry as _pkg_func + from aish_freekey import register_free_key_with_retry as _pkg_func # noqa: F401 setup_module._HAS_FREEKEY_PYTHON_PACKAGE = True except ImportError: pass From 6ac1dbc050d9b5d5ee429b759fb8fad32667e3fa Mon Sep 17 00:00:00 2001 From: xueZhixin Date: Wed, 18 Mar 2026 11:16:52 +0800 Subject: [PATCH 3/3] fix registration prompt issue --- src/aish/config.py | 14 ++++++++++ src/aish/i18n/en-US.yaml | 2 ++ src/aish/i18n/zh-CN.yaml | 2 ++ src/aish/shell.py | 48 ++++++++++++++++++++++++++++----- src/aish/wizard/setup_wizard.py | 8 ++++-- 5 files changed, 66 insertions(+), 8 deletions(-) diff --git a/src/aish/config.py b/src/aish/config.py index 714dc18..7a82825 100644 --- a/src/aish/config.py +++ b/src/aish/config.py @@ -253,6 +253,11 @@ class ConfigModel(BaseModel): ), ) + is_free_key: bool = Field( + default=False, + description="Whether the current configuration uses a free API key", + ) + @field_validator("tool_arg_preview", mode="before") @classmethod def normalize_tool_arg_preview(cls, v: Any) -> dict[str, ToolArgPreviewSettings]: @@ -567,6 +572,15 @@ def set_api_key(self, api_key: Optional[str]) -> None: self.config_model.api_key = api_key self.save_config() + def is_free_key(self) -> bool: + """Check if using free API key""" + return self.config_model.is_free_key + + def set_is_free_key(self, is_free_key: bool) -> None: + """Set whether using free API key""" + self.config_model.is_free_key = is_free_key + self.save_config() + @property def model_config(self) -> ConfigModel: """Get the underlying Pydantic model""" diff --git a/src/aish/i18n/en-US.yaml b/src/aish/i18n/en-US.yaml index 92d27bc..207c485 100644 --- a/src/aish/i18n/en-US.yaml +++ b/src/aish/i18n/en-US.yaml @@ -64,6 +64,8 @@ cli: • Data is used solely for this service and will not be shared with third parties free_key_quota_exhausted: "Free quota exhausted" free_key_quota_exhausted_hint: "Your free API Key quota has been exhausted. Please configure your own API Key to continue." + free_key_service_unavailable: "Service temporarily unavailable" + free_key_service_unavailable_hint: "Server resources have reached the limit or the service is temporarily unavailable. Please contact the administrator to increase the quota, or configure your own API Key to continue." provider_header: "Choose a provider (type to filter)" provider_filter_prompt: "Filter: " provider_filter_hint: "Type to filter, use ↑/↓ to select, Enter to confirm, Esc to cancel." diff --git a/src/aish/i18n/zh-CN.yaml b/src/aish/i18n/zh-CN.yaml index 05c9cd9..b0e0d54 100644 --- a/src/aish/i18n/zh-CN.yaml +++ b/src/aish/i18n/zh-CN.yaml @@ -64,6 +64,8 @@ cli: • 数据仅用于本服务，不会与第三方共享 free_key_quota_exhausted: "免费额度已用完" free_key_quota_exhausted_hint: "您的免费 API Key 额度已耗尽，请配置您自己的 API Key 继续使用。" + free_key_service_unavailable: "服务暂时不可用" + free_key_service_unavailable_hint: "服务器资源已达上限或服务暂时不可用，请联系管理员增加配额，或配置您自己的 API Key 继续使用。" provider_header: "选择 Provider" provider_filter_prompt: "过滤：" provider_filter_hint: "输入关键词实时过滤，↑/↓ 选择，回车确认，Esc 取消" diff --git a/src/aish/shell.py b/src/aish/shell.py index 74274fc..a800eda 100644 --- a/src/aish/shell.py +++ b/src/aish/shell.py @@ -1440,6 +1440,7 @@ async def handle_setup_command(self, user_input: str) -> None: new_model = new_config.model new_api_base = new_config.api_base new_api_key = new_config.api_key + new_is_free_key = new_config.is_free_key self.llm_session.update_model( new_model, @@ -1450,6 +1451,7 @@ async def handle_setup_command(self, user_input: str) -> None: self.config.model = new_model self.config.api_base = new_api_base self.config.api_key = new_api_key + self.config.is_free_key = new_is_free_key if self.session_record is not None: self.session_record.model = new_model @@ -3094,12 +3096,33 @@ def handle_error_event(self, event: LLMEvent): error_type = event.data.get("error_type", "general") error_details = event.data.get("error_details") - # Check for quota/rate limit errors - error_lower = str(error_message).lower() - quota_exhausted = any( - keyword in error_lower - for keyword in ["rate limit", "quota", "insufficient", "429", "exceeded"] - ) + # Check for quota/rate limit errors (only for free key mode) + # self.config is ConfigModel, use .is_free_key attribute + is_free_key = getattr(self.config, "is_free_key", False) + quota_exhausted = False + service_unavailable = False + if is_free_key: + error_lower = str(error_message).lower() + # Check for quota exhaustion (rate limit, quota exceeded) + quota_exhausted = any( + keyword in error_lower + for keyword in [ + "rate limit", "quota", "insufficient", "429", "exceeded", + "too many requests" + ] + ) + # Check for service unavailability (auth errors, not found, server errors) + # These may indicate the free key is invalid or server resources exhausted + if not quota_exhausted: + service_unavailable = any( + keyword in error_lower + for keyword in [ + "401", "403", "404", "500", "502", "503", "504", + "authentication", "unauthorized", "forbidden", + "not found", "internal server error", "bad gateway", + "service unavailable", "gateway timeout" + ] + ) if error_type == "streaming_error": self.console.print(f"❌ Streaming Error: {error_message}", style="red") @@ -3132,6 +3155,19 @@ def handle_error_event(self, event: LLMEvent): f"💡 {t('shell.hint.run_setup')}", style="dim" ) + # Show service unavailable hint (for free key) + elif service_unavailable: + self.console.print( + Panel( + t("cli.setup.free_key_service_unavailable_hint"), + title=t("cli.setup.free_key_service_unavailable"), + border_style="yellow", + ) + ) + self.console.print( + f"💡 {t('shell.hint.run_setup')}", + style="dim" + ) else: self.console.print(f"❌ Error: {error_message}", style="red") diff --git a/src/aish/wizard/setup_wizard.py b/src/aish/wizard/setup_wizard.py index 0a2c2cd..5fc50a8 100644 --- a/src/aish/wizard/setup_wizard.py +++ b/src/aish/wizard/setup_wizard.py @@ -1503,7 +1503,8 @@ def _handle_free_key_registration() -> tuple[str, str, str] | object | None: Returns: tuple[str, str, str]: (api_key, api_base, model) on success FALLBACK_MANUAL_SETUP: user chose to fallback to manual setup - None: user cancelled + or disagreed with privacy notice + None: user cancelled during registration """ if not HAS_FREE_KEY_MODULE: return FALLBACK_MANUAL_SETUP @@ -1534,7 +1535,7 @@ def _handle_free_key_registration() -> tuple[str, str, str] | object | None: ) if consent == "disagree": - return None # Return to entry selection + return FALLBACK_MANUAL_SETUP # Fallback to manual setup # Detect geo location console.print(t("cli.setup.free_key_detecting_location"), style="dim") @@ -1702,10 +1703,12 @@ def _persist_setup_config( api_base: Optional[str], api_key: str, model: str, + is_free_key: bool = False, ) -> None: config.set_api_base(api_base) config.set_api_key(api_key) config.set_model(model) + config.set_is_free_key(is_free_key) def _interactive_setup(config: Config) -> Optional[ConfigModel]: @@ -1774,6 +1777,7 @@ def _interactive_setup(config: Config) -> Optional[ConfigModel]: api_base=provider.api_base, api_key=api_key, model=model, + is_free_key=True, ) return config.model_config else: