| 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 | llm_api_key, |
| 19 | llm_base_url, |
| 20 | llm_model, |
| 21 | llm_model_provider, |
| 22 | load_agent_config, |
| 23 | reranker_api_key, |
| 24 | reranker_base_url, |
| 25 | reranker_model, |
| 26 | video_api_key, |
| 27 | video_base_url, |
| 28 | video_model, |
| 29 | video_provider, |
| 30 | ) |
| 31 | |
| 32 | |
| 33 | class AgentConfigTests(unittest.TestCase): |
| 34 | def setUp(self): |
| 35 | load_agent_config.cache_clear() |
| 36 | |
| 37 | def tearDown(self): |
| 38 | load_agent_config.cache_clear() |
| 39 | |
| 40 | def test_reads_agent_local_config(self): |
| 41 | with tempfile.TemporaryDirectory() as tmp: |
| 42 | config_dir = Path(tmp) / "configs" |
| 43 | config_dir.mkdir() |
| 44 | (config_dir / "agent.local.yaml").write_text(yaml.safe_dump({ |
| 45 | "llm": {"model_provider": "openai", "model": "config-llm", "base_url": "https://config.test/v1", "api_key": "config-key"}, |
| 46 | "image": {"model": "config-image", "base_url": "https://image.test", "api_key": "image-key"}, |
| 47 | "video": {"model": "config-video", "base_url": "https://openrouter.ai/api/v1", "api_key": "video-key"}, |
| 48 | "embedding": {"model_provider": "openai", "model": "config-embedding", "base_url": "https://embedding.test/v1", "api_key": "embedding-key"}, |
| 49 | "reranker": {"model": "config-reranker", "base_url": "https://reranker.test", "api_key": "reranker-key"}, |
| 50 | }), encoding="utf-8") |
| 51 | with patch.dict(os.environ, {}, clear=True): |
| 52 | self.assertEqual(llm_model(tmp), "config-llm") |
| 53 | self.assertEqual(llm_model_provider(tmp), "openai") |
| 54 | self.assertEqual(llm_base_url(tmp), "https://config.test/v1") |
| 55 | self.assertEqual(llm_api_key(tmp), "config-key") |
| 56 | self.assertEqual(image_model(tmp), "config-image") |
| 57 | self.assertEqual(image_base_url(tmp), "https://image.test") |
| 58 | self.assertEqual(image_api_key(tmp), "image-key") |
| 59 | self.assertEqual(video_model(tmp), "config-video") |
| 60 | self.assertEqual(video_provider(tmp), "openrouter") |
| 61 | self.assertEqual(video_base_url(tmp), "https://openrouter.ai/api/v1") |
| 62 | self.assertEqual(video_api_key(tmp), "video-key") |
| 63 | self.assertEqual(embedding_model_provider(tmp), "openai") |
| 64 | self.assertEqual(embedding_model(tmp), "config-embedding") |
| 65 | self.assertEqual(embedding_base_url(tmp), "https://embedding.test/v1") |
| 66 | self.assertEqual(embedding_api_key(tmp), "embedding-key") |
| 67 | self.assertEqual(reranker_model(tmp), "config-reranker") |
| 68 | self.assertEqual(reranker_base_url(tmp), "https://reranker.test") |
| 69 | self.assertEqual(reranker_api_key(tmp), "reranker-key") |
| 70 | |
| 71 | def test_environment_overrides_agent_local_config(self): |
| 72 | with tempfile.TemporaryDirectory() as tmp: |
| 73 | config_dir = Path(tmp) / "configs" |
| 74 | config_dir.mkdir() |
| 75 | (config_dir / "agent.local.yaml").write_text(yaml.safe_dump({"llm": {"model": "config-llm", "api_key": "config-key"}}), encoding="utf-8") |
| 76 | 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): |
| 77 | self.assertEqual(llm_model(tmp), "env-llm") |
| 78 | self.assertEqual(llm_model_provider(tmp), "openai") |
| 79 | self.assertEqual(llm_api_key(tmp), "env-key") |
| 80 | self.assertEqual(video_provider(tmp), "openrouter") |
| 81 | self.assertEqual(video_base_url(tmp), "https://openrouter.ai/api/v1") |
| 82 | self.assertEqual(embedding_model(tmp), "env-embedding") |
| 83 | self.assertEqual(embedding_base_url(tmp), "https://env-embedding.test/v1") |
| 84 | self.assertEqual(embedding_api_key(tmp), "env-embedding-key") |
| 85 | self.assertEqual(reranker_model(tmp), "env-reranker") |
| 86 | self.assertEqual(reranker_base_url(tmp), "https://env-reranker.test") |
| 87 | self.assertEqual(reranker_api_key(tmp), "env-reranker-key") |
| 88 | |
| 89 | def test_image_and_video_keys_fall_back_to_llm_key(self): |
| 90 | with tempfile.TemporaryDirectory() as tmp: |
| 91 | config_dir = Path(tmp) / "configs" |
| 92 | config_dir.mkdir() |
| 93 | (config_dir / "agent.local.yaml").write_text(yaml.safe_dump({"llm": {"api_key": "shared-key"}}), encoding="utf-8") |
| 94 | with patch.dict(os.environ, {}, clear=True): |
| 95 | self.assertEqual(image_api_key(tmp), "shared-key") |
| 96 | self.assertEqual(video_api_key(tmp), "shared-key") |
| 97 | |
| 98 | def test_video_provider_is_inferred_from_base_url(self): |
| 99 | self.assertEqual(api_provider_from_base_url("https://openrouter.ai/api/v1"), "openrouter") |
| 100 | self.assertEqual(api_provider_from_base_url("https://yunwu.ai/v1"), "yunwu") |
| 101 | self.assertEqual(api_provider_from_base_url("https://example.com/v1"), "") |
| 102 | |
| 103 | |
| 104 | if __name__ == "__main__": |
| 105 | unittest.main() |
| 106 |