From 25b241faa5fe42ab32cda82da0166d1b4ac55467 Mon Sep 17 00:00:00 2001 From: binyangzhu000-sudo <224954946+binyangzhu000-sudo@users.noreply.github.com> Date: Fri, 24 Jul 2026 02:51:36 +0800 Subject: [PATCH] Add Atlas Cloud LLM provider --- Personality_changes/change_test.py | 13 +++- Personality_changes/llm_link.py | 54 +++++++++++++-- Personality_test/llm_link.py | 64 ++++++++++++++++-- para.env.example | 10 ++- tests/test_atlascloud_provider.py | 105 +++++++++++++++++++++++++++++ 5 files changed, 233 insertions(+), 13 deletions(-) create mode 100644 tests/test_atlascloud_provider.py diff --git a/Personality_changes/change_test.py b/Personality_changes/change_test.py index 0eed60a..2980c2e 100644 --- a/Personality_changes/change_test.py +++ b/Personality_changes/change_test.py @@ -3,10 +3,16 @@ from Class import * compensation_weight = 0.05 -save_load = {'QWEN' : 'data/QWEN_result.txt', 'OPENAI' : 'data/OPENAI_result.txt', 'LLAMA' : 'data/LLAMA_result.txt'} +save_load = {'QWEN' : 'data/QWEN_result.txt', 'OPENAI' : 'data/OPENAI_result.txt', 'LLAMA' : 'data/LLAMA_result.txt', 'ATLASCLOUD' : 'data/ATLASCLOUD_result.txt'} +model_aliases = {'ATLAS' : 'ATLASCLOUD', 'ATLAS_CLOUD' : 'ATLASCLOUD'} MBTI_LIST = ["INTP", "INTJ", "INFP", "INFJ", "ISTP", "ISTJ", "ISFP", "ISFJ", "ENFP", "ENFJ", "ENTP", "ENTJ", "ESFP", "ESFJ", "ESTP", "ESTJ"] +def normalize_model_name(model: str): + normalized = model.replace("-", "_").upper() + return model_aliases.get(normalized, normalized) + + def one_func_test(mbti : str, function : str, model : str): agent = Agent(mbti) with open('data/weight.json', encoding='UTF-8') as file: @@ -57,6 +63,7 @@ def one_func_test(mbti : str, function : str, model : str): return all_log, all_false_rsp, result, weight_data def one_mbti_test(mbti : str, model : str, scene : list): + model = normalize_model_name(model) save_msg = '\nmbti:' + mbti + '\n' all_result = [] @@ -82,10 +89,10 @@ def one_mbti_test(mbti : str, model : str, scene : list): parser.add_argument('--method', type=str, default = 'all_scene') # all_scene / single_scene parser.add_argument('--scene', type=str, default = 'Se') parser.add_argument('--mbti', type=str, default= 'INTJ') - parser.add_argument('--model', type=str, default= 'OPENAI')# QWEN, OPENAI, LLAMA + parser.add_argument('--model', type=str, default= 'OPENAI')# QWEN, OPENAI, LLAMA, ATLASCLOUD args = parser.parse_args() if args.method == 'all_scene': scene = ["Fi", "Fe", "Ti", "Te", "Si", "Se", "Ni", "Ne"] elif args.method == 'single_scene': scene = [args.scene] - one_mbti_test(args.mbti, args.model, scene) \ No newline at end of file + one_mbti_test(args.mbti, args.model, scene) diff --git a/Personality_changes/llm_link.py b/Personality_changes/llm_link.py index 6943dbe..57f9def 100644 --- a/Personality_changes/llm_link.py +++ b/Personality_changes/llm_link.py @@ -3,6 +3,22 @@ import os +ATLAS_CLOUD_BASE_URL = "https://api.atlascloud.ai/v1" +ATLAS_CLOUD_DEFAULT_MODEL = "qwen/qwen3.5-flash" + + +def _get_env(*names: str, default=None): + for name in names: + value = os.getenv(name) + if value: + return value + return default + + +def _normalize_model_name(model: str) -> str: + return model.replace("-", "_").upper() + + def get_openai_rsp(content : str, sys_prompt : str, history_msg : list): load_dotenv("para.env") @@ -55,6 +71,32 @@ def get_llama_rsp(content : str, sys_prompt : str, history_msg : list): return rsp +def get_atlascloud_rsp(content : str, sys_prompt : str, history_msg : list): + load_dotenv("para.env") + api = _get_env("ATLASCLOUD_API_KEY", "ATLAS_CLOUD_API_KEY") + url = _get_env( + "ATLASCLOUD_API_BASE", + "ATLAS_CLOUD_API_BASE", + "ATLASCLOUD_BASE_URL", + "ATLAS_CLOUD_BASE_URL", + default=ATLAS_CLOUD_BASE_URL, + ) + model = _get_env( + "ATLASCLOUD_MODEL", + "ATLAS_CLOUD_MODEL", + default=ATLAS_CLOUD_DEFAULT_MODEL, + ) + + client = OpenAI( + api_key=api, + base_url=url, + ) + + rsp = one_dialogue(content, client, model, sys_prompt, history_msg) + + return rsp + + # 单对话处理 def one_dialogue(content: str, client : OpenAI, model : str, sys_prompt : str, history_msg : list): @@ -73,16 +115,20 @@ def one_dialogue(content: str, client : OpenAI, model : str, sys_prompt : str, h def get_rsp(content : str, model : str, sys_prompt : str, history_msg : list): + normalized_model = _normalize_model_name(model) - if model == "OPENAI": + if normalized_model == "OPENAI": return get_openai_rsp(content, sys_prompt, history_msg) - if model == "QWEN": + if normalized_model == "QWEN": return get_qwen_rsp(content, sys_prompt, history_msg) - if model == "LLAMA": + if normalized_model == "LLAMA": return get_llama_rsp(content, sys_prompt, history_msg) - return None \ No newline at end of file + if normalized_model in {"ATLASCLOUD", "ATLAS_CLOUD", "ATLAS"}: + return get_atlascloud_rsp(content, sys_prompt, history_msg) + + return None diff --git a/Personality_test/llm_link.py b/Personality_test/llm_link.py index df467e4..c299601 100644 --- a/Personality_test/llm_link.py +++ b/Personality_test/llm_link.py @@ -3,6 +3,39 @@ import os +ATLAS_CLOUD_BASE_URL = "https://api.atlascloud.ai/v1" +ATLAS_CLOUD_DEFAULT_MODEL = "qwen/qwen3.5-flash" + + +def _get_env(*names: str, default=None): + for name in names: + value = os.getenv(name) + if value: + return value + return default + + +def _normalize_model_name(model: str) -> str: + return model.replace("-", "_").upper() + + +def _chat(content: str, api_key, base_url, model: str): + client = OpenAI( + api_key=api_key, + base_url=base_url, + ) + + completion = client.chat.completions.create( + model=model, + messages=[ + {"role": "user", "content": content} + ], + temperature=0.6, + ) + + return completion.choices[0].message.content + + def get_openai_rsp(content : str): load_dotenv("para.env") @@ -76,16 +109,39 @@ def get_llama_rsp(content :str): return rsp +def get_atlascloud_rsp(content : str): + load_dotenv("para.env") + api = _get_env("ATLASCLOUD_API_KEY", "ATLAS_CLOUD_API_KEY") + url = _get_env( + "ATLASCLOUD_API_BASE", + "ATLAS_CLOUD_API_BASE", + "ATLASCLOUD_BASE_URL", + "ATLAS_CLOUD_BASE_URL", + default=ATLAS_CLOUD_BASE_URL, + ) + model = _get_env( + "ATLASCLOUD_MODEL", + "ATLAS_CLOUD_MODEL", + default=ATLAS_CLOUD_DEFAULT_MODEL, + ) + + return _chat(content, api, url, model) + + def get_rsp(content : str, model : str): + normalized_model = _normalize_model_name(model) - if model == "OPENAI": + if normalized_model == "OPENAI": return get_openai_rsp(content) - if model == "QWEN": + if normalized_model == "QWEN": return get_qwen_rsp(content) - if model == "LLAMA": - return get_llama_rsp(content) \ No newline at end of file + if normalized_model == "LLAMA": + return get_llama_rsp(content) + + if normalized_model in {"ATLASCLOUD", "ATLAS_CLOUD", "ATLAS"}: + return get_atlascloud_rsp(content) diff --git a/para.env.example b/para.env.example index 6e3af0e..d32fe94 100644 --- a/para.env.example +++ b/para.env.example @@ -1,5 +1,5 @@ # LLM config -# options: "OPENAI", "QWEN", "LLAMA" +# options: "OPENAI", "QWEN", "LLAMA", "ATLASCLOUD" LLM_MODEL="QWEN" # OpenAI config @@ -15,4 +15,10 @@ QWEN_MODEL="qwen3-235b-a22b-instruct-2507" # Llama config LLAMA_API_KEY="your_llama_api_key_here" LLAMA_BASE_URL="your_llama_base_url_here" -LLAMA_MODEL="meta-llama/llama-4-maverick" \ No newline at end of file +LLAMA_MODEL="meta-llama/llama-4-maverick" + +# Atlas Cloud config (OpenAI-compatible) +ATLASCLOUD_API_KEY="your_atlascloud_api_key_here" +# ATLAS_CLOUD_API_KEY is also supported as an alias. +ATLASCLOUD_API_BASE="https://api.atlascloud.ai/v1" +ATLASCLOUD_MODEL="qwen/qwen3.5-flash" diff --git a/tests/test_atlascloud_provider.py b/tests/test_atlascloud_provider.py new file mode 100644 index 0000000..f052002 --- /dev/null +++ b/tests/test_atlascloud_provider.py @@ -0,0 +1,105 @@ +import importlib.util +import os +import sys +import types +import unittest +from pathlib import Path +from unittest.mock import patch + + +ROOT = Path(__file__).resolve().parents[1] + + +class FakeCompletions: + def __init__(self, client): + self.client = client + + def create(self, **kwargs): + self.client.requests.append(kwargs) + message = types.SimpleNamespace(content="atlas response") + choice = types.SimpleNamespace(message=message) + return types.SimpleNamespace(choices=[choice]) + + +class FakeChat: + def __init__(self, client): + self.completions = FakeCompletions(client) + + +class FakeOpenAI: + calls = [] + + def __init__(self, api_key=None, base_url=None): + self.api_key = api_key + self.base_url = base_url + self.requests = [] + self.chat = FakeChat(self) + FakeOpenAI.calls.append(self) + + +def load_llm_link(module_name: str, relative_path: str): + openai_module = types.ModuleType("openai") + openai_module.OpenAI = FakeOpenAI + dotenv_module = types.ModuleType("dotenv") + dotenv_module.load_dotenv = lambda *args, **kwargs: None + spec = importlib.util.spec_from_file_location(module_name, ROOT / relative_path) + + with patch.dict(sys.modules, {"openai": openai_module, "dotenv": dotenv_module}): + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + + return module + + +class AtlasCloudProviderTests(unittest.TestCase): + def setUp(self): + FakeOpenAI.calls.clear() + + def test_personality_test_routes_atlascloud_alias(self): + module = load_llm_link("personality_test_llm_link", "Personality_test/llm_link.py") + + with patch.dict( + os.environ, + { + "ATLAS_CLOUD_API_KEY": "atlas-key", + "ATLAS_CLOUD_API_BASE": "https://atlas.example/v1", + "ATLAS_CLOUD_MODEL": "deepseek-ai/deepseek-v4-pro", + }, + clear=True, + ): + response = module.get_rsp("hello", "atlas-cloud") + + self.assertEqual(response, "atlas response") + self.assertEqual(FakeOpenAI.calls[0].api_key, "atlas-key") + self.assertEqual(FakeOpenAI.calls[0].base_url, "https://atlas.example/v1") + self.assertEqual(FakeOpenAI.calls[0].requests[0]["model"], "deepseek-ai/deepseek-v4-pro") + self.assertEqual(FakeOpenAI.calls[0].requests[0]["messages"], [{"role": "user", "content": "hello"}]) + + def test_personality_changes_uses_atlascloud_defaults(self): + module = load_llm_link("personality_changes_llm_link", "Personality_changes/llm_link.py") + + with patch.dict(os.environ, {"ATLASCLOUD_API_KEY": "atlas-key"}, clear=True): + response = module.get_rsp( + "next question", + "ATLASCLOUD", + "system prompt", + [{"role": "assistant", "content": "previous answer"}], + ) + + self.assertEqual(response, "atlas response") + self.assertEqual(FakeOpenAI.calls[0].api_key, "atlas-key") + self.assertEqual(FakeOpenAI.calls[0].base_url, "https://api.atlascloud.ai/v1") + self.assertEqual(FakeOpenAI.calls[0].requests[0]["model"], "qwen/qwen3.5-flash") + self.assertEqual( + FakeOpenAI.calls[0].requests[0]["messages"], + [ + {"role": "system", "content": "system prompt"}, + {"role": "assistant", "content": "previous answer"}, + {"role": "user", "content": "next question"}, + ], + ) + + +if __name__ == "__main__": + unittest.main()