| 1 | import os |
| 2 | import tempfile |
| 3 | import unittest |
| 4 | from pathlib import Path |
| 5 | from unittest.mock import patch |
| 6 | |
| 7 | import yaml |
| 8 | |
| 9 | from agent_runtime.config import ( |
| 10 | api_provider_from_base_url, |
| 11 | embedding_api_key, |
| 12 | embedding_base_url, |
| 13 | embedding_model, |
| 14 | embedding_model_provider, |
| 15 | image_api_key, |
| 16 | image_base_url, |
| 17 | image_model, |
| 18 | image_num_candidates, |
| 19 | llm_api_key, |
| 20 | llm_base_url, |
| 21 | llm_model, |
| 22 | llm_model_provider, |
| 23 | load_agent_config, |
| 24 | reranker_api_key, |
| 25 | reranker_base_url, |
| 26 | reranker_model, |
| 27 | video_api_key, |
| 28 | video_base_url, |
| 29 | video_model, |
| 30 | video_provider, |
| 31 | ) |
| 32 | |
| 33 | |
| 34 | class AgentConfigTests(unittest.TestCase): |
| 35 | def setUp(self): |
| 36 | load_agent_config.cache_clear() |
| 37 | |
| 38 | def tearDown(self): |
| 39 | load_agent_config.cache_clear() |
| 40 | |
| 41 | def test_reads_agent_local_config(self): |
| 42 | with tempfile.TemporaryDirectory() as tmp: |
| 43 | config_dir = Path(tmp) / "configs" |
| 44 | config_dir.mkdir() |
| 45 | (config_dir / "agent.local.yaml").write_text(yaml.safe_dump({ |
| 46 | "llm": {"model_provider": "openai", "model": "config-llm", "base_url": "https://config.test/v1", "api_key": "config-key"}, |
| 47 | "image": {"model": "config-image", "base_url": "https://image.test", "api_key": "image-key"}, |
| 48 | "video": {"model": "config-video", "base_url": "https://openrouter.ai/api/v1", "api_key": "video-key"}, |
| 49 | "embedding": {"model_provider": "openai", "model": "config-embedding", "base_url": "https://embedding.test/v1", "api_key": "embedding-key"}, |
| 50 | "reranker": {"model": "config-reranker", "base_url": "https://reranker.test", "api_key": "reranker-key"}, |
| 51 | }), encoding="utf-8") |
| 52 | with patch.dict(os.environ, {}, clear=True): |
| 53 | self.assertEqual(llm_model(tmp), "config-llm") |
| 54 | self.assertEqual(llm_model_provider(tmp), "openai") |
| 55 | self.assertEqual(llm_base_url(tmp), "https://config.test/v1") |
| 56 | self.assertEqual(llm_api_key(tmp), "config-key") |
| 57 | self.assertEqual(image_model(tmp), "config-image") |
| 58 | self.assertEqual(image_base_url(tmp), "https://image.test") |
| 59 | self.assertEqual(image_api_key(tmp), "image-key") |
| 60 | self.assertEqual(video_model(tmp), "config-video") |
| 61 | self.assertEqual(video_provider(tmp), "openrouter") |
| 62 | self.assertEqual(video_base_url(tmp), "https://openrouter.ai/api/v1") |
| 63 | self.assertEqual(video_api_key(tmp), "video-key") |
| 64 | self.assertEqual(embedding_model_provider(tmp), "openai") |
| 65 | self.assertEqual(embedding_model(tmp), "config-embedding") |
| 66 | self.assertEqual(embedding_base_url(tmp), "https://embedding.test/v1") |
| 67 | self.assertEqual(embedding_api_key(tmp), "embedding-key") |
| 68 | self.assertEqual(reranker_model(tmp), "config-reranker") |
| 69 | self.assertEqual(reranker_base_url(tmp), "https://reranker.test") |
| 70 | self.assertEqual(reranker_api_key(tmp), "reranker-key") |
| 71 | |
| 72 | def test_environment_overrides_agent_local_config(self): |
| 73 | with tempfile.TemporaryDirectory() as tmp: |
| 74 | config_dir = Path(tmp) / "configs" |
| 75 | config_dir.mkdir() |
| 76 | (config_dir / "agent.local.yaml").write_text(yaml.safe_dump({"llm": {"model": "config-llm", "api_key": "config-key"}}), encoding="utf-8") |
| 77 | with patch.dict(os.environ, {"VIMAX_LLM_MODEL": "env-llm", "VIMAX_LLM_MODEL_PROVIDER": "openai", "VIMAX_LLM_API_KEY": "env-key", "VIMAX_VIDEO_BASE_URL": "https://openrouter.ai/api/v1", "VIMAX_EMBEDDING_MODEL": "env-embedding", "VIMAX_EMBEDDING_BASE_URL": "https://env-embedding.test/v1", "VIMAX_EMBEDDING_API_KEY": "env-embedding-key", "VIMAX_RERANKER_MODEL": "env-reranker", "VIMAX_RERANKER_BASE_URL": "https://env-reranker.test", "VIMAX_RERANKER_API_KEY": "env-reranker-key"}, clear=True): |
| 78 | self.assertEqual(llm_model(tmp), "env-llm") |
| 79 | self.assertEqual(llm_model_provider(tmp), "openai") |
| 80 | self.assertEqual(llm_api_key(tmp), "env-key") |
| 81 | self.assertEqual(video_provider(tmp), "openrouter") |
| 82 | self.assertEqual(video_base_url(tmp), "https://openrouter.ai/api/v1") |
| 83 | self.assertEqual(embedding_model(tmp), "env-embedding") |
| 84 | self.assertEqual(embedding_base_url(tmp), "https://env-embedding.test/v1") |
| 85 | self.assertEqual(embedding_api_key(tmp), "env-embedding-key") |
| 86 | self.assertEqual(reranker_model(tmp), "env-reranker") |
| 87 | self.assertEqual(reranker_base_url(tmp), "https://env-reranker.test") |
| 88 | self.assertEqual(reranker_api_key(tmp), "env-reranker-key") |
| 89 | |
| 90 | def test_image_and_video_keys_fall_back_to_llm_key(self): |
| 91 | with tempfile.TemporaryDirectory() as tmp: |
| 92 | config_dir = Path(tmp) / "configs" |
| 93 | config_dir.mkdir() |
| 94 | (config_dir / "agent.local.yaml").write_text(yaml.safe_dump({"llm": {"api_key": "shared-key"}}), encoding="utf-8") |
| 95 | with patch.dict(os.environ, {}, clear=True): |
| 96 | self.assertEqual(image_api_key(tmp), "shared-key") |
| 97 | self.assertEqual(video_api_key(tmp), "shared-key") |
| 98 | |
| 99 | def test_video_provider_is_inferred_from_base_url(self): |
| 100 | self.assertEqual(api_provider_from_base_url("https://openrouter.ai/api/v1"), "openrouter") |
| 101 | self.assertEqual(api_provider_from_base_url("https://yunwu.ai/v1"), "yunwu") |
| 102 | self.assertEqual(api_provider_from_base_url("https://example.com/v1"), "") |
| 103 | |
| 104 | def test_image_candidates_default_and_environment_override(self): |
| 105 | with tempfile.TemporaryDirectory() as tmp, patch.dict(os.environ, {}, clear=True): |
| 106 | self.assertEqual(image_num_candidates(tmp), 2) |
| 107 | with patch.dict(os.environ, {"VIMAX_IMAGE_NUM_CANDIDATES": "1"}): |
| 108 | self.assertEqual(image_num_candidates(tmp), 1) |
| 109 | |
| 110 | def test_image_candidates_use_workspace_config(self): |
| 111 | with tempfile.TemporaryDirectory() as tmp, patch.dict(os.environ, {}, clear=True): |
| 112 | config_dir = Path(tmp) / "configs" |
| 113 | config_dir.mkdir() |
| 114 | (config_dir / "agent.local.yaml").write_text("image_selection:\n num_candidates: 3\n", encoding="utf-8") |
| 115 | self.assertEqual(image_num_candidates(tmp), 3) |
| 116 | with patch.dict(os.environ, {"VIMAX_IMAGE_NUM_CANDIDATES": "1"}): |
| 117 | self.assertEqual(image_num_candidates(tmp), 1) |
| 118 | |
| 119 | def test_invalid_image_candidate_counts_are_rejected(self): |
| 120 | with tempfile.TemporaryDirectory() as tmp: |
| 121 | for value in ("0", "-1", "1.5", "abc", ""): |
| 122 | with self.subTest(value=value), patch.dict(os.environ, {"VIMAX_IMAGE_NUM_CANDIDATES": value}, clear=True): |
| 123 | with self.assertRaisesRegex(ValueError, "positive integer"): |
| 124 | image_num_candidates(tmp) |
| 125 | |
| 126 | |
| 127 | if __name__ == "__main__": |
| 128 | unittest.main() |
| 129 |