| from __future__ import annotations |
|
|
| import json |
| import tempfile |
| import unittest |
| from pathlib import Path |
| from unittest import mock |
|
|
| from fastapi import FastAPI |
| from fastapi.testclient import TestClient |
|
|
| from api import ai as ai_module |
| from services.chat_model_service import resolve_chat_models |
| from services.config import ConfigStore, _normalize_chat_model_overrides |
|
|
|
|
| def _candidate(model_id: str, **extra: object) -> dict[str, object]: |
| return {"id": model_id, "object": "model", "owned_by": "chatgpt", **extra} |
|
|
|
|
| def _third_party(model_id: str, **extra: object) -> dict[str, object]: |
| """第三方接入模型的条目形态（group_id + capabilities 由第三方服务写死）。""" |
|
|
| return { |
| "id": model_id, |
| "object": "model", |
| "owned_by": "aichimo", |
| "group_id": "g1", |
| "capabilities": ["txt2img", "img2img"], |
| **extra, |
| } |
|
|
|
|
| class ChatModelOverrideNormalizationTests(unittest.TestCase): |
| def test_trims_dedupes_and_drops_blanks(self) -> None: |
| normalized = _normalize_chat_model_overrides({ |
| "include": [" gpt-5-6 ", "", None, "gpt-5-6", " ", "research"], |
| "exclude": ["gpt-5-5", "gpt-5-5", None], |
| }) |
| self.assertEqual(normalized, {"include": ["gpt-5-6", "research"], "exclude": ["gpt-5-5"]}) |
|
|
| def test_include_wins_when_a_model_is_listed_in_both(self) -> None: |
| normalized = _normalize_chat_model_overrides({ |
| "include": ["gpt-5-6"], |
| "exclude": ["gpt-5-6", "gpt-5-5"], |
| }) |
| self.assertEqual(normalized, {"include": ["gpt-5-6"], "exclude": ["gpt-5-5"]}) |
|
|
| def test_invalid_payloads_fall_back_to_empty(self) -> None: |
| for value in (None, "gpt-5-6", 42, ["gpt-5-6"]): |
| self.assertEqual(_normalize_chat_model_overrides(value), {"include": [], "exclude": []}) |
|
|
| def test_overrides_round_trip_through_config_store(self) -> None: |
| with tempfile.TemporaryDirectory() as tmp_dir: |
| path = Path(tmp_dir) / "config.json" |
| path.write_text(json.dumps({"auth-key": "test-auth"}), encoding="utf-8") |
| store = ConfigStore(path) |
|
|
| self.assertEqual(store.chat_models, {"include": [], "exclude": []}) |
| saved = store.update({"chat_models": {"include": [" gpt-5-6 "], "exclude": ["gpt-5-5", "gpt-5-5"]}}) |
| self.assertEqual(saved["chat_models"], {"include": ["gpt-5-6"], "exclude": ["gpt-5-5"]}) |
| self.assertEqual(store.get()["chat_models"], {"include": ["gpt-5-6"], "exclude": ["gpt-5-5"]}) |
|
|
| |
| store.update({"proxy": "http://proxy.example:8080"}) |
| self.assertEqual(store.get()["chat_models"], {"include": ["gpt-5-6"], "exclude": ["gpt-5-5"]}) |
|
|
| |
| reloaded = ConfigStore(path) |
| self.assertEqual(reloaded.chat_models, {"include": ["gpt-5-6"], "exclude": ["gpt-5-5"]}) |
|
|
|
|
| class ResolveChatModelsTests(unittest.TestCase): |
| def setUp(self) -> None: |
| self.items = [ |
| _candidate("auto", title="Auto"), |
| _candidate("gpt-5-6", title="GPT-5.6 Luna"), |
| _candidate("gpt-5-6-mini", title="GPT-5.6 Luna"), |
| _candidate("research"), |
| _candidate("python"), |
| _candidate("gpt-image-2"), |
| |
| _third_party("Nano Banana 2-4k"), |
| _third_party("gemini-3-pro-image"), |
| |
| _third_party("qwen3-max", capabilities=[]), |
| ] |
|
|
| def resolved(self, overrides: dict[str, object] | None = None) -> dict[str, dict[str, object]]: |
| return {item["id"]: item for item in resolve_chat_models(self.items, overrides)} |
|
|
| def test_third_party_models_are_kept_as_candidates(self) -> None: |
| """第三方模型不能被硬剔除，否则管理员连看都看不到、更开不了。""" |
|
|
| resolved = self.resolved() |
| for model_id in ("Nano Banana 2-4k", "gemini-3-pro-image", "qwen3-max"): |
| self.assertIn(model_id, resolved, f"{model_id} 必须出现在候选里") |
| self.assertEqual(resolved["Nano Banana 2-4k"]["chat_source"], "third_party") |
| self.assertEqual(resolved["qwen3-max"]["chat_source"], "third_party") |
| self.assertEqual(resolved["gpt-5-6"]["chat_source"], "upstream") |
|
|
| def test_third_party_image_models_default_hidden(self) -> None: |
| resolved = self.resolved() |
| for model_id in ("Nano Banana 2-4k", "gemini-3-pro-image"): |
| self.assertFalse(resolved[model_id]["chat_default_enabled"]) |
| self.assertTrue(resolved[model_id]["chat_is_image"]) |
| self.assertEqual(resolved[model_id]["chat_default_reason"], "image_model") |
| enabled = {item["id"] for item in resolve_chat_models(self.items) if item["chat_enabled"]} |
| self.assertNotIn("Nano Banana 2-4k", enabled) |
| self.assertNotIn("gemini-3-pro-image", enabled) |
|
|
| def test_third_party_image_models_can_be_forced_visible(self) -> None: |
| """默认隐藏只是默认值：管理员可以选择把它放进对话下拉。""" |
|
|
| resolved = self.resolved({"include": ["Nano Banana 2-4k"], "exclude": []}) |
| self.assertTrue(resolved["Nano Banana 2-4k"]["chat_enabled"]) |
| |
| self.assertFalse(resolved["gemini-3-pro-image"]["chat_enabled"]) |
|
|
| def test_third_party_non_image_models_default_to_chat(self) -> None: |
| """第三方接入的文本 / 多模态模型应当默认可聊（不能因为"来自第三方"就被排除）。""" |
|
|
| resolved = self.resolved() |
| self.assertTrue(resolved["qwen3-max"]["chat_default_enabled"]) |
| self.assertFalse(resolved["qwen3-max"]["chat_is_image"]) |
| self.assertEqual(resolved["qwen3-max"]["chat_default_reason"], "") |
|
|
| def test_third_party_model_without_capabilities_but_image_name_is_hidden(self) -> None: |
| resolved = { |
| item["id"]: item |
| for item in resolve_chat_models([_third_party("flux-dev-image", capabilities=[])]) |
| } |
| self.assertFalse(resolved["flux-dev-image"]["chat_default_enabled"]) |
| self.assertEqual(resolved["flux-dev-image"]["chat_default_reason"], "image_model") |
|
|
| def test_default_rules_keep_chat_models_and_hide_pseudo_models(self) -> None: |
| enabled = {item["id"] for item in resolve_chat_models(self.items) if item["chat_enabled"]} |
| self.assertEqual(enabled, {"auto", "gpt-5-6", "gpt-5-6-mini", "qwen3-max"}) |
| defaults = {item["id"]: item["chat_default_enabled"] for item in resolve_chat_models(self.items)} |
| self.assertFalse(defaults["research"]) |
| self.assertFalse(defaults["gpt-image-2"]) |
| self.assertTrue(defaults["gpt-5-6"]) |
|
|
| def test_default_reasons_explain_why_a_model_is_hidden(self) -> None: |
| resolved = self.resolved() |
| self.assertEqual(resolved["research"]["chat_default_reason"], "tool_model") |
| self.assertEqual(resolved["python"]["chat_default_reason"], "tool_model") |
| self.assertEqual(resolved["gpt-image-2"]["chat_default_reason"], "image_model") |
| self.assertEqual(resolved["gpt-5-6"]["chat_default_reason"], "") |
|
|
| def test_overrides_apply_on_top_of_defaults(self) -> None: |
| resolved = { |
| item["id"]: item["chat_enabled"] |
| for item in resolve_chat_models( |
| self.items, |
| {"include": ["research"], "exclude": ["gpt-5-6"]}, |
| ) |
| } |
| self.assertTrue(resolved["research"]) |
| self.assertFalse(resolved["gpt-5-6"]) |
| |
| self.assertTrue(resolved["gpt-5-6-mini"]) |
| self.assertFalse(resolved["python"]) |
|
|
| def test_enabled_models_sort_first(self) -> None: |
| resolved = resolve_chat_models(self.items) |
| flags = [item["chat_enabled"] for item in resolved] |
| self.assertEqual(flags, sorted(flags, reverse=True)) |
|
|
|
|
| class ChatModelsEndpointTests(unittest.TestCase): |
| def _get( |
| self, |
| payload: list[dict[str, object]], |
| overrides: dict[str, list[str]] | None = None, |
| third_party: list[dict[str, object]] | None = None, |
| ): |
| with ( |
| mock.patch.object(ai_module, "require_identity", return_value={"role": "admin"}), |
| mock.patch.object( |
| ai_module.openai_v1_models, |
| "list_models", |
| return_value={"object": "list", "data": payload}, |
| ), |
| mock.patch.object( |
| ai_module.third_party_image_service, |
| "public_models_for_identity", |
| return_value=third_party or [], |
| ), |
| mock.patch.object( |
| type(ai_module.config), |
| "chat_models", |
| new_callable=mock.PropertyMock, |
| return_value=overrides or {"include": [], "exclude": []}, |
| ), |
| ): |
| app = FastAPI() |
| app.include_router(ai_module.create_router()) |
| return TestClient(app).get("/api/chat/models") |
|
|
| def test_returns_chat_flags_for_upstream_models(self) -> None: |
| response = self._get([ |
| _candidate("auto"), |
| _candidate("gpt-5-6"), |
| _candidate("research"), |
| ]) |
|
|
| self.assertEqual(response.status_code, 200) |
| payload = response.json() |
| self.assertEqual(payload["object"], "list") |
| self.assertEqual(payload["include"], []) |
| self.assertEqual(payload["exclude"], []) |
| flags = {item["id"]: item["chat_enabled"] for item in payload["data"]} |
| self.assertEqual(flags, {"auto": True, "gpt-5-6": True, "research": False}) |
|
|
| def test_merges_third_party_models_into_candidates(self) -> None: |
| """回归：/api/chat/models 之前只取上游模型，第三方模型一份都拿不到。""" |
|
|
| response = self._get( |
| [_candidate("auto"), _candidate("gpt-5-6")], |
| third_party=[_third_party("Nano Banana 2-4k"), _third_party("qwen3-max", capabilities=[])], |
| ) |
|
|
| payload = response.json() |
| flags = {item["id"]: item["chat_enabled"] for item in payload["data"]} |
| self.assertIn("Nano Banana 2-4k", flags, "第三方生图模型必须出现在候选里") |
| self.assertFalse(flags["Nano Banana 2-4k"]) |
| self.assertTrue(flags["gpt-5-6"]) |
| self.assertTrue(flags["qwen3-max"], "第三方非生图模型应默认可聊") |
| sources = {item["id"]: item.get("chat_source") for item in payload["data"]} |
| self.assertEqual(sources["Nano Banana 2-4k"], "third_party") |
| self.assertEqual(sources["gpt-5-6"], "upstream") |
|
|
| def test_third_party_models_are_deduplicated_against_upstream(self) -> None: |
| response = self._get( |
| [_candidate("auto"), _third_party("Nano Banana 2-4k")], |
| third_party=[_third_party("Nano Banana 2-4k")], |
| ) |
|
|
| ids = [item["id"] for item in response.json()["data"]] |
| self.assertEqual(ids.count("Nano Banana 2-4k"), 1) |
|
|
| def test_echoes_overrides_and_applies_them(self) -> None: |
| response = self._get( |
| [_candidate("gpt-5-6"), _candidate("research")], |
| {"include": ["research"], "exclude": ["gpt-5-6"]}, |
| ) |
|
|
| payload = response.json() |
| self.assertEqual(payload["include"], ["research"]) |
| self.assertEqual(payload["exclude"], ["gpt-5-6"]) |
| flags = {item["id"]: item["chat_enabled"] for item in payload["data"]} |
| self.assertTrue(flags["research"]) |
| self.assertFalse(flags["gpt-5-6"]) |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
