From 626de190cee949ef932a14ae231ddff6d0cc1aad Mon Sep 17 00:00:00 2001 From: binyangzhu000-sudo <224954946+binyangzhu000-sudo@users.noreply.github.com> Date: Wed, 22 Jul 2026 23:55:15 +0800 Subject: [PATCH] Add Atlas Cloud model aliases --- .env.example | 2 + quantmind/configs/__init__.py | 17 +++++- quantmind/configs/base.py | 81 +++++++++++++++++++++++++++++ quantmind/flows/_paper_summary.py | 6 +-- quantmind/flows/paper/_structure.py | 4 +- quantmind/mind/retrieval.py | 4 +- tests/configs/test_base.py | 73 +++++++++++++++++++++++++- tests/flows/test_structure.py | 33 +++++++++++- 8 files changed, 210 insertions(+), 10 deletions(-) diff --git a/.env.example b/.env.example index c10140c..0ebc2c3 100644 --- a/.env.example +++ b/.env.example @@ -4,6 +4,7 @@ # API Keys LLAMA_CLOUD_API_KEY=your_llama_cloud_api_key_here OPENAI_API_KEY=your_openai_api_key_here +ATLASCLOUD_API_KEY=your_atlascloud_api_key_here # QuantMind Settings QUANTMIND_LOG_LEVEL=INFO @@ -21,6 +22,7 @@ QUANTMIND_ARXIV_MAX_RESULTS=100 # Optional: Override specific service endpoints # LLAMA_CLOUD_BASE_URL=https://api.llamaindex.ai # OPENAI_BASE_URL=https://api.openai.com +# ATLASCLOUD_BASE_URL=https://api.atlascloud.ai/v1 # Development Settings # QUANTMIND_DEBUG=true diff --git a/quantmind/configs/__init__.py b/quantmind/configs/__init__.py index 3676f6f..b5a97f4 100644 --- a/quantmind/configs/__init__.py +++ b/quantmind/configs/__init__.py @@ -8,7 +8,16 @@ - the magic-input resolver (PR5) has one introspection target. """ -from quantmind.configs.base import BaseFlowCfg, BaseInput +from quantmind.configs.base import ( + ATLASCLOUD_BASE_URL, + ATLASCLOUD_DEFAULT_CHAT_MODEL, + ATLASCLOUD_DEFAULT_REASONING_MODEL, + BaseFlowCfg, + BaseInput, + atlascloud_model, + is_atlascloud_model, + resolve_agent_model, +) from quantmind.configs.earnings import EarningsFlowCfg, EarningsInput from quantmind.configs.news import NewsCollectionCfg, NewsWindow from quantmind.configs.paper import PaperFlowCfg, PaperInput @@ -18,6 +27,9 @@ __all__ = [ "BaseFlowCfg", "BaseInput", + "ATLASCLOUD_BASE_URL", + "ATLASCLOUD_DEFAULT_CHAT_MODEL", + "ATLASCLOUD_DEFAULT_REASONING_MODEL", "EarningsFlowCfg", "EarningsInput", "NewsCollectionCfg", @@ -26,4 +38,7 @@ "PaperInput", "PaperStructureCfg", "RetrievalCfg", + "atlascloud_model", + "is_atlascloud_model", + "resolve_agent_model", ] diff --git a/quantmind/configs/base.py b/quantmind/configs/base.py index 05b5ef5..a8251d3 100644 --- a/quantmind/configs/base.py +++ b/quantmind/configs/base.py @@ -6,9 +6,22 @@ discriminated-union member; subclasses set a `Literal` discriminator field. """ +from __future__ import annotations + +import os +from typing import TYPE_CHECKING + from agents import ModelSettings from pydantic import BaseModel, ConfigDict +if TYPE_CHECKING: + from agents.models.interface import Model + +ATLASCLOUD_BASE_URL = "https://api.atlascloud.ai/v1" +ATLASCLOUD_DEFAULT_CHAT_MODEL = "qwen/qwen3.5-flash" +ATLASCLOUD_DEFAULT_REASONING_MODEL = "deepseek-ai/deepseek-v4-pro" +ATLASCLOUD_MODEL_PREFIXES = ("atlascloud/", "atlas-cloud/", "atlas/") + class BaseFlowCfg(BaseModel): """Base configuration shared by all flows.""" @@ -47,3 +60,71 @@ class BaseInput(BaseModel): """Parent of every flow's discriminated-union input member.""" model_config = ConfigDict(extra="forbid") + + +def atlascloud_model(model: str = ATLASCLOUD_DEFAULT_CHAT_MODEL) -> str: + """Return a QuantMind model value routed through Atlas Cloud. + + The returned value can be used anywhere ``BaseFlowCfg.model`` is accepted. + Runtime credentials are read from ``ATLASCLOUD_API_KEY`` (or + ``ATLAS_CLOUD_API_KEY``) when the flow actually builds its SDK agent. + """ + value = model.strip() + if not value: + raise ValueError("atlascloud model name must not be blank") + if is_atlascloud_model(value): + return value + return f"atlascloud/{value}" + + +def is_atlascloud_model(model: str) -> bool: + """Whether a model string uses one of QuantMind's Atlas Cloud aliases.""" + return model.startswith(ATLASCLOUD_MODEL_PREFIXES) + + +def resolve_agent_model(model: str) -> str | Model: + """Resolve a config model string into an Agents SDK model. + + Non-Atlas values are returned unchanged. ``atlascloud/...`` aliases are + backed by the Agents SDK LiteLLM adapter with Atlas Cloud's + OpenAI-compatible endpoint and API key environment variables. + """ + if not is_atlascloud_model(model): + return model + + from agents.extensions.models.litellm_model import LitellmModel + + return LitellmModel( + model=f"openai/{_strip_atlascloud_prefix(model)}", + base_url=_atlascloud_base_url(), + api_key=_atlascloud_api_key(), + ) + + +def _strip_atlascloud_prefix(model: str) -> str: + for prefix in ATLASCLOUD_MODEL_PREFIXES: + if model.startswith(prefix): + value = model[len(prefix) :].strip() + if not value: + raise ValueError("atlascloud model name must not be blank") + return value + raise ValueError("model is not an Atlas Cloud alias") + + +def _atlascloud_api_key() -> str: + value = os.getenv("ATLASCLOUD_API_KEY") or os.getenv("ATLAS_CLOUD_API_KEY") + if not value: + raise ValueError( + "Set ATLASCLOUD_API_KEY before using atlascloud/... models" + ) + return value + + +def _atlascloud_base_url() -> str: + return ( + os.getenv("ATLASCLOUD_API_BASE") + or os.getenv("ATLASCLOUD_BASE_URL") + or os.getenv("ATLAS_CLOUD_API_BASE") + or os.getenv("ATLAS_CLOUD_BASE_URL") + or ATLASCLOUD_BASE_URL + ).rstrip("/") diff --git a/quantmind/flows/_paper_summary.py b/quantmind/flows/_paper_summary.py index c7738bd..f011958 100644 --- a/quantmind/flows/_paper_summary.py +++ b/quantmind/flows/_paper_summary.py @@ -24,7 +24,7 @@ from agents import Agent, ModelSettings from pydantic import BaseModel, ConfigDict, Field, field_validator -from quantmind.configs import PaperFlowCfg +from quantmind.configs import PaperFlowCfg, resolve_agent_model from quantmind.flows._runner import run_with_observability from quantmind.knowledge import PaperChunkSet, PaperSourceRevision @@ -293,7 +293,7 @@ async def summarize( researcher: Agent[Any] = Agent( name="paper_chunk_researcher", instructions=_RESEARCH_INSTRUCTIONS, - model=cfg.model, + model=resolve_agent_model(cfg.model), model_settings=model_settings, output_type=PaperResearchDraft, ) @@ -317,7 +317,7 @@ async def study(group: _ChunkGroup) -> PaperResearchDraft: reducer: Agent[Any] = Agent( name="paper_summary_reducer", instructions=_summary_instructions(cfg), - model=cfg.model, + model=resolve_agent_model(cfg.model), model_settings=model_settings, output_type=PaperSummaryDraft, ) diff --git a/quantmind/flows/paper/_structure.py b/quantmind/flows/paper/_structure.py index 23a3c8b..f110775 100644 --- a/quantmind/flows/paper/_structure.py +++ b/quantmind/flows/paper/_structure.py @@ -8,7 +8,7 @@ from agents import Agent, ModelSettings -from quantmind.configs import PaperStructureCfg +from quantmind.configs import PaperStructureCfg, resolve_agent_model from quantmind.flows._runner import run_with_observability from quantmind.knowledge import PaperSourceRevision, PaperStructureTreeDraft from quantmind.preprocess import OutlineSignals @@ -131,7 +131,7 @@ def build_agent(json_object: bool) -> Agent[Any]: model_settings = _structure_model_settings(cfg) kwargs: dict[str, Any] = { "name": "paper_structure_builder", - "model": cfg.model, + "model": resolve_agent_model(cfg.model), } if json_object: kwargs["instructions"] = json_object_instructions( diff --git a/quantmind/mind/retrieval.py b/quantmind/mind/retrieval.py index a6886cc..5e4ecd6 100644 --- a/quantmind/mind/retrieval.py +++ b/quantmind/mind/retrieval.py @@ -28,7 +28,7 @@ from agents.exceptions import MaxTurnsExceeded from pydantic import BaseModel, ConfigDict, Field -from quantmind.configs import RetrievalCfg +from quantmind.configs import RetrievalCfg, resolve_agent_model from quantmind.knowledge import ( ArtifactLocator, Citation, @@ -344,7 +344,7 @@ def build_agent(json_object: bool) -> Agent[Any]: model_settings = cfg.model_settings or ModelSettings() kwargs: dict[str, Any] = { "name": "structure_agentic_retriever", - "model": cfg.model, + "model": resolve_agent_model(cfg.model), "tools": [get_document_structure, get_node_content], } if json_object: diff --git a/tests/configs/test_base.py b/tests/configs/test_base.py index 164e237..6f0e9b6 100644 --- a/tests/configs/test_base.py +++ b/tests/configs/test_base.py @@ -1,11 +1,23 @@ """Tests for configs.base.""" +import os import unittest +from unittest.mock import patch from agents import ModelSettings +from agents.extensions.models.litellm_model import LitellmModel from pydantic import ValidationError -from quantmind.configs.base import BaseFlowCfg, BaseInput +from quantmind.configs.base import ( + ATLASCLOUD_BASE_URL, + ATLASCLOUD_DEFAULT_CHAT_MODEL, + ATLASCLOUD_DEFAULT_REASONING_MODEL, + BaseFlowCfg, + BaseInput, + atlascloud_model, + is_atlascloud_model, + resolve_agent_model, +) class BaseFlowCfgTests(unittest.TestCase): @@ -40,8 +52,66 @@ class _SampleInput(BaseInput): _SampleInput(x=1, y=2) # type: ignore[call-arg] +class AtlasCloudModelTests(unittest.TestCase): + def test_atlascloud_model_defaults_and_preserves_aliases(self) -> None: + self.assertEqual( + atlascloud_model(), + f"atlascloud/{ATLASCLOUD_DEFAULT_CHAT_MODEL}", + ) + self.assertEqual( + atlascloud_model(ATLASCLOUD_DEFAULT_REASONING_MODEL), + f"atlascloud/{ATLASCLOUD_DEFAULT_REASONING_MODEL}", + ) + self.assertEqual( + atlascloud_model("atlas-cloud/qwen/qwen3.5-flash"), + "atlas-cloud/qwen/qwen3.5-flash", + ) + self.assertTrue(is_atlascloud_model("atlas/qwen/qwen3.5-flash")) + + def test_atlascloud_model_rejects_blank_values(self) -> None: + with self.assertRaisesRegex(ValueError, "must not be blank"): + atlascloud_model(" ") + + def test_resolve_agent_model_uses_atlascloud_environment(self) -> None: + env = { + "ATLASCLOUD_API_KEY": "atlas-test-key", + "ATLASCLOUD_BASE_URL": "https://atlas.test/v1/", + } + with patch.dict(os.environ, env, clear=False): + model = resolve_agent_model("atlascloud/qwen/qwen3.5-flash") + + self.assertIsInstance(model, LitellmModel) + assert isinstance(model, LitellmModel) + self.assertEqual(model.model, "openai/qwen/qwen3.5-flash") + self.assertEqual(model.api_key, "atlas-test-key") + self.assertEqual(model.base_url, "https://atlas.test/v1") + + def test_resolve_agent_model_requires_api_key(self) -> None: + with patch.dict( + os.environ, + { + "ATLASCLOUD_API_KEY": "", + "ATLAS_CLOUD_API_KEY": "", + }, + clear=False, + ): + with self.assertRaisesRegex(ValueError, "ATLASCLOUD_API_KEY"): + resolve_agent_model("atlascloud/qwen/qwen3.5-flash") + + def test_resolve_agent_model_leaves_non_atlas_values_unchanged( + self, + ) -> None: + self.assertEqual( + resolve_agent_model("litellm/anthropic/claude-test"), + "litellm/anthropic/claude-test", + ) + + class PackageExportTests(unittest.TestCase): def test_top_level_imports(self): + from quantmind.configs import ( + ATLASCLOUD_BASE_URL as ATLASCLOUD_BASE_URL_EXPORT, + ) from quantmind.configs import ( BaseFlowCfg as BaseFlowCfgExport, ) @@ -58,6 +128,7 @@ def test_top_level_imports(self): self.assertTrue(issubclass(NewsCollectionCfg, BaseFlowCfgExport)) self.assertTrue(issubclass(EarningsFlowCfg, BaseFlowCfgExport)) self.assertEqual(BaseInputExport.__name__, "BaseInput") + self.assertEqual(ATLASCLOUD_BASE_URL_EXPORT, ATLASCLOUD_BASE_URL) if __name__ == "__main__": diff --git a/tests/flows/test_structure.py b/tests/flows/test_structure.py index 9dc9173..17429bb 100644 --- a/tests/flows/test_structure.py +++ b/tests/flows/test_structure.py @@ -1,15 +1,17 @@ """Offline tests for the config-bound ``PaperFlow.build`` structure shape.""" import asyncio +import os import unittest from pathlib import Path from unittest.mock import AsyncMock, patch import httpx from agents import ModelSettings +from agents.extensions.models.litellm_model import LitellmModel from openai import BadRequestError -from quantmind.configs import PaperFlowCfg, PaperStructureCfg +from quantmind.configs import PaperFlowCfg, PaperStructureCfg, atlascloud_model from quantmind.configs.paper import LocalFilePath from quantmind.flows import PaperFlow, PaperStructureError from quantmind.flows.paper._structure import ( @@ -194,6 +196,35 @@ async def test_strict_structured_output_is_the_default_for_every_model( self.assertIs(agent.output_type, PaperStructureTreeDraft) self.assertIsNone(agent.model_settings.extra_body) + async def test_atlascloud_model_alias_uses_openai_compatible_adapter( + self, + ) -> None: + result = build_paper_result() + cfg = PaperStructureCfg(model=atlascloud_model()) + run_mock = AsyncMock(return_value=_fixture_draft()) + with patch.dict( + os.environ, + {"ATLASCLOUD_API_KEY": "atlas-test-key"}, + clear=False, + ): + with patch( + "quantmind.flows.paper._structure.run_with_observability", + new=run_mock, + ): + draft = await _AgentsPaperStructureProvider().structure( + signals=_empty_signals(), + source=result.source_revision, + cfg=cfg, + ) + + self.assertEqual(draft, _fixture_draft()) + agent = run_mock.await_args.args[0] + self.assertIsInstance(agent.model, LitellmModel) + assert isinstance(agent.model, LitellmModel) + self.assertEqual(agent.model.model, "openai/qwen/qwen3.5-flash") + self.assertEqual(agent.model.base_url, "https://api.atlascloud.ai/v1") + self.assertEqual(agent.model.api_key, "atlas-test-key") + async def test_falls_back_to_json_object_when_json_schema_is_rejected( self, ) -> None: