From 2a78cd0a1e54dd10267c628b3e47ba9d4565794f Mon Sep 17 00:00:00 2001 From: aiapienthusiast <327629069+aiapienthusiast@users.noreply.github.com> Date: Thu, 1 Oct 2026 20:00:53 +0400 Subject: [PATCH 1/3] feat: add Cheaper Inference as a provider option --- README.md | 2 +- .../4_basic_chatbot_different_providers.py | 12 +++++++++++- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 270f5588..3a2de045 100644 --- a/README.md +++ b/README.md @@ -418,7 +418,7 @@ uv sync --all-packages ## Provider & Model Compatibility -Atomic Agents depends on the [Instructor](https://github.com/jxnl/instructor) package. This means that in all examples where OpenAI is used, any other API supported by Instructor can also be used, such as Ollama, Groq, Mistral, Cohere, Anthropic, Gemini, OpenAI and any OpenAI-compatible API such as OpenRouter, MiniMax, EdenAI, and more. For a complete list, please refer to the Instructor documentation on its [GitHub page](https://github.com/jxnl/instructor). +Atomic Agents depends on the [Instructor](https://github.com/jxnl/instructor) package. This means that in all examples where OpenAI is used, any other API supported by Instructor can also be used, such as Ollama, Groq, Mistral, Cohere, Anthropic, Gemini, OpenAI and any OpenAI-compatible API such as OpenRouter, MiniMax, EdenAI, Cheaper Inference, and more. For a complete list, please refer to the Instructor documentation on its [GitHub page](https://github.com/jxnl/instructor). ## Support diff --git a/atomic-examples/quickstart/quickstart/4_basic_chatbot_different_providers.py b/atomic-examples/quickstart/quickstart/4_basic_chatbot_different_providers.py index 7f8d10ed..6b57326d 100644 --- a/atomic-examples/quickstart/quickstart/4_basic_chatbot_different_providers.py +++ b/atomic-examples/quickstart/quickstart/4_basic_chatbot_different_providers.py @@ -90,6 +90,16 @@ def setup_client(provider): model = "openai/gpt-4o-mini" model_api_parameters = {"max_tokens": 2048} assistant_role = "assistant" + elif provider == "9" or provider == "cheaperinference": + from openai import OpenAI as CheaperInferenceClient + + api_key = os.getenv("CHEAPER_INFERENCE_API_KEY") + client = instructor.from_openai( + CheaperInferenceClient(base_url="https://api.cheaperinference.com/v1", api_key=api_key) + ) + model = "gpt-5.4-mini" + model_api_parameters = {"max_tokens": 2048} + assistant_role = "assistant" else: raise ValueError(f"Unsupported provider: {provider}") @@ -97,7 +107,7 @@ def setup_client(provider): # Prompt the user to choose a provider from one in the list below. -providers_list = ["openai", "anthropic", "groq", "ollama", "gemini", "openrouter", "minimax", "edenai"] +providers_list = ["openai", "anthropic", "groq", "ollama", "gemini", "openrouter", "minimax", "edenai", "cheaperinference"] y = "bold yellow" b = "bold blue" g = "bold green" From 2f28aa636d6259745b5a935294f9fbe2cdda4e68 Mon Sep 17 00:00:00 2001 From: aiapienthusiast <327629069+aiapienthusiast@users.noreply.github.com> Date: Mon, 5 Oct 2026 18:37:16 +0400 Subject: [PATCH 2/3] fix: validate Cheaper Inference key, split provider setup, add mocked tests --- .../agents/test_cheaperinference_provider.py | 260 ++++++++++++++ .../4_basic_chatbot_different_providers.py | 332 ++++++++++-------- 2 files changed, 443 insertions(+), 149 deletions(-) create mode 100644 atomic-agents/tests/agents/test_cheaperinference_provider.py diff --git a/atomic-agents/tests/agents/test_cheaperinference_provider.py b/atomic-agents/tests/agents/test_cheaperinference_provider.py new file mode 100644 index 00000000..d3fea4c2 --- /dev/null +++ b/atomic-agents/tests/agents/test_cheaperinference_provider.py @@ -0,0 +1,260 @@ +"""Tests for the Cheaper Inference branch of the multi-provider quickstart example. + +All HTTP traffic is mocked. Real network transports raise if any test reaches them. +""" + +import importlib.util +import json +import os + +import httpx +import instructor +import openai +import pytest + +from atomic_agents import AtomicAgent, AgentConfig, BasicChatInputSchema, BasicChatOutputSchema + +EXAMPLE_PATH = os.path.join( + os.path.dirname(__file__), + "..", + "..", + "..", + "atomic-examples", + "quickstart", + "quickstart", + "4_basic_chatbot_different_providers.py", +) + +VENDOR_KEY = "ci-test-vendor-key" +UNRELATED_OPENAI_KEY = "sk-unrelated-dummy-openai-key" +BASE_URL = "https://api.cheaperinference.com/v1" + + +def _load_example(): + spec = importlib.util.spec_from_file_location("quickstart_providers_example", EXAMPLE_PATH) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@pytest.fixture +def example(): + return _load_example() + + +@pytest.fixture(autouse=True) +def block_network(monkeypatch): + """Fail any request that reaches a real transport.""" + sent = [] + + def _blocked(self, request): + sent.append(request) + raise AssertionError(f"Unexpected network request: {request.method} {request.url}") + + monkeypatch.setattr(httpx.HTTPTransport, "handle_request", _blocked) + return sent + + +@pytest.fixture +def unrelated_openai_key(monkeypatch): + monkeypatch.setenv("OPENAI_API_KEY", UNRELATED_OPENAI_KEY) + monkeypatch.delenv("CHEAPER_INFERENCE_API_KEY", raising=False) + + +class TestSelector: + def test_cheaperinference_is_menu_option_9(self, example): + assert example.providers_list[8] == "cheaperinference" + assert example.PROVIDER_NUMBERS["9"] == "cheaperinference" + + def test_existing_menu_numbers_unchanged(self, example): + assert example.providers_list[:8] == [ + "openai", + "anthropic", + "groq", + "ollama", + "gemini", + "openrouter", + "minimax", + "edenai", + ] + + @pytest.mark.parametrize("choice", ["9", "cheaperinference"]) + def test_selector_dispatches_to_cheaperinference(self, example, monkeypatch, choice): + sentinel = ("client", "model", {}, "assistant") + monkeypatch.setitem(example.PROVIDER_SETUPS, "cheaperinference", lambda: sentinel) + assert example.setup_client(choice) is sentinel + + @pytest.mark.parametrize("choice", ["0", "10", "cheaper", ""]) + def test_unsupported_provider_raises(self, example, choice): + with pytest.raises(ValueError, match="Unsupported provider"): + example.setup_client(choice) + + +class TestCredentials: + @pytest.mark.parametrize("vendor_value", [None, "", " "]) + @pytest.mark.parametrize("choice", ["9", "cheaperinference"]) + def test_missing_key_never_builds_client(self, example, monkeypatch, block_network, vendor_value, choice): + monkeypatch.setenv("OPENAI_API_KEY", UNRELATED_OPENAI_KEY) + if vendor_value is None: + monkeypatch.delenv("CHEAPER_INFERENCE_API_KEY", raising=False) + else: + monkeypatch.setenv("CHEAPER_INFERENCE_API_KEY", vendor_value) + + constructed = [] + monkeypatch.setattr(openai, "OpenAI", lambda *a, **kw: constructed.append(kw)) + sends = [] + monkeypatch.setattr(httpx.Client, "send", lambda self, *a, **kw: sends.append(a)) + + with pytest.raises(ValueError, match="CHEAPER_INFERENCE_API_KEY") as excinfo: + example.setup_client(choice) + + assert constructed == [] + assert sends == [] + assert block_network == [] + assert UNRELATED_OPENAI_KEY not in str(excinfo.value) + + def test_vendor_key_is_passed_explicitly(self, example, monkeypatch, unrelated_openai_key): + monkeypatch.setenv("CHEAPER_INFERENCE_API_KEY", VENDOR_KEY) + constructed = [] + + class _RecordingOpenAI(openai.OpenAI): + def __init__(self, **kwargs): + constructed.append(kwargs) + super().__init__(**kwargs) + + monkeypatch.setattr(openai, "OpenAI", _RecordingOpenAI) + + client, _, _, _ = example.setup_client("cheaperinference") + + assert constructed == [{"base_url": BASE_URL, "api_key": VENDOR_KEY}] + assert client.client.api_key == VENDOR_KEY + + +def _tool_call_response(arguments): + return { + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 0, + "model": "gpt-5.4-mini", + "choices": [ + { + "index": 0, + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_test", + "type": "function", + "function": {"name": "BasicChatOutputSchema", "arguments": json.dumps(arguments)}, + } + ], + }, + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + } + + +@pytest.fixture +def mocked_gateway(example, monkeypatch, unrelated_openai_key): + """Build the Cheaper Inference client with an httpx MockTransport instead of the network.""" + monkeypatch.setenv("CHEAPER_INFERENCE_API_KEY", VENDOR_KEY) + requests = [] + + def _handler(request): + requests.append(request) + return httpx.Response(200, json=_tool_call_response({"chat_message": "Hello from the mock"})) + + class _MockedOpenAI(openai.OpenAI): + def __init__(self, **kwargs): + super().__init__(http_client=httpx.Client(transport=httpx.MockTransport(_handler)), max_retries=0, **kwargs) + + monkeypatch.setattr(openai, "OpenAI", _MockedOpenAI) + return example.setup_client("cheaperinference"), requests + + +class TestRequestShape: + def test_setup_returns_expected_configuration(self, mocked_gateway): + (client, model, model_api_parameters, assistant_role), _ = mocked_gateway + assert isinstance(client, instructor.Instructor) + assert client.mode == instructor.Mode.TOOLS + assert str(client.client.base_url) == BASE_URL + "/" + assert model == "gpt-5.4-mini" + assert model_api_parameters == {"max_tokens": 2048} + assert assistant_role == "assistant" + + def test_request_sent_to_gateway_with_vendor_key(self, mocked_gateway): + (client, model, model_api_parameters, assistant_role), requests = mocked_gateway + agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema]( + AgentConfig( + client=client, + model=model, + model_api_parameters=model_api_parameters, + assistant_role=assistant_role, + ) + ) + + agent.run(BasicChatInputSchema(chat_message="Hi")) + + assert len(requests) == 1 + request = requests[0] + assert request.method == "POST" + assert str(request.url) == BASE_URL + "/chat/completions" + assert request.headers["authorization"] == f"Bearer {VENDOR_KEY}" + assert UNRELATED_OPENAI_KEY not in str(request.headers) + + body = json.loads(request.content) + assert body["model"] == "gpt-5.4-mini" + assert body["max_tokens"] == 2048 + assert body["tools"][0]["function"]["name"] == "BasicChatOutputSchema" + assert body["tool_choice"] == {"type": "function", "function": {"name": "BasicChatOutputSchema"}} + assert body["messages"][-1]["role"] == "user" + assert json.loads(body["messages"][-1]["content"]) == {"chat_message": "Hi"} + + +class TestStructuredResponse: + def test_tool_call_parsed_into_output_schema(self, mocked_gateway): + (client, model, model_api_parameters, assistant_role), requests = mocked_gateway + agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema]( + AgentConfig( + client=client, + model=model, + model_api_parameters=model_api_parameters, + assistant_role=assistant_role, + ) + ) + + response = agent.run(BasicChatInputSchema(chat_message="Hi")) + + assert isinstance(response, BasicChatOutputSchema) + assert response.chat_message == "Hello from the mock" + assert len(requests) == 1 + history = agent.history.get_history() + assert [message["role"] for message in history] == ["user", "assistant"] + + +class TestOtherProvidersUnchanged: + """The dispatch refactor keeps the OpenAI-SDK providers on the same settings.""" + + @pytest.mark.parametrize( + "choice, env_var, base_url, model, params, mode", + [ + ("1", "OPENAI_API_KEY", None, "gpt-5-mini", {"reasoning_effort": "low", "max_tokens": 2048}, "TOOLS"), + ("4", None, "http://localhost:11434/v1", "llama3", {"max_tokens": 2048}, "JSON"), + ("6", "OPENROUTER_API_KEY", "https://openrouter.ai/api/v1", "mistral/ministral-8b", {"max_tokens": 2048}, "TOOLS"), + ("7", "MINIMAX_API_KEY", "https://api.minimax.io/v1", "MiniMax-M3", {"max_tokens": 2048}, "JSON"), + ("8", "EDENAI_API_KEY", "https://api.edenai.run/v3", "openai/gpt-4o-mini", {"max_tokens": 2048}, "TOOLS"), + ], + ) + def test_openai_sdk_providers(self, example, monkeypatch, choice, env_var, base_url, model, params, mode): + if env_var: + monkeypatch.setenv(env_var, "dummy-key") + client, got_model, got_params, role = example.setup_client(choice) + assert got_model == model + assert got_params == params + assert role == "assistant" + assert client.mode == getattr(instructor.Mode, mode) + if base_url: + assert str(client.client.base_url) == base_url + "/" diff --git a/atomic-examples/quickstart/quickstart/4_basic_chatbot_different_providers.py b/atomic-examples/quickstart/quickstart/4_basic_chatbot_different_providers.py index 6b57326d..19c94e05 100644 --- a/atomic-examples/quickstart/quickstart/4_basic_chatbot_different_providers.py +++ b/atomic-examples/quickstart/quickstart/4_basic_chatbot_different_providers.py @@ -7,169 +7,203 @@ from atomic_agents import AtomicAgent, AgentConfig, BasicChatInputSchema, BasicChatOutputSchema from dotenv import load_dotenv -load_dotenv() - # Initialize a Rich Console for pretty console outputs console = Console() -# History setup -history = ChatHistory() +CHEAPER_INFERENCE_BASE_URL = "https://api.cheaperinference.com/v1" + + +def _require_api_key(env_var): + """Return a nonempty API key from the environment, or raise before any client exists.""" + api_key = (os.getenv(env_var) or "").strip() + if not api_key: + raise ValueError(f"{env_var} is not set. Set it to a nonempty API key to use this provider.") + return api_key + + +# Each setup function returns (client, model, model_api_parameters, assistant_role) +def _setup_openai(): + from openai import OpenAI + + api_key = os.getenv("OPENAI_API_KEY") + client = instructor.from_openai(OpenAI(api_key=api_key)) + return client, "gpt-5-mini", {"reasoning_effort": "low", "max_tokens": 2048}, "assistant" + + +def _setup_anthropic(): + from anthropic import Anthropic + + api_key = os.getenv("ANTHROPIC_API_KEY") + client = instructor.from_anthropic(Anthropic(api_key=api_key)) + return client, "claude-3-5-haiku-20241022", {"max_tokens": 2048}, "assistant" + + +def _setup_groq(): + from groq import Groq + + api_key = os.getenv("GROQ_API_KEY") + client = instructor.from_groq(Groq(api_key=api_key), mode=instructor.Mode.JSON) + return client, "mixtral-8x7b-32768", {"max_tokens": 2048}, "assistant" + + +def _setup_ollama(): + from openai import OpenAI as OllamaClient + + client = instructor.from_openai( + OllamaClient(base_url="http://localhost:11434/v1", api_key="ollama"), mode=instructor.Mode.JSON + ) + return client, "llama3", {"max_tokens": 2048}, "assistant" + + +def _setup_gemini(): + import google.genai + + api_key = os.getenv("GEMINI_API_KEY") + client = instructor.from_genai( + google.genai.Client(api_key=api_key), + mode=instructor.Mode.GENAI_TOOLS, + ) + return client, "gemini-2.5-flash", {}, "model" + +def _setup_openrouter(): + from openai import OpenAI as OpenRouterClient -# Function to set up the client based on the chosen provider + api_key = os.getenv("OPENROUTER_API_KEY") + client = instructor.from_openai(OpenRouterClient(base_url="https://openrouter.ai/api/v1", api_key=api_key)) + return client, "mistral/ministral-8b", {"max_tokens": 2048}, "assistant" + + +def _setup_minimax(): + from openai import OpenAI as MiniMaxClient + + api_key = os.getenv("MINIMAX_API_KEY") + client = instructor.from_openai( + MiniMaxClient(base_url="https://api.minimax.io/v1", api_key=api_key), + mode=instructor.Mode.JSON, + ) + return client, "MiniMax-M3", {"max_tokens": 2048}, "assistant" + + +def _setup_edenai(): + from openai import OpenAI as EdenAIClient + + api_key = os.getenv("EDENAI_API_KEY") + client = instructor.from_openai(EdenAIClient(base_url="https://api.edenai.run/v3", api_key=api_key)) + return client, "openai/gpt-4o-mini", {"max_tokens": 2048}, "assistant" + + +def _setup_cheaperinference(): + # Validate first: with api_key=None the OpenAI SDK falls back to OPENAI_API_KEY. + api_key = _require_api_key("CHEAPER_INFERENCE_API_KEY") + + from openai import OpenAI as CheaperInferenceClient + + client = instructor.from_openai( + CheaperInferenceClient(base_url=CHEAPER_INFERENCE_BASE_URL, api_key=api_key), + mode=instructor.Mode.TOOLS, + ) + return client, "gpt-5.4-mini", {"max_tokens": 2048}, "assistant" + + +# Provider name -> setup function. The order defines the menu numbers (1, 2, ...). +PROVIDER_SETUPS = { + "openai": _setup_openai, + "anthropic": _setup_anthropic, + "groq": _setup_groq, + "ollama": _setup_ollama, + "gemini": _setup_gemini, + "openrouter": _setup_openrouter, + "minimax": _setup_minimax, + "edenai": _setup_edenai, + "cheaperinference": _setup_cheaperinference, +} +providers_list = list(PROVIDER_SETUPS) +PROVIDER_NUMBERS = {str(i): name for i, name in enumerate(providers_list, start=1)} + + +# Function to set up the client based on the chosen provider (menu number or name) def setup_client(provider): console.log(f"provider: {provider}") - if provider == "1" or provider == "openai": - from openai import OpenAI - - api_key = os.getenv("OPENAI_API_KEY") - client = instructor.from_openai(OpenAI(api_key=api_key)) - model = "gpt-5-mini" - model_api_parameters = {"reasoning_effort": "low", "max_tokens": 2048} - assistant_role = "assistant" - elif provider == "2" or provider == "anthropic": - from anthropic import Anthropic - - api_key = os.getenv("ANTHROPIC_API_KEY") - client = instructor.from_anthropic(Anthropic(api_key=api_key)) - model = "claude-3-5-haiku-20241022" - model_api_parameters = {"max_tokens": 2048} - assistant_role = "assistant" - elif provider == "3" or provider == "groq": - from groq import Groq - - api_key = os.getenv("GROQ_API_KEY") - client = instructor.from_groq(Groq(api_key=api_key), mode=instructor.Mode.JSON) - model = "mixtral-8x7b-32768" - model_api_parameters = {"max_tokens": 2048} - assistant_role = "assistant" - elif provider == "4" or provider == "ollama": - from openai import OpenAI as OllamaClient - - client = instructor.from_openai( - OllamaClient(base_url="http://localhost:11434/v1", api_key="ollama"), mode=instructor.Mode.JSON - ) - model = "llama3" - model_api_parameters = {"max_tokens": 2048} - assistant_role = "assistant" - elif provider == "5" or provider == "gemini": - import google.genai - - api_key = os.getenv("GEMINI_API_KEY") - client = instructor.from_genai( - google.genai.Client(api_key=api_key), - mode=instructor.Mode.GENAI_TOOLS, - ) - model = "gemini-2.5-flash" - model_api_parameters = {} - assistant_role = "model" - elif provider == "6" or provider == "openrouter": - from openai import OpenAI as OpenRouterClient - - api_key = os.getenv("OPENROUTER_API_KEY") - client = instructor.from_openai(OpenRouterClient(base_url="https://openrouter.ai/api/v1", api_key=api_key)) - model = "mistral/ministral-8b" - model_api_parameters = {"max_tokens": 2048} - assistant_role = "assistant" - elif provider == "7" or provider == "minimax": - from openai import OpenAI as MiniMaxClient - - api_key = os.getenv("MINIMAX_API_KEY") - client = instructor.from_openai( - MiniMaxClient(base_url="https://api.minimax.io/v1", api_key=api_key), - mode=instructor.Mode.JSON, - ) - model = "MiniMax-M3" - model_api_parameters = {"max_tokens": 2048} - assistant_role = "assistant" - elif provider == "8" or provider == "edenai": - from openai import OpenAI as EdenAIClient - - api_key = os.getenv("EDENAI_API_KEY") - client = instructor.from_openai(EdenAIClient(base_url="https://api.edenai.run/v3", api_key=api_key)) - model = "openai/gpt-4o-mini" - model_api_parameters = {"max_tokens": 2048} - assistant_role = "assistant" - elif provider == "9" or provider == "cheaperinference": - from openai import OpenAI as CheaperInferenceClient - - api_key = os.getenv("CHEAPER_INFERENCE_API_KEY") - client = instructor.from_openai( - CheaperInferenceClient(base_url="https://api.cheaperinference.com/v1", api_key=api_key) - ) - model = "gpt-5.4-mini" - model_api_parameters = {"max_tokens": 2048} - assistant_role = "assistant" - else: + setup = PROVIDER_SETUPS.get(PROVIDER_NUMBERS.get(provider, provider)) + if setup is None: raise ValueError(f"Unsupported provider: {provider}") + return setup() - return client, model, model_api_parameters, assistant_role +def main(): + load_dotenv() -# Prompt the user to choose a provider from one in the list below. -providers_list = ["openai", "anthropic", "groq", "ollama", "gemini", "openrouter", "minimax", "edenai", "cheaperinference"] -y = "bold yellow" -b = "bold blue" -g = "bold green" -provider_inner_str = ( - f"{' / '.join(f'[[{g}]{i + 1}[/{g}]]. [{b}]{provider}[/{b}]' for i, provider in enumerate(providers_list))}" -) -providers_str = f"[{y}]Choose a provider ({provider_inner_str}): [/{y}]" + # History setup + history = ChatHistory() -provider = console.input(providers_str).lower() + # Prompt the user to choose a provider from one in the list below. + y = "bold yellow" + b = "bold blue" + g = "bold green" + provider_inner_str = ( + f"{' / '.join(f'[[{g}]{i + 1}[/{g}]]. [{b}]{provider}[/{b}]' for i, provider in enumerate(providers_list))}" + ) + providers_str = f"[{y}]Choose a provider ({provider_inner_str}): [/{y}]" + + provider = console.input(providers_str).lower() -# Set up the client and model based on the chosen provider -client, model, model_api_parameters, assistant_role = setup_client(provider) + # Set up the client and model based on the chosen provider + client, model, model_api_parameters, assistant_role = setup_client(provider) -# Initialize history with an initial message from the assistant -initial_message = BasicChatOutputSchema(chat_message="Hello! How can I assist you today?") -history.add_message(assistant_role, initial_message) + # Initialize history with an initial message from the assistant + initial_message = BasicChatOutputSchema(chat_message="Hello! How can I assist you today?") + history.add_message(assistant_role, initial_message) -# Agent setup with specified configuration -agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema]( - config=AgentConfig( - client=client, - model=model, - history=history, - assistant_role=assistant_role, - model_api_parameters=model_api_parameters, + # Agent setup with specified configuration + agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema]( + config=AgentConfig( + client=client, + model=model, + history=history, + assistant_role=assistant_role, + model_api_parameters=model_api_parameters, + ) ) -) - -# Generate the default system prompt for the agent -default_system_prompt = agent.system_prompt_generator.generate_prompt() -# Display the system prompt in a styled panel -console.print(Panel(default_system_prompt, width=console.width, style="bold cyan"), style="bold cyan") - -# Display the initial message from the assistant -console.print(Text("Agent:", style="bold green"), end=" ") -console.print(Text(initial_message.chat_message, style="bold green")) - -# Start an infinite loop to handle user inputs and agent responses -while True: - # Prompt the user for input with a styled prompt - user_input = console.input("[bold blue]You:[/bold blue] ") - # Check if the user wants to exit the chat - if user_input.lower() in ["/exit", "/quit"]: - console.print("Exiting chat...") - break - # Check if the user wants to see token count (works with any provider!) - if user_input.lower() == "/tokens": - token_info = agent.get_context_token_count() - console.print(f"[bold magenta]Token Usage ({model}):[/bold magenta]") - console.print(f" Total: {token_info.total} tokens") - console.print(f" System prompt: {token_info.system_prompt} tokens") - console.print(f" History: {token_info.history} tokens") - if token_info.max_tokens: - console.print(f" Max context: {token_info.max_tokens} tokens") - if token_info.utilization: - console.print(f" Context utilization: {token_info.utilization:.1%}") - continue - - # Process the user's input through the agent and get the response - input_schema = BasicChatInputSchema(chat_message=user_input) - response = agent.run(input_schema) - - agent_message = Text(response.chat_message, style="bold green") + + # Generate the default system prompt for the agent + default_system_prompt = agent.system_prompt_generator.generate_prompt() + # Display the system prompt in a styled panel + console.print(Panel(default_system_prompt, width=console.width, style="bold cyan"), style="bold cyan") + + # Display the initial message from the assistant console.print(Text("Agent:", style="bold green"), end=" ") - console.print(agent_message) + console.print(Text(initial_message.chat_message, style="bold green")) + + # Start an infinite loop to handle user inputs and agent responses + while True: + # Prompt the user for input with a styled prompt + user_input = console.input("[bold blue]You:[/bold blue] ") + # Check if the user wants to exit the chat + if user_input.lower() in ["/exit", "/quit"]: + console.print("Exiting chat...") + break + # Check if the user wants to see token count (works with any provider!) + if user_input.lower() == "/tokens": + token_info = agent.get_context_token_count() + console.print(f"[bold magenta]Token Usage ({model}):[/bold magenta]") + console.print(f" Total: {token_info.total} tokens") + console.print(f" System prompt: {token_info.system_prompt} tokens") + console.print(f" History: {token_info.history} tokens") + if token_info.max_tokens: + console.print(f" Max context: {token_info.max_tokens} tokens") + if token_info.utilization: + console.print(f" Context utilization: {token_info.utilization:.1%}") + continue + + # Process the user's input through the agent and get the response + input_schema = BasicChatInputSchema(chat_message=user_input) + response = agent.run(input_schema) + + agent_message = Text(response.chat_message, style="bold green") + console.print(Text("Agent:", style="bold green"), end=" ") + console.print(agent_message) + + +if __name__ == "__main__": + main() From fad71da62aead512e75ea821c8c647dd7ea6e51d Mon Sep 17 00:00:00 2001 From: aiapienthusiast <327629069+aiapienthusiast@users.noreply.github.com> Date: Fri, 9 Oct 2026 20:00:20 +0400 Subject: [PATCH 3/3] refactor: split quickstart provider main into small helpers --- .../agents/test_cheaperinference_provider.py | 106 +++++++++++++++++ .../4_basic_chatbot_different_providers.py | 110 ++++++++++-------- 2 files changed, 169 insertions(+), 47 deletions(-) diff --git a/atomic-agents/tests/agents/test_cheaperinference_provider.py b/atomic-agents/tests/agents/test_cheaperinference_provider.py index d3fea4c2..8f131c19 100644 --- a/atomic-agents/tests/agents/test_cheaperinference_provider.py +++ b/atomic-agents/tests/agents/test_cheaperinference_provider.py @@ -4,13 +4,16 @@ """ import importlib.util +import io import json import os +from types import SimpleNamespace import httpx import instructor import openai import pytest +from rich.console import Console from atomic_agents import AtomicAgent, AgentConfig, BasicChatInputSchema, BasicChatOutputSchema @@ -258,3 +261,106 @@ def test_openai_sdk_providers(self, example, monkeypatch, choice, env_var, base_ assert client.mode == getattr(instructor.Mode, mode) if base_url: assert str(client.client.base_url) == base_url + "/" + + +@pytest.fixture +def captured_console(example, monkeypatch): + """Replace the example console with one that records output and reads scripted inputs.""" + console = Console(file=io.StringIO(), width=200, force_terminal=False) + inputs = [] + prompts = [] + + def _input(prompt=""): + prompts.append(prompt) + return inputs.pop(0) + + monkeypatch.setattr(console, "input", _input) + monkeypatch.setattr(example, "console", console) + return console, inputs, prompts + + +class _FakeAgent: + def __init__(self, token_info): + self.token_info = token_info + self.runs = [] + + def get_context_token_count(self): + return self.token_info + + def run(self, input_schema): + self.runs.append(input_schema) + return BasicChatOutputSchema(chat_message=f"echo: {input_schema.chat_message}") + + +class TestInteractiveFlow: + def test_provider_prompt_lists_numbered_providers(self, example): + prompt = example._provider_prompt() + for number, name in example.PROVIDER_NUMBERS.items(): + assert f"[[bold green]{number}[/bold green]]. [bold blue]{name}[/bold blue]" in prompt + + def test_choose_provider_lowercases_input(self, example, captured_console): + _, inputs, prompts = captured_console + inputs.append("CheaperInference") + assert example._choose_provider() == "cheaperinference" + assert prompts == [example._provider_prompt()] + + @pytest.mark.parametrize("command", ["/exit", "/quit", "/EXIT"]) + def test_exit_commands_stop_the_loop(self, example, captured_console, command): + console, _, _ = captured_console + agent = _FakeAgent(token_info=None) + assert example._handle_turn(agent, "m", command) is False + assert agent.runs == [] + assert "Exiting chat..." in console.file.getvalue() + + def test_tokens_command_prints_usage_without_model_call(self, example, captured_console): + console, _, _ = captured_console + token_info = SimpleNamespace(total=30, system_prompt=10, history=20, max_tokens=100, utilization=0.3) + agent = _FakeAgent(token_info) + assert example._handle_turn(agent, "gpt-5.4-mini", "/tokens") is True + output = console.file.getvalue() + assert agent.runs == [] + assert "Token Usage (gpt-5.4-mini):" in output + assert "Total: 30 tokens" in output + assert "Max context: 100 tokens" in output + assert "Context utilization: 30.0%" in output + + def test_tokens_command_skips_unknown_limits(self, example, captured_console): + console, _, _ = captured_console + token_info = SimpleNamespace(total=30, system_prompt=10, history=20, max_tokens=None, utilization=None) + example._handle_turn(_FakeAgent(token_info), "m", "/tokens") + output = console.file.getvalue() + assert "Max context" not in output + assert "Context utilization" not in output + + def test_message_is_sent_to_agent_and_printed(self, example, captured_console): + console, _, _ = captured_console + agent = _FakeAgent(token_info=None) + assert example._handle_turn(agent, "m", "Hi") is True + assert [schema.chat_message for schema in agent.runs] == ["Hi"] + assert "Agent: echo: Hi" in console.file.getvalue() + + def test_build_agent_starts_history_with_greeting(self, example, mocked_gateway): + (client, model, model_api_parameters, assistant_role), requests = mocked_gateway + agent, initial_message = example._build_agent(client, model, model_api_parameters, assistant_role) + assert initial_message.chat_message == "Hello! How can I assist you today?" + assert agent.model == "gpt-5.4-mini" + assert agent.model_api_parameters == {"max_tokens": 2048} + assert [message["role"] for message in agent.history.get_history()] == ["assistant"] + assert requests == [] + + def test_main_runs_one_turn_then_exits(self, example, monkeypatch, mocked_gateway, captured_console): + setup, requests = mocked_gateway + console, inputs, prompts = captured_console + monkeypatch.setattr(example, "load_dotenv", lambda: None) + monkeypatch.setitem(example.PROVIDER_SETUPS, "cheaperinference", lambda: setup) + inputs.extend(["9", "Hi", "/quit"]) + + example.main() + + output = console.file.getvalue() + assert inputs == [] + assert len(prompts) == 3 + assert len(requests) == 1 + assert "Agent: Hello! How can I assist you today?" in output + assert "Agent: Hello from the mock" in output + assert output.rstrip().endswith("Exiting chat...") diff --git a/atomic-examples/quickstart/quickstart/4_basic_chatbot_different_providers.py b/atomic-examples/quickstart/quickstart/4_basic_chatbot_different_providers.py index 19c94e05..7bd88378 100644 --- a/atomic-examples/quickstart/quickstart/4_basic_chatbot_different_providers.py +++ b/atomic-examples/quickstart/quickstart/4_basic_chatbot_different_providers.py @@ -131,31 +131,33 @@ def setup_client(provider): return setup() -def main(): - load_dotenv() - - # History setup - history = ChatHistory() - - # Prompt the user to choose a provider from one in the list below. +def _provider_prompt(): + """Build the styled provider menu, e.g. "[1]. openai / [2]. anthropic / ...".""" y = "bold yellow" b = "bold blue" g = "bold green" provider_inner_str = ( f"{' / '.join(f'[[{g}]{i + 1}[/{g}]]. [{b}]{provider}[/{b}]' for i, provider in enumerate(providers_list))}" ) - providers_str = f"[{y}]Choose a provider ({provider_inner_str}): [/{y}]" + return f"[{y}]Choose a provider ({provider_inner_str}): [/{y}]" - provider = console.input(providers_str).lower() - # Set up the client and model based on the chosen provider - client, model, model_api_parameters, assistant_role = setup_client(provider) +def _choose_provider(): + """Prompt the user to choose a provider by menu number or name.""" + return console.input(_provider_prompt()).lower() + - # Initialize history with an initial message from the assistant +def _print_agent_message(message): + console.print(Text("Agent:", style="bold green"), end=" ") + console.print(Text(message, style="bold green")) + + +def _build_agent(client, model, model_api_parameters, assistant_role): + """Create the agent with a history that starts with the assistant greeting.""" + history = ChatHistory() initial_message = BasicChatOutputSchema(chat_message="Hello! How can I assist you today?") history.add_message(assistant_role, initial_message) - # Agent setup with specified configuration agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema]( config=AgentConfig( client=client, @@ -165,44 +167,58 @@ def main(): model_api_parameters=model_api_parameters, ) ) + return agent, initial_message + + +def _print_token_usage(agent, model): + """Show the context token count (works with any provider).""" + token_info = agent.get_context_token_count() + console.print(f"[bold magenta]Token Usage ({model}):[/bold magenta]") + console.print(f" Total: {token_info.total} tokens") + console.print(f" System prompt: {token_info.system_prompt} tokens") + console.print(f" History: {token_info.history} tokens") + if token_info.max_tokens: + console.print(f" Max context: {token_info.max_tokens} tokens") + if token_info.utilization: + console.print(f" Context utilization: {token_info.utilization:.1%}") + + +def _handle_turn(agent, model, user_input): + """Handle one user input. Return False when the user wants to exit.""" + command = user_input.lower() + if command in ["/exit", "/quit"]: + console.print("Exiting chat...") + return False + if command == "/tokens": + _print_token_usage(agent, model) + return True + + response = agent.run(BasicChatInputSchema(chat_message=user_input)) + _print_agent_message(response.chat_message) + return True + + +def _chat_loop(agent, model): + """Read user inputs and print agent responses until the user exits.""" + while True: + user_input = console.input("[bold blue]You:[/bold blue] ") + if not _handle_turn(agent, model, user_input): + break + + +def main(): + load_dotenv() + + provider = _choose_provider() + client, model, model_api_parameters, assistant_role = setup_client(provider) + agent, initial_message = _build_agent(client, model, model_api_parameters, assistant_role) - # Generate the default system prompt for the agent + # Display the default system prompt in a styled panel, then the greeting default_system_prompt = agent.system_prompt_generator.generate_prompt() - # Display the system prompt in a styled panel console.print(Panel(default_system_prompt, width=console.width, style="bold cyan"), style="bold cyan") + _print_agent_message(initial_message.chat_message) - # Display the initial message from the assistant - console.print(Text("Agent:", style="bold green"), end=" ") - console.print(Text(initial_message.chat_message, style="bold green")) - - # Start an infinite loop to handle user inputs and agent responses - while True: - # Prompt the user for input with a styled prompt - user_input = console.input("[bold blue]You:[/bold blue] ") - # Check if the user wants to exit the chat - if user_input.lower() in ["/exit", "/quit"]: - console.print("Exiting chat...") - break - # Check if the user wants to see token count (works with any provider!) - if user_input.lower() == "/tokens": - token_info = agent.get_context_token_count() - console.print(f"[bold magenta]Token Usage ({model}):[/bold magenta]") - console.print(f" Total: {token_info.total} tokens") - console.print(f" System prompt: {token_info.system_prompt} tokens") - console.print(f" History: {token_info.history} tokens") - if token_info.max_tokens: - console.print(f" Max context: {token_info.max_tokens} tokens") - if token_info.utilization: - console.print(f" Context utilization: {token_info.utilization:.1%}") - continue - - # Process the user's input through the agent and get the response - input_schema = BasicChatInputSchema(chat_message=user_input) - response = agent.run(input_schema) - - agent_message = Text(response.chat_message, style="bold green") - console.print(Text("Agent:", style="bold green"), end=" ") - console.print(agent_message) + _chat_loop(agent, model) if __name__ == "__main__":