mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
e4a41faa33
- Fix E402: Move module-level imports to top of file - Fix F401: Add noqa for availability check imports - Fix F402: Rename loop variables shadowing imports - Fix E722: Replace bare except with except Exception - Fix B904: Add exception chaining (from e) - Fix F811: Remove duplicate imports - Fix B027: Add noqa for empty close() method - Fix E741: Rename ambiguous variable l -> label - Fix I001: Import sorting issues - Apply ruff format to all 106 files All 902 tests pass.
294 lines
9.9 KiB
Python
294 lines
9.9 KiB
Python
"""Tests for universal provider support.
|
|
|
|
Tests OpenAICompatibleProvider, GoogleProvider, and LiteLLMProvider.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from headroom.providers import (
|
|
GoogleProvider,
|
|
ModelCapabilities,
|
|
OpenAICompatibleProvider,
|
|
create_groq_provider,
|
|
create_lmstudio_provider,
|
|
create_ollama_provider,
|
|
create_together_provider,
|
|
create_vllm_provider,
|
|
is_litellm_available,
|
|
)
|
|
|
|
|
|
def _transformers_available() -> bool:
|
|
"""Check if transformers is available."""
|
|
try:
|
|
import transformers # noqa: F401
|
|
|
|
return True
|
|
except ImportError:
|
|
return False
|
|
|
|
|
|
class TestOpenAICompatibleProvider:
|
|
"""Tests for OpenAICompatibleProvider."""
|
|
|
|
def test_init_default(self):
|
|
"""Test initialization with defaults."""
|
|
provider = OpenAICompatibleProvider()
|
|
assert provider.name == "openai_compatible"
|
|
assert provider.base_url is None
|
|
|
|
def test_init_with_config(self):
|
|
"""Test initialization with configuration."""
|
|
provider = OpenAICompatibleProvider(
|
|
name="custom",
|
|
base_url="http://localhost:8080/v1",
|
|
api_key="test-key",
|
|
)
|
|
assert provider.name == "custom"
|
|
assert provider.base_url == "http://localhost:8080/v1"
|
|
assert provider.api_key == "test-key"
|
|
|
|
def test_supports_any_model(self):
|
|
"""Test that provider supports any model."""
|
|
provider = OpenAICompatibleProvider()
|
|
assert provider.supports_model("any-model") is True
|
|
assert provider.supports_model("llama-3") is True
|
|
assert provider.supports_model("custom-finetuned") is True
|
|
|
|
@pytest.mark.skipif(
|
|
not _transformers_available(),
|
|
reason="transformers not installed - needed for HuggingFace tokenizer",
|
|
)
|
|
def test_get_token_counter(self):
|
|
"""Test getting token counter."""
|
|
provider = OpenAICompatibleProvider()
|
|
counter = provider.get_token_counter("llama-3-8b")
|
|
assert counter is not None
|
|
# Should be able to count tokens
|
|
count = counter.count_text("Hello, world!")
|
|
assert count > 0
|
|
|
|
def test_get_context_limit_known_model(self):
|
|
"""Test context limit for known models."""
|
|
provider = OpenAICompatibleProvider()
|
|
# Llama 3.1 has 128K context
|
|
limit = provider.get_context_limit("llama-3.1-8b")
|
|
assert limit == 128000
|
|
|
|
def test_get_context_limit_unknown_model(self):
|
|
"""Test context limit for unknown models (defaults to 128K)."""
|
|
provider = OpenAICompatibleProvider()
|
|
limit = provider.get_context_limit("unknown-model")
|
|
assert limit == 128000
|
|
|
|
def test_register_model(self):
|
|
"""Test registering a custom model."""
|
|
provider = OpenAICompatibleProvider()
|
|
provider.register_model(
|
|
"my-model",
|
|
context_window=64000,
|
|
max_output_tokens=8192,
|
|
input_cost_per_1m=1.0,
|
|
output_cost_per_1m=2.0,
|
|
)
|
|
assert provider.get_context_limit("my-model") == 64000
|
|
|
|
def test_estimate_cost_registered_model(self):
|
|
"""Test cost estimation for registered model."""
|
|
provider = OpenAICompatibleProvider()
|
|
provider.register_model(
|
|
"priced-model",
|
|
input_cost_per_1m=1.0,
|
|
output_cost_per_1m=2.0,
|
|
)
|
|
cost = provider.estimate_cost(
|
|
input_tokens=1000000,
|
|
output_tokens=500000,
|
|
model="priced-model",
|
|
)
|
|
assert cost == 2.0 # 1.0 + 1.0
|
|
|
|
def test_estimate_cost_unknown_model(self):
|
|
"""Test cost estimation returns None for unknown model."""
|
|
provider = OpenAICompatibleProvider()
|
|
cost = provider.estimate_cost(
|
|
input_tokens=1000,
|
|
output_tokens=500,
|
|
model="unknown-model",
|
|
)
|
|
assert cost is None
|
|
|
|
|
|
class TestModelCapabilities:
|
|
"""Tests for ModelCapabilities dataclass."""
|
|
|
|
def test_default_values(self):
|
|
"""Test default capability values."""
|
|
caps = ModelCapabilities(model="test-model")
|
|
assert caps.context_window == 128000
|
|
assert caps.max_output_tokens == 4096
|
|
assert caps.supports_tools is True
|
|
assert caps.supports_vision is False
|
|
assert caps.supports_streaming is True
|
|
|
|
def test_custom_values(self):
|
|
"""Test custom capability values."""
|
|
caps = ModelCapabilities(
|
|
model="custom-model",
|
|
context_window=32000,
|
|
max_output_tokens=16384,
|
|
supports_tools=False,
|
|
supports_vision=True,
|
|
input_cost_per_1m=0.5,
|
|
output_cost_per_1m=1.5,
|
|
)
|
|
assert caps.context_window == 32000
|
|
assert caps.max_output_tokens == 16384
|
|
assert caps.supports_tools is False
|
|
assert caps.supports_vision is True
|
|
assert caps.input_cost_per_1m == 0.5
|
|
assert caps.output_cost_per_1m == 1.5
|
|
|
|
|
|
class TestGoogleProvider:
|
|
"""Tests for GoogleProvider."""
|
|
|
|
@pytest.fixture
|
|
def provider(self):
|
|
"""Create Google provider."""
|
|
return GoogleProvider()
|
|
|
|
def test_name(self, provider):
|
|
"""Test provider name."""
|
|
assert provider.name == "google"
|
|
|
|
def test_supports_gemini_models(self, provider):
|
|
"""Test support for Gemini models."""
|
|
assert provider.supports_model("gemini-2.0-flash") is True
|
|
assert provider.supports_model("gemini-1.5-pro") is True
|
|
assert provider.supports_model("gemini-1.5-flash") is True
|
|
|
|
def test_not_supports_other_models(self, provider):
|
|
"""Test non-support for other models."""
|
|
assert provider.supports_model("gpt-4o") is False
|
|
assert provider.supports_model("claude-3") is False
|
|
|
|
def test_get_token_counter(self, provider):
|
|
"""Test getting token counter."""
|
|
counter = provider.get_token_counter("gemini-2.0-flash")
|
|
assert counter is not None
|
|
count = counter.count_text("Hello, world!")
|
|
assert count > 0
|
|
|
|
def test_get_context_limit_gemini_2(self, provider):
|
|
"""Test context limit for Gemini 2.0."""
|
|
limit = provider.get_context_limit("gemini-2.0-flash")
|
|
assert limit == 1000000 # 1M tokens
|
|
|
|
def test_get_context_limit_gemini_1_5_pro(self, provider):
|
|
"""Test context limit for Gemini 1.5 Pro (2M!)."""
|
|
limit = provider.get_context_limit("gemini-1.5-pro")
|
|
assert limit == 2000000 # 2M tokens!
|
|
|
|
def test_estimate_cost(self, provider):
|
|
"""Test cost estimation."""
|
|
cost = provider.estimate_cost(
|
|
input_tokens=1000000,
|
|
output_tokens=500000,
|
|
model="gemini-2.0-flash",
|
|
)
|
|
assert cost is not None
|
|
# 1M input * $0.10 + 0.5M output * $0.40 = $0.10 + $0.20 = $0.30
|
|
assert abs(cost - 0.30) < 0.01
|
|
|
|
def test_openai_compatible_url(self):
|
|
"""Test OpenAI-compatible URL."""
|
|
url = GoogleProvider.get_openai_compatible_url("test-key")
|
|
assert "generativelanguage.googleapis.com" in url
|
|
|
|
|
|
class TestProviderFactoryFunctions:
|
|
"""Tests for provider factory functions."""
|
|
|
|
def test_create_ollama_provider(self):
|
|
"""Test creating Ollama provider."""
|
|
provider = create_ollama_provider()
|
|
assert provider.name == "ollama"
|
|
assert provider.base_url == "http://localhost:11434/v1"
|
|
|
|
def test_create_ollama_provider_custom_url(self):
|
|
"""Test creating Ollama provider with custom URL."""
|
|
provider = create_ollama_provider("http://192.168.1.100:11434/v1")
|
|
assert provider.base_url == "http://192.168.1.100:11434/v1"
|
|
|
|
def test_create_together_provider(self):
|
|
"""Test creating Together provider."""
|
|
provider = create_together_provider()
|
|
assert provider.name == "together"
|
|
assert "together.xyz" in provider.base_url
|
|
|
|
def test_create_groq_provider(self):
|
|
"""Test creating Groq provider."""
|
|
provider = create_groq_provider()
|
|
assert provider.name == "groq"
|
|
assert "groq.com" in provider.base_url
|
|
|
|
def test_create_vllm_provider(self):
|
|
"""Test creating vLLM provider."""
|
|
provider = create_vllm_provider("http://localhost:8000/v1")
|
|
assert provider.name == "vllm"
|
|
assert provider.base_url == "http://localhost:8000/v1"
|
|
|
|
def test_create_lmstudio_provider(self):
|
|
"""Test creating LM Studio provider."""
|
|
provider = create_lmstudio_provider()
|
|
assert provider.name == "lmstudio"
|
|
assert provider.base_url == "http://localhost:1234/v1"
|
|
|
|
|
|
class TestLiteLLMProvider:
|
|
"""Tests for LiteLLM provider."""
|
|
|
|
def test_is_litellm_available(self):
|
|
"""Test checking LiteLLM availability."""
|
|
result = is_litellm_available()
|
|
assert isinstance(result, bool)
|
|
|
|
@pytest.mark.skipif(
|
|
not is_litellm_available(),
|
|
reason="LiteLLM not installed",
|
|
)
|
|
def test_create_litellm_provider(self):
|
|
"""Test creating LiteLLM provider."""
|
|
from headroom.providers import create_litellm_provider
|
|
|
|
provider = create_litellm_provider()
|
|
assert provider.name == "litellm"
|
|
|
|
@pytest.mark.skipif(
|
|
not is_litellm_available(),
|
|
reason="LiteLLM not installed",
|
|
)
|
|
def test_litellm_supports_any_model(self):
|
|
"""Test LiteLLM supports any model."""
|
|
from headroom.providers import create_litellm_provider
|
|
|
|
provider = create_litellm_provider()
|
|
assert provider.supports_model("gpt-4o") is True
|
|
assert provider.supports_model("claude-3-sonnet") is True
|
|
assert provider.supports_model("any-model") is True
|
|
|
|
@pytest.mark.skipif(
|
|
not is_litellm_available(),
|
|
reason="LiteLLM not installed",
|
|
)
|
|
def test_litellm_list_providers(self):
|
|
"""Test listing LiteLLM providers."""
|
|
from headroom.providers import LiteLLMProvider
|
|
|
|
providers = LiteLLMProvider.list_supported_providers()
|
|
assert "openai" in providers
|
|
assert "anthropic" in providers
|
|
assert "ollama" in providers
|