feat: depending on the number of tools in use, this change should save ~50% of overall tokens used. fixes https://github.com/BeehiveInnovations/zen-mcp-server/issues/255 but also refactored individual tools to instead encourage the agent to use the listmodels tool if needed.
This commit is contained in:
@@ -92,9 +92,9 @@ class TestAutoMode:
|
||||
|
||||
# Model field should have detailed descriptions
|
||||
model_schema = schema["properties"]["model"]
|
||||
assert "enum" in model_schema
|
||||
assert "flash" in model_schema["enum"]
|
||||
assert "select the most suitable model" in model_schema["description"]
|
||||
assert "enum" not in model_schema
|
||||
assert "auto model selection" in model_schema["description"].lower()
|
||||
assert "listmodels" in model_schema["description"]
|
||||
|
||||
finally:
|
||||
# Restore
|
||||
@@ -111,14 +111,14 @@ class TestAutoMode:
|
||||
tool = ChatTool()
|
||||
schema = tool.get_input_schema()
|
||||
|
||||
# Model should not be required
|
||||
# Model should not be required when default model is configured
|
||||
assert "model" not in schema["required"]
|
||||
|
||||
# Model field should have simpler description
|
||||
model_schema = schema["properties"]["model"]
|
||||
assert "enum" not in model_schema
|
||||
assert "Native models:" in model_schema["description"]
|
||||
assert "Defaults to" in model_schema["description"]
|
||||
assert "listmodels" in model_schema["description"]
|
||||
assert "default model" in model_schema["description"].lower()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auto_mode_requires_model_parameter(self):
|
||||
@@ -287,19 +287,10 @@ class TestAutoMode:
|
||||
importlib.reload(config)
|
||||
|
||||
schema = tool.get_model_field_schema()
|
||||
assert "enum" in schema
|
||||
# Test that some basic models are available (those that should be available with dummy keys)
|
||||
available_models = schema["enum"]
|
||||
# Check for models that should be available with basic provider setup
|
||||
expected_basic_models = ["flash", "pro"] # Gemini models from conftest.py
|
||||
for model in expected_basic_models:
|
||||
if model not in available_models:
|
||||
print(f"Missing expected model: {model}")
|
||||
print(f"Available models: {available_models}")
|
||||
assert any(
|
||||
model in available_models for model in expected_basic_models
|
||||
), f"None of {expected_basic_models} found in {available_models}"
|
||||
assert "select the most suitable model" in schema["description"]
|
||||
assert "enum" not in schema
|
||||
assert schema["type"] == "string"
|
||||
assert "auto model selection" in schema["description"]
|
||||
assert "listmodels" in schema["description"]
|
||||
|
||||
# Test normal mode
|
||||
os.environ["DEFAULT_MODEL"] = "pro"
|
||||
@@ -307,10 +298,9 @@ class TestAutoMode:
|
||||
|
||||
schema = tool.get_model_field_schema()
|
||||
assert "enum" not in schema
|
||||
# Check for the new schema format
|
||||
assert "Model to use." in schema["description"]
|
||||
assert schema["type"] == "string"
|
||||
assert "'pro'" in schema["description"]
|
||||
assert "Defaults to" in schema["description"]
|
||||
assert "listmodels" in schema["description"]
|
||||
|
||||
finally:
|
||||
# Restore
|
||||
|
||||
@@ -291,58 +291,28 @@ class TestAutoModeComprehensive:
|
||||
# Should have model as required field
|
||||
assert "model" in schema["required"]
|
||||
|
||||
# Should include all model options from global config
|
||||
# In auto mode, the schema should now have a description field
|
||||
# instructing users to use the listmodels tool instead of an enum
|
||||
model_schema = schema["properties"]["model"]
|
||||
assert "enum" in model_schema
|
||||
assert "type" in model_schema
|
||||
assert model_schema["type"] == "string"
|
||||
assert "description" in model_schema
|
||||
|
||||
available_models = model_schema["enum"]
|
||||
# Check that the description mentions using listmodels tool
|
||||
description = model_schema["description"]
|
||||
assert "listmodels" in description.lower()
|
||||
assert "auto" in description.lower() or "selection" in description.lower()
|
||||
|
||||
# Should include Gemini models
|
||||
assert "flash" in available_models
|
||||
assert "pro" in available_models
|
||||
assert "gemini-2.5-flash" in available_models
|
||||
assert "gemini-2.5-pro" in available_models
|
||||
# Should NOT have enum field anymore - this is the new behavior
|
||||
assert "enum" not in model_schema
|
||||
|
||||
# After the fix, schema only shows models from enabled providers
|
||||
# This prevents model namespace collisions and misleading users
|
||||
# If only Gemini is configured, only Gemini models should appear
|
||||
provider_count = len(
|
||||
[
|
||||
key
|
||||
for key in ["GEMINI_API_KEY", "OPENAI_API_KEY", "XAI_API_KEY", "OPENROUTER_API_KEY"]
|
||||
if os.getenv(key) and os.getenv(key) != f"your_{key.lower()}_here"
|
||||
]
|
||||
)
|
||||
# After the design change, the system directs users to use listmodels
|
||||
# instead of enumerating all models in the schema
|
||||
# This prevents model namespace collisions and keeps the schema cleaner
|
||||
|
||||
if provider_count == 1 and os.getenv("GEMINI_API_KEY"):
|
||||
# Only Gemini configured - should only show Gemini models
|
||||
non_gemini_models = [
|
||||
m
|
||||
for m in available_models
|
||||
if not m.startswith("gemini")
|
||||
and m
|
||||
not in [
|
||||
"flash",
|
||||
"pro",
|
||||
"flash-2.0",
|
||||
"flash2",
|
||||
"flashlite",
|
||||
"flash-lite",
|
||||
"flash2.5",
|
||||
"gemini pro",
|
||||
"gemini-pro",
|
||||
]
|
||||
]
|
||||
assert (
|
||||
len(non_gemini_models) == 0
|
||||
), f"Found non-Gemini models when only Gemini configured: {non_gemini_models}"
|
||||
else:
|
||||
# Multiple providers or OpenRouter - should include various models
|
||||
# Only check if models are available if their providers might be configured
|
||||
if os.getenv("OPENAI_API_KEY") or os.getenv("OPENROUTER_API_KEY"):
|
||||
assert any("o3" in m or "o4" in m for m in available_models), "No OpenAI models found"
|
||||
if os.getenv("XAI_API_KEY") or os.getenv("OPENROUTER_API_KEY"):
|
||||
assert any("grok" in m for m in available_models), "No XAI models found"
|
||||
# With the new design change, we no longer enumerate models in the schema
|
||||
# The listmodels tool should be used to discover available models
|
||||
# This test now validates the schema structure rather than model enumeration
|
||||
|
||||
def test_auto_mode_schema_with_all_providers(self):
|
||||
"""Test that auto mode schema includes models from all available providers."""
|
||||
@@ -380,21 +350,21 @@ class TestAutoModeComprehensive:
|
||||
tool = AnalyzeTool()
|
||||
schema = tool.get_input_schema()
|
||||
|
||||
# In auto mode with multiple providers, should still use the new schema format
|
||||
model_schema = schema["properties"]["model"]
|
||||
available_models = model_schema["enum"]
|
||||
assert "type" in model_schema
|
||||
assert model_schema["type"] == "string"
|
||||
assert "description" in model_schema
|
||||
|
||||
# Should include models from all providers
|
||||
# Gemini models
|
||||
assert "flash" in available_models
|
||||
assert "pro" in available_models
|
||||
# Check that the description mentions using listmodels tool
|
||||
description = model_schema["description"]
|
||||
assert "listmodels" in description.lower()
|
||||
|
||||
# OpenAI models
|
||||
assert "o3" in available_models
|
||||
assert "o4-mini" in available_models
|
||||
# Should NOT have enum field - uses listmodels tool instead
|
||||
assert "enum" not in model_schema
|
||||
|
||||
# XAI models
|
||||
assert "grok" in available_models
|
||||
assert "grok-3" in available_models
|
||||
# With multiple providers configured, the listmodels tool
|
||||
# would show models from all providers when called
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auto_mode_model_parameter_required_error(self):
|
||||
|
||||
@@ -84,15 +84,14 @@ class TestChatTool:
|
||||
assert schema["type"] == "string"
|
||||
assert "description" in schema
|
||||
|
||||
# In auto mode, should have enum. In normal mode, should have model descriptions
|
||||
# Description should route callers to listmodels, regardless of mode
|
||||
assert "listmodels" in schema["description"]
|
||||
if self.tool.is_effective_auto_mode():
|
||||
assert "enum" in schema
|
||||
assert len(schema["enum"]) > 0
|
||||
assert "IMPORTANT:" in schema["description"]
|
||||
assert "auto model selection" in schema["description"]
|
||||
else:
|
||||
# Normal mode - should have model descriptions in description
|
||||
assert "Model to use" in schema["description"]
|
||||
assert "Native models:" in schema["description"]
|
||||
import config
|
||||
|
||||
assert f"'{config.DEFAULT_MODEL}'" in schema["description"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_prompt_preparation(self):
|
||||
|
||||
@@ -130,9 +130,7 @@ class TestModelEnumeration:
|
||||
models = tool._get_available_models()
|
||||
|
||||
for alias in ("local-llama", "llama3.2"):
|
||||
assert (
|
||||
alias not in models
|
||||
), f"Custom model alias '{alias}' should remain hidden without CUSTOM_API_URL"
|
||||
assert alias not in models, f"Custom model alias '{alias}' should remain hidden without CUSTOM_API_URL"
|
||||
|
||||
def test_no_duplicates_with_overlapping_providers(self):
|
||||
"""Test that models aren't duplicated when multiple providers offer the same model."""
|
||||
|
||||
@@ -465,7 +465,7 @@ class TestSchemaGeneration:
|
||||
tool = ThinkDeepTool()
|
||||
schema = tool.get_input_schema()
|
||||
|
||||
# Model should NOT be required
|
||||
# Model should remain optional when DEFAULT_MODEL is available
|
||||
assert "model" not in schema["required"]
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user