{"repo": "channinglua/prax-agent", "files": {"src/prax/core/provider_setup.py": {"57f7e20c942489300b883e5999942cdbfca5642a": {"t": 1785674160.5509858, "n": 12655, "h": "5aa67a4a0c1b2ac7", "c": "\"\"\"Helpers for provider setup templates and local model config generation.\"\"\"\n\nfrom __future__ import annotations\n\nfrom copy import deepcopy\nfrom pathlib import Path\nfrom typing import Any\n\nimport yaml\n\n\nFLOW_TEMPLATES: dict[str, dict[str, Any]] = {\n    \"glm\": {\n        \"default_model\": \"glm-5\",\n        \"upgrade_chain\": [\"glm-4-flash\", \"glm-4\", \"glm-5\"],\n        \"providers\": {\n            \"zhipu\": {\n                \"base_url\": \"https://open.bigmodel.cn/api/paas/v4\",\n                \"api_key_env\": \"ZHIPU_API_KEY\",\n                \"format\": \"openai\",\n                \"models\": [\n                    {\n                        \"name\": \"glm-4-flash\",\n                        \"aliases\": [\"glm-flash\"],\n                        \"request_mode\": \"chat_completions\",\n                        \"tier\": \"low\",\n                        \"cost_per_1m_tokens\": 0.1,\n                        \"supports_tools\": True,\n                        \"supports_streaming\": True,\n                    },\n                    {\n                        \"name\": \"glm-4\",\n                        \"aliases\": [\"glm\"],\n                        \"request_mode\": \"chat_completions\",\n                        \"tier\": \"standard\",\n                        \"cost_per_1m_tokens\": 1.0,\n                        \"supports_tools\": True,\n                        \"supports_streaming\": True,\n                    },\n                    {\n                        \"name\": \"glm-5\",\n                        \"aliases\": [\"glm-pro\"],\n                        \"request_mode\": \"chat_completions\",\n                        \"tier\": \"high\",\n                        \"cost_per_1m_tokens\": 2.0,\n                        \"supports_tools\": True,\n                        \"supports_streaming\": True,\n                    },\n                ],\n            }\n        },\n    },\n    \"codex\": {\n        \"default_model\": \"codex\",\n        \"upgrade_chain\": [\"codex\", \"gpt-5.5\", \"claude-sonnet-4-8\"],\n        \"providers\": {\n            \"openai\": {\n                \"base_url\": \"https://api.openai.com/v1\",\n                \"api_key_env\": \"OPENAI_API_KEY\",\n                \"format\": \"openai\",\n                \"models\": [\n                    {\n                        \"name\": \"gpt-5.5\",\n                        \"api_model\": \"gpt-5.5\",\n                        \"aliases\": [\"gpt\", \"gpt4\"],\n                        \"request_mode\": \"chat_completions\",\n                        \"tier\": \"standard\",\n                        \"cost_per_1m_tokens\": 5.0,\n                        \"supports_tools\": True,\n                        \"supports_streaming\": True,\n                    },\n                    {\n                        \"name\": \"codex\",\n                        \"api_model\": \"<replace-with-codex-model>\",\n                        \"aliases\": [\"codex\", \"responses\"],\n                        \"request_mode\": \"responses\",\n                        \"tier\": \"high\",\n                        \"cost_per_1m_tokens\": 0.0,\n                        \"supports_tools\": True,\n                        \"supports_streaming\": True,\n                        \"supports_reasoning_effort\": True,\n                        \"default_reasoning_effort\": \"medium\",\n                    },\n                ],\n            }\n        },\n    },\n    \"claude\": {\n        \"default_model\": \"claude-sonnet-4-8\",\n        \"upgrade_chain\": [\"claude-sonnet-4-8\", \"claude-opus-4-8\"],\n        \"providers\": {\n            \"anthropic\": {\n                \"base_url\": \"https://api.anthropic.com\",\n                \"api_key_env\": [\"ANTHROPIC_API_KEY\", \"ANTHROPIC_AUTH_TOKEN\"],\n                \"format\": \"anthropic\",\n                \"models\": [\n                    {\n                        \"name\": \"claude-sonnet-4-8\",\n                        \"aliases\": [\"sonnet\"],\n                        \"request_mode\": \"chat_completions\",\n                        \"tier\": \"premium\",\n                        \"cost_per_1m_tokens\": 15.0,\n                        \"supports_tools\": True,\n                        \"supports_streaming\": True,\n                        \"supports_thinking\": True,\n                        \"default_thinking_budget_tokens\": 12000,\n                    },\n                    {\n                        \"name\": \"claude-opus-4-8\",\n                        \"aliases\": [\"opus\"],\n                        \"request_mode\": \"chat_completions\",\n                        \"tier\": \"ultra\",\n                        \"cost_per_1m_tokens\": 75.0,\n                        \"supports_tools\": True,\n                        \"supports_streaming\": True,\n                        \"supports_thinking\": True,\n                        \"default_thinking_budget_tokens\": 32000,\n                    },\n                ],\n            }\n        },\n    },\n}\n\n\ndef flow_names(target: str) -> list[str]:\n    if target == \"all\":\n        return [\"glm\", \"codex\", \"claude\"]\n    if target not in FLOW_TEMPLATES:\n        raise ValueError(\"flow must be glm|codex|claude|all\")\n    return [target]\n\n\ndef build_flow_template(flow: str) -> dict[str, Any]:\n    if flow == \"all\":\n        return merge_flow_templates(flow_names(flow))\n    if flow not in FLOW_TEMPLATES:\n        raise ValueError(\"flow must be glm|codex|claude|all\")\n    return deepcopy(FLOW_TEMPLATES[flow])\n\n\ndef merge_flow_templates(flows: list[str]) -> dict[str, Any]:\n    result: dict[str, Any] = {\"providers\": {}, \"upgrade_chain\": []}\n    default_model = None\n    for flow in flows:\n        template = deepcopy(FLOW_TEMPLATES[flow])\n        if default_model is None:\n            default_model = template.get(\"default_model\")\n        result[\"providers\"] = _merge_providers(result[\"providers\"], template.get(\"providers\", {}))\n        result[\"upgrade_chain\"] = _merge_upgrade_chain(result[\"upgrade_chain\"], template.get(\"upgrade_chain\", []))\n    result[\"default_model\"] = default_model or \"gpt-5.5\"\n    return result\n\n\ndef merge_into_local_config(\n    existing: dict[str, Any] | None,\n    flow: str,\n    *,\n    overwrite: bool = False,\n    set_default: bool = False,\n    full: bool = False,\n) -> dict[str, Any]:\n    \"\"\"Build the workspace yaml content.\n\n    Default behaviour as of 0.5.5: write an empty skeleton (`providers: {}`)\n    so the user-level `~/.prax/models.yaml` and bundled defaults can flow\n    through unchanged. Without this, the auto-generated workspace yaml\n    silently overrode user-level `base_url` / `api_key_env` settings —\n    every field looked configured but was just a copy of the bundled\n    default. See https://github.com/.../issues/<id-of-init-models-bug>.\n\n    - `full=True` reproduces the old behaviour (full template merge),\n      kept for users who want the schema fully expanded as a starting\n      point and for back-compat with scripts.\n    - `overwrite=True` (the `--force` flag) still writes the full\n      template — explicit overwrite always wins.\n    \"\"\"\n    if overwrite:\n        config = build_flow_template(flow)\n        if set_default and flow != \"all\":\n            config[\"default_model\"] = FLOW_TEMPLATES[flow][\"default_model\"]\n        return config\n\n    if not full:\n        # Skeleton mode: do not seed any provider fields. If existing\n        # workspace yaml already had hand-written content, preserve it\n        # verbatim — we don't touch what the user explicitly authored.\n        if existing:\n            return deepcopy(existing)\n        return {\"providers\": {}}\n\n    base = deepcopy(existing or {})\n    template = build_flow_template(flow)\n    merged = {\n        \"providers\": _merge_providers(base.get(\"providers\", {}), template.get(\"providers\", {})),\n        \"upgrade_chain\": _merge_upgrade_chain(base.get(\"upgrade_chain\", []), template.get(\"upgrade_chain\", [])),\n        \"default_model\": base.get(\"default_model\") or template.get(\"default_model\"),\n    }\n    if set_default and flow != \"all\":\n        merged[\"default_model\"] = FLOW_TEMPLATES[flow][\"default_model\"]\n    return merged\n\n\nSKELETON_HEADER = \"\"\"\\\n# .prax/models.yaml — workspace-level model overrides for this project.\n#\n# YOU DON'T HAVE TO EDIT THIS FILE. prax already ships with working\n# defaults for OpenAI / Anthropic / Zhipu (GLM). To start using prax:\n#\n#   1. Set the API key for the provider you want, e.g.:\n#        export OPENAI_API_KEY=sk-...        # OpenAI / compatible relays\n#        export ANTHROPIC_API_KEY=sk-ant-... # Claude\n#        export ZHIPU_API_KEY=...            # GLM\n#   2. Run `prax prompt \"hi\"` to test.\n#\n# This file is for OVERRIDING those defaults — only fields you write\n# here will take effect; the rest inherit from ~/.prax/models.yaml\n# (your cross-project user-level config) and bundled defaults.\n#\n# Common overrides:\n#\n#   providers:\n#     openai:\n#       base_url: https://my-relay.example.com/openai   # route through a relay\n#   default_model: claude-sonnet-4-8                    # this project uses Claude\n#\n# Want the full schema as a starting point (all providers, every field,\n# placeholders to replace)? Run:\n#\n#   prax /init-models --full\n#\n# Want to see what's currently in effect (bundled + user + workspace\n# merged)? Run:\n#\n#   prax /providers\n\"\"\"\n\n\ndef load_user_models_config() -> dict[str, Any] | None:\n    \"\"\"Read `~/.prax/models.yaml` if present. Used to detect overlap when\n    the user is about to write workspace yaml that would override their\n    user-level config (e.g. running /init-models repeatedly).\n    \"\"\"\n    path = Path.home() / \".prax\" / \"models.yaml\"\n    if not path.exists():\n        return None\n    try:\n        return yaml.safe_load(path.read_text(encoding=\"utf-8\")) or {}\n    except yaml.YAMLError:\n        return None\n\n\ndef detect_user_level_overlap(flow: str) -> list[str]:\n    \"\"\"Return user-level provider names that the requested flow would\n    cover. Empty list = no risk of unintended override.\n    \"\"\"\n    user = load_user_models_config()\n    if not user:\n        return []\n    user_providers = set((user.get(\"providers\") or {}).keys())\n    template_providers: set[str] = set()\n    for fname in flow_names(flow):\n        template_providers.update(FLOW_TEMPLATES[fname].get(\"providers\", {}).keys())\n    return sorted(user_providers & template_providers)\n\n\ndef render_yaml(config: dict[str, Any]) -> str:\n    return yaml.safe_dump(config, sort_keys=False, allow_unicode=True)\n\n\ndef load_local_models_config(cwd: str) -> dict[str, Any] | None:\n    path = Path(cwd) / \".prax\" / \"models.yaml\"\n    if not path.exists():\n        return None\n    return yaml.safe_load(path.read_text(encoding=\"utf-8\"))\n\n\ndef write_local_models_config(cwd: str, config: dict[str, Any], *, header: str = \"\") -> Path:\n    path = Path(cwd) / \".prax\" / \"models.yaml\"\n    path.parent.mkdir(parents=True, exist_ok=True)\n    body = render_yaml(config)\n    text = (header + \"\\n\" + body) if header else body\n    path.write_text(text, encoding=\"utf-8\")\n    return path\n\n\ndef render_env_example(flow: str) -> str:\n    lines = [\"# Generated by prax provider setup\"]\n    keys: list[str] = []\n    for name in flow_names(flow):\n        template = FLOW_TEMPLATES[name]\n        for provider_cfg in template.get(\"providers\", {}).values():\n            raw = provider_cfg.get(\"api_key_env\", \"\")\n            if isinstance(raw, list):\n                keys.extend(str(item) for item in raw)\n            elif raw:\n                keys.append(str(raw))\n    unique_keys: list[str] = []\n    for key in keys:\n        if key not in unique_keys:\n            unique_keys.append(key)\n    for key in unique_keys:\n        lines.append(f\"{key}=<your-{key.lower()}>\")\n    return \"\\n\".join(lines) + \"\\n\"\n\n\ndef write_env_example(cwd: str, flow: str) -> Path:\n    path = Path(cwd) / \".prax\" / \".env.example\"\n    path.parent.mkdir(parents=True, exist_ok=True)\n    path.write_text(render_env_example(flow), encoding=\"utf-8\")\n    return path\n\n\ndef _merge_providers(base: dict[str, Any], extra: dict[str, Any]) -> dict[str, Any]:\n    merged = deepcopy(base)\n    for provider_name, provider_cfg in extra.items():\n        if provider_name not in merged:\n            merged[provider_name] = deepcopy(provider_cfg)\n            continue\n        existing_provider = merged[provider_name]\n        for key, value in provider_cfg.items():\n            if key != \"models\":\n                existing_provider[key] = deepcopy(value)\n                continue\n            existing_models = {model[\"name\"]: deepcopy(model) for model in existing_provider.get(\"models\", [])}\n            for model in value:\n                existing_models[model[\"name\"]] = deepcopy(model)\n            existing_provider[\"models\"] = list(existing_models.values())\n    return merged\n\n\ndef _merge_upgrade_chain(base: list[str], extra: list[str]) -> list[str]:\n    result: list[str] = []\n    for item in [*base, *extra]:\n        if item not in result:\n            result.append(item)\n    return result\n"}}}}