diff --git a/archytas/agent.py b/archytas/agent.py index e4c3443..12ed656 100644 --- a/archytas/agent.py +++ b/archytas/agent.py @@ -180,6 +180,17 @@ def __init__( # make_tool_dict runs). Map of tool_name -> tool callable. self.statetools: dict[str, Callable] = {} + def set_model(self, model: BaseArchytasModel) -> None: + """Swap the active model at runtime, keeping the chat history in sync. + + The chat history keeps its own reference to the model (used for token + budgeting and the model metadata it serializes), so assigning + ``agent.model`` alone leaves that reference stale. Always route runtime + model changes through this method so both stay consistent. + """ + self.model = model + self.chat_history.set_model(model) + def get_prompt_sections(self) -> list[PromptSection]: """Return the ordered list of ``PromptSection``s composing the system prompt. diff --git a/archytas/models/azure.py b/archytas/models/azure.py index e7f29f3..334a5c4 100644 --- a/archytas/models/azure.py +++ b/archytas/models/azure.py @@ -18,12 +18,14 @@ def auth(self, **kwargs) -> None: else: auth_token = self.config.api_key if not auth_token: - auth_token = DEFERRED_TOKEN_VALUE - set_env_auth(AZURE_OPENAI_API_KEY=auth_token) - - # Replace local auth token from value from environment variables to allow fetching preset auth variables in the - # environment. - auth_token = os.environ.get('AZURE_OPENAI_API_KEY', DEFERRED_TOKEN_VALUE) + # No key configured: defer to the environment (e.g. a user-set + # AZURE_OPENAI_API_KEY), without clobbering it. + set_env_auth(AZURE_OPENAI_API_KEY=DEFERRED_TOKEN_VALUE) + auth_token = os.environ.get('AZURE_OPENAI_API_KEY', DEFERRED_TOKEN_VALUE) + else: + # An explicit key wins over the environment -- including a value an + # earlier model instance wrote there -- so a new key takes effect now. + os.environ['AZURE_OPENAI_API_KEY'] = auth_token if auth_token != DEFERRED_TOKEN_VALUE: self.config.api_key = auth_token diff --git a/archytas/models/openai.py b/archytas/models/openai.py index 6a663d7..43887e0 100644 --- a/archytas/models/openai.py +++ b/archytas/models/openai.py @@ -51,12 +51,14 @@ def auth(self, **kwargs) -> None: else: auth_token = self.config.api_key if not auth_token: - auth_token = DEFERRED_TOKEN_VALUE - set_env_auth(OPENAI_API_KEY=auth_token) - - # Replace local auth token from value from environment variables to allow fetching preset auth variables in the - # environment. - auth_token = os.environ.get('OPENAI_API_KEY', DEFERRED_TOKEN_VALUE) + # No key configured: defer to the environment (e.g. a user-set + # OPENAI_API_KEY), without clobbering it. + set_env_auth(OPENAI_API_KEY=DEFERRED_TOKEN_VALUE) + auth_token = os.environ.get('OPENAI_API_KEY', DEFERRED_TOKEN_VALUE) + else: + # An explicit key wins over the environment -- including a value an + # earlier model instance wrote there -- so a new key takes effect now. + os.environ['OPENAI_API_KEY'] = auth_token if auth_token != DEFERRED_TOKEN_VALUE: self.config.api_key = auth_token diff --git a/tests/test_agent_set_model.py b/tests/test_agent_set_model.py new file mode 100644 index 0000000..9f3dd81 --- /dev/null +++ b/tests/test_agent_set_model.py @@ -0,0 +1,22 @@ +"""Regression test for Agent.set_model keeping the chat history in sync. + +The chat history holds its own model reference (used for token budgeting and +the model metadata it serializes for UIs). Assigning ``agent.model`` alone +leaves that reference stale; ``Agent.set_model`` must update both. +""" +from archytas.agent import Agent +from archytas.models.openai import OpenAIModel + + +def test_set_model_updates_agent_and_chat_history(monkeypatch): + monkeypatch.setenv("OPENAI_API_KEY", "sk-dummy") + first = OpenAIModel({"api_key": "sk-a", "model_name": "gpt-4o-mini"}) + second = OpenAIModel({"api_key": "sk-b", "model_name": "gpt-4o"}) + + agent = Agent(model=first, spinner=None, rich_print=False) + assert agent.model is first + assert agent.chat_history.model is first + + agent.set_model(second) + assert agent.model is second + assert agent.chat_history.model is second diff --git a/tests/test_auth_env_precedence.py b/tests/test_auth_env_precedence.py new file mode 100644 index 0000000..86b1343 --- /dev/null +++ b/tests/test_auth_env_precedence.py @@ -0,0 +1,31 @@ +"""Regression tests for API-key precedence in OpenAI/Azure model auth. + +These guard the subtle bug where an explicitly-configured key was shadowed by +a stale value left in ``os.environ`` (written there by an earlier model +instance via ``set_env_auth``'s ``setdefault``), so a freshly-entered key only +took effect after a process restart. + +No network calls: constructing the model is enough to resolve the effective key. +""" +import pytest + +from archytas.models.openai import OpenAIModel + + +def _client_key(model) -> str: + secret = model.model.openai_api_key + return secret.get_secret_value() if hasattr(secret, "get_secret_value") else str(secret) + + +def test_explicit_key_overrides_stale_env(monkeypatch): + # Simulate a stale/bad key left in the environment by an earlier build. + monkeypatch.setenv("OPENAI_API_KEY", "sk-STALE-must-not-win") + model = OpenAIModel({"api_key": "sk-explicit-config-key", "model_name": "gpt-4o-mini"}) + assert _client_key(model) == "sk-explicit-config-key" + + +def test_defers_to_env_when_no_key_configured(monkeypatch): + # With no key configured, a user-provided environment variable should win. + monkeypatch.setenv("OPENAI_API_KEY", "sk-from-user-environment") + model = OpenAIModel({"model_name": "gpt-4o-mini"}) + assert _client_key(model) == "sk-from-user-environment"