From 375af054b7164c3de107f8594fae4f8c77c98078 Mon Sep 17 00:00:00 2001 From: Yijia-Xiao Date: Wed, 16 Sep 2026 21:45:08 +0000 Subject: [PATCH] feat(cli): offer the previous run's selections as defaults (#1236, #920) - analysts, depth, provider, models and language prefill; prompts still shown - values no longer offered by the current catalog are dropped - environment variables keep skipping their step --- cli/prefs.py | 91 +++++++++++++++++++++++++ cli/utils.py | 85 +++++++++++++++--------- tests/test_cli_prefs.py | 143 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 288 insertions(+), 31 deletions(-) create mode 100644 cli/prefs.py create mode 100644 tests/test_cli_prefs.py diff --git a/cli/prefs.py b/cli/prefs.py new file mode 100644 index 000000000..711756f85 --- /dev/null +++ b/cli/prefs.py @@ -0,0 +1,91 @@ +"""What the last run chose, offered back as the next run's defaults. + +The interactive flow asks the same questions every time, and only some of them +have an environment variable to skip them (the analyst set has none). Remembered +answers prefill the prompts so Enter accepts them; they never skip a step, so a +run always starts on choices the user has seen. + +Only answers that are stable between runs are kept. The ticker and the analysis +date are not: they change every run, and a remembered date would quietly offer a +stale one. + +Every value is checked against the current choices on the way out, because +models and providers are added and retired between versions. A remembered model +that is no longer offered is dropped rather than shown. +""" + +from __future__ import annotations + +import json +import os +from pathlib import Path + +from cli.models import AnalystType, AssetType +from cli.utils import _llm_provider_table, filter_analysts_for_asset_type +from tradingagents.llm_clients.model_catalog import get_model_options + +_PREFS_PATH = Path(os.path.expanduser("~")) / ".tradingagents" / "cli_prefs.json" + +REMEMBERED = ( + "output_language", "analysts", "research_depth", "llm_provider", + "quick_think_llm", "deep_think_llm", "backend_url", +) + + +def load_last_run() -> dict: + """The previous run's answers, or an empty dict when there is nothing usable. + + Convenience state: an unreadable or corrupt file means no defaults, never an + error in the user's way. + """ + try: + data = json.loads(_PREFS_PATH.read_text(encoding="utf-8")) + return data if isinstance(data, dict) else {} + except (OSError, ValueError): + return {} + + +def save_last_run(selections: dict) -> None: + """Record the answers worth offering next time; failure is never fatal.""" + kept = {k: v for k, v in selections.items() if k in REMEMBERED and v not in (None, "", [])} + kept["analysts"] = [getattr(a, "value", a) for a in kept.get("analysts", [])] or None + kept = {k: v for k, v in kept.items() if v is not None} + try: + _PREFS_PATH.parent.mkdir(parents=True, exist_ok=True) + temp = _PREFS_PATH.with_suffix(".tmp") + temp.write_text(json.dumps(kept, indent=2), encoding="utf-8") + os.replace(temp, _PREFS_PATH) # a concurrent run reads one file or the other + except OSError: + return + + +def sanitize(prefs: dict, asset_type) -> dict: + """Keep only the remembered answers that are still choosable now.""" + kept: dict = {} + if isinstance(prefs.get("output_language"), str): + kept["output_language"] = prefs["output_language"] + if prefs.get("research_depth") in (1, 3, 5): + kept["research_depth"] = prefs["research_depth"] + + known = {a.value for a in AnalystType} + analysts = [a for a in prefs.get("analysts") or [] if a in known] + allowed = filter_analysts_for_asset_type([AnalystType(a) for a in analysts], AssetType(asset_type)) + if allowed: + kept["analysts"] = [a.value for a in allowed] + + provider = prefs.get("llm_provider") + # Region-specific providers (qwen-cn) are picked in a second prompt, so the + # base key is what the provider menu matches. + base = (provider or "").split("-cn")[0] + if base and base in {key for _, key, _ in _llm_provider_table()}: + kept["llm_provider"] = provider + if isinstance(prefs.get("backend_url"), str) and prefs["backend_url"]: + kept["backend_url"] = prefs["backend_url"] + for field, mode in (("quick_think_llm", "quick"), ("deep_think_llm", "deep")): + try: + offered = {model for _, model in get_model_options(base, mode)} + except KeyError: + continue + if prefs.get(field) in offered: + kept[field] = prefs[field] + return kept diff --git a/cli/utils.py b/cli/utils.py index 987ffa496..534f9afa8 100644 --- a/cli/utils.py +++ b/cli/utils.py @@ -132,8 +132,16 @@ def get_analysis_date() -> str: return date.strip() -def select_analysts(asset_type: AssetType = AssetType.STOCK) -> list[AnalystType]: - """Select analysts using an interactive checkbox.""" +def _matching_choice(options, default): + """The option value equal to ``default``, or None to leave the menu as is.""" + return next((value for _, value in options if value == default), None) + + +def select_analysts(asset_type: AssetType = AssetType.STOCK, default=None) -> list[AnalystType]: + """Select analysts using an interactive checkbox. + + ``default`` pre-checks the previous run's analysts; the prompt still shows. + """ available_analysts = filter_analysts_for_asset_type( [value for _, value in ANALYST_ORDER], asset_type, @@ -141,7 +149,7 @@ def select_analysts(asset_type: AssetType = AssetType.STOCK) -> list[AnalystType choices = questionary.checkbox( "Select Your [Analysts Team]:", choices=[ - questionary.Choice(display, value=value) + questionary.Choice(display, value=value, checked=value.value in (default or [])) for display, value in ANALYST_ORDER if value in available_analysts ], @@ -164,7 +172,7 @@ def select_analysts(asset_type: AssetType = AssetType.STOCK) -> list[AnalystType return choices -def select_research_depth() -> int: +def select_research_depth(default=None) -> int: """Select research depth using an interactive selection.""" # Define research depth options with their corresponding values @@ -179,6 +187,7 @@ def select_research_depth() -> int: choices=[ questionary.Choice(display, value=value) for display, value in DEPTH_OPTIONS ], + default=_matching_choice(DEPTH_OPTIONS, default), instruction="\n- Use arrow keys to navigate\n- Press Enter to select", style=questionary.Style( [ @@ -289,7 +298,7 @@ def _prompt_custom_model_id() -> str: return _require_text("Enter model ID:", "Please enter a model ID.") -def _select_model(provider: str, mode: str) -> str: +def _select_model(provider: str, mode: str, default=None) -> str: """Select a model for the given provider and mode (quick/deep).""" if provider.lower() == "openrouter": return select_openrouter_model(mode) @@ -306,6 +315,7 @@ def _select_model(provider: str, mode: str) -> str: questionary.Choice(display, value=value) for display, value in get_model_options(provider, mode) ], + default=_matching_choice(get_model_options(provider, mode), default), instruction="\n- Use arrow keys to navigate\n- Press Enter to select", style=questionary.Style( [ @@ -326,14 +336,14 @@ def _select_model(provider: str, mode: str) -> str: return choice -def select_shallow_thinking_agent(provider) -> str: +def select_shallow_thinking_agent(provider, default=None) -> str: """Select shallow thinking llm engine using an interactive selection.""" - return _select_model(provider, "quick") + return _select_model(provider, "quick", default) -def select_deep_thinking_agent(provider) -> str: +def select_deep_thinking_agent(provider, default=None) -> str: """Select deep thinking llm engine using an interactive selection.""" - return _select_model(provider, "deep") + return _select_model(provider, "deep", default) def _llm_provider_table() -> list[tuple[str, str, str | None]]: """(display_name, provider_key, base_url) for every supported provider. @@ -402,9 +412,15 @@ def prompt_openai_compatible_url() -> str: return url.strip() -def select_llm_provider() -> tuple[str, str | None]: +def select_llm_provider(default=None) -> tuple[str, str | None]: """Select the LLM provider and its API endpoint.""" PROVIDERS = _llm_provider_table() + # A region-specific key (qwen-cn) is chosen in a later prompt; the menu + # lists the base provider. + base = (default or "").split("-cn")[0] + preselected = next( + ((key, url) for _, key, url in PROVIDERS if key == base), None + ) choice = questionary.select( "Select your LLM Provider:", @@ -412,6 +428,7 @@ def select_llm_provider() -> tuple[str, str | None]: questionary.Choice(display, value=(provider_key, url)) for display, provider_key, url in PROVIDERS ], + default=preselected, instruction="\n- Use arrow keys to navigate\n- Press Enter to select", style=questionary.Style( [ @@ -458,9 +475,9 @@ def ask_anthropic_effort() -> str | None: return questionary.select( "Select Effort Level:", choices=[ - questionary.Choice("High (recommended)", "high"), - questionary.Choice("Medium (balanced)", "medium"), - questionary.Choice("Low (faster, cheaper)", "low"), + questionary.Choice("High (recommended)", "high"), + questionary.Choice("Medium (balanced)", "medium"), + questionary.Choice("Low (faster, cheaper)", "low"), ], style=questionary.Style([ ("selected", "fg:cyan noinherit"), @@ -479,8 +496,8 @@ def ask_gemini_thinking_config() -> str | None: return questionary.select( "Select Thinking Mode:", choices=[ - questionary.Choice("Enable Thinking (recommended)", "high"), - questionary.Choice("Minimal/Disable Thinking", "minimal"), + questionary.Choice("Enable Thinking (recommended)", "high"), + questionary.Choice("Minimal/Disable Thinking", "minimal"), ], style=questionary.Style([ ("selected", "fg:green noinherit"), @@ -654,24 +671,30 @@ def ensure_api_key(provider: str) -> str | None: return key -def ask_output_language() -> str: - """Ask for report output language.""" +def ask_output_language(default=None) -> str: + """Ask for report output language. + + ``default`` is offered only when it is one of the listed languages: a custom + one entered last time is free text, which the menu cannot preselect. + """ + choices = [ + questionary.Choice("English (default)", "English"), + questionary.Choice("Chinese (中文)", "Chinese"), + questionary.Choice("Japanese (日本語)", "Japanese"), + questionary.Choice("Korean (한국어)", "Korean"), + questionary.Choice("Hindi (हिन्दी)", "Hindi"), + questionary.Choice("Spanish (Español)", "Spanish"), + questionary.Choice("Portuguese (Português)", "Portuguese"), + questionary.Choice("French (Français)", "French"), + questionary.Choice("German (Deutsch)", "German"), + questionary.Choice("Arabic (العربية)", "Arabic"), + questionary.Choice("Russian (Русский)", "Russian"), + questionary.Choice("Custom language", "custom"), + ] choice = questionary.select( "Select Output Language:", - choices=[ - questionary.Choice("English (default)", "English"), - questionary.Choice("Chinese (中文)", "Chinese"), - questionary.Choice("Japanese (日本語)", "Japanese"), - questionary.Choice("Korean (한국어)", "Korean"), - questionary.Choice("Hindi (हिन्दी)", "Hindi"), - questionary.Choice("Spanish (Español)", "Spanish"), - questionary.Choice("Portuguese (Português)", "Portuguese"), - questionary.Choice("French (Français)", "French"), - questionary.Choice("German (Deutsch)", "German"), - questionary.Choice("Arabic (العربية)", "Arabic"), - questionary.Choice("Russian (Русский)", "Russian"), - questionary.Choice("Custom language", "custom"), - ], + choices=choices, + default=_matching_choice([(c.title, c.value) for c in choices], default), style=questionary.Style([ ("selected", "fg:yellow noinherit"), ("highlighted", "fg:yellow noinherit"), diff --git a/tests/test_cli_prefs.py b/tests/test_cli_prefs.py new file mode 100644 index 000000000..a0426f4ab --- /dev/null +++ b/tests/test_cli_prefs.py @@ -0,0 +1,143 @@ +"""The CLI remembers what you chose last time and offers it back. + +Prefill only: every prompt still appears, so a run never starts on a choice the +user did not see. Environment variables keep skipping their step outright and +win over anything remembered. Remembered values are validated against the +current choices each time, since models and providers come and go between +versions and a stale one must not be offered. +""" + +from __future__ import annotations + +from unittest import mock + +import pytest + +from cli.models import AnalystType +from cli.prefs import load_last_run, sanitize, save_last_run + +SAVED = { + "output_language": "English", + "analysts": ["market", "fundamentals"], + "research_depth": 3, + "llm_provider": "openai", + "quick_think_llm": "gpt-5.6-mini", + "deep_think_llm": "gpt-5.6", +} + + +@pytest.fixture(autouse=True) +def _home(tmp_path, monkeypatch): + monkeypatch.setattr("cli.prefs._PREFS_PATH", tmp_path / "cli_prefs.json") + return tmp_path + + +@pytest.mark.unit +def test_round_trip(): + save_last_run(SAVED) + assert load_last_run() == SAVED + + +@pytest.mark.unit +def test_missing_file_is_not_an_error(): + assert load_last_run() == {} + + +@pytest.mark.unit +def test_a_corrupt_file_degrades_to_no_memory(_home): + (_home / "cli_prefs.json").write_text("{not json") + assert load_last_run() == {} + + +@pytest.mark.unit +def test_a_half_written_file_cannot_be_observed(_home): + """Two runs finishing together must never leave a torn file behind.""" + save_last_run(SAVED) + save_last_run({**SAVED, "research_depth": 5}) + assert load_last_run()["research_depth"] == 5 + assert list((_home).glob("*.tmp*")) == [] + + +# --- validation against the current choices --------------------------------- + +@pytest.mark.unit +def test_a_model_that_no_longer_exists_is_dropped(): + # gpt-5.4 is still accepted by config, but is no longer in the picker's list. + kept = sanitize({**SAVED, "quick_think_llm": "gpt-5.4"}, "stock") + assert "quick_think_llm" not in kept + assert kept["deep_think_llm"] == "gpt-5.6" # the valid sibling survives + + +@pytest.mark.unit +def test_an_unknown_provider_drops_itself_and_its_models(): + kept = sanitize({**SAVED, "llm_provider": "no-such-provider"}, "stock") + assert "llm_provider" not in kept + assert "quick_think_llm" not in kept and "deep_think_llm" not in kept + + +@pytest.mark.unit +def test_analysts_are_narrowed_to_the_asset_type(): + kept = sanitize(SAVED, "crypto") + assert AnalystType.FUNDAMENTALS.value not in kept["analysts"] + assert AnalystType.MARKET.value in kept["analysts"] + + +@pytest.mark.unit +def test_junk_values_are_dropped_rather_than_offered(): + kept = sanitize({"research_depth": 99, "analysts": ["astrology"], "output_language": 5}, "stock") + assert kept == {} + + +@pytest.mark.unit +def test_a_region_specific_provider_survives(): + kept = sanitize({**SAVED, "llm_provider": "qwen-cn", "quick_think_llm": None}, "stock") + assert kept["llm_provider"] == "qwen-cn" + + +# --- wiring ------------------------------------------------------------------ + +def _answer_every_prompt(monkeypatch): + """Drive the real selection flow, answering each prompt with a fixed value.""" + import cli.main as m + + monkeypatch.setattr(m, "fetch_announcements", lambda: []) + monkeypatch.setattr(m, "display_announcements", lambda *a: None) + monkeypatch.setattr(m, "get_ticker", lambda: "NVDA") + monkeypatch.setattr(m, "get_analysis_date", lambda: "2026-09-01") + monkeypatch.setattr(m, "ask_output_language", lambda default=None: "English") + monkeypatch.setattr(m, "select_analysts", lambda asset_type, default=None: [AnalystType.MARKET]) + monkeypatch.setattr(m, "select_research_depth", lambda default=None: 3) + monkeypatch.setattr(m, "select_llm_provider", lambda default=None: ("openai", None)) + monkeypatch.setattr(m, "select_shallow_thinking_agent", lambda p, default=None: "gpt-5.6-mini") + monkeypatch.setattr(m, "select_deep_thinking_agent", lambda p, default=None: "gpt-5.6") + monkeypatch.setattr(m, "ask_openai_reasoning_effort", lambda: "medium") + return m + + +@pytest.mark.unit +def test_selections_are_remembered_after_a_run(monkeypatch): + """Drives the real flow: a stubbed selections dict would hide a key mismatch.""" + m = _answer_every_prompt(monkeypatch) + + m.get_user_selections() + + remembered = load_last_run() + assert remembered["analysts"] == ["market"] + assert remembered["quick_think_llm"] == "gpt-5.6-mini" + assert remembered["deep_think_llm"] == "gpt-5.6" + assert remembered["llm_provider"] == "openai" + assert "ticker" not in remembered # changes every run; never remembered + assert "analysis_date" not in remembered # a stale date must not be offered + + +@pytest.mark.unit +def test_a_custom_language_is_remembered_without_breaking_the_next_run(): + """A free-text answer is not one of the menu's choices, and questionary + rejects a default it cannot find, so offering it back would crash startup.""" + from cli.utils import ask_output_language + + save_last_run({"output_language": "Turkish"}) + with mock.patch("cli.utils.questionary.select") as select: + select.return_value.ask.return_value = "English" + ask_output_language(load_last_run()["output_language"]) + assert select.call_args.kwargs["default"] is None