From e9f15a2678391dc0ab5621f746ff7b2258227551 Mon Sep 17 00:00:00 2001 From: daavoo Date: Fri, 13 Feb 2026 16:45:23 +0100 Subject: [PATCH 1/2] feat: Add `any-llm-args` to config. --- src/any_agent/config.py | 7 +++++ src/any_agent/frameworks/llama_index.py | 25 ++++++++++++------ src/any_agent/frameworks/openai.py | 22 +++++++++++----- src/any_agent/frameworks/tinyagent.py | 2 +- tests/unit/frameworks/test_llama_index.py | 31 +++++++++++++++++++++++ tests/unit/frameworks/test_openai.py | 25 ++++++++++++++++++ tests/unit/frameworks/test_tinyagent.py | 15 +++++++++++ 7 files changed, 112 insertions(+), 15 deletions(-) diff --git a/src/any_agent/config.py b/src/any_agent/config.py index 4b346c23..48962012 100644 --- a/src/any_agent/config.py +++ b/src/any_agent/config.py @@ -236,6 +236,13 @@ class CalendarEvent(BaseModel): Refer to [any-llm Completion API Docs](https://mozilla-ai.github.io/any-llm/api/completion/) for more info. """ + any_llm_args: MutableMapping[str, Any] | None = None + """Pass arguments to `AnyLLM.create()` when using integrations backed by any-llm. + + Use this for provider/client initialization options that are not completion-time + generation params (which should be passed via `model_args`). + """ + output_type: type[BaseModel] | None = None """Control the output schema from calling `run`. By default, the agent will return a type str. diff --git a/src/any_agent/frameworks/llama_index.py b/src/any_agent/frameworks/llama_index.py index 7b24c577..3617145d 100644 --- a/src/any_agent/frameworks/llama_index.py +++ b/src/any_agent/frameworks/llama_index.py @@ -102,6 +102,7 @@ def __init__( max_retries: int = 10, api_key: str | None = None, api_base: str | None = None, + any_llm_args: dict[str, Any] | None = None, **kwargs: Any, ) -> None: additional_kwargs = additional_kwargs or {} @@ -116,10 +117,14 @@ def __init__( ) self._parse_model(model) + llm_create_kwargs: dict[str, Any] = dict(any_llm_args or {}) + if api_key is not None: + llm_create_kwargs["api_key"] = api_key + if api_base is not None: + llm_create_kwargs["api_base"] = api_base self._client = AnyLLM.create( provider=self._provider, - api_key=api_key, - api_base=api_base, + **llm_create_kwargs, ) def _parse_model(self, model: str) -> None: @@ -512,14 +517,18 @@ def _get_model(self, agent_config: AgentConfig) -> "LLM": model_id = agent_config.model_id + model_kwargs: dict[str, Any] = { + "model": model_id, + "api_key": agent_config.api_key, + "api_base": agent_config.api_base, + "additional_kwargs": additional_kwargs, + } + if model_type is DEFAULT_MODEL_TYPE and agent_config.any_llm_args is not None: + model_kwargs["any_llm_args"] = agent_config.any_llm_args + return cast( "LLM", - model_type( - model=model_id, - api_key=agent_config.api_key, - api_base=agent_config.api_base, - additional_kwargs=additional_kwargs, # type: ignore[arg-type] - ), + model_type(**model_kwargs), # type: ignore[arg-type] ) async def _load_agent(self) -> None: diff --git a/src/any_agent/frameworks/openai.py b/src/any_agent/frameworks/openai.py index 3aefcc18..beab72f8 100644 --- a/src/any_agent/frameworks/openai.py +++ b/src/any_agent/frameworks/openai.py @@ -80,12 +80,19 @@ def __init__( model: str, base_url: str | None = None, api_key: str | None = None, + any_llm_args: dict[str, Any] | None = None, ): provider, model_id = AnyLLM.split_model_provider(model) self.model = model self.base_url = base_url self.api_key = api_key - self.llm = AnyLLM.create(provider=provider, api_key=api_key, api_base=base_url) + llm_create_kwargs: dict[str, Any] = dict(any_llm_args or {}) + if api_key is not None: + llm_create_kwargs["api_key"] = api_key + if base_url is not None: + llm_create_kwargs["api_base"] = base_url + + self.llm = AnyLLM.create(provider=provider, **llm_create_kwargs) self.model_id = model_id async def get_response( @@ -399,11 +406,14 @@ def _get_model( base_url = agent_config.api_base or cast( "str | None", model_args.get("api_base") ) - return model_type( - model=agent_config.model_id, - base_url=base_url, - api_key=agent_config.api_key, - ) + model_kwargs: dict[str, Any] = { + "model": agent_config.model_id, + "base_url": base_url, + "api_key": agent_config.api_key, + } + if model_type is DEFAULT_MODEL_TYPE and agent_config.any_llm_args is not None: + model_kwargs["any_llm_args"] = agent_config.any_llm_args + return model_type(**model_kwargs) async def _load_agent(self) -> None: """Load the OpenAI agent with the given configuration.""" diff --git a/src/any_agent/frameworks/tinyagent.py b/src/any_agent/frameworks/tinyagent.py index c45f6dc9..a7b255b7 100644 --- a/src/any_agent/frameworks/tinyagent.py +++ b/src/any_agent/frameworks/tinyagent.py @@ -101,7 +101,7 @@ def __init__(self, config: AgentConfig) -> None: self.uses_openai = provider_name == LLMProvider.OPENAI # Create the LLM instance using the AnyLLM class pattern - llm_kwargs: dict[str, Any] = {} + llm_kwargs: dict[str, Any] = dict(self.config.any_llm_args or {}) if self.config.api_key: llm_kwargs["api_key"] = self.config.api_key if self.config.api_base: diff --git a/tests/unit/frameworks/test_llama_index.py b/tests/unit/frameworks/test_llama_index.py index a4e0e41f..1f3065de 100644 --- a/tests/unit/frameworks/test_llama_index.py +++ b/tests/unit/frameworks/test_llama_index.py @@ -70,3 +70,34 @@ def test_run_llama_index_agent_custom_args() -> None: ) agent.run("foo", timeout=10) agent_mock.run.assert_called_once_with("foo", timeout=10) + + +def test_load_llama_index_agent_forwards_any_llm_args() -> None: + model_mock = MagicMock() + create_mock = MagicMock() + create_mock.return_value = MagicMock() + any_llm_args = {"timeout": 17, "max_retries": 3} + + from llama_index.core.tools import FunctionTool + + with ( + patch("any_agent.frameworks.llama_index.DEFAULT_AGENT_TYPE", create_mock), + patch("any_agent.frameworks.llama_index.DEFAULT_MODEL_TYPE", model_mock), + patch.object(FunctionTool, "from_defaults"), + ): + AnyAgent.create( + AgentFramework.LLAMA_INDEX, + AgentConfig( + model_id="gemini/gemini-2.0-flash", + instructions="You are a helpful assistant", + any_llm_args=any_llm_args, + ), + ) + + model_mock.assert_called_once_with( + model="gemini/gemini-2.0-flash", + api_key=None, + api_base=None, + additional_kwargs={}, + any_llm_args=any_llm_args, + ) diff --git a/tests/unit/frameworks/test_openai.py b/tests/unit/frameworks/test_openai.py index b2b605b9..267dc1df 100644 --- a/tests/unit/frameworks/test_openai.py +++ b/tests/unit/frameworks/test_openai.py @@ -84,6 +84,31 @@ def test_openai_with_api_key() -> None: ) +def test_openai_forwards_any_llm_args() -> None: + mock_agent = MagicMock() + mock_model = MagicMock() + any_llm_args = {"timeout": 42, "headers": {"x-test": "1"}} + + with ( + patch("any_agent.frameworks.openai.Agent", mock_agent), + patch("any_agent.frameworks.openai.DEFAULT_MODEL_TYPE", mock_model), + ): + AnyAgent.create( + AgentFramework.OPENAI, + AgentConfig( + model_id="mistral:mistral-small-latest", + any_llm_args=any_llm_args, + ), + ) + + mock_model.assert_called_once_with( + model="mistral:mistral-small-latest", + base_url=None, + api_key=None, + any_llm_args=any_llm_args, + ) + + def test_load_openai_with_mcp_server() -> None: mock_agent = MagicMock() mock_function_tool = MagicMock() diff --git a/tests/unit/frameworks/test_tinyagent.py b/tests/unit/frameworks/test_tinyagent.py index f02639f9..8cf49ef9 100644 --- a/tests/unit/frameworks/test_tinyagent.py +++ b/tests/unit/frameworks/test_tinyagent.py @@ -268,6 +268,21 @@ def test_uses_openai_handles_gateway_provider( assert agent.uses_openai is expected_uses_openai +def test_tinyagent_forwards_any_llm_args_to_anyllm_create() -> None: + any_llm_args = {"timeout": 99, "organization": "test-org"} + provider, _ = AnyLLM.split_model_provider(DEFAULT_SMALL_MODEL_ID) + + with patch("any_agent.frameworks.tinyagent.AnyLLM.create") as mock_create: + TinyAgent( + AgentConfig( + model_id=DEFAULT_SMALL_MODEL_ID, + any_llm_args=any_llm_args, + ) + ) + + mock_create.assert_called_once_with(provider, **any_llm_args) + + @pytest.mark.asyncio async def test_tool_result_appended_when_tool_not_found() -> None: """Test that tool_result message is appended when a tool is not found. From d4afea43613ac7ce2ea46cb6b85a732b7f7b105f Mon Sep 17 00:00:00 2001 From: daavoo Date: Fri, 13 Feb 2026 16:52:05 +0100 Subject: [PATCH 2/2] fix lint --- src/any_agent/frameworks/llama_index.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/any_agent/frameworks/llama_index.py b/src/any_agent/frameworks/llama_index.py index 3617145d..144c5063 100644 --- a/src/any_agent/frameworks/llama_index.py +++ b/src/any_agent/frameworks/llama_index.py @@ -528,7 +528,7 @@ def _get_model(self, agent_config: AgentConfig) -> "LLM": return cast( "LLM", - model_type(**model_kwargs), # type: ignore[arg-type] + model_type(**model_kwargs), ) async def _load_agent(self) -> None: