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-agents/tests/agents/test_cheaperinference_provider.py b/atomic-agents/tests/agents/test_cheaperinference_provider.py new file mode 100644 index 00000000..8f131c19 --- /dev/null +++ b/atomic-agents/tests/agents/test_cheaperinference_provider.py @@ -0,0 +1,366 @@ +"""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 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 + +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 + "/" + + +@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 7f8d10ed..7bd88378 100644 --- a/atomic-examples/quickstart/quickstart/4_basic_chatbot_different_providers.py +++ b/atomic-examples/quickstart/quickstart/4_basic_chatbot_different_providers.py @@ -7,159 +7,219 @@ 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 + + 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" -# Function to set up the client based on the chosen provider + +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" - 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 _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))}" + ) + return f"[{y}]Choose a provider ({provider_inner_str}): [/{y}]" -# Prompt the user to choose a provider from one in the list below. -providers_list = ["openai", "anthropic", "groq", "ollama", "gemini", "openrouter", "minimax", "edenai"] -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) +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 -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, +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 = 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"]: + 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...") - 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) + 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) + + # Display the default system prompt in a styled panel, then the greeting + default_system_prompt = agent.system_prompt_generator.generate_prompt() + console.print(Panel(default_system_prompt, width=console.width, style="bold cyan"), style="bold cyan") + _print_agent_message(initial_message.chat_message) + + _chat_loop(agent, model) + + +if __name__ == "__main__": + main()