mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
175746cc26
This commit prepares Headroom for public open source release with comprehensive documentation, licensing, and community infrastructure. License & Legal: - Add Apache 2.0 LICENSE file - Add NOTICE file with third-party attributions - Add SECURITY.md for vulnerability reporting Community: - Add CONTRIBUTING.md with contribution guidelines - Add CODE_OF_CONDUCT.md (Contributor Covenant) - Add GitHub issue templates (bug report, feature request) - Add pull request template Documentation: - Update README.md with compelling value proposition - Add docs/getting-started.md - Add docs/proxy.md for proxy server documentation - Add docs/transforms.md for transform reference - Add docs/api.md for API reference - Add examples/README.md Package Infrastructure: - Add headroom/py.typed for PEP 561 compliance - Add headroom/cli.py for CLI entry point - Add .github/workflows/ci.yml for CI pipeline - Add .github/workflows/publish.yml for PyPI publishing - Update pyproject.toml with proper metadata New Features: - Add multi-provider support (Google, Cohere, LiteLLM, OpenAI-compatible) - Add universal tokenizer registry with multiple backends - Add model registry with pricing and context limits - Add production proxy server with caching and rate limiting Code Quality: - Fix 83 lint issues via ruff auto-fix - Fix version consistency (benchmarks 0.1.0 → 0.2.0) - Add skip decorators for optional dependency tests
124 lines
4.2 KiB
Python
124 lines
4.2 KiB
Python
"""Tests for Cohere provider."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from headroom.providers import CohereProvider
|
|
|
|
|
|
class TestCohereProvider:
|
|
"""Tests for CohereProvider."""
|
|
|
|
@pytest.fixture
|
|
def provider(self):
|
|
"""Create Cohere provider without client (estimation mode)."""
|
|
return CohereProvider()
|
|
|
|
def test_name(self, provider):
|
|
"""Test provider name."""
|
|
assert provider.name == "cohere"
|
|
|
|
def test_supports_command_models(self, provider):
|
|
"""Test support for Command models."""
|
|
assert provider.supports_model("command-r-plus") is True
|
|
assert provider.supports_model("command-r") is True
|
|
assert provider.supports_model("command-a") is True
|
|
assert provider.supports_model("command") 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
|
|
assert provider.supports_model("gemini-2.0") is False
|
|
|
|
def test_get_token_counter(self, provider):
|
|
"""Test getting token counter."""
|
|
counter = provider.get_token_counter("command-r-plus")
|
|
assert counter is not None
|
|
count = counter.count_text("Hello, world!")
|
|
assert count > 0
|
|
|
|
def test_get_context_limit_command_a(self, provider):
|
|
"""Test context limit for Command A (256K)."""
|
|
limit = provider.get_context_limit("command-a")
|
|
assert limit == 256000
|
|
|
|
def test_get_context_limit_command_r_plus(self, provider):
|
|
"""Test context limit for Command R+."""
|
|
limit = provider.get_context_limit("command-r-plus")
|
|
assert limit == 128000
|
|
|
|
def test_get_context_limit_command_r(self, provider):
|
|
"""Test context limit for Command R."""
|
|
limit = provider.get_context_limit("command-r")
|
|
assert limit == 128000
|
|
|
|
def test_get_context_limit_legacy_command(self, provider):
|
|
"""Test context limit for legacy Command."""
|
|
limit = provider.get_context_limit("command")
|
|
assert limit == 4096
|
|
|
|
def test_estimate_cost_command_r_plus(self, provider):
|
|
"""Test cost estimation for Command R+."""
|
|
cost = provider.estimate_cost(
|
|
input_tokens=1000000,
|
|
output_tokens=500000,
|
|
model="command-r-plus",
|
|
)
|
|
assert cost is not None
|
|
# 1M input * $2.50/1M + 0.5M output * $10.00/1M = $2.50 + $5.00 = $7.50
|
|
assert abs(cost - 7.50) < 0.01
|
|
|
|
def test_estimate_cost_command_r(self, provider):
|
|
"""Test cost estimation for Command R."""
|
|
cost = provider.estimate_cost(
|
|
input_tokens=1000000,
|
|
output_tokens=500000,
|
|
model="command-r",
|
|
)
|
|
assert cost is not None
|
|
# 1M input * $0.15/1M + 0.5M output * $0.60/1M = $0.15 + $0.30 = $0.45
|
|
assert abs(cost - 0.45) < 0.01
|
|
|
|
def test_estimate_cost_unknown_model(self, provider):
|
|
"""Test cost estimation returns None for unknown model."""
|
|
cost = provider.estimate_cost(
|
|
input_tokens=1000,
|
|
output_tokens=500,
|
|
model="unknown-model",
|
|
)
|
|
assert cost is None
|
|
|
|
|
|
class TestCohereTokenCounter:
|
|
"""Tests for CohereTokenCounter."""
|
|
|
|
@pytest.fixture
|
|
def counter(self):
|
|
"""Create token counter without client."""
|
|
provider = CohereProvider()
|
|
return provider.get_token_counter("command-r-plus")
|
|
|
|
def test_count_text_empty(self, counter):
|
|
"""Test counting empty text."""
|
|
assert counter.count_text("") == 0
|
|
|
|
def test_count_text_simple(self, counter):
|
|
"""Test counting simple text."""
|
|
count = counter.count_text("Hello, world!")
|
|
assert count > 0
|
|
assert count < 20 # Should be a few tokens
|
|
|
|
def test_count_messages(self, counter):
|
|
"""Test counting messages."""
|
|
messages = [
|
|
{"role": "user", "content": "Hello!"},
|
|
{"role": "assistant", "content": "Hi there!"},
|
|
]
|
|
count = counter.count_messages(messages)
|
|
assert count > 0
|
|
|
|
def test_count_messages_empty(self, counter):
|
|
"""Test counting empty messages."""
|
|
assert counter.count_messages([]) == 0
|