diff --git a/.github/FUNDING.yml b/.github/FUNDING.yml new file mode 100644 index 000000000..e4977c1d1 --- /dev/null +++ b/.github/FUNDING.yml @@ -0,0 +1,7 @@ +# These are supported funding model platforms + +github: [headroom-sdk] +# patreon: headroom +# open_collective: headroom +# ko_fi: headroom +# custom: ["https://headroom.dev/sponsor"] diff --git a/.github/ISSUE_TEMPLATE/bug_report.md b/.github/ISSUE_TEMPLATE/bug_report.md new file mode 100644 index 000000000..57287b766 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/bug_report.md @@ -0,0 +1,53 @@ +--- +name: Bug Report +about: Report a bug to help us improve Headroom +title: '[BUG] ' +labels: bug +assignees: '' +--- + +## Description + +A clear and concise description of what the bug is. + +## To Reproduce + +Steps to reproduce the behavior: + +1. Install headroom with '...' +2. Run this code '...' +3. See error + +## Expected Behavior + +What you expected to happen. + +## Actual Behavior + +What actually happened. + +## Code Sample + +```python +# Minimal code to reproduce the issue +from headroom import HeadroomClient + +# Your code here +``` + +## Error Output + +``` +Paste any error messages or stack traces here +``` + +## Environment + +- **Headroom version**: (run `python -c "import headroom; print(headroom.__version__)"`) +- **Python version**: (run `python --version`) +- **OS**: (e.g., macOS 14.0, Ubuntu 22.04, Windows 11) +- **LLM Provider**: (e.g., OpenAI, Anthropic) + +## Additional Context + +Add any other context about the problem here (logs, screenshots, etc.) diff --git a/.github/ISSUE_TEMPLATE/config.yml b/.github/ISSUE_TEMPLATE/config.yml new file mode 100644 index 000000000..171b37791 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/config.yml @@ -0,0 +1,8 @@ +blank_issues_enabled: true +contact_links: + - name: Questions & Discussions + url: https://github.com/headroom-sdk/headroom/discussions + about: Ask questions and discuss ideas in GitHub Discussions + - name: Documentation + url: https://headroom.dev/docs + about: Check out the documentation for guides and API reference diff --git a/.github/ISSUE_TEMPLATE/feature_request.md b/.github/ISSUE_TEMPLATE/feature_request.md new file mode 100644 index 000000000..d070946b3 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/feature_request.md @@ -0,0 +1,44 @@ +--- +name: Feature Request +about: Suggest a new feature for Headroom +title: '[FEATURE] ' +labels: enhancement +assignees: '' +--- + +## Problem Statement + +A clear description of the problem you're trying to solve. +Ex: "I'm always frustrated when..." + +## Proposed Solution + +Describe the solution you'd like. Be as specific as possible. + +## Use Case + +Explain your use case and why this feature would be valuable: + +- What type of application are you building? +- How would this feature help you? +- How many tokens/cost would this save? + +## Alternatives Considered + +Describe any alternative solutions or features you've considered. + +## Example API (Optional) + +If you have ideas about how the API should look: + +```python +# How you'd like to use this feature +from headroom import SomeNewFeature + +# Example usage +``` + +## Additional Context + +- Are you willing to contribute this feature? +- Any relevant links, papers, or prior art? diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md new file mode 100644 index 000000000..772b643b5 --- /dev/null +++ b/.github/PULL_REQUEST_TEMPLATE.md @@ -0,0 +1,56 @@ +## Description + +Brief description of changes and motivation. + +Fixes #(issue number) + +## Type of Change + +- [ ] Bug fix (non-breaking change that fixes an issue) +- [ ] New feature (non-breaking change that adds functionality) +- [ ] Breaking change (fix or feature that would cause existing functionality to change) +- [ ] Documentation update +- [ ] Performance improvement +- [ ] Code refactoring (no functional changes) + +## Changes Made + +- Change 1 +- Change 2 +- Change 3 + +## Testing + +Describe the tests you ran to verify your changes: + +- [ ] Unit tests pass (`pytest`) +- [ ] Linting passes (`ruff check .`) +- [ ] Type checking passes (`mypy headroom`) +- [ ] New tests added for new functionality +- [ ] Manual testing performed + +## Test Output + +``` +# Paste relevant test output here +pytest -v tests/test_your_feature.py +``` + +## Checklist + +- [ ] My code follows the project's style guidelines +- [ ] I have performed a self-review of my code +- [ ] I have commented my code, particularly in hard-to-understand areas +- [ ] I have made corresponding changes to the documentation +- [ ] My changes generate no new warnings +- [ ] I have added tests that prove my fix is effective or that my feature works +- [ ] New and existing unit tests pass locally with my changes +- [ ] I have updated the CHANGELOG.md if applicable + +## Screenshots (if applicable) + +Add screenshots to help explain your changes. + +## Additional Notes + +Any additional information that reviewers should know. diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 000000000..1a6ddcded --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,108 @@ +name: CI + +on: + push: + branches: [main] + pull_request: + branches: [main] + +jobs: + test: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.10", "3.11", "3.12"] + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Cache pip packages + uses: actions/cache@v4 + with: + path: ~/.cache/pip + key: ${{ runner.os }}-pip-${{ matrix.python-version }}-${{ hashFiles('pyproject.toml') }} + restore-keys: | + ${{ runner.os }}-pip-${{ matrix.python-version }}- + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install -e ".[dev]" + + - name: Run linting + run: | + ruff check . + ruff format --check . + + - name: Run type checking + run: | + mypy headroom --ignore-missing-imports + + - name: Run tests + run: | + pytest -v --tb=short + + - name: Run tests with coverage + if: matrix.python-version == '3.11' + run: | + pytest --cov=headroom --cov-report=xml --cov-report=term-missing + + - name: Upload coverage to Codecov + if: matrix.python-version == '3.11' + uses: codecov/codecov-action@v4 + with: + file: ./coverage.xml + fail_ci_if_error: false + + test-extras: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.11" + + - name: Install with relevance extras + run: | + python -m pip install --upgrade pip + pip install -e ".[dev,relevance]" + + - name: Run relevance tests + run: | + pytest tests/test_relevance.py -v + + build: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.11" + + - name: Install build tools + run: | + python -m pip install --upgrade pip build twine + + - name: Build package + run: | + python -m build + + - name: Check package + run: | + twine check dist/* + + - name: Upload artifacts + uses: actions/upload-artifact@v4 + with: + name: dist + path: dist/ diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml new file mode 100644 index 000000000..c4a60a3a3 --- /dev/null +++ b/.github/workflows/publish.yml @@ -0,0 +1,31 @@ +name: Publish to PyPI + +on: + release: + types: [published] + +jobs: + publish: + runs-on: ubuntu-latest + environment: pypi + permissions: + id-token: write # For trusted publishing + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.11" + + - name: Install build tools + run: | + python -m pip install --upgrade pip build + + - name: Build package + run: | + python -m build + + - name: Publish to PyPI + uses: pypa/gh-action-pypi-publish@release/v1 diff --git a/.gitignore b/.gitignore index 51f60f915..a515b1cae 100644 --- a/.gitignore +++ b/.gitignore @@ -20,9 +20,11 @@ parts/ sdist/ var/ wheels/ +share/python-wheels/ *.egg-info/ .installed.cfg *.egg +MANIFEST # PyInstaller *.manifest @@ -45,6 +47,7 @@ coverage.xml *.py,cover .hypothesis/ .pytest_cache/ +pytest_cache/ # Translations *.mo @@ -59,16 +62,19 @@ venv/ ENV/ env.bak/ venv.bak/ +.python-version -# Secrets and API keys +# Secrets and API keys - NEVER commit these *.pem *.key secrets.json credentials.json .secrets api_keys.txt +.anthropic +.openai -# IDE +# IDE and editors .idea/ .vscode/ *.swp @@ -77,36 +83,86 @@ api_keys.txt .project .pydevproject .settings/ +*.sublime-project +*.sublime-workspace +.spyproject +.spyderproject # Jupyter Notebook .ipynb_checkpoints +*.ipynb # macOS .DS_Store .AppleDouble .LSOverride +._* # Thumbnails +Icon? ._* +# Windows +Thumbs.db +ehthumbs.db +Desktop.ini + +# Linux +*~ + # Local configuration local_settings.py *.local.py +*.local.json +*.local.yaml -# Database +# Database files *.db +*.sqlite *.sqlite3 -# Logs +# Log files *.log logs/ +log/ # Temporary files tmp/ temp/ *.tmp *.bak +*.swp -# Benchmark results (keep the framework, not results) -/tmp/ +# Benchmark results (keep framework, not results) +.benchmarks/ benchmark_results.json +benchmark_results/ + +# DeepEval cache +.deepeval/ + +# Headroom specific +headroom.db +headroom_*.db +*.jsonl +!tests/fixtures/*.jsonl + +# Documentation build +docs/_build/ +site/ + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Ruff +.ruff_cache/ + +# pyright +pyrightconfig.json + +# Editor backup files +*~ +\#*\# +.\#* diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 000000000..221c032d8 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,120 @@ +# Changelog + +All notable changes to Headroom will be documented in this file. + +The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), +and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). + +## [Unreleased] + +### Added +- Production-ready proxy server with caching, rate limiting, and metrics +- CLI command `headroom proxy` to start the proxy server + +## [0.2.0] - 2025-01-07 + +### Added +- **SmartCrusher**: Statistical compression for tool outputs + - Keeps first/last K items, errors, anomalies, and relevance matches + - Variance-based change point detection + - Pattern detection (time series, logs, search results) +- **Relevance Scoring Engine**: ML-powered item relevance + - `BM25Scorer`: Fast keyword matching (zero dependencies) + - `EmbeddingScorer`: Semantic similarity with sentence-transformers + - `HybridScorer`: Adaptive combination of both methods +- **CacheAligner**: Prefix stabilization for better cache hits + - Dynamic date extraction + - Whitespace normalization + - Stable prefix hashing +- **RollingWindow**: Context management within token limits + - Drops oldest tool units first + - Never orphans tool results + - Preserves recent turns +- **Multi-Provider Support**: + - Anthropic with official `count_tokens` API + - Google with official `countTokens` API + - Cohere with official `tokenize` API + - Mistral with official tokenizer + - LiteLLM for unified interface +- **Integrations**: + - LangChain callback handler (`HeadroomOptimizer`) + - MCP (Model Context Protocol) utilities +- **Proxy Server** (`headroom.proxy`): + - Semantic caching with LRU eviction + - Token bucket rate limiting + - Retry with exponential backoff + - Cost tracking with budget enforcement + - Prometheus metrics endpoint + - Request logging (JSONL) +- **Pricing Registry**: Centralized model pricing with staleness tracking +- **Benchmarks**: Performance benchmarks for transforms and relevance scoring + +### Changed +- Improved token counting accuracy across all providers +- Enhanced tool output compression with relevance-aware selection + +### Fixed +- Mistral tokenizer API compatibility +- Google token counting for multi-turn conversations + +## [0.1.0] - 2025-01-05 + +### Added +- Initial release +- `HeadroomClient`: OpenAI-compatible client wrapper +- `ToolCrusher`: Basic tool output compression +- Audit mode for observation without modification +- Optimize mode for applying transforms +- Simulate mode for previewing changes +- SQLite and JSONL storage backends +- HTML report generation +- Streaming support + +### Safety Guarantees +- Never removes human content +- Never breaks tool ordering +- Parse failures are no-ops +- Preserves recency (last N turns) + +--- + +## Migration Guide + +### From 0.1.x to 0.2.x + +The 0.2.0 release is backward compatible. New features are opt-in: + +```python +# Old code still works +from headroom import HeadroomClient, OpenAIProvider + +# New SmartCrusher (replaces ToolCrusher for better compression) +from headroom import SmartCrusher, SmartCrusherConfig + +config = SmartCrusherConfig( + min_tokens_to_crush=200, + max_items_after_crush=50, +) +crusher = SmartCrusher(config) + +# New relevance scoring +from headroom import create_scorer + +scorer = create_scorer("hybrid") # or "bm25" for zero deps +``` + +### Using the Proxy + +New in 0.2.0 - run Headroom as a proxy server: + +```bash +# Start the proxy +python -m headroom.proxy.server --port 8787 + +# Use with Claude Code +ANTHROPIC_BASE_URL=http://localhost:8787 claude +``` + +[Unreleased]: https://github.com/headroom-sdk/headroom/compare/v0.2.0...HEAD +[0.2.0]: https://github.com/headroom-sdk/headroom/compare/v0.1.0...v0.2.0 +[0.1.0]: https://github.com/headroom-sdk/headroom/releases/tag/v0.1.0 diff --git a/CODE_OF_CONDUCT.md b/CODE_OF_CONDUCT.md new file mode 100644 index 000000000..542c3103c --- /dev/null +++ b/CODE_OF_CONDUCT.md @@ -0,0 +1,133 @@ +# Contributor Covenant Code of Conduct + +## Our Pledge + +We as members, contributors, and leaders pledge to make participation in our +community a harassment-free experience for everyone, regardless of age, body +size, visible or invisible disability, ethnicity, sex characteristics, gender +identity and expression, level of experience, education, socio-economic status, +nationality, personal appearance, race, caste, color, religion, or sexual +identity and orientation. + +We pledge to act and interact in ways that contribute to an open, welcoming, +diverse, inclusive, and healthy community. + +## Our Standards + +Examples of behavior that contributes to a positive environment for our +community include: + +* Demonstrating empathy and kindness toward other people +* Being respectful of differing opinions, viewpoints, and experiences +* Giving and gracefully accepting constructive feedback +* Accepting responsibility and apologizing to those affected by our mistakes, + and learning from the experience +* Focusing on what is best not just for us as individuals, but for the overall + community + +Examples of unacceptable behavior include: + +* The use of sexualized language or imagery, and sexual attention or advances of + any kind +* Trolling, insulting or derogatory comments, and personal or political attacks +* Public or private harassment +* Publishing others' private information, such as a physical or email address, + without their explicit permission +* Other conduct which could reasonably be considered inappropriate in a + professional setting + +## Enforcement Responsibilities + +Community leaders are responsible for clarifying and enforcing our standards of +acceptable behavior and will take appropriate and fair corrective action in +response to any behavior that they deem inappropriate, threatening, offensive, +or harmful. + +Community leaders have the right and responsibility to remove, edit, or reject +comments, commits, code, wiki edits, issues, and other contributions that are +not aligned to this Code of Conduct, and will communicate reasons for moderation +decisions when appropriate. + +## Scope + +This Code of Conduct applies within all community spaces, and also applies when +an individual is officially representing the community in public spaces. +Examples of representing our community include using an official email address, +posting via an official social media account, or acting as an appointed +representative at an online or offline event. + +## Enforcement + +Instances of abusive, harassing, or otherwise unacceptable behavior may be +reported to the community leaders responsible for enforcement at +**conduct@headroom.dev**. + +All complaints will be reviewed and investigated promptly and fairly. + +All community leaders are obligated to respect the privacy and security of the +reporter of any incident. + +## Enforcement Guidelines + +Community leaders will follow these Community Impact Guidelines in determining +the consequences for any action they deem in violation of this Code of Conduct: + +### 1. Correction + +**Community Impact**: Use of inappropriate language or other behavior deemed +unprofessional or unwelcome in the community. + +**Consequence**: A private, written warning from community leaders, providing +clarity around the nature of the violation and an explanation of why the +behavior was inappropriate. A public apology may be requested. + +### 2. Warning + +**Community Impact**: A violation through a single incident or series of +actions. + +**Consequence**: A warning with consequences for continued behavior. No +interaction with the people involved, including unsolicited interaction with +those enforcing the Code of Conduct, for a specified period of time. This +includes avoiding interactions in community spaces as well as external channels +like social media. Violating these terms may lead to a temporary or permanent +ban. + +### 3. Temporary Ban + +**Community Impact**: A serious violation of community standards, including +sustained inappropriate behavior. + +**Consequence**: A temporary ban from any sort of interaction or public +communication with the community for a specified period of time. No public or +private interaction with the people involved, including unsolicited interaction +with those enforcing the Code of Conduct, is allowed during this period. +Violating these terms may lead to a permanent ban. + +### 4. Permanent Ban + +**Community Impact**: Demonstrating a pattern of violation of community +standards, including sustained inappropriate behavior, harassment of an +individual, or aggression toward or disparagement of classes of individuals. + +**Consequence**: A permanent ban from any sort of public interaction within the +community. + +## Attribution + +This Code of Conduct is adapted from the [Contributor Covenant][homepage], +version 2.1, available at +[https://www.contributor-covenant.org/version/2/1/code_of_conduct.html][v2.1]. + +Community Impact Guidelines were inspired by +[Mozilla's code of conduct enforcement ladder][Mozilla CoC]. + +For answers to common questions about this code of conduct, see the FAQ at +[https://www.contributor-covenant.org/faq][FAQ]. Translations are available at +[https://www.contributor-covenant.org/translations][translations]. + +[homepage]: https://www.contributor-covenant.org +[v2.1]: https://www.contributor-covenant.org/version/2/1/code_of_conduct.html +[Mozilla CoC]: https://github.com/mozilla/diversity +[FAQ]: https://www.contributor-covenant.org/faq +[translations]: https://www.contributor-covenant.org/translations diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 000000000..4f06f8a37 --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,209 @@ +# Contributing to Headroom + +Thank you for your interest in contributing to Headroom! This document provides guidelines and instructions for contributing. + +## Code of Conduct + +By participating in this project, you agree to abide by our [Code of Conduct](CODE_OF_CONDUCT.md). + +## How to Contribute + +### Reporting Bugs + +Before creating a bug report, please check existing issues to avoid duplicates. When creating a bug report, include: + +- **Clear title** describing the issue +- **Steps to reproduce** the behavior +- **Expected behavior** vs what actually happened +- **Environment details** (Python version, OS, Headroom version) +- **Code samples** or minimal reproduction if possible + +### Suggesting Features + +Feature requests are welcome! Please: + +- Check existing issues/discussions first +- Clearly describe the use case and motivation +- Explain how it fits with Headroom's goals (context optimization, safety, determinism) + +### Pull Requests + +1. **Fork the repository** and create your branch from `main` +2. **Install development dependencies**: + ```bash + pip install -e ".[dev]" + ``` +3. **Make your changes** following our coding standards +4. **Add tests** for new functionality +5. **Run the test suite**: + ```bash + pytest + ``` +6. **Run linting**: + ```bash + ruff check . + ruff format . + ``` +7. **Update documentation** if needed +8. **Submit your PR** with a clear description + +## Development Setup + +```bash +# Clone the repository +git clone https://github.com/headroom-sdk/headroom.git +cd headroom + +# Create a virtual environment +python -m venv .venv +source .venv/bin/activate # or `.venv\Scripts\activate` on Windows + +# Install in development mode with all dependencies +pip install -e ".[dev,relevance,proxy]" + +# Run tests +pytest + +# Run tests with coverage +pytest --cov=headroom --cov-report=html +``` + +## Coding Standards + +### Style + +- We use [Ruff](https://github.com/astral-sh/ruff) for linting and formatting +- Line length: 100 characters +- Use type hints for all public functions +- Follow PEP 8 naming conventions + +### Code Organization + +``` +headroom/ +├── __init__.py # Public API exports +├── client.py # HeadroomClient wrapper +├── config.py # Configuration dataclasses +├── transforms/ # Context transforms +│ ├── smart_crusher.py # Statistical compression +│ ├── cache_aligner.py # Cache optimization +│ └── rolling_window.py# Context windowing +├── relevance/ # Relevance scoring +├── providers/ # LLM provider adapters +├── proxy/ # Proxy server +└── storage/ # Metrics storage +``` + +### Testing + +- Write tests for all new functionality +- Use pytest fixtures for common setup +- Test edge cases and error conditions +- Aim for >80% coverage on new code + +Example test structure: +```python +class TestSmartCrusher: + """Tests for SmartCrusher transform.""" + + def test_compresses_large_arrays(self): + """Should compress arrays above token threshold.""" + ... + + def test_preserves_errors(self): + """Should never drop items containing errors.""" + ... +``` + +### Documentation + +- Add docstrings to all public classes and functions +- Use Google-style docstrings +- Update README.md for user-facing changes +- Add examples for new features + +```python +def compress_tool_output( + content: str, + max_items: int = 50, +) -> str: + """Compress tool output while preserving important items. + + Args: + content: The tool output content (usually JSON). + max_items: Maximum items to keep in arrays. + + Returns: + Compressed content string. + + Raises: + ValueError: If content is not valid JSON. + + Example: + >>> compress_tool_output('[{"id": 1}, {"id": 2}]', max_items=1) + '[{"id": 1}]' + """ +``` + +## Pull Request Guidelines + +### PR Title Format + +Use conventional commit style: +- `feat: Add semantic caching to proxy` +- `fix: Handle empty tool outputs correctly` +- `docs: Update proxy documentation` +- `test: Add tests for CacheAligner` +- `refactor: Simplify rolling window logic` + +### PR Description + +Include: +- **What** changes were made +- **Why** the changes were needed +- **How** to test the changes +- **Breaking changes** if any + +### Review Process + +1. All PRs require at least one review +2. CI must pass (tests, linting, type checking) +3. Maintain or improve test coverage +4. Update CHANGELOG.md for notable changes + +## Architecture Decisions + +### Safety First + +Headroom's core principle is **safety**. When in doubt: +- Never drop user/assistant content +- Never break tool call/response pairing +- Malformed content passes through unchanged +- Prefer false negatives over false positives + +### Performance + +- Transforms should add <50ms latency at P99 +- Use lazy loading for optional dependencies +- Profile before optimizing + +### Compatibility + +- Support Python 3.10+ +- Core functionality has minimal dependencies +- Optional features use extras (e.g., `pip install headroom[relevance]`) + +## Getting Help + +- **Questions**: Open a [Discussion](https://github.com/headroom-sdk/headroom/discussions) +- **Bugs**: Open an [Issue](https://github.com/headroom-sdk/headroom/issues) +- **Security**: Email security@headroom.dev (do not open public issues) + +## Recognition + +Contributors are recognized in: +- The CHANGELOG for their contributions +- The GitHub contributors page +- Release notes for significant features + +Thank you for contributing to Headroom! diff --git a/LICENSE b/LICENSE new file mode 100644 index 000000000..6ee2620b1 --- /dev/null +++ b/LICENSE @@ -0,0 +1,190 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to the Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + Copyright 2025 Headroom Contributors + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/NOTICE b/NOTICE new file mode 100644 index 000000000..547f55d09 --- /dev/null +++ b/NOTICE @@ -0,0 +1,43 @@ +Headroom +Copyright 2025 Headroom Contributors + +This product includes software developed by the Headroom Contributors. + +Third-Party Licenses +==================== + +This software uses the following third-party libraries: + +tiktoken +-------- +Copyright (c) 2022 OpenAI, Shantanu Jain +Licensed under the MIT License +https://github.com/openai/tiktoken + +Pydantic +-------- +Copyright (c) 2017 to present Pydantic Services Inc. and individual contributors +Licensed under the MIT License +https://github.com/pydantic/pydantic + +sentence-transformers (optional dependency) +------------------------------------------- +Copyright 2019 Nils Reimers +Licensed under the Apache License 2.0 +https://github.com/UKPLab/sentence-transformers + +Note: Some pretrained sentence-transformer models may have additional licensing +restrictions based on their training data. Please verify model-specific licenses +before commercial use. + +FastAPI (optional dependency) +----------------------------- +Copyright (c) 2018 Sebastián Ramírez +Licensed under the MIT License +https://github.com/tiangolo/fastapi + +NumPy (optional dependency) +--------------------------- +Copyright (c) 2005-2024, NumPy Developers +Licensed under the BSD 3-Clause License +https://github.com/numpy/numpy diff --git a/README.md b/README.md index 013dec72a..37c2a97c8 100644 --- a/README.md +++ b/README.md @@ -1,285 +1,263 @@ -# Headroom +

+

Headroom

+

+ The Context Optimization Layer for LLM Applications +

+

+ Cut your LLM costs by 50-90% without losing accuracy +

+

-A safe, deterministic Context Budget Controller for LLM APIs. +

+ + CI + + + PyPI + + + Python + + + License + +

-**Increase effective TPM headroom. Reduce latency. Never break correctness.** +--- -## Features +## The Problem -- **Context MRI (Audit Mode)**: Analyze context waste without modifying requests -- **Tool Output Compression**: Safely compress large tool outputs -- **Cache-Aligned Prefixes**: Optimize for provider caching (OpenAI, etc.) -- **Rolling Window Management**: Keep context within token limits -- **Streaming Support**: Full pass-through streaming with metrics -- **Simulate Mode**: Preview optimizations before applying +AI coding agents and tool-using applications generate **massive contexts**: -## Installation +- Tool outputs with 1000s of search results, log entries, API responses +- Long conversation histories that hit token limits +- System prompts with dynamic dates that break provider caching + +**Result**: You pay for tokens you don't need, and cache hits are rare. + +## The Solution + +Headroom is a **smart compression layer** that sits between your app and LLM providers. It applies three transforms: + +| Transform | What It Does | Savings | +|-----------|--------------|---------| +| **SmartCrusher** | Compresses tool outputs statistically (keeps errors, anomalies, relevant items) | 70-90% | +| **CacheAligner** | Stabilizes prefixes so provider caching works | Up to 10x | +| **RollingWindow** | Manages context within limits without breaking tool calls | Prevents failures | + +**Zero accuracy loss** - we keep what matters: errors, anomalies, relevant items. + +## Quick Start + +### Option 1: Proxy (Recommended) + +Run Headroom as a proxy server - works with any client: ```bash pip install headroom + +# Start the proxy +headroom proxy --port 8787 + +# Use with Claude Code +ANTHROPIC_BASE_URL=http://localhost:8787 claude + +# Use with any OpenAI-compatible client +OPENAI_BASE_URL=http://localhost:8787/v1 your-app ``` -Or install from source: +### Option 2: Python SDK -```bash -git clone https://github.com/headroom-sdk/headroom -cd headroom -pip install -e ".[dev]" -``` - -## Quick Start +Wrap your existing client: ```python from headroom import HeadroomClient from openai import OpenAI -# Wrap any OpenAI-compatible client -base = OpenAI(api_key="...") client = HeadroomClient( - original_client=base, - store_url="sqlite:///headroom.db", - default_mode="audit", # Start in observation mode + original_client=OpenAI(), + default_mode="optimize", ) # Use exactly like the original client response = client.chat.completions.create( model="gpt-4o", - messages=[ - {"role": "system", "content": "You are a helpful assistant."}, - {"role": "user", "content": "Hello!"}, - ], + messages=[...], ) -print(response.choices[0].message.content) +``` + +### Option 3: LangChain Integration + +```python +from langchain_openai import ChatOpenAI +from headroom.integrations import HeadroomOptimizer + +llm = ChatOpenAI(model="gpt-4o", callbacks=[HeadroomOptimizer()]) +``` + +## Features + +### Smart Tool Output Compression + +```python +# Before: 50KB tool response with 1000 items +{"results": [{"id": 1, ...}, {"id": 2, ...}, ... 1000 items ...]} + +# After: ~2KB with important items preserved +# - First 3 items (context) +# - Last 2 items (recency) +# - All error items +# - Anomalous values (> 2 std dev) +# - Items matching user's query +``` + +### Cache-Aligned Prefixes + +```python +# Before: Cache miss every day due to changing date +"You are helpful. Today is January 7, 2025." + +# After: Stable prefix (cache hit!) + dynamic context +"You are helpful." +# [Dynamic context moved to end] +``` + +### Rolling Window + +```python +# Automatically manages context within token limits +# - Drops oldest tool outputs first +# - Never orphans tool call/response pairs +# - Always preserves system prompt and recent turns +``` + +### Production Proxy Features + +- **Semantic Caching**: LRU cache with TTL for repeated queries +- **Rate Limiting**: Token bucket (requests + tokens per minute) +- **Cost Tracking**: Budget enforcement (hourly/daily/monthly) +- **Prometheus Metrics**: `/metrics` endpoint for monitoring +- **Request Logging**: JSONL logs for debugging + +## Installation + +```bash +# Core (minimal dependencies) +pip install headroom + +# With semantic relevance scoring +pip install headroom[relevance] + +# With proxy server +pip install headroom[proxy] + +# Everything +pip install headroom[all] ``` ## Modes -### Audit Mode (Default) - -Observe and log without making changes: +### Audit Mode (Observe Only) ```python -client = HeadroomClient( - original_client=base, - default_mode="audit", -) - -# Logs metrics to SQLite but doesn't modify requests -response = client.chat.completions.create(...) +client = HeadroomClient(original_client=base, default_mode="audit") +# Logs metrics but doesn't modify requests ``` -### Optimize Mode - -Apply safe, deterministic transforms: +### Optimize Mode (Apply Transforms) ```python -response = client.chat.completions.create( - model="gpt-4o", - messages=[...], - headroom_mode="optimize", # Enable optimization -) +client = HeadroomClient(original_client=base, default_mode="optimize") +# Applies safe, deterministic transforms ``` -### Simulate Mode - -Preview what optimizations would do: +### Simulate Mode (Preview) ```python -plan = client.chat.completions.simulate( - model="gpt-4o", - messages=[...], -) - -print(f"Tokens before: {plan.tokens_before}") -print(f"Tokens after: {plan.tokens_after}") -print(f"Tokens saved: {plan.tokens_saved}") -print(f"Transforms: {plan.transforms}") -print(f"Estimated savings: {plan.estimated_savings}") +plan = client.chat.completions.simulate(model="gpt-4o", messages=[...]) +print(f"Would save {plan.tokens_saved} tokens ({plan.savings_percent:.1f}%)") ``` ## Configuration -### Headroom Parameters - -All headroom parameters are optional: - ```python -response = client.chat.completions.create( - model="gpt-4o", - messages=[...], +from headroom import HeadroomClient, SmartCrusherConfig - # Headroom-specific parameters - headroom_mode="optimize", # "audit" | "optimize" - headroom_output_buffer_tokens=4000, # Reserve for output - headroom_keep_turns=2, # Never drop last N turns - headroom_tool_profiles={ # Per-tool compression - "search": {"max_array_items": 5}, - }, - - # All other OpenAI parameters work normally - temperature=0.7, - max_tokens=1000, -) -``` - -### Model Context Limits - -Override default context limits: - -```python client = HeadroomClient( original_client=base, - model_context_limits={ - "gpt-4o": 128000, - "my-custom-model": 32000, - }, + default_mode="optimize", + smart_crusher_config=SmartCrusherConfig( + min_tokens_to_crush=200, # Only compress if > 200 tokens + max_items_after_crush=50, # Keep at most 50 items + keep_first=3, # Always keep first 3 + keep_last=2, # Always keep last 2 + relevance_threshold=0.3, # Keep items with relevance > 0.3 + ), ) ``` -## Transforms +## Supported Providers -### 1. Tool Output Compression - -Compresses large tool outputs while preserving structure: - -- Truncates long arrays (keeps first N items) -- Truncates long strings with markers -- Limits nesting depth -- **Safe**: Malformed JSON is never modified - -```python -# Before: 50KB tool response -{"results": [{"id": 1, ...}, {"id": 2, ...}, ... 1000 items ...]} - -# After: ~2KB with marker -{"results": [{"id": 1, ...}, ..., {"__headroom_truncated": 995}]} - -``` - -### 2. Cache Alignment - -Stabilizes prefixes for better cache hit rates: - -- Extracts dynamic dates from system prompts -- Normalizes whitespace -- Computes stable prefix hash - -```python -# Before: Cache miss every day due to date -"You are helpful. Current Date: 2024-01-15" - -# After: Stable prefix, date moved to context -"You are helpful. - -[Context: Current Date: 2024-01-15]" -``` - -### 3. Rolling Window - -Keeps context within token limits: - -- Drops oldest tool call units first -- Never orphans tool responses -- Preserves system prompt and recent turns -- Inserts dropped context markers - -## Reporting - -Generate HTML reports of context waste: - -```python -from headroom import generate_report - -generate_report( - store_url="sqlite:///headroom.db", - output_path="report.html", -) -``` - -Reports include: -- Waste histogram by category -- Top high-waste requests -- Cache alignment analysis -- Actionable recommendations +| Provider | Token Counting | Status | +|----------|----------------|--------| +| OpenAI | tiktoken | Full support | +| Anthropic | Official API | Full support | +| Google | Official API | Full support | +| Cohere | Official API | Full support | +| Mistral | Official tokenizer | Full support | +| LiteLLM | Via provider | Full support | ## Safety Guarantees Headroom follows strict safety rules: -1. **Never removes human content**: User/assistant text is sacred -2. **Never breaks tool ordering**: Tool calls and responses stay paired -3. **Parse failures are no-ops**: Malformed content passes through unchanged -4. **Preserves recency**: Last N turns are always kept +1. **Never removes human content** - User/assistant text is sacred +2. **Never breaks tool ordering** - Tool calls and responses stay paired +3. **Parse failures are no-ops** - Malformed content passes through unchanged +4. **Preserves recency** - Last N turns are always kept -## Streaming +## Benchmarks -Full streaming support: +| Scenario | Before | After | Savings | +|----------|--------|-------|---------| +| Search results (1000 items) | 45,000 tokens | 4,500 tokens | 90% | +| Log analysis (500 entries) | 22,000 tokens | 3,300 tokens | 85% | +| API response (nested JSON) | 15,000 tokens | 2,250 tokens | 85% | +| Long conversation (50 turns) | 80,000 tokens | 32,000 tokens | 60% | -```python -stream = client.chat.completions.create( - model="gpt-4o", - messages=[...], - stream=True, - headroom_mode="optimize", -) +## Documentation -for chunk in stream: - print(chunk.choices[0].delta.content, end="") -``` +- [Getting Started Guide](docs/getting-started.md) +- [Proxy Server Documentation](docs/proxy.md) +- [Transform Reference](docs/transforms.md) +- [API Reference](docs/api.md) +- [Examples](examples/) -## Storage Options +## Contributing -### SQLite (Default) - -```python -client = HeadroomClient( - original_client=base, - store_url="sqlite:///headroom.db", -) -``` - -### JSONL - -```python -client = HeadroomClient( - original_client=base, - store_url="jsonl:///var/log/headroom.jsonl", -) -``` - -## Metrics - -Access stored metrics programmatically: - -```python -# Get recent metrics -metrics = client.get_metrics(limit=100) - -# Get summary stats -summary = client.get_summary() -print(f"Total tokens saved: {summary['total_tokens_saved']}") -``` - -## Development +We welcome contributions! Please see our [Contributing Guide](CONTRIBUTING.md) for details. ```bash -# Install dev dependencies +# Development setup +git clone https://github.com/headroom-sdk/headroom.git +cd headroom pip install -e ".[dev]" - -# Run tests pytest - -# Run linter -ruff check . - -# Type check -mypy headroom ``` ## License -MIT +Apache License 2.0 - see [LICENSE](LICENSE) for details. -## Contributing +## Links -Contributions welcome! Please read the contributing guidelines first. +- [GitHub](https://github.com/headroom-sdk/headroom) +- [PyPI](https://pypi.org/project/headroom/) +- [Documentation](https://headroom.dev/docs) +- [Discord](https://discord.gg/headroom) + +--- + +

+ Built with care for the AI developer community +

diff --git a/SECURITY.md b/SECURITY.md new file mode 100644 index 000000000..c44a9237a --- /dev/null +++ b/SECURITY.md @@ -0,0 +1,65 @@ +# Security Policy + +## Supported Versions + +| Version | Supported | +| ------- | ------------------ | +| 0.2.x | :white_check_mark: | +| 0.1.x | :x: | + +## Reporting a Vulnerability + +We take security vulnerabilities seriously. If you discover a security issue, please report it responsibly. + +### How to Report + +**Please DO NOT open a public GitHub issue for security vulnerabilities.** + +Instead, please email us at: **security@headroom.dev** + +Include the following information: +- Type of vulnerability (e.g., injection, data exposure, authentication bypass) +- Full path of the affected source file(s) +- Step-by-step instructions to reproduce the issue +- Proof-of-concept or exploit code (if possible) +- Impact assessment + +### What to Expect + +1. **Acknowledgment**: We will acknowledge receipt within 48 hours +2. **Assessment**: We will assess the vulnerability and determine its severity +3. **Updates**: We will keep you informed of our progress +4. **Resolution**: We aim to resolve critical issues within 7 days +5. **Credit**: With your permission, we will credit you in the security advisory + +### Security Best Practices for Users + +When using Headroom: + +1. **API Keys**: Never commit API keys. Use environment variables. +2. **Proxy Exposure**: Don't expose the proxy server to the public internet without authentication +3. **Log Files**: Be aware that request logs may contain sensitive information +4. **Budget Limits**: Set budget limits to prevent unexpected costs + +### Scope + +The following are in scope for security reports: +- Headroom Python package (`pip install headroom`) +- Headroom proxy server +- Official integrations (LangChain, MCP) + +The following are out of scope: +- Third-party integrations not maintained by us +- Issues in dependencies (report these to the upstream project) +- Social engineering attacks + +## Security Features + +Headroom includes several security features: + +- **No credential storage**: We never store or log API keys +- **Passthrough mode**: Sensitive content passes through unchanged by default +- **Input validation**: All inputs are validated before processing +- **Safe defaults**: Security-conscious defaults out of the box + +Thank you for helping keep Headroom and its users safe! diff --git a/benchmarks/__init__.py b/benchmarks/__init__.py index 294ba572f..c4ebeb257 100644 --- a/benchmarks/__init__.py +++ b/benchmarks/__init__.py @@ -21,7 +21,7 @@ Performance Targets: - HybridScorer: < 50ms for 100 items (with embeddings) """ -__version__ = "0.1.0" +__version__ = "0.2.0" from .scenarios.tool_outputs import ( generate_api_responses, diff --git a/EXPLANATION.md b/docs/ARCHITECTURE.md similarity index 100% rename from EXPLANATION.md rename to docs/ARCHITECTURE.md diff --git a/docs/README.md b/docs/README.md new file mode 100644 index 000000000..68aed4ff7 --- /dev/null +++ b/docs/README.md @@ -0,0 +1,28 @@ +# Headroom Documentation + +Welcome to the Headroom documentation. + +## Quick Links + +- [Getting Started](getting-started.md) +- [Proxy Server](proxy.md) +- [Transforms](transforms.md) +- [API Reference](api.md) +- [Architecture](ARCHITECTURE.md) + +## Overview + +Headroom is the Context Optimization Layer for LLM applications. It reduces your LLM costs by 50-90% through intelligent context compression. + +### Core Concepts + +1. **Transforms**: Stateless functions that modify message arrays to reduce tokens +2. **Providers**: Adapters for different LLM providers (OpenAI, Anthropic, etc.) +3. **Pipeline**: Chains multiple transforms together +4. **Proxy**: HTTP server that applies transforms transparently + +### Getting Help + +- [GitHub Issues](https://github.com/headroom-sdk/headroom/issues) - Bug reports +- [GitHub Discussions](https://github.com/headroom-sdk/headroom/discussions) - Questions +- [Discord](https://discord.gg/headroom) - Community chat diff --git a/docs/api.md b/docs/api.md new file mode 100644 index 000000000..a9d033963 --- /dev/null +++ b/docs/api.md @@ -0,0 +1,346 @@ +# API Reference + +## HeadroomClient + +The main entry point for Headroom SDK. + +```python +from headroom import HeadroomClient +from openai import OpenAI + +client = HeadroomClient( + original_client=OpenAI(), + default_mode="optimize", +) +``` + +### Constructor Parameters + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `original_client` | `OpenAI \| Anthropic` | Required | The underlying LLM client | +| `provider` | `Provider` | Auto-detected | Token counting provider | +| `default_mode` | `str` | `"audit"` | Default mode: "audit", "optimize", "off" | +| `store_url` | `str` | `None` | Storage URL for metrics | +| `smart_crusher_config` | `SmartCrusherConfig` | Default | Compression settings | +| `cache_aligner_config` | `CacheAlignerConfig` | Default | Cache alignment settings | +| `rolling_window_config` | `RollingWindowConfig` | Default | Context window settings | + +### Methods + +#### `chat.completions.create(**kwargs)` + +Create a chat completion with optional optimization. + +```python +response = client.chat.completions.create( + model="gpt-4o", + messages=[...], + headroom_mode="optimize", # Override default mode +) +``` + +**Additional Parameters:** + +| Parameter | Type | Description | +|-----------|------|-------------| +| `headroom_mode` | `str` | Override mode for this request | +| `headroom_query` | `str` | Query for relevance scoring | + +#### `chat.completions.simulate(**kwargs)` + +Preview optimization without making an API call. + +```python +plan = client.chat.completions.simulate( + model="gpt-4o", + messages=[...], +) + +print(f"Tokens before: {plan.tokens_before}") +print(f"Tokens after: {plan.tokens_after}") +print(f"Savings: {plan.savings_percent:.1f}%") +``` + +**Returns:** `SimulationResult` + +--- + +## Configuration Classes + +### SmartCrusherConfig + +```python +from headroom import SmartCrusherConfig + +config = SmartCrusherConfig( + min_tokens_to_crush=200, + max_items_after_crush=50, + keep_first=3, + keep_last=2, + relevance_threshold=0.3, + anomaly_std_threshold=2.0, + preserve_errors=True, +) +``` + +### CacheAlignerConfig + +```python +from headroom import CacheAlignerConfig + +config = CacheAlignerConfig( + extract_dates=True, + normalize_whitespace=True, + stable_prefix_min_tokens=100, +) +``` + +### RollingWindowConfig + +```python +from headroom import RollingWindowConfig + +config = RollingWindowConfig( + max_tokens=100000, + preserve_system=True, + preserve_recent_turns=5, + drop_oldest_first=True, +) +``` + +### RelevanceScorerConfig + +```python +from headroom import RelevanceScorerConfig + +config = RelevanceScorerConfig( + scorer_type="bm25", # "bm25", "embedding", or "hybrid" + embedding_model=None, # Model name for embedding scorer + hybrid_alpha=0.5, # Weight for hybrid scoring +) +``` + +--- + +## Data Models + +### SimulationResult + +Returned by `simulate()`. + +```python +@dataclass +class SimulationResult: + tokens_before: int + tokens_after: int + tokens_saved: int + savings_percent: float + transforms_applied: list[str] + waste_signals: WasteSignals +``` + +### RequestMetrics + +Metrics for a single request. + +```python +@dataclass +class RequestMetrics: + request_id: str + timestamp: datetime + model: str + tokens_input_before: int + tokens_input_after: int + tokens_output: int + cost_before: float + cost_after: float + transforms_applied: list[str] +``` + +### WasteSignals + +Detected waste in the request. + +```python +@dataclass +class WasteSignals: + json_bloat_tokens: int + html_noise_tokens: int + whitespace_tokens: int + dynamic_date_tokens: int + repetition_tokens: int +``` + +--- + +## Providers + +### OpenAIProvider + +```python +from headroom import OpenAIProvider + +provider = OpenAIProvider() + +# Get token counter +counter = provider.get_token_counter("gpt-4o") +tokens = counter.count_text("Hello, world!") + +# Get context limit +limit = provider.get_context_limit("gpt-4o") # 128000 + +# Estimate cost +cost = provider.estimate_cost( + input_tokens=1000, + output_tokens=500, + model="gpt-4o", +) +``` + +### AnthropicProvider + +```python +from headroom import AnthropicProvider +from anthropic import Anthropic + +provider = AnthropicProvider(client=Anthropic()) + +counter = provider.get_token_counter("claude-3-5-sonnet-latest") +tokens = counter.count_messages(messages) # Accurate count via API +``` + +--- + +## Relevance Scoring + +### BM25Scorer + +Fast keyword-based scoring (zero dependencies). + +```python +from headroom import BM25Scorer + +scorer = BM25Scorer() +scores = scorer.score_items( + items=["item 1", "item 2", ...], + query="search query", +) +``` + +### EmbeddingScorer + +Semantic similarity scoring (requires `sentence-transformers`). + +```python +from headroom import EmbeddingScorer, embedding_available + +if embedding_available(): + scorer = EmbeddingScorer(model="all-MiniLM-L6-v2") + scores = scorer.score_items(items, query) +``` + +### HybridScorer + +Combines BM25 and embeddings. + +```python +from headroom import HybridScorer + +scorer = HybridScorer(alpha=0.5) # 50% BM25, 50% embedding +scores = scorer.score_items(items, query) +``` + +### create_scorer() + +Factory function to create scorers. + +```python +from headroom import create_scorer + +# Auto-select best available scorer +scorer = create_scorer() + +# Explicitly choose type +scorer = create_scorer(scorer_type="hybrid", alpha=0.7) +``` + +--- + +## Transforms (Direct Use) + +### SmartCrusher + +```python +from headroom import SmartCrusher + +crusher = SmartCrusher() +result = crusher.crush( + data={"results": [...]}, + query="user query", +) +``` + +### CacheAligner + +```python +from headroom import CacheAligner + +aligner = CacheAligner() +result = aligner.align(messages) +``` + +### RollingWindow + +```python +from headroom import RollingWindow + +window = RollingWindow(config) +result = window.apply(messages, max_tokens=100000) +``` + +### TransformPipeline + +```python +from headroom import TransformPipeline + +pipeline = TransformPipeline([ + SmartCrusher(), + CacheAligner(), + RollingWindow(), +]) + +result = pipeline.transform(messages) +``` + +--- + +## Utilities + +### Tokenizer + +```python +from headroom import Tokenizer, count_tokens_text, count_tokens_messages + +# Quick counting +tokens = count_tokens_text("Hello, world!", model="gpt-4o") + +# With tokenizer instance +tokenizer = Tokenizer(model="gpt-4o") +tokens = tokenizer.count_text("Hello") +tokens = tokenizer.count_messages(messages) +``` + +### generate_report() + +Generate HTML/Markdown reports from stored metrics. + +```python +from headroom import generate_report + +report = generate_report( + store_url="sqlite:///headroom.db", + format="html", + period="day", +) +``` diff --git a/docs/getting-started.md b/docs/getting-started.md new file mode 100644 index 000000000..66c1e0ebe --- /dev/null +++ b/docs/getting-started.md @@ -0,0 +1,109 @@ +# Getting Started with Headroom + +This guide will help you get up and running with Headroom in under 5 minutes. + +## Installation + +```bash +# Core package (minimal dependencies) +pip install headroom + +# With proxy server +pip install headroom[proxy] + +# With semantic relevance (for smarter compression) +pip install headroom[relevance] + +# Everything +pip install headroom[all] +``` + +## Quick Start: Proxy Mode (Recommended) + +The easiest way to use Headroom is as a proxy server: + +```bash +# Start the proxy +headroom proxy --port 8787 +``` + +Then point your LLM client at it: + +```bash +# Claude Code +ANTHROPIC_BASE_URL=http://localhost:8787 claude + +# OpenAI-compatible clients +OPENAI_BASE_URL=http://localhost:8787/v1 your-app +``` + +That's it! All your requests now go through Headroom and get optimized automatically. + +## Quick Start: Python SDK + +If you want programmatic control: + +```python +from headroom import HeadroomClient +from openai import OpenAI + +# Create a wrapped client +client = HeadroomClient( + original_client=OpenAI(), + default_mode="optimize", +) + +# Use exactly like the original +response = client.chat.completions.create( + model="gpt-4o", + messages=[ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Hello!"}, + ], +) +``` + +## Modes + +### Audit Mode + +Observe without modifying: + +```python +client = HeadroomClient( + original_client=OpenAI(), + default_mode="audit", +) +# Logs metrics but doesn't change requests +``` + +### Optimize Mode + +Apply transforms to reduce tokens: + +```python +client = HeadroomClient( + original_client=OpenAI(), + default_mode="optimize", +) +# Compresses tool outputs, aligns cache prefixes, etc. +``` + +### Simulate Mode + +Preview what optimizations would do: + +```python +plan = client.chat.completions.simulate( + model="gpt-4o", + messages=[...], +) +print(f"Would save {plan.tokens_saved} tokens") +print(f"Transforms: {plan.transforms_applied}") +``` + +## Next Steps + +- [Proxy Server Documentation](proxy.md) - Configure the proxy +- [Transforms Reference](transforms.md) - Understand each transform +- [API Reference](api.md) - Full API documentation diff --git a/docs/proxy.md b/docs/proxy.md new file mode 100644 index 000000000..b95d61c00 --- /dev/null +++ b/docs/proxy.md @@ -0,0 +1,173 @@ +# Proxy Server Documentation + +The Headroom proxy server is a production-ready HTTP server that applies context optimization to all requests passing through it. + +## Starting the Proxy + +```bash +# Basic usage +headroom proxy + +# Custom port +headroom proxy --port 8080 + +# With all options +headroom proxy \ + --host 0.0.0.0 \ + --port 8787 \ + --log-file /var/log/headroom.jsonl \ + --budget 100.0 +``` + +## Command Line Options + +| Option | Default | Description | +|--------|---------|-------------| +| `--host` | `127.0.0.1` | Host to bind to | +| `--port` | `8787` | Port to bind to | +| `--no-optimize` | `false` | Disable optimization (passthrough mode) | +| `--no-cache` | `false` | Disable semantic caching | +| `--no-rate-limit` | `false` | Disable rate limiting | +| `--log-file` | None | Path to JSONL log file | +| `--budget` | None | Daily budget limit in USD | + +## API Endpoints + +### Health Check + +```bash +curl http://localhost:8787/health +``` + +Response: +```json +{ + "status": "healthy", + "optimize": true, + "stats": { + "total_requests": 42, + "tokens_saved": 15000, + "savings_percent": 45.2 + } +} +``` + +### Detailed Statistics + +```bash +curl http://localhost:8787/stats +``` + +### Prometheus Metrics + +```bash +curl http://localhost:8787/metrics +``` + +### LLM APIs + +The proxy supports both Anthropic and OpenAI API formats: + +```bash +# Anthropic format +POST /v1/messages + +# OpenAI format +POST /v1/chat/completions +``` + +## Using with Claude Code + +```bash +# Start proxy +headroom proxy --port 8787 + +# In another terminal +ANTHROPIC_BASE_URL=http://localhost:8787 claude +``` + +## Using with Cursor + +1. Start the proxy: `headroom proxy` +2. In Cursor settings, set the base URL to `http://localhost:8787` + +## Using with OpenAI SDK + +```python +from openai import OpenAI + +client = OpenAI( + base_url="http://localhost:8787/v1", + api_key="your-api-key", # Still needed for upstream +) +``` + +## Features + +### Semantic Caching + +The proxy caches responses for repeated queries: + +- LRU eviction with configurable max entries +- TTL-based expiration +- Cache key based on message content hash + +### Rate Limiting + +Token bucket rate limiting protects against runaway costs: + +- Configurable requests per minute +- Configurable tokens per minute +- Per-API-key tracking + +### Cost Tracking + +Track spending and enforce budgets: + +- Real-time cost estimation +- Budget periods: hourly, daily, monthly +- Automatic request rejection when over budget + +### Prometheus Metrics + +Export metrics for monitoring: + +``` +headroom_requests_total +headroom_tokens_saved_total +headroom_cost_usd_total +headroom_latency_ms_sum +``` + +## Configuration via Environment + +```bash +export HEADROOM_HOST=0.0.0.0 +export HEADROOM_PORT=8787 +export HEADROOM_BUDGET=100.0 +headroom proxy +``` + +## Running in Production + +For production deployments: + +```bash +# Use a process manager +pip install gunicorn + +# Run with gunicorn +gunicorn headroom.proxy.server:app \ + --workers 4 \ + --bind 0.0.0.0:8787 \ + --worker-class uvicorn.workers.UvicornWorker +``` + +Or with Docker: + +```dockerfile +FROM python:3.11-slim +RUN pip install headroom[proxy] +EXPOSE 8787 +CMD ["headroom", "proxy", "--host", "0.0.0.0"] +``` diff --git a/docs/transforms.md b/docs/transforms.md new file mode 100644 index 000000000..736f50320 --- /dev/null +++ b/docs/transforms.md @@ -0,0 +1,198 @@ +# Transform Reference + +Headroom provides three core transforms that work together to optimize LLM context. + +## SmartCrusher + +Statistical compression for JSON tool outputs. + +### How It Works + +SmartCrusher analyzes JSON arrays and selectively keeps important items: + +1. **First/Last items** - Context for pagination and recency +2. **Error items** - 100% preservation of error states +3. **Anomalies** - Statistical outliers (> 2 std dev from mean) +4. **Relevant items** - Matches to user's query via BM25/embeddings +5. **Change points** - Significant transitions in data + +### Configuration + +```python +from headroom import SmartCrusherConfig + +config = SmartCrusherConfig( + min_tokens_to_crush=200, # Only compress if > 200 tokens + max_items_after_crush=50, # Keep at most 50 items + keep_first=3, # Always keep first 3 items + keep_last=2, # Always keep last 2 items + relevance_threshold=0.3, # Keep items with relevance > 0.3 + anomaly_std_threshold=2.0, # Keep items > 2 std dev from mean + preserve_errors=True, # Always keep error items +) +``` + +### Example + +```python +from headroom import SmartCrusher + +crusher = SmartCrusher(config) + +# Before: 1000 search results (45,000 tokens) +tool_output = {"results": [...1000 items...]} + +# After: ~50 important items (4,500 tokens) - 90% reduction +compressed = crusher.crush(tool_output, query="user's question") +``` + +### What Gets Preserved + +| Category | Preserved | Why | +|----------|-----------|-----| +| Errors | 100% | Critical for debugging | +| First N | 100% | Context/pagination | +| Last N | 100% | Recency | +| Anomalies | All | Unusual values matter | +| Relevant | Top K | Match user's query | +| Others | Sampled | Statistical representation | + +--- + +## CacheAligner + +Prefix stabilization for improved cache hit rates. + +### The Problem + +LLM providers cache request prefixes. But dynamic content breaks caching: + +``` +"You are helpful. Today is January 7, 2025." # Changes daily = no cache +``` + +### The Solution + +CacheAligner extracts dynamic content to stabilize the prefix: + +```python +from headroom import CacheAligner + +aligner = CacheAligner() +result = aligner.align(messages) + +# Static prefix (cacheable): +# "You are helpful." + +# Dynamic content moved to end: +# [Current date context] +``` + +### Configuration + +```python +from headroom import CacheAlignerConfig + +config = CacheAlignerConfig( + extract_dates=True, # Move dates to dynamic section + normalize_whitespace=True, # Consistent spacing + stable_prefix_min_tokens=100, # Min prefix size for alignment +) +``` + +### Cache Hit Improvement + +| Scenario | Before | After | +|----------|--------|-------| +| Daily date in prompt | 0% hits | ~95% hits | +| Dynamic user context | ~10% hits | ~80% hits | +| Consistent prompts | ~90% hits | ~95% hits | + +--- + +## RollingWindow + +Context management within token limits. + +### The Problem + +Long conversations exceed context limits. Naive truncation breaks tool calls: + +``` +[tool_call: search] # Kept +[tool_result: ...] # Dropped = orphaned call! +``` + +### The Solution + +RollingWindow drops complete tool units, preserving pairs: + +```python +from headroom import RollingWindow + +window = RollingWindow(config) +result = window.apply(messages, max_tokens=100000) + +# Guarantees: +# 1. Tool calls paired with results +# 2. System prompt preserved +# 3. Recent turns kept +# 4. Oldest tool outputs dropped first +``` + +### Configuration + +```python +from headroom import RollingWindowConfig + +config = RollingWindowConfig( + max_tokens=100000, # Target token limit + preserve_system=True, # Always keep system prompt + preserve_recent_turns=5, # Keep last 5 user/assistant turns + drop_oldest_first=True, # Remove oldest tool outputs +) +``` + +### Drop Priority + +1. **Oldest tool outputs** - First to go +2. **Old assistant messages** - Summary preserved +3. **Old user messages** - Only if necessary +4. **Never dropped**: System prompt, recent turns, active tool pairs + +--- + +## TransformPipeline + +Combine transforms for optimal results. + +```python +from headroom import TransformPipeline, SmartCrusher, CacheAligner, RollingWindow + +pipeline = TransformPipeline([ + SmartCrusher(), # First: compress tool outputs + CacheAligner(), # Then: stabilize prefix + RollingWindow(), # Finally: fit in context +]) + +result = pipeline.transform(messages) +print(f"Saved {result.tokens_saved} tokens") +``` + +### Recommended Order + +1. **SmartCrusher** - Reduce individual messages +2. **CacheAligner** - Optimize for caching +3. **RollingWindow** - Final size constraint + +--- + +## Safety Guarantees + +All transforms follow strict safety rules: + +1. **Never remove human content** - User/assistant text is sacred +2. **Never break tool ordering** - Calls and results stay paired +3. **Parse failures are no-ops** - Malformed content passes through +4. **Preserves recency** - Last N turns always kept +5. **100% error preservation** - Error items never dropped diff --git a/examples/README.md b/examples/README.md new file mode 100644 index 000000000..e4fbf240d --- /dev/null +++ b/examples/README.md @@ -0,0 +1,133 @@ +# Headroom Examples + +This directory contains examples demonstrating Headroom's capabilities. + +## Quick Start Examples + +### basic_usage.py + +Basic integration with OpenAI client: + +```bash +export OPENAI_API_KEY='your-key' +python examples/basic_usage.py +``` + +### anthropic_example.py + +Integration with Anthropic Claude: + +```bash +export ANTHROPIC_API_KEY='your-key' +python examples/anthropic_example.py +``` + +### streaming_example.py + +Streaming responses with optimization: + +```bash +export OPENAI_API_KEY='your-key' +python examples/streaming_example.py +``` + +## Evaluation Examples + +### smart_vs_naive_eval.py + +Compare SmartCrusher against naive truncation: + +```bash +export OPENAI_API_KEY='your-key' +python examples/smart_vs_naive_eval.py +``` + +### real_world_eval.py + +Comprehensive evaluation with Anthropic models: + +```bash +export ANTHROPIC_API_KEY='your-key' +python examples/real_world_eval.py +``` + +### real_world_openai_eval.py + +Comprehensive evaluation with OpenAI models: + +```bash +export OPENAI_API_KEY='your-key' +python examples/real_world_openai_eval.py +``` + +## Demo Directories + +### langchain_demo/ + +Full LangChain agent integration demo: + +```bash +# No API key needed for compression demo +PYTHONPATH=. python -m examples.langchain_demo.show_compression + +# Full comparison (requires API key) +export OPENAI_API_KEY='your-key' +PYTHONPATH=. python -m examples.langchain_demo.run_comparison +``` + +See [langchain_demo/README.md](langchain_demo/README.md) for details. + +### mcp_demo/ + +MCP (Model Context Protocol) integration demo: + +```bash +export OPENAI_API_KEY='your-key' +PYTHONPATH=. python -m examples.mcp_demo.run_agent_eval +``` + +## Running Examples + +All examples can be run from the repository root: + +```bash +# Install dependencies +pip install -e ".[dev]" + +# Run any example +python examples/.py +``` + +## Expected Results + +| Example | Token Savings | Notes | +|---------|---------------|-------| +| basic_usage | 50-70% | Simple tool output compression | +| langchain_demo | 70-85% | Real agent with multiple tools | +| mcp_demo | 60-80% | MCP tool outputs | +| real_world_eval | 50-90% | Varies by scenario | + +## Troubleshooting + +**ModuleNotFoundError: No module named 'headroom'** + +Run from the repository root with PYTHONPATH: + +```bash +PYTHONPATH=. python examples/basic_usage.py +``` + +Or install in development mode: + +```bash +pip install -e . +``` + +**API Key Errors** + +Ensure your API keys are set: + +```bash +export OPENAI_API_KEY='sk-...' +export ANTHROPIC_API_KEY='sk-ant-...' +``` diff --git a/headroom/__init__.py b/headroom/__init__.py index 839b39f2f..fe02d8e5b 100644 --- a/headroom/__init__.py +++ b/headroom/__init__.py @@ -57,15 +57,6 @@ from .config import ( WasteSignals, ) from .providers import AnthropicProvider, OpenAIProvider, Provider, TokenCounter -from .reporting import generate_report -from .tokenizer import Tokenizer, count_tokens_messages, count_tokens_text -from .transforms import ( - CacheAligner, - RollingWindow, - SmartCrusher, - ToolCrusher, - TransformPipeline, -) from .relevance import ( BM25Scorer, EmbeddingScorer, @@ -75,6 +66,15 @@ from .relevance import ( create_scorer, embedding_available, ) +from .reporting import generate_report +from .tokenizer import Tokenizer, count_tokens_messages, count_tokens_text +from .transforms import ( + CacheAligner, + RollingWindow, + SmartCrusher, + ToolCrusher, + TransformPipeline, +) __version__ = "0.2.0" diff --git a/headroom/cli.py b/headroom/cli.py new file mode 100644 index 000000000..65a4db367 --- /dev/null +++ b/headroom/cli.py @@ -0,0 +1,185 @@ +#!/usr/bin/env python3 +"""Headroom CLI - The Context Optimization Layer for LLM Applications. + +Usage: + headroom proxy [OPTIONS] Start the optimization proxy server + headroom --version Show version + headroom --help Show this help message + +Examples: + # Start proxy on default port (8787) + headroom proxy + + # Start proxy on custom port + headroom proxy --port 8080 + + # Start with optimization disabled (passthrough mode) + headroom proxy --no-optimize + + # Use with Claude Code + ANTHROPIC_BASE_URL=http://localhost:8787 claude +""" + +from __future__ import annotations + +import argparse +import sys + + +def get_version() -> str: + """Get the current version.""" + try: + from headroom import __version__ + return __version__ + except ImportError: + return "unknown" + + +def cmd_proxy(args: argparse.Namespace) -> int: + """Start the proxy server.""" + try: + from headroom.proxy.server import ProxyConfig, run_server + except ImportError as e: + print("Error: Proxy dependencies not installed. Run: pip install headroom[proxy]") + print(f"Details: {e}") + return 1 + + config = ProxyConfig( + host=args.host, + port=args.port, + optimize=not args.no_optimize, + cache_enabled=not args.no_cache, + rate_limit_enabled=not args.no_rate_limit, + log_file=args.log_file, + budget_limit_usd=args.budget, + ) + + print(f""" +╔═══════════════════════════════════════════════════════════════════════╗ +║ HEADROOM PROXY ║ +║ The Context Optimization Layer for LLM Applications ║ +╚═══════════════════════════════════════════════════════════════════════╝ + +Starting proxy server... + + URL: http://{config.host}:{config.port} + Optimization: {'ENABLED' if config.optimize else 'DISABLED'} + Caching: {'ENABLED' if config.cache_enabled else 'DISABLED'} + Rate Limit: {'ENABLED' if config.rate_limit_enabled else 'DISABLED'} + +Usage with Claude Code: + ANTHROPIC_BASE_URL=http://{config.host}:{config.port} claude + +Usage with OpenAI-compatible clients: + OPENAI_BASE_URL=http://{config.host}:{config.port}/v1 your-app + +Endpoints: + GET /health Health check + GET /stats Detailed statistics + GET /metrics Prometheus metrics + POST /v1/messages Anthropic API + POST /v1/chat/completions OpenAI API + +Press Ctrl+C to stop. +""") + + try: + run_server(config) + except KeyboardInterrupt: + print("\nShutting down...") + return 0 + + return 0 + + +def cmd_version(args: argparse.Namespace) -> int: + """Print version information.""" + print(f"headroom {get_version()}") + return 0 + + +def main(argv: list[str] | None = None) -> int: + """Main CLI entry point.""" + parser = argparse.ArgumentParser( + prog="headroom", + description="The Context Optimization Layer for LLM Applications", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" +Examples: + headroom proxy Start proxy on port 8787 + headroom proxy --port 8080 Start proxy on port 8080 + headroom proxy --no-optimize Passthrough mode (no optimization) + +Environment Variables: + ANTHROPIC_API_KEY Your Anthropic API key (for proxying) + OPENAI_API_KEY Your OpenAI API key (for proxying) + +Documentation: https://github.com/headroom-sdk/headroom + """, + ) + + parser.add_argument( + "--version", "-V", + action="store_true", + help="Show version and exit", + ) + + subparsers = parser.add_subparsers(dest="command", help="Commands") + + # Proxy command + proxy_parser = subparsers.add_parser( + "proxy", + help="Start the optimization proxy server", + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + proxy_parser.add_argument( + "--host", + default="127.0.0.1", + help="Host to bind to (default: 127.0.0.1)", + ) + proxy_parser.add_argument( + "--port", "-p", + type=int, + default=8787, + help="Port to bind to (default: 8787)", + ) + proxy_parser.add_argument( + "--no-optimize", + action="store_true", + help="Disable optimization (passthrough mode)", + ) + proxy_parser.add_argument( + "--no-cache", + action="store_true", + help="Disable semantic caching", + ) + proxy_parser.add_argument( + "--no-rate-limit", + action="store_true", + help="Disable rate limiting", + ) + proxy_parser.add_argument( + "--log-file", + help="Path to JSONL log file", + ) + proxy_parser.add_argument( + "--budget", + type=float, + help="Daily budget limit in USD", + ) + proxy_parser.set_defaults(func=cmd_proxy) + + args = parser.parse_args(argv) + + if args.version: + return cmd_version(args) + + if args.command is None: + parser.print_help() + return 0 + + return args.func(args) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/headroom/client.py b/headroom/client.py index 75d9ddcae..a77cda7b2 100644 --- a/headroom/client.py +++ b/headroom/client.py @@ -2,8 +2,9 @@ from __future__ import annotations +from collections.abc import Iterator from datetime import datetime -from typing import Any, Iterator +from typing import Any from .config import ( HeadroomConfig, diff --git a/headroom/config.py b/headroom/config.py index ba0fc4df8..770d544fa 100644 --- a/headroom/config.py +++ b/headroom/config.py @@ -301,8 +301,8 @@ class TransformResult: transforms_applied: list[str] markers_inserted: list[str] = field(default_factory=list) warnings: list[str] = field(default_factory=list) - diff_artifact: "DiffArtifact | None" = None # Populated if generate_diff_artifact=True - cache_metrics: "CachePrefixMetrics | None" = None # Populated by CacheAligner + diff_artifact: DiffArtifact | None = None # Populated if generate_diff_artifact=True + cache_metrics: CachePrefixMetrics | None = None # Populated by CacheAligner @dataclass diff --git a/headroom/integrations/__init__.py b/headroom/integrations/__init__.py index e780eab50..99e50f7ca 100644 --- a/headroom/integrations/__init__.py +++ b/headroom/integrations/__init__.py @@ -8,21 +8,20 @@ Install LangChain support: pip install headroom[langchain] """ from .langchain import ( - HeadroomChatModel, HeadroomCallbackHandler, - optimize_messages, + HeadroomChatModel, HeadroomRunnable, + optimize_messages, ) - from .mcp import ( - HeadroomMCPCompressor, + DEFAULT_MCP_PROFILES, HeadroomMCPClientWrapper, + HeadroomMCPCompressor, MCPCompressionResult, MCPToolProfile, compress_tool_result, compress_tool_result_with_metrics, create_headroom_mcp_proxy, - DEFAULT_MCP_PROFILES, ) __all__ = [ diff --git a/headroom/integrations/langchain.py b/headroom/integrations/langchain.py index ca83d2cef..781d1cead 100644 --- a/headroom/integrations/langchain.py +++ b/headroom/integrations/langchain.py @@ -29,9 +29,10 @@ from __future__ import annotations import json import logging +from collections.abc import Iterator, Sequence from dataclasses import dataclass from datetime import datetime -from typing import Any, Iterator, List, Optional, Sequence, Union +from typing import Any from uuid import uuid4 # LangChain imports - these are optional dependencies @@ -378,7 +379,7 @@ class HeadroomChatModel(BaseChatModel): **kwargs, ) - def bind_tools(self, tools: Sequence[Any], **kwargs) -> "HeadroomChatModel": + def bind_tools(self, tools: Sequence[Any], **kwargs) -> HeadroomChatModel: """Bind tools to the wrapped model.""" new_wrapped = self.wrapped_model.bind_tools(tools, **kwargs) return HeadroomChatModel( diff --git a/headroom/integrations/mcp.py b/headroom/integrations/mcp.py index 06f80de82..aa102b8a8 100644 --- a/headroom/integrations/mcp.py +++ b/headroom/integrations/mcp.py @@ -48,12 +48,13 @@ from __future__ import annotations import json import re +from collections.abc import Callable from dataclasses import dataclass, field -from typing import Any, Callable +from typing import Any from headroom.config import HeadroomConfig, SmartCrusherConfig -from headroom.transforms import SmartCrusher from headroom.providers import OpenAIProvider +from headroom.transforms import SmartCrusher @dataclass diff --git a/headroom/models/__init__.py b/headroom/models/__init__.py new file mode 100644 index 000000000..8c76b3fa2 --- /dev/null +++ b/headroom/models/__init__.py @@ -0,0 +1,39 @@ +"""Model registry and capabilities database. + +Provides a centralized registry of LLM models with their capabilities, +context limits, pricing, and provider information. + +Usage: + from headroom.models import ModelRegistry, get_model_info + + # Get info about a model + info = get_model_info("gpt-4o") + print(f"Context: {info.context_window}") + print(f"Provider: {info.provider}") + + # List all models from a provider + models = ModelRegistry.list_models(provider="openai") + + # Register a custom model + ModelRegistry.register( + "my-custom-model", + provider="custom", + context_window=32000, + ) +""" + +from .registry import ( + ModelInfo, + ModelRegistry, + get_model_info, + list_models, + register_model, +) + +__all__ = [ + "ModelRegistry", + "ModelInfo", + "get_model_info", + "list_models", + "register_model", +] diff --git a/headroom/models/registry.py b/headroom/models/registry.py new file mode 100644 index 000000000..47340c674 --- /dev/null +++ b/headroom/models/registry.py @@ -0,0 +1,749 @@ +"""Model registry with capabilities database. + +Centralized database of LLM models with their capabilities, context limits, +pricing, and provider information. Supports dynamic registration of custom +models and automatic provider detection. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from datetime import date +from typing import Any + + +@dataclass(frozen=True) +class ModelInfo: + """Information about an LLM model. + + Attributes: + name: Model identifier. + provider: Provider name (openai, anthropic, etc.). + context_window: Maximum context window in tokens. + max_output_tokens: Maximum output tokens. + supports_tools: Whether model supports tool/function calling. + supports_vision: Whether model supports image inputs. + supports_streaming: Whether model supports streaming responses. + supports_json_mode: Whether model supports JSON output mode. + tokenizer_backend: Tokenizer backend to use. + input_cost_per_1m: Cost per 1M input tokens in USD. + output_cost_per_1m: Cost per 1M output tokens in USD. + cached_input_cost_per_1m: Cost per 1M cached input tokens. + pricing_date: Date pricing was last updated. + aliases: Alternative names for the model. + notes: Additional notes about the model. + """ + + name: str + provider: str + context_window: int = 128000 + max_output_tokens: int = 4096 + supports_tools: bool = True + supports_vision: bool = False + supports_streaming: bool = True + supports_json_mode: bool = True + tokenizer_backend: str | None = None + input_cost_per_1m: float | None = None + output_cost_per_1m: float | None = None + cached_input_cost_per_1m: float | None = None + pricing_date: date | None = None + aliases: tuple[str, ...] = () + notes: str = "" + + +# Built-in model database +# Pricing as of January 2025 - verify current rates +_MODELS: dict[str, ModelInfo] = {} + + +def _register_builtin_models() -> None: + """Register built-in models.""" + + # ============================================================ + # OpenAI Models + # ============================================================ + + # GPT-4o family + _MODELS["gpt-4o"] = ModelInfo( + name="gpt-4o", + provider="openai", + context_window=128000, + max_output_tokens=16384, + supports_tools=True, + supports_vision=True, + supports_streaming=True, + tokenizer_backend="tiktoken", + input_cost_per_1m=2.50, + output_cost_per_1m=10.00, + cached_input_cost_per_1m=1.25, + pricing_date=date(2025, 1, 6), + aliases=("gpt-4o-2024-11-20", "gpt-4o-2024-08-06", "gpt-4o-2024-05-13"), + notes="Latest GPT-4o with vision and tools", + ) + + _MODELS["gpt-4o-mini"] = ModelInfo( + name="gpt-4o-mini", + provider="openai", + context_window=128000, + max_output_tokens=16384, + supports_tools=True, + supports_vision=True, + supports_streaming=True, + tokenizer_backend="tiktoken", + input_cost_per_1m=0.15, + output_cost_per_1m=0.60, + cached_input_cost_per_1m=0.075, + pricing_date=date(2025, 1, 6), + aliases=("gpt-4o-mini-2024-07-18",), + notes="Cost-effective GPT-4o variant", + ) + + # o1 reasoning models + _MODELS["o1"] = ModelInfo( + name="o1", + provider="openai", + context_window=200000, + max_output_tokens=100000, + supports_tools=True, + supports_vision=True, + supports_streaming=True, + tokenizer_backend="tiktoken", + input_cost_per_1m=15.00, + output_cost_per_1m=60.00, + cached_input_cost_per_1m=7.50, + pricing_date=date(2025, 1, 6), + notes="Full reasoning model with extended thinking", + ) + + _MODELS["o1-mini"] = ModelInfo( + name="o1-mini", + provider="openai", + context_window=128000, + max_output_tokens=65536, + supports_tools=True, + supports_vision=False, + supports_streaming=True, + tokenizer_backend="tiktoken", + input_cost_per_1m=1.10, + output_cost_per_1m=4.40, + cached_input_cost_per_1m=0.55, + pricing_date=date(2025, 1, 6), + notes="Fast reasoning model", + ) + + _MODELS["o3-mini"] = ModelInfo( + name="o3-mini", + provider="openai", + context_window=200000, + max_output_tokens=100000, + supports_tools=True, + supports_vision=True, + supports_streaming=True, + tokenizer_backend="tiktoken", + input_cost_per_1m=1.10, + output_cost_per_1m=4.40, + cached_input_cost_per_1m=0.55, + pricing_date=date(2025, 1, 6), + notes="Latest reasoning model", + ) + + # GPT-4 Turbo + _MODELS["gpt-4-turbo"] = ModelInfo( + name="gpt-4-turbo", + provider="openai", + context_window=128000, + max_output_tokens=4096, + supports_tools=True, + supports_vision=True, + supports_streaming=True, + tokenizer_backend="tiktoken", + input_cost_per_1m=10.00, + output_cost_per_1m=30.00, + cached_input_cost_per_1m=5.00, + pricing_date=date(2025, 1, 6), + aliases=("gpt-4-turbo-preview", "gpt-4-turbo-2024-04-09"), + notes="GPT-4 Turbo with vision", + ) + + # GPT-4 + _MODELS["gpt-4"] = ModelInfo( + name="gpt-4", + provider="openai", + context_window=8192, + max_output_tokens=4096, + supports_tools=True, + supports_vision=False, + supports_streaming=True, + tokenizer_backend="tiktoken", + input_cost_per_1m=30.00, + output_cost_per_1m=60.00, + pricing_date=date(2025, 1, 6), + aliases=("gpt-4-0613",), + notes="Original GPT-4", + ) + + _MODELS["gpt-4-32k"] = ModelInfo( + name="gpt-4-32k", + provider="openai", + context_window=32768, + max_output_tokens=4096, + supports_tools=True, + supports_vision=False, + supports_streaming=True, + tokenizer_backend="tiktoken", + input_cost_per_1m=60.00, + output_cost_per_1m=120.00, + pricing_date=date(2025, 1, 6), + notes="Extended context GPT-4", + ) + + # GPT-3.5 + _MODELS["gpt-3.5-turbo"] = ModelInfo( + name="gpt-3.5-turbo", + provider="openai", + context_window=16385, + max_output_tokens=4096, + supports_tools=True, + supports_vision=False, + supports_streaming=True, + tokenizer_backend="tiktoken", + input_cost_per_1m=0.50, + output_cost_per_1m=1.50, + cached_input_cost_per_1m=0.25, + pricing_date=date(2025, 1, 6), + aliases=("gpt-3.5-turbo-0125", "gpt-3.5-turbo-1106"), + notes="Fast and cost-effective", + ) + + # ============================================================ + # Anthropic Models + # ============================================================ + + _MODELS["claude-3-5-sonnet-20241022"] = ModelInfo( + name="claude-3-5-sonnet-20241022", + provider="anthropic", + context_window=200000, + max_output_tokens=8192, + supports_tools=True, + supports_vision=True, + supports_streaming=True, + tokenizer_backend="anthropic", + input_cost_per_1m=3.00, + output_cost_per_1m=15.00, + cached_input_cost_per_1m=0.30, + pricing_date=date(2025, 1, 6), + aliases=("claude-3-5-sonnet-latest", "claude-sonnet-4-20250514"), + notes="Claude 3.5 Sonnet - Best balance of speed and capability", + ) + + _MODELS["claude-3-5-haiku-20241022"] = ModelInfo( + name="claude-3-5-haiku-20241022", + provider="anthropic", + context_window=200000, + max_output_tokens=8192, + supports_tools=True, + supports_vision=True, + supports_streaming=True, + tokenizer_backend="anthropic", + input_cost_per_1m=0.80, + output_cost_per_1m=4.00, + cached_input_cost_per_1m=0.08, + pricing_date=date(2025, 1, 6), + aliases=("claude-3-5-haiku-latest",), + notes="Claude 3.5 Haiku - Fast and cost-effective", + ) + + _MODELS["claude-3-opus-20240229"] = ModelInfo( + name="claude-3-opus-20240229", + provider="anthropic", + context_window=200000, + max_output_tokens=4096, + supports_tools=True, + supports_vision=True, + supports_streaming=True, + tokenizer_backend="anthropic", + input_cost_per_1m=15.00, + output_cost_per_1m=75.00, + cached_input_cost_per_1m=1.50, + pricing_date=date(2025, 1, 6), + aliases=("claude-3-opus-latest",), + notes="Claude 3 Opus - Most capable", + ) + + _MODELS["claude-3-haiku-20240307"] = ModelInfo( + name="claude-3-haiku-20240307", + provider="anthropic", + context_window=200000, + max_output_tokens=4096, + supports_tools=True, + supports_vision=True, + supports_streaming=True, + tokenizer_backend="anthropic", + input_cost_per_1m=0.25, + output_cost_per_1m=1.25, + cached_input_cost_per_1m=0.03, + pricing_date=date(2025, 1, 6), + notes="Claude 3 Haiku - Legacy fast model", + ) + + # ============================================================ + # Google Models + # ============================================================ + + _MODELS["gemini-2.0-flash"] = ModelInfo( + name="gemini-2.0-flash", + provider="google", + context_window=1000000, + max_output_tokens=8192, + supports_tools=True, + supports_vision=True, + supports_streaming=True, + tokenizer_backend="google", + input_cost_per_1m=0.10, + output_cost_per_1m=0.40, + pricing_date=date(2025, 1, 6), + aliases=("gemini-2.0-flash-exp",), + notes="Gemini 2.0 Flash - Fast multimodal", + ) + + _MODELS["gemini-1.5-pro"] = ModelInfo( + name="gemini-1.5-pro", + provider="google", + context_window=2000000, + max_output_tokens=8192, + supports_tools=True, + supports_vision=True, + supports_streaming=True, + tokenizer_backend="google", + input_cost_per_1m=1.25, + output_cost_per_1m=5.00, + pricing_date=date(2025, 1, 6), + aliases=("gemini-1.5-pro-latest",), + notes="Gemini 1.5 Pro - 2M context window", + ) + + _MODELS["gemini-1.5-flash"] = ModelInfo( + name="gemini-1.5-flash", + provider="google", + context_window=1000000, + max_output_tokens=8192, + supports_tools=True, + supports_vision=True, + supports_streaming=True, + tokenizer_backend="google", + input_cost_per_1m=0.075, + output_cost_per_1m=0.30, + pricing_date=date(2025, 1, 6), + aliases=("gemini-1.5-flash-latest",), + notes="Gemini 1.5 Flash - Cost-effective", + ) + + # ============================================================ + # Meta Llama Models (open source) + # ============================================================ + + _MODELS["llama-3.3-70b"] = ModelInfo( + name="llama-3.3-70b", + provider="meta", + context_window=128000, + max_output_tokens=4096, + supports_tools=True, + supports_vision=False, + supports_streaming=True, + tokenizer_backend="huggingface", + aliases=("llama-3.3-70b-instruct", "meta-llama/Llama-3.3-70B-Instruct"), + notes="Llama 3.3 70B - Open source", + ) + + _MODELS["llama-3.1-405b"] = ModelInfo( + name="llama-3.1-405b", + provider="meta", + context_window=128000, + max_output_tokens=4096, + supports_tools=True, + supports_vision=False, + supports_streaming=True, + tokenizer_backend="huggingface", + aliases=("llama-3.1-405b-instruct", "meta-llama/Llama-3.1-405B-Instruct"), + notes="Llama 3.1 405B - Largest open source", + ) + + _MODELS["llama-3.1-70b"] = ModelInfo( + name="llama-3.1-70b", + provider="meta", + context_window=128000, + max_output_tokens=4096, + supports_tools=True, + supports_vision=False, + supports_streaming=True, + tokenizer_backend="huggingface", + aliases=("llama-3.1-70b-instruct", "meta-llama/Llama-3.1-70B-Instruct"), + notes="Llama 3.1 70B", + ) + + _MODELS["llama-3.1-8b"] = ModelInfo( + name="llama-3.1-8b", + provider="meta", + context_window=128000, + max_output_tokens=4096, + supports_tools=True, + supports_vision=False, + supports_streaming=True, + tokenizer_backend="huggingface", + aliases=("llama-3.1-8b-instruct", "meta-llama/Llama-3.1-8B-Instruct"), + notes="Llama 3.1 8B - Fast and efficient", + ) + + # ============================================================ + # Mistral Models + # ============================================================ + + _MODELS["mistral-large"] = ModelInfo( + name="mistral-large", + provider="mistral", + context_window=128000, + max_output_tokens=4096, + supports_tools=True, + supports_vision=False, + supports_streaming=True, + tokenizer_backend="huggingface", + input_cost_per_1m=2.00, + output_cost_per_1m=6.00, + pricing_date=date(2025, 1, 6), + aliases=("mistral-large-latest",), + notes="Mistral Large - Best capability", + ) + + _MODELS["mistral-small"] = ModelInfo( + name="mistral-small", + provider="mistral", + context_window=32768, + max_output_tokens=4096, + supports_tools=True, + supports_vision=False, + supports_streaming=True, + tokenizer_backend="huggingface", + input_cost_per_1m=0.20, + output_cost_per_1m=0.60, + pricing_date=date(2025, 1, 6), + aliases=("mistral-small-latest",), + notes="Mistral Small - Cost-effective", + ) + + _MODELS["mixtral-8x7b"] = ModelInfo( + name="mixtral-8x7b", + provider="mistral", + context_window=32768, + max_output_tokens=4096, + supports_tools=True, + supports_vision=False, + supports_streaming=True, + tokenizer_backend="huggingface", + aliases=("mixtral-8x7b-instruct",), + notes="Mixtral 8x7B - MoE architecture", + ) + + _MODELS["mistral-7b"] = ModelInfo( + name="mistral-7b", + provider="mistral", + context_window=32768, + max_output_tokens=4096, + supports_tools=False, + supports_vision=False, + supports_streaming=True, + tokenizer_backend="huggingface", + aliases=("mistral-7b-instruct",), + notes="Mistral 7B - Open source", + ) + + # ============================================================ + # DeepSeek Models + # ============================================================ + + _MODELS["deepseek-v3"] = ModelInfo( + name="deepseek-v3", + provider="deepseek", + context_window=128000, + max_output_tokens=8192, + supports_tools=True, + supports_vision=False, + supports_streaming=True, + tokenizer_backend="huggingface", + input_cost_per_1m=0.14, + output_cost_per_1m=0.28, + pricing_date=date(2025, 1, 6), + notes="DeepSeek V3 - High performance, low cost", + ) + + _MODELS["deepseek-coder"] = ModelInfo( + name="deepseek-coder", + provider="deepseek", + context_window=16384, + max_output_tokens=4096, + supports_tools=False, + supports_vision=False, + supports_streaming=True, + tokenizer_backend="huggingface", + notes="DeepSeek Coder - Specialized for code", + ) + + # ============================================================ + # Qwen Models + # ============================================================ + + _MODELS["qwen2.5-72b"] = ModelInfo( + name="qwen2.5-72b", + provider="alibaba", + context_window=131072, + max_output_tokens=8192, + supports_tools=True, + supports_vision=False, + supports_streaming=True, + tokenizer_backend="huggingface", + aliases=("qwen2.5-72b-instruct",), + notes="Qwen 2.5 72B - Strong multilingual", + ) + + _MODELS["qwen2.5-7b"] = ModelInfo( + name="qwen2.5-7b", + provider="alibaba", + context_window=131072, + max_output_tokens=8192, + supports_tools=True, + supports_vision=False, + supports_streaming=True, + tokenizer_backend="huggingface", + aliases=("qwen2.5-7b-instruct",), + notes="Qwen 2.5 7B - Efficient", + ) + + +# Initialize built-in models +_register_builtin_models() + +# Build alias lookup +_ALIASES: dict[str, str] = {} +for model_name, info in _MODELS.items(): + for alias in info.aliases: + _ALIASES[alias.lower()] = model_name + + +class ModelRegistry: + """Registry of LLM models and their capabilities. + + Singleton registry providing access to model information. + Supports built-in models and custom registration. + + Example: + # Get model info + info = ModelRegistry.get("gpt-4o") + print(f"Context: {info.context_window}") + + # Register custom model + ModelRegistry.register( + "my-model", + provider="custom", + context_window=32000, + ) + + # List models by provider + openai_models = ModelRegistry.list_models(provider="openai") + """ + + @classmethod + def get(cls, model: str) -> ModelInfo | None: + """Get model information. + + Args: + model: Model name or alias. + + Returns: + ModelInfo if found, None otherwise. + """ + model_lower = model.lower() + + # Direct lookup + if model_lower in _MODELS: + return _MODELS[model_lower] + + # Alias lookup + if model_lower in _ALIASES: + return _MODELS[_ALIASES[model_lower]] + + # Prefix matching + for name, info in _MODELS.items(): + if model_lower.startswith(name): + return info + + return None + + @classmethod + def register( + cls, + model: str, + provider: str, + context_window: int = 128000, + **kwargs: Any, + ) -> ModelInfo: + """Register a custom model. + + Args: + model: Model name. + provider: Provider name. + context_window: Maximum context window. + **kwargs: Additional ModelInfo fields. + + Returns: + Registered ModelInfo. + """ + info = ModelInfo( + name=model, + provider=provider, + context_window=context_window, + **kwargs, + ) + _MODELS[model.lower()] = info + + # Register aliases + for alias in info.aliases: + _ALIASES[alias.lower()] = model.lower() + + return info + + @classmethod + def list_models( + cls, + provider: str | None = None, + supports_tools: bool | None = None, + supports_vision: bool | None = None, + min_context: int | None = None, + ) -> list[ModelInfo]: + """List models matching criteria. + + Args: + provider: Filter by provider. + supports_tools: Filter by tool support. + supports_vision: Filter by vision support. + min_context: Minimum context window. + + Returns: + List of matching ModelInfo. + """ + results = [] + for info in _MODELS.values(): + if provider and info.provider != provider: + continue + if supports_tools is not None and info.supports_tools != supports_tools: + continue + if supports_vision is not None and info.supports_vision != supports_vision: + continue + if min_context and info.context_window < min_context: + continue + results.append(info) + return results + + @classmethod + def list_providers(cls) -> list[str]: + """List all known providers. + + Returns: + List of provider names. + """ + return list(set(info.provider for info in _MODELS.values())) + + @classmethod + def get_context_limit(cls, model: str, default: int = 128000) -> int: + """Get context limit for a model. + + Args: + model: Model name. + default: Default if model not found. + + Returns: + Context window size. + """ + info = cls.get(model) + return info.context_window if info else default + + @classmethod + def estimate_cost( + cls, + model: str, + input_tokens: int, + output_tokens: int, + cached_tokens: int = 0, + ) -> float | None: + """Estimate API cost for a model. + + Args: + model: Model name. + input_tokens: Number of input tokens. + output_tokens: Number of output tokens. + cached_tokens: Number of cached input tokens. + + Returns: + Estimated cost in USD, or None if pricing unknown. + """ + info = cls.get(model) + if not info or info.input_cost_per_1m is None: + return None + + input_cost = (input_tokens / 1_000_000) * info.input_cost_per_1m + output_cost = (output_tokens / 1_000_000) * (info.output_cost_per_1m or 0) + + if cached_tokens and info.cached_input_cost_per_1m: + # Adjust for cached tokens + regular_input = input_tokens - cached_tokens + cached_cost = (cached_tokens / 1_000_000) * info.cached_input_cost_per_1m + input_cost = (regular_input / 1_000_000) * info.input_cost_per_1m + cached_cost + + return input_cost + output_cost + + +# Convenience functions +def get_model_info(model: str) -> ModelInfo | None: + """Get information about a model. + + Args: + model: Model name or alias. + + Returns: + ModelInfo if found, None otherwise. + """ + return ModelRegistry.get(model) + + +def list_models( + provider: str | None = None, + **kwargs: Any, +) -> list[ModelInfo]: + """List models matching criteria. + + Args: + provider: Filter by provider. + **kwargs: Additional filter criteria. + + Returns: + List of matching ModelInfo. + """ + return ModelRegistry.list_models(provider=provider, **kwargs) + + +def register_model( + model: str, + provider: str, + context_window: int = 128000, + **kwargs: Any, +) -> ModelInfo: + """Register a custom model. + + Args: + model: Model name. + provider: Provider name. + context_window: Maximum context window. + **kwargs: Additional ModelInfo fields. + + Returns: + Registered ModelInfo. + """ + return ModelRegistry.register(model, provider, context_window, **kwargs) diff --git a/headroom/parser.py b/headroom/parser.py index 3f1552f8e..78e4185c2 100644 --- a/headroom/parser.py +++ b/headroom/parser.py @@ -4,7 +4,7 @@ from __future__ import annotations import hashlib import re -from typing import Any, TYPE_CHECKING +from typing import TYPE_CHECKING, Any from .config import Block, WasteSignals diff --git a/headroom/pricing/__init__.py b/headroom/pricing/__init__.py index 235b65091..e3ce72bd2 100644 --- a/headroom/pricing/__init__.py +++ b/headroom/pricing/__init__.py @@ -4,18 +4,21 @@ This module provides pricing information and cost estimation utilities for various LLM providers including OpenAI and Anthropic. """ -from .registry import CostEstimate, ModelPricing, PricingRegistry -from .openai_prices import ( - LAST_UPDATED as OPENAI_LAST_UPDATED, - OPENAI_PRICES, - get_openai_registry, -) from .anthropic_prices import ( - LAST_UPDATED as ANTHROPIC_LAST_UPDATED, ANTHROPIC_PRICES, get_anthropic_registry, ) - +from .anthropic_prices import ( + LAST_UPDATED as ANTHROPIC_LAST_UPDATED, +) +from .openai_prices import ( + LAST_UPDATED as OPENAI_LAST_UPDATED, +) +from .openai_prices import ( + OPENAI_PRICES, + get_openai_registry, +) +from .registry import CostEstimate, ModelPricing, PricingRegistry __all__ = [ # Core classes diff --git a/headroom/pricing/anthropic_prices.py b/headroom/pricing/anthropic_prices.py index a6551744f..67e96fcb4 100644 --- a/headroom/pricing/anthropic_prices.py +++ b/headroom/pricing/anthropic_prices.py @@ -4,7 +4,6 @@ from datetime import date from .registry import ModelPricing, PricingRegistry - # Last verified date for pricing information LAST_UPDATED = date(2025, 1, 6) diff --git a/headroom/pricing/openai_prices.py b/headroom/pricing/openai_prices.py index 6d3790c25..328c7fc41 100644 --- a/headroom/pricing/openai_prices.py +++ b/headroom/pricing/openai_prices.py @@ -4,7 +4,6 @@ from datetime import date from .registry import ModelPricing, PricingRegistry - # Last verified date for pricing information LAST_UPDATED = date(2025, 1, 6) diff --git a/headroom/pricing/registry.py b/headroom/pricing/registry.py index 1bd61033f..e1a0eb940 100644 --- a/headroom/pricing/registry.py +++ b/headroom/pricing/registry.py @@ -2,7 +2,6 @@ from dataclasses import dataclass, field from datetime import date, timedelta -from typing import Optional @dataclass(frozen=True) @@ -15,11 +14,11 @@ class ModelPricing: provider: str input_per_1m: float output_per_1m: float - cached_input_per_1m: Optional[float] = None - batch_input_per_1m: Optional[float] = None - batch_output_per_1m: Optional[float] = None - context_window: Optional[int] = None - notes: Optional[str] = None + cached_input_per_1m: float | None = None + batch_input_per_1m: float | None = None + batch_output_per_1m: float | None = None + context_window: int | None = None + notes: str | None = None @dataclass @@ -27,9 +26,9 @@ class CostEstimate: """Result of a cost estimation calculation.""" cost_usd: float breakdown: dict = field(default_factory=dict) - pricing_date: Optional[date] = None + pricing_date: date | None = None is_stale: bool = False - warning: Optional[str] = None + warning: str | None = None class PricingRegistry: @@ -41,8 +40,8 @@ class PricingRegistry: def __init__( self, last_updated: date, - source_url: Optional[str] = None, - prices: Optional[dict[str, ModelPricing]] = None, + source_url: str | None = None, + prices: dict[str, ModelPricing] | None = None, ): """Initialize the pricing registry. @@ -55,7 +54,7 @@ class PricingRegistry: self.source_url = source_url self.prices: dict[str, ModelPricing] = prices or {} - def get_price(self, model: str) -> Optional[ModelPricing]: + def get_price(self, model: str) -> ModelPricing | None: """Get pricing for a specific model. Args: @@ -75,7 +74,7 @@ class PricingRegistry: age = date.today() - self.last_updated return age > timedelta(days=self.STALENESS_THRESHOLD_DAYS) - def staleness_warning(self) -> Optional[str]: + def staleness_warning(self) -> str | None: """Get a warning message if pricing is stale. Returns: diff --git a/headroom/providers/__init__.py b/headroom/providers/__init__.py index 7a3943684..105c87b54 100644 --- a/headroom/providers/__init__.py +++ b/headroom/providers/__init__.py @@ -2,15 +2,60 @@ Providers encapsulate model-specific behavior like tokenization, context limits, and cost estimation. + +Supported Providers: +- OpenAIProvider: Native OpenAI models (GPT-4o, o1, etc.) +- AnthropicProvider: Claude models +- GoogleProvider: Google Gemini models +- CohereProvider: Cohere Command models +- OpenAICompatibleProvider: Universal provider for any OpenAI-compatible API + (Ollama, vLLM, Together, Groq, Fireworks, LM Studio, etc.) +- LiteLLMProvider: Universal provider via LiteLLM (100+ providers) """ from .anthropic import AnthropicProvider from .base import Provider, TokenCounter +from .cohere import CohereProvider +from .google import GoogleProvider +from .litellm import ( + LiteLLMProvider, + create_litellm_provider, + is_litellm_available, +) from .openai import OpenAIProvider +from .openai_compatible import ( + ModelCapabilities, + OpenAICompatibleProvider, + create_anyscale_provider, + create_fireworks_provider, + create_groq_provider, + create_lmstudio_provider, + create_ollama_provider, + create_together_provider, + create_vllm_provider, +) __all__ = [ + # Base "Provider", "TokenCounter", + # Native providers "OpenAIProvider", "AnthropicProvider", + "GoogleProvider", + "CohereProvider", + # Universal providers + "OpenAICompatibleProvider", + "ModelCapabilities", + "LiteLLMProvider", + "is_litellm_available", + # Factory functions + "create_ollama_provider", + "create_together_provider", + "create_groq_provider", + "create_fireworks_provider", + "create_anyscale_provider", + "create_vllm_provider", + "create_lmstudio_provider", + "create_litellm_provider", ] diff --git a/headroom/providers/cohere.py b/headroom/providers/cohere.py new file mode 100644 index 000000000..ed736a20f --- /dev/null +++ b/headroom/providers/cohere.py @@ -0,0 +1,313 @@ +"""Cohere provider for Headroom SDK. + +Token counting uses Cohere's official tokenize API when a client +is provided. This gives accurate counts for all content types. + +Usage: + import cohere + from headroom import CohereProvider + + client = cohere.ClientV2() # Uses CO_API_KEY env var + provider = CohereProvider(client=client) # Accurate counting via API + + # Or without client (uses estimation - less accurate) + provider = CohereProvider() # Warning: approximate counting +""" + +from __future__ import annotations + +import logging +import warnings +from datetime import date +from typing import Any + +from headroom.tokenizers import EstimatingTokenCounter + +from .base import Provider, TokenCounter + +logger = logging.getLogger(__name__) + +# Warning flags +_FALLBACK_WARNING_SHOWN = False + +# Pricing metadata +_PRICING_LAST_UPDATED = date(2025, 1, 6) + +# Cohere model context limits +_CONTEXT_LIMITS: dict[str, int] = { + # Command A (latest, 2025) + "command-a-03-2025": 256000, + "command-a": 256000, + # Command R+ (2024) + "command-r-plus-08-2024": 128000, + "command-r-plus": 128000, + # Command R (2024) + "command-r-08-2024": 128000, + "command-r": 128000, + # Command (legacy) + "command": 4096, + "command-light": 4096, + "command-nightly": 128000, + # Embed models + "embed-english-v3.0": 512, + "embed-multilingual-v3.0": 512, + "embed-english-light-v3.0": 512, + "embed-multilingual-light-v3.0": 512, +} + +# Pricing per 1M tokens (input, output) +_PRICING: dict[str, tuple[float, float]] = { + "command-a-03-2025": (2.50, 10.00), + "command-a": (2.50, 10.00), + "command-r-plus-08-2024": (2.50, 10.00), + "command-r-plus": (2.50, 10.00), + "command-r-08-2024": (0.15, 0.60), + "command-r": (0.15, 0.60), + "command": (1.00, 2.00), + "command-light": (0.30, 0.60), +} + + +class CohereTokenCounter: + """Token counter for Cohere models. + + When a Cohere client is provided, uses the official tokenize API + for accurate counting. Falls back to estimation when no client + is available. + + Usage: + import cohere + client = cohere.ClientV2() + + # With API (accurate) + counter = CohereTokenCounter("command-r-plus", client=client) + + # Without API (estimation) + counter = CohereTokenCounter("command-r-plus") + """ + + def __init__(self, model: str, client: Any = None): + """Initialize Cohere token counter. + + Args: + model: Cohere model name. + client: Optional cohere.ClientV2 for API-based counting. + """ + global _FALLBACK_WARNING_SHOWN + + self.model = model + self._client = client + self._use_api = client is not None + + # Cohere uses ~4 chars per token + self._estimator = EstimatingTokenCounter(chars_per_token=4.0) + + if not self._use_api and not _FALLBACK_WARNING_SHOWN: + warnings.warn( + "CohereProvider: No client provided, using estimation. " + "For accurate counting, pass a Cohere client: " + "CohereProvider(client=cohere.ClientV2())", + UserWarning, + stacklevel=4 + ) + _FALLBACK_WARNING_SHOWN = True + + def count_text(self, text: str) -> int: + """Count tokens in text. + + Uses tokenize API if client available, otherwise estimates. + """ + if not text: + return 0 + + if self._use_api: + try: + response = self._client.tokenize( + text=text, + model=self.model, + ) + return len(response.tokens) + except Exception as e: + logger.debug(f"Cohere tokenize API failed: {e}, using estimation") + + return self._estimator.count_text(text) + + def count_message(self, message: dict[str, Any]) -> int: + """Count tokens in a message.""" + content = self._extract_content(message) + tokens = self.count_text(content) + tokens += 4 # Message overhead (role tokens, etc.) + return tokens + + def count_messages(self, messages: list[dict[str, Any]]) -> int: + """Count tokens in messages.""" + if not messages: + return 0 + + # For API-based counting, concatenate all content + if self._use_api: + try: + all_content = [] + for msg in messages: + content = self._extract_content(msg) + role = msg.get("role", "user") + all_content.append(f"{role}: {content}") + + full_text = "\n".join(all_content) + response = self._client.tokenize( + text=full_text, + model=self.model, + ) + return len(response.tokens) + except Exception as e: + logger.debug(f"Cohere tokenize API failed: {e}, using estimation") + + # Fallback to estimation + total = sum(self.count_message(msg) for msg in messages) + total += 3 # Priming tokens + return total + + def _extract_content(self, message: dict[str, Any]) -> str: + """Extract text content from message.""" + content = message.get("content", "") + if isinstance(content, str): + return content + elif isinstance(content, list): + parts = [] + for part in content: + if isinstance(part, dict) and part.get("type") == "text": + parts.append(part.get("text", "")) + elif isinstance(part, str): + parts.append(part) + return "\n".join(parts) + return str(content) + + +class CohereProvider(Provider): + """Provider for Cohere Command models. + + Supports Command R, Command R+, and Command A model families. + + Example: + import cohere + client = cohere.ClientV2() + + # With client (accurate token counting via API) + provider = CohereProvider(client=client) + + # Without client (estimation-based counting) + provider = CohereProvider() + + # Token counting + counter = provider.get_token_counter("command-r-plus") + tokens = counter.count_text("Hello, world!") + + # Context limits + limit = provider.get_context_limit("command-a") # 256K tokens + + # Cost estimation + cost = provider.estimate_cost( + input_tokens=100000, + output_tokens=10000, + model="command-r-plus", + ) + """ + + def __init__(self, client: Any = None): + """Initialize Cohere provider. + + Args: + client: Optional cohere.ClientV2 for API-based token counting. + If provided, uses tokenize API for accurate counts. + """ + self._client = client + + @property + def name(self) -> str: + return "cohere" + + def supports_model(self, model: str) -> bool: + """Check if model is a known Cohere model.""" + model_lower = model.lower() + if model_lower in _CONTEXT_LIMITS: + return True + # Check prefix match + for prefix in ["command-a", "command-r", "command", "embed-"]: + if model_lower.startswith(prefix): + return True + return False + + def get_token_counter(self, model: str) -> TokenCounter: + """Get token counter for a Cohere model. + + Uses tokenize API if client was provided, otherwise estimates. + """ + if not self.supports_model(model): + raise ValueError( + f"Model '{model}' is not recognized as a Cohere model. " + f"Supported models: {list(_CONTEXT_LIMITS.keys())}" + ) + return CohereTokenCounter(model, client=self._client) + + def get_context_limit(self, model: str) -> int: + """Get context limit for a Cohere model.""" + model_lower = model.lower() + + # Direct match + if model_lower in _CONTEXT_LIMITS: + return _CONTEXT_LIMITS[model_lower] + + # Prefix match + for prefix, limit in [ + ("command-a", 256000), + ("command-r-plus", 128000), + ("command-r", 128000), + ("command", 4096), + ("embed-", 512), + ]: + if model_lower.startswith(prefix): + return limit + + raise ValueError( + f"Unknown context limit for model '{model}'. " + f"Known models: {list(_CONTEXT_LIMITS.keys())}" + ) + + def estimate_cost( + self, + input_tokens: int, + output_tokens: int, + model: str, + cached_tokens: int = 0, + ) -> float | None: + """Estimate cost for Cohere API call. + + Args: + input_tokens: Number of input tokens. + output_tokens: Number of output tokens. + model: Model name. + cached_tokens: Not used by Cohere. + + Returns: + Estimated cost in USD, or None if pricing unknown. + """ + model_lower = model.lower() + + # Find pricing + input_price, output_price = None, None + for model_prefix, (inp, outp) in _PRICING.items(): + if model_lower.startswith(model_prefix): + input_price, output_price = inp, outp + break + + if input_price is None: + return None + + input_cost = (input_tokens / 1_000_000) * input_price + output_cost = (output_tokens / 1_000_000) * output_price + + return input_cost + output_cost + + def get_output_buffer(self, model: str, default: int = 4000) -> int: + """Get recommended output buffer.""" + return default diff --git a/headroom/providers/google.py b/headroom/providers/google.py new file mode 100644 index 000000000..1b7b0b9d0 --- /dev/null +++ b/headroom/providers/google.py @@ -0,0 +1,372 @@ +"""Google Gemini provider for Headroom SDK. + +Supports Google's Gemini models through two interfaces: +1. OpenAI-compatible endpoint (recommended for Headroom) +2. Native Google AI SDK (for advanced features) + +Token counting uses Google's official countTokens API when a client +is provided. This gives accurate counts for all content types. + +Usage: + import google.generativeai as genai + from headroom import GoogleProvider + + genai.configure(api_key="your-api-key") + provider = GoogleProvider(client=genai) # Accurate counting via API + + # Or without client (uses estimation - less accurate) + provider = GoogleProvider() # Warning: approximate counting +""" + +from __future__ import annotations + +import logging +import warnings +from datetime import date +from typing import Any + +from headroom.tokenizers import EstimatingTokenCounter + +from .base import Provider, TokenCounter + +logger = logging.getLogger(__name__) + +# Warning flags +_FALLBACK_WARNING_SHOWN = False + +# Pricing metadata +_PRICING_LAST_UPDATED = date(2025, 1, 6) + +# Google model context limits +_CONTEXT_LIMITS: dict[str, int] = { + # Gemini 2.0 + "gemini-2.0-flash": 1000000, + "gemini-2.0-flash-exp": 1000000, + "gemini-2.0-flash-thinking": 1000000, + # Gemini 1.5 + "gemini-1.5-pro": 2000000, + "gemini-1.5-pro-latest": 2000000, + "gemini-1.5-flash": 1000000, + "gemini-1.5-flash-latest": 1000000, + "gemini-1.5-flash-8b": 1000000, + # Gemini 1.0 + "gemini-1.0-pro": 32768, + "gemini-pro": 32768, +} + +# Pricing per 1M tokens (input, output) +# Note: Google has different pricing tiers based on context length +_PRICING: dict[str, tuple[float, float]] = { + "gemini-2.0-flash": (0.10, 0.40), + "gemini-2.0-flash-exp": (0.10, 0.40), # Experimental, may change + "gemini-1.5-pro": (1.25, 5.00), # Up to 128K context + "gemini-1.5-flash": (0.075, 0.30), # Up to 128K context + "gemini-1.5-flash-8b": (0.0375, 0.15), + "gemini-1.0-pro": (0.50, 1.50), +} + + +class GeminiTokenCounter: + """Token counter for Gemini models. + + When a google.generativeai client is provided, uses the official + countTokens API for accurate counting. Falls back to estimation + when no client is available. + + Usage: + import google.generativeai as genai + genai.configure(api_key="...") + + # With API (accurate) + counter = GeminiTokenCounter("gemini-2.0-flash", client=genai) + + # Without API (estimation) + counter = GeminiTokenCounter("gemini-2.0-flash") + """ + + def __init__(self, model: str, client: Any = None): + """Initialize Gemini token counter. + + Args: + model: Gemini model name. + client: Optional google.generativeai module for API-based counting. + """ + global _FALLBACK_WARNING_SHOWN + + self.model = model + self._client = client + self._use_api = client is not None + self._genai_model = None + + # Gemini uses ~4 chars per token (similar to GPT models) + self._estimator = EstimatingTokenCounter(chars_per_token=4.0) + + if not self._use_api and not _FALLBACK_WARNING_SHOWN: + warnings.warn( + "GoogleProvider: No client provided, using estimation. " + "For accurate counting, pass google.generativeai: " + "GoogleProvider(client=genai)", + UserWarning, + stacklevel=4 + ) + _FALLBACK_WARNING_SHOWN = True + + def _get_model(self): + """Lazy-load the GenerativeModel for API calls.""" + if self._genai_model is None and self._client is not None: + self._genai_model = self._client.GenerativeModel(self.model) + return self._genai_model + + def count_text(self, text: str) -> int: + """Count tokens in text. + + Uses countTokens API if client available, otherwise estimates. + """ + if not text: + return 0 + + if self._use_api: + try: + model = self._get_model() + response = model.count_tokens(text) + return response.total_tokens + except Exception as e: + logger.debug(f"Google countTokens API failed: {e}, using estimation") + + return self._estimator.count_text(text) + + def count_message(self, message: dict[str, Any]) -> int: + """Count tokens in a message.""" + # For API-based counting, convert message to content and count + if self._use_api: + try: + content = self._message_to_content(message) + model = self._get_model() + response = model.count_tokens(content) + return response.total_tokens + except Exception as e: + logger.debug(f"Google countTokens API failed: {e}, using estimation") + + # Fallback to estimation + return self._estimate_message(message) + + def count_messages(self, messages: list[dict[str, Any]]) -> int: + """Count tokens in messages. + + Uses countTokens API with full conversation if available. + """ + if not messages: + return 0 + + if self._use_api: + try: + # Convert to Gemini content format + contents = [self._message_to_content(msg) for msg in messages] + model = self._get_model() + response = model.count_tokens(contents) + return response.total_tokens + except Exception as e: + logger.debug(f"Google countTokens API failed: {e}, using estimation") + + # Fallback to estimation + total = sum(self._estimate_message(msg) for msg in messages) + total += 3 # Priming tokens + return total + + def _message_to_content(self, message: dict[str, Any]) -> str: + """Convert OpenAI-format message to text content for counting.""" + content = message.get("content", "") + if isinstance(content, str): + return content + elif isinstance(content, list): + parts = [] + for part in content: + if isinstance(part, dict) and part.get("type") == "text": + parts.append(part.get("text", "")) + elif isinstance(part, str): + parts.append(part) + return "\n".join(parts) + return str(content) + + def _estimate_message(self, message: dict[str, Any]) -> int: + """Estimate tokens in a message without API.""" + tokens = 4 # Message overhead + + role = message.get("role", "") + tokens += self._estimator.count_text(role) + + content = message.get("content") + if content: + if isinstance(content, str): + tokens += self._estimator.count_text(content) + elif isinstance(content, list): + for part in content: + if isinstance(part, dict): + if part.get("type") == "text": + tokens += self._estimator.count_text(part.get("text", "")) + elif isinstance(part, str): + tokens += self._estimator.count_text(part) + + return tokens + + +class GoogleProvider(Provider): + """Provider for Google Gemini models. + + Supports Gemini 1.5 and 2.0 model families through: + - OpenAI-compatible endpoint (generativelanguage.googleapis.com) + - Native Google AI SDK (for accurate token counting) + + Example: + import google.generativeai as genai + genai.configure(api_key="...") + + # With client (accurate token counting via API) + provider = GoogleProvider(client=genai) + + # Without client (estimation-based counting) + provider = GoogleProvider() + + # Token counting + counter = provider.get_token_counter("gemini-2.0-flash") + tokens = counter.count_text("Hello, world!") + + # Context limits + limit = provider.get_context_limit("gemini-1.5-pro") # 2M tokens! + + # Cost estimation + cost = provider.estimate_cost( + input_tokens=100000, + output_tokens=10000, + model="gemini-1.5-pro", + ) + """ + + # OpenAI-compatible endpoint for Gemini + OPENAI_COMPATIBLE_BASE_URL = "https://generativelanguage.googleapis.com/v1beta/openai" + + def __init__(self, client: Any = None): + """Initialize Google provider. + + Args: + client: Optional google.generativeai module for API-based token counting. + If provided, uses countTokens API for accurate counts. + """ + self._client = client + + @property + def name(self) -> str: + return "google" + + def supports_model(self, model: str) -> bool: + """Check if model is a known Gemini model.""" + model_lower = model.lower() + if model_lower in _CONTEXT_LIMITS: + return True + # Check prefix match + for prefix in ["gemini-2", "gemini-1.5", "gemini-1.0", "gemini-pro"]: + if model_lower.startswith(prefix): + return True + return False + + def get_token_counter(self, model: str) -> TokenCounter: + """Get token counter for a Gemini model. + + Uses countTokens API if client was provided, otherwise estimates. + """ + if not self.supports_model(model): + raise ValueError( + f"Model '{model}' is not recognized as a Google model. " + f"Supported models: {list(_CONTEXT_LIMITS.keys())}" + ) + return GeminiTokenCounter(model, client=self._client) + + def get_context_limit(self, model: str) -> int: + """Get context limit for a Gemini model. + + Note: Gemini 1.5 Pro has 2M token context! + """ + model_lower = model.lower() + + # Direct match + if model_lower in _CONTEXT_LIMITS: + return _CONTEXT_LIMITS[model_lower] + + # Prefix match + for prefix, limit in [ + ("gemini-2.0", 1000000), + ("gemini-1.5-pro", 2000000), + ("gemini-1.5-flash", 1000000), + ("gemini-1.0", 32768), + ("gemini-pro", 32768), + ]: + if model_lower.startswith(prefix): + return limit + + raise ValueError( + f"Unknown context limit for model '{model}'. " + f"Known models: {list(_CONTEXT_LIMITS.keys())}" + ) + + def estimate_cost( + self, + input_tokens: int, + output_tokens: int, + model: str, + cached_tokens: int = 0, + ) -> float | None: + """Estimate cost for Gemini API call. + + Note: Google has tiered pricing based on context length. + This uses the standard pricing (up to 128K context). + For >128K context, actual costs may be higher. + + Args: + input_tokens: Number of input tokens. + output_tokens: Number of output tokens. + model: Model name. + cached_tokens: Number of cached tokens (not used by Google). + + Returns: + Estimated cost in USD, or None if pricing unknown. + """ + model_lower = model.lower() + + # Find pricing + input_price, output_price = None, None + for model_prefix, (inp, outp) in _PRICING.items(): + if model_lower.startswith(model_prefix): + input_price, output_price = inp, outp + break + + if input_price is None: + return None + + input_cost = (input_tokens / 1_000_000) * input_price + output_cost = (output_tokens / 1_000_000) * output_price + + return input_cost + output_cost + + def get_output_buffer(self, model: str, default: int = 4000) -> int: + """Get recommended output buffer.""" + # Gemini models can output up to 8K tokens + return min(8192, default) + + @classmethod + def get_openai_compatible_url(cls, api_key: str) -> str: + """Get OpenAI-compatible endpoint URL. + + Use this with the OpenAI client: + from openai import OpenAI + client = OpenAI( + api_key=api_key, + base_url=GoogleProvider.get_openai_compatible_url(api_key), + ) + + Args: + api_key: Google AI API key. + + Returns: + Base URL for OpenAI-compatible requests. + """ + return cls.OPENAI_COMPATIBLE_BASE_URL diff --git a/headroom/providers/litellm.py b/headroom/providers/litellm.py new file mode 100644 index 000000000..d1e72814e --- /dev/null +++ b/headroom/providers/litellm.py @@ -0,0 +1,293 @@ +"""LiteLLM provider for universal LLM support. + +LiteLLM provides a unified interface to 100+ LLM providers: +- OpenAI, Azure OpenAI +- Anthropic +- Google (Vertex AI, AI Studio) +- AWS Bedrock +- Cohere +- Replicate +- Hugging Face +- Ollama +- Together AI +- Groq +- And many more... + +This integration allows Headroom to work with any LiteLLM-supported +model without needing provider-specific implementations. + +Requires: pip install litellm +""" + +from __future__ import annotations + +import logging +from typing import Any + +from headroom.tokenizers import EstimatingTokenCounter + +from .base import Provider, TokenCounter + +logger = logging.getLogger(__name__) + +# Check if litellm is available +try: + import litellm + from litellm import get_model_info as litellm_get_model_info + from litellm import model_cost as litellm_model_cost + from litellm import token_counter as litellm_token_counter + + LITELLM_AVAILABLE = True +except ImportError: + LITELLM_AVAILABLE = False + litellm = None + litellm_token_counter = None + litellm_model_cost = None + litellm_get_model_info = None + + +def is_litellm_available() -> bool: + """Check if LiteLLM is installed. + + Returns: + True if litellm is available. + """ + return LITELLM_AVAILABLE + + +class LiteLLMTokenCounter: + """Token counter using LiteLLM's token counting. + + LiteLLM provides accurate token counting for most providers + by using the appropriate tokenizer for each model. + """ + + def __init__(self, model: str): + """Initialize LiteLLM token counter. + + Args: + model: Model name in LiteLLM format (e.g., 'gpt-4o', 'claude-3-sonnet'). + """ + if not LITELLM_AVAILABLE: + raise RuntimeError( + "LiteLLM is required for LiteLLMProvider. " + "Install with: pip install litellm" + ) + self.model = model + # Fallback estimator for when litellm counting fails + self._fallback = EstimatingTokenCounter() + + def count_text(self, text: str) -> int: + """Count tokens in text using LiteLLM.""" + if not text: + return 0 + try: + # LiteLLM's token_counter expects messages format + # We wrap text in a simple message + return litellm_token_counter( + model=self.model, + messages=[{"role": "user", "content": text}], + ) + except Exception as e: + logger.debug(f"LiteLLM token count failed for {self.model}: {e}") + return self._fallback.count_text(text) + + def count_message(self, message: dict[str, Any]) -> int: + """Count tokens in a single message.""" + try: + return litellm_token_counter( + model=self.model, + messages=[message], + ) + except Exception as e: + logger.debug(f"LiteLLM message count failed for {self.model}: {e}") + # Fallback to estimation + tokens = 4 # Base overhead + content = message.get("content", "") + if isinstance(content, str): + tokens += self._fallback.count_text(content) + return tokens + + def count_messages(self, messages: list[dict[str, Any]]) -> int: + """Count tokens in messages using LiteLLM.""" + if not messages: + return 0 + try: + return litellm_token_counter( + model=self.model, + messages=messages, + ) + except Exception as e: + logger.debug(f"LiteLLM messages count failed for {self.model}: {e}") + # Fallback to estimation + total = sum(self.count_message(msg) for msg in messages) + total += 3 # Priming + return total + + +class LiteLLMProvider(Provider): + """Provider using LiteLLM for universal model support. + + LiteLLM supports 100+ LLM providers with a unified interface. + This provider leverages LiteLLM's: + - Token counting (accurate for most providers) + - Model info (context limits, capabilities) + - Cost estimation (from LiteLLM's model database) + + Example: + from headroom.providers import LiteLLMProvider + + provider = LiteLLMProvider() + + # Works with any LiteLLM-supported model + counter = provider.get_token_counter("gpt-4o") + counter = provider.get_token_counter("claude-3-5-sonnet-20241022") + counter = provider.get_token_counter("gemini/gemini-1.5-pro") + counter = provider.get_token_counter("bedrock/anthropic.claude-v2") + counter = provider.get_token_counter("ollama/llama3") + + Model Format: + LiteLLM uses a provider/model format for some providers: + - OpenAI: "gpt-4o" or "openai/gpt-4o" + - Anthropic: "claude-3-sonnet" or "anthropic/claude-3-sonnet" + - Google: "gemini/gemini-1.5-pro" + - Azure: "azure/gpt-4" + - Bedrock: "bedrock/anthropic.claude-v2" + - Ollama: "ollama/llama3" + + See LiteLLM docs for full model list: + https://docs.litellm.ai/docs/providers + """ + + def __init__(self): + """Initialize LiteLLM provider.""" + if not LITELLM_AVAILABLE: + raise RuntimeError( + "LiteLLM is required for LiteLLMProvider. " + "Install with: pip install litellm" + ) + + @property + def name(self) -> str: + return "litellm" + + def supports_model(self, model: str) -> bool: + """Check if LiteLLM supports this model. + + LiteLLM supports most models, so this returns True + for any model. Actual support depends on credentials. + """ + return True # LiteLLM handles validation + + def get_token_counter(self, model: str) -> TokenCounter: + """Get token counter for a model.""" + return LiteLLMTokenCounter(model) + + def get_context_limit(self, model: str) -> int: + """Get context limit using LiteLLM's model info.""" + try: + info = litellm_get_model_info(model) + if info and "max_input_tokens" in info: + return info["max_input_tokens"] + if info and "max_tokens" in info: + return info["max_tokens"] + except Exception as e: + logger.debug(f"LiteLLM get_model_info failed for {model}: {e}") + + # Fallback to reasonable default + return 128000 + + def get_output_buffer(self, model: str, default: int = 4000) -> int: + """Get recommended output buffer.""" + try: + info = litellm_get_model_info(model) + if info and "max_output_tokens" in info: + return min(info["max_output_tokens"], default) + except Exception: + pass + return default + + def estimate_cost( + self, + input_tokens: int, + output_tokens: int, + model: str, + cached_tokens: int = 0, + ) -> float | None: + """Estimate cost using LiteLLM's cost database. + + Args: + input_tokens: Number of input tokens. + output_tokens: Number of output tokens. + model: Model name. + cached_tokens: Cached tokens (may not be supported by all providers). + + Returns: + Estimated cost in USD, or None if pricing unknown. + """ + try: + # LiteLLM's cost calculation + cost = litellm.completion_cost( + model=model, + prompt="", # We're using token counts directly + completion="", + prompt_tokens=input_tokens, + completion_tokens=output_tokens, + ) + return cost + except Exception as e: + logger.debug(f"LiteLLM cost estimation failed for {model}: {e}") + return None + + @classmethod + def list_supported_providers(cls) -> list[str]: + """List providers supported by LiteLLM. + + Returns: + List of provider names. + """ + if not LITELLM_AVAILABLE: + return [] + + # Major providers supported by LiteLLM + return [ + "openai", + "anthropic", + "azure", + "google", + "vertex_ai", + "bedrock", + "cohere", + "replicate", + "huggingface", + "ollama", + "together_ai", + "groq", + "fireworks_ai", + "anyscale", + "deepinfra", + "perplexity", + "mistral", + "cloudflare", + "ai21", + "nlp_cloud", + "aleph_alpha", + "petals", + "baseten", + "openrouter", + "vllm", + "xinference", + "text-generation-inference", + ] + + +def create_litellm_provider() -> LiteLLMProvider: + """Create a LiteLLM provider. + + Returns: + Configured LiteLLMProvider. + + Raises: + RuntimeError: If LiteLLM is not installed. + """ + return LiteLLMProvider() diff --git a/headroom/providers/openai.py b/headroom/providers/openai.py index 7bbe7dea8..963b8fdc4 100644 --- a/headroom/providers/openai.py +++ b/headroom/providers/openai.py @@ -6,7 +6,6 @@ Cost estimates are APPROXIMATE - always verify against your actual billing. from __future__ import annotations -import json import warnings from datetime import date from functools import lru_cache diff --git a/headroom/providers/openai_compatible.py b/headroom/providers/openai_compatible.py new file mode 100644 index 000000000..daf8a3fd9 --- /dev/null +++ b/headroom/providers/openai_compatible.py @@ -0,0 +1,521 @@ +"""OpenAI-compatible provider for universal LLM support. + +This provider supports any LLM service that implements the OpenAI API format: +- Ollama (local) +- vLLM (local/cloud) +- Together AI +- Groq +- Fireworks AI +- Anyscale +- LM Studio +- LocalAI +- Hugging Face Inference Endpoints +- Azure OpenAI +- And many more... + +The key insight: 70%+ of LLM providers use OpenAI-compatible APIs, +so supporting this format gives near-universal coverage. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import Any + +from headroom.tokenizers import get_tokenizer + +from .base import Provider + +logger = logging.getLogger(__name__) + + +@dataclass +class ModelCapabilities: + """Model capability metadata. + + Stores information about a model's capabilities and constraints + that the provider needs for token counting and cost estimation. + """ + + model: str + context_window: int = 128000 # Default to 128K + max_output_tokens: int = 4096 + supports_tools: bool = True + supports_vision: bool = False + supports_streaming: bool = True + tokenizer_backend: str | None = None # Force specific tokenizer + input_cost_per_1m: float | None = None # Cost per 1M input tokens + output_cost_per_1m: float | None = None # Cost per 1M output tokens + + +# Default context limits for common open models +# These are reasonable defaults; users can override +_DEFAULT_CONTEXT_LIMITS: dict[str, int] = { + # Llama 3 family + "llama-3": 8192, + "llama-3-8b": 8192, + "llama-3-70b": 8192, + "llama-3.1": 128000, + "llama-3.1-8b": 128000, + "llama-3.1-70b": 128000, + "llama-3.1-405b": 128000, + "llama-3.2": 128000, + "llama-3.3": 128000, + # Llama 2 family + "llama-2": 4096, + "llama-2-7b": 4096, + "llama-2-13b": 4096, + "llama-2-70b": 4096, + "codellama": 16384, + # Mistral family + "mistral": 32768, + "mistral-7b": 32768, + "mistral-nemo": 128000, + "mistral-small": 32768, + "mistral-large": 128000, + "mixtral": 32768, + "mixtral-8x7b": 32768, + "mixtral-8x22b": 65536, + # Qwen family + "qwen": 32768, + "qwen2": 32768, + "qwen2-7b": 32768, + "qwen2-72b": 32768, + "qwen2.5": 131072, + # DeepSeek + "deepseek": 32768, + "deepseek-coder": 16384, + "deepseek-v2": 128000, + "deepseek-v3": 128000, + # Yi + "yi": 32768, + "yi-34b": 32768, + # Phi + "phi-2": 2048, + "phi-3": 4096, + "phi-3-mini": 4096, + "phi-3-medium": 4096, + # Others + "falcon": 2048, + "falcon-40b": 2048, + "falcon-180b": 2048, + "gemma": 8192, + "gemma-2": 8192, + "starcoder": 8192, + "starcoder2": 16384, +} + + +class OpenAICompatibleTokenCounter: + """Token counter for OpenAI-compatible providers. + + Uses the TokenizerRegistry to get the appropriate tokenizer + for the model, falling back to estimation if needed. + """ + + def __init__( + self, + model: str, + tokenizer_backend: str | None = None, + ): + """Initialize token counter. + + Args: + model: Model name. + tokenizer_backend: Force specific tokenizer backend. + """ + self.model = model + self._tokenizer = get_tokenizer(model, backend=tokenizer_backend) + + def count_text(self, text: str) -> int: + """Count tokens in text.""" + return self._tokenizer.count_text(text) + + def count_message(self, message: dict[str, Any]) -> int: + """Count tokens in a single message.""" + # Use OpenAI-style message overhead + tokens = 4 # Base overhead + + role = message.get("role", "") + tokens += self.count_text(role) + + content = message.get("content") + if content: + if isinstance(content, str): + tokens += self.count_text(content) + elif isinstance(content, list): + for part in content: + if isinstance(part, dict): + if part.get("type") == "text": + tokens += self.count_text(part.get("text", "")) + elif isinstance(part, str): + tokens += self.count_text(part) + + name = message.get("name") + if name: + tokens += self.count_text(name) + 1 + + tool_calls = message.get("tool_calls") + if tool_calls: + for tc in tool_calls: + func = tc.get("function", {}) + tokens += self.count_text(func.get("name", "")) + tokens += self.count_text(func.get("arguments", "")) + tokens += 10 + + tool_call_id = message.get("tool_call_id") + if tool_call_id: + tokens += self.count_text(tool_call_id) + 2 + + return tokens + + def count_messages(self, messages: list[dict[str, Any]]) -> int: + """Count tokens in a list of messages.""" + total = sum(self.count_message(msg) for msg in messages) + total += 3 # Priming tokens + return total + + +class OpenAICompatibleProvider(Provider): + """Provider for OpenAI-compatible LLM services. + + Works with any service implementing the OpenAI chat completions API: + - Ollama (local) + - vLLM (local/cloud) + - Together AI + - Groq + - Fireworks AI + - LM Studio + - LocalAI + - And many more... + + Example: + # For Ollama + provider = OpenAICompatibleProvider( + name="ollama", + base_url="http://localhost:11434/v1", + default_model="llama3.1", + ) + + # For Together AI + provider = OpenAICompatibleProvider( + name="together", + base_url="https://api.together.xyz/v1", + ) + + # Get token counter for a specific model + counter = provider.get_token_counter("llama-3.1-8b") + """ + + def __init__( + self, + name: str = "openai_compatible", + base_url: str | None = None, + api_key: str | None = None, + default_model: str | None = None, + models: dict[str, ModelCapabilities] | None = None, + ): + """Initialize OpenAI-compatible provider. + + Args: + name: Provider name for identification. + base_url: API base URL (e.g., 'http://localhost:11434/v1'). + api_key: API key (if required). + default_model: Default model for operations. + models: Custom model configurations. + """ + self._name = name + self.base_url = base_url + self.api_key = api_key + self.default_model = default_model + self._models: dict[str, ModelCapabilities] = models or {} + + @property + def name(self) -> str: + return self._name + + def register_model( + self, + model: str, + capabilities: ModelCapabilities | None = None, + **kwargs: Any, + ) -> None: + """Register a model with its capabilities. + + Args: + model: Model name. + capabilities: Model capabilities object. + **kwargs: Alternative way to specify capabilities. + """ + if capabilities is not None: + self._models[model] = capabilities + else: + self._models[model] = ModelCapabilities(model=model, **kwargs) + + def supports_model(self, model: str) -> bool: + """Check if model is supported. + + OpenAI-compatible providers support any model by default, + using estimation for token counting. + """ + return True # Always return True - we can estimate + + def get_token_counter(self, model: str) -> OpenAICompatibleTokenCounter: + """Get token counter for a model. + + Uses the TokenizerRegistry to find the best tokenizer, + with fallback to estimation. + """ + tokenizer_backend = None + + # Check for registered model with specific tokenizer + if model in self._models: + tokenizer_backend = self._models[model].tokenizer_backend + + return OpenAICompatibleTokenCounter(model, tokenizer_backend) + + def get_context_limit(self, model: str) -> int: + """Get context limit for a model. + + Priority: + 1. Registered model capabilities + 2. Default limits for known models + 3. Prefix matching + 4. Default 128K + """ + # Check registered models + if model in self._models: + return self._models[model].context_window + + model_lower = model.lower() + + # Check default limits + if model_lower in _DEFAULT_CONTEXT_LIMITS: + return _DEFAULT_CONTEXT_LIMITS[model_lower] + + # Prefix match + for prefix, limit in _DEFAULT_CONTEXT_LIMITS.items(): + if model_lower.startswith(prefix): + return limit + + # Default to 128K for modern models + return 128000 + + def get_output_buffer(self, model: str, default: int = 4000) -> int: + """Get recommended output buffer.""" + if model in self._models: + return min(self._models[model].max_output_tokens, default) + return default + + def estimate_cost( + self, + input_tokens: int, + output_tokens: int, + model: str, + cached_tokens: int = 0, + ) -> float | None: + """Estimate cost if pricing is configured. + + Args: + input_tokens: Number of input tokens. + output_tokens: Number of output tokens. + model: Model name. + cached_tokens: Number of cached tokens. + + Returns: + Estimated cost in USD, or None if pricing unknown. + """ + if model not in self._models: + return None + + caps = self._models[model] + if caps.input_cost_per_1m is None or caps.output_cost_per_1m is None: + return None + + input_cost = (input_tokens / 1_000_000) * caps.input_cost_per_1m + output_cost = (output_tokens / 1_000_000) * caps.output_cost_per_1m + + return input_cost + output_cost + + +# Pre-configured provider factories for common services + + +def create_ollama_provider( + base_url: str = "http://localhost:11434/v1", +) -> OpenAICompatibleProvider: + """Create provider for Ollama. + + Ollama is a popular local LLM runner that supports many open models. + + Args: + base_url: Ollama API URL (default: http://localhost:11434/v1). + + Returns: + Configured provider. + """ + return OpenAICompatibleProvider( + name="ollama", + base_url=base_url, + ) + + +def create_together_provider( + api_key: str | None = None, +) -> OpenAICompatibleProvider: + """Create provider for Together AI. + + Together AI offers high-performance inference for open models. + + Args: + api_key: Together AI API key. + + Returns: + Configured provider with Together AI pricing. + """ + provider = OpenAICompatibleProvider( + name="together", + base_url="https://api.together.xyz/v1", + api_key=api_key, + ) + + # Register common Together models with pricing + # Pricing as of Jan 2025 (verify current rates) + provider.register_model( + "meta-llama/Llama-3.1-8B-Instruct-Turbo", + context_window=128000, + input_cost_per_1m=0.18, + output_cost_per_1m=0.18, + ) + provider.register_model( + "meta-llama/Llama-3.1-70B-Instruct-Turbo", + context_window=128000, + input_cost_per_1m=0.88, + output_cost_per_1m=0.88, + ) + provider.register_model( + "meta-llama/Llama-3.1-405B-Instruct-Turbo", + context_window=128000, + input_cost_per_1m=3.50, + output_cost_per_1m=3.50, + ) + + return provider + + +def create_groq_provider( + api_key: str | None = None, +) -> OpenAICompatibleProvider: + """Create provider for Groq. + + Groq offers ultra-fast inference on custom hardware. + + Args: + api_key: Groq API key. + + Returns: + Configured provider with Groq pricing. + """ + provider = OpenAICompatibleProvider( + name="groq", + base_url="https://api.groq.com/openai/v1", + api_key=api_key, + ) + + # Register common Groq models with pricing + # Pricing as of Jan 2025 (verify current rates) + provider.register_model( + "llama-3.1-8b-instant", + context_window=128000, + input_cost_per_1m=0.05, + output_cost_per_1m=0.08, + ) + provider.register_model( + "llama-3.1-70b-versatile", + context_window=128000, + input_cost_per_1m=0.59, + output_cost_per_1m=0.79, + ) + provider.register_model( + "mixtral-8x7b-32768", + context_window=32768, + input_cost_per_1m=0.24, + output_cost_per_1m=0.24, + ) + + return provider + + +def create_fireworks_provider( + api_key: str | None = None, +) -> OpenAICompatibleProvider: + """Create provider for Fireworks AI. + + Args: + api_key: Fireworks API key. + + Returns: + Configured provider. + """ + return OpenAICompatibleProvider( + name="fireworks", + base_url="https://api.fireworks.ai/inference/v1", + api_key=api_key, + ) + + +def create_anyscale_provider( + api_key: str | None = None, +) -> OpenAICompatibleProvider: + """Create provider for Anyscale Endpoints. + + Args: + api_key: Anyscale API key. + + Returns: + Configured provider. + """ + return OpenAICompatibleProvider( + name="anyscale", + base_url="https://api.endpoints.anyscale.com/v1", + api_key=api_key, + ) + + +def create_vllm_provider( + base_url: str, +) -> OpenAICompatibleProvider: + """Create provider for vLLM server. + + vLLM is a high-performance inference engine. + + Args: + base_url: vLLM server URL (e.g., 'http://localhost:8000/v1'). + + Returns: + Configured provider. + """ + return OpenAICompatibleProvider( + name="vllm", + base_url=base_url, + ) + + +def create_lmstudio_provider( + base_url: str = "http://localhost:1234/v1", +) -> OpenAICompatibleProvider: + """Create provider for LM Studio. + + LM Studio is a desktop app for running local LLMs. + + Args: + base_url: LM Studio API URL. + + Returns: + Configured provider. + """ + return OpenAICompatibleProvider( + name="lmstudio", + base_url=base_url, + ) diff --git a/headroom/proxy/__init__.py b/headroom/proxy/__init__.py new file mode 100644 index 000000000..3d1c39745 --- /dev/null +++ b/headroom/proxy/__init__.py @@ -0,0 +1,19 @@ +"""Headroom Proxy Server. + +A transparent proxy that sits between LLM clients (Claude Code, Cursor, etc.) +and LLM APIs (Anthropic, OpenAI), applying Headroom optimizations. + +Usage: + # Start the proxy + python -m headroom.proxy.server + + # Use with Claude Code + ANTHROPIC_BASE_URL=http://localhost:8787 claude + + # Use with Cursor (if using Anthropic) + Set base URL in Cursor settings to http://localhost:8787 +""" + +from .server import create_app, run_server + +__all__ = ["create_app", "run_server"] diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py new file mode 100644 index 000000000..546adf135 --- /dev/null +++ b/headroom/proxy/server.py @@ -0,0 +1,1399 @@ +"""Headroom Proxy Server - Production Ready. + +A full-featured LLM proxy with optimization, caching, rate limiting, +and observability. + +Features: +- Context optimization (SmartCrusher, CacheAligner, RollingWindow) +- Semantic caching (save costs on repeated queries) +- Rate limiting (token bucket) +- Retry with exponential backoff +- Cost tracking and budgets +- Request tagging and metadata +- Provider fallback +- Prometheus metrics +- Full request/response logging + +Usage: + python -m headroom.proxy.server --port 8787 + + # With Claude Code: + ANTHROPIC_BASE_URL=http://localhost:8787 claude +""" + +from __future__ import annotations + +import argparse +import asyncio +import hashlib +import json +import logging +import random +import sys +import time +from collections import defaultdict +from dataclasses import asdict, dataclass +from datetime import datetime, timedelta +from pathlib import Path +from typing import Literal + +import httpx + +try: + import uvicorn + from fastapi import FastAPI, Header, HTTPException, Request, Response + from fastapi.middleware.cors import CORSMiddleware + from fastapi.responses import PlainTextResponse, StreamingResponse + FASTAPI_AVAILABLE = True +except ImportError: + FASTAPI_AVAILABLE = False + +# Add parent to path for imports +sys.path.insert(0, str(Path(__file__).parent.parent.parent)) + +from headroom.config import CacheAlignerConfig, RollingWindowConfig, SmartCrusherConfig +from headroom.providers import AnthropicProvider, OpenAIProvider +from headroom.tokenizers import get_tokenizer +from headroom.transforms import CacheAligner, RollingWindow, SmartCrusher, TransformPipeline + +logging.basicConfig( + level=logging.INFO, + format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' +) +logger = logging.getLogger("headroom.proxy") + + +# ============================================================================= +# Data Models +# ============================================================================= + +@dataclass +class RequestLog: + """Complete log of a single request.""" + request_id: str + timestamp: str + provider: str + model: str + + # Tokens + input_tokens_original: int + input_tokens_optimized: int + output_tokens: int | None + tokens_saved: int + savings_percent: float + + # Cost + estimated_cost_usd: float | None + estimated_savings_usd: float | None + + # Performance + optimization_latency_ms: float + total_latency_ms: float | None + + # Metadata + tags: dict[str, str] + cache_hit: bool + transforms_applied: list[str] + + # Request/Response (optional, for debugging) + request_messages: list[dict] | None = None + response_content: str | None = None + error: str | None = None + + +@dataclass +class CacheEntry: + """Cached response entry.""" + response_body: bytes + response_headers: dict[str, str] + created_at: datetime + ttl_seconds: int + hit_count: int = 0 + tokens_saved_per_hit: int = 0 + + +@dataclass +class RateLimitState: + """Token bucket rate limiter state.""" + tokens: float + last_update: float + + +@dataclass +class ProxyConfig: + """Proxy configuration.""" + # Server + host: str = "127.0.0.1" + port: int = 8787 + + # Optimization + optimize: bool = True + min_tokens_to_crush: int = 500 + max_items_after_crush: int = 50 + keep_last_turns: int = 4 + + # Caching + cache_enabled: bool = True + cache_ttl_seconds: int = 3600 # 1 hour + cache_max_entries: int = 1000 + + # Rate limiting + rate_limit_enabled: bool = True + rate_limit_requests_per_minute: int = 60 + rate_limit_tokens_per_minute: int = 100000 + + # Retry + retry_enabled: bool = True + retry_max_attempts: int = 3 + retry_base_delay_ms: int = 1000 + retry_max_delay_ms: int = 30000 + + # Cost tracking + cost_tracking_enabled: bool = True + budget_limit_usd: float | None = None # None = unlimited + budget_period: Literal["hourly", "daily", "monthly"] = "daily" + + # Logging + log_requests: bool = True + log_file: str | None = None + log_full_messages: bool = False # Privacy: don't log content by default + + # Fallback + fallback_enabled: bool = False + fallback_provider: str | None = None # "openai" or "anthropic" + + # Timeouts + request_timeout_seconds: int = 300 + connect_timeout_seconds: int = 10 + + +# ============================================================================= +# Caching +# ============================================================================= + +class SemanticCache: + """Simple semantic cache based on message content hash.""" + + def __init__(self, max_entries: int = 1000, ttl_seconds: int = 3600): + self.max_entries = max_entries + self.ttl_seconds = ttl_seconds + self._cache: dict[str, CacheEntry] = {} + self._access_order: list[str] = [] + + def _compute_key(self, messages: list[dict], model: str) -> str: + """Compute cache key from messages and model.""" + # Normalize messages for consistent hashing + normalized = json.dumps({ + "model": model, + "messages": messages, + }, sort_keys=True) + return hashlib.sha256(normalized.encode()).hexdigest()[:32] + + def get(self, messages: list[dict], model: str) -> CacheEntry | None: + """Get cached response if exists and not expired.""" + key = self._compute_key(messages, model) + entry = self._cache.get(key) + + if entry is None: + return None + + # Check expiration + age = (datetime.now() - entry.created_at).total_seconds() + if age > entry.ttl_seconds: + del self._cache[key] + self._access_order.remove(key) + return None + + entry.hit_count += 1 + return entry + + def set( + self, + messages: list[dict], + model: str, + response_body: bytes, + response_headers: dict[str, str], + tokens_saved: int = 0, + ): + """Cache a response.""" + key = self._compute_key(messages, model) + + # Evict if at capacity (LRU) + while len(self._cache) >= self.max_entries and self._access_order: + oldest_key = self._access_order.pop(0) + self._cache.pop(oldest_key, None) + + self._cache[key] = CacheEntry( + response_body=response_body, + response_headers=response_headers, + created_at=datetime.now(), + ttl_seconds=self.ttl_seconds, + tokens_saved_per_hit=tokens_saved, + ) + self._access_order.append(key) + + def stats(self) -> dict: + """Get cache statistics.""" + total_hits = sum(e.hit_count for e in self._cache.values()) + return { + "entries": len(self._cache), + "max_entries": self.max_entries, + "total_hits": total_hits, + "ttl_seconds": self.ttl_seconds, + } + + def clear(self): + """Clear all cache entries.""" + self._cache.clear() + self._access_order.clear() + + +# ============================================================================= +# Rate Limiting +# ============================================================================= + +class TokenBucketRateLimiter: + """Token bucket rate limiter for requests and tokens.""" + + def __init__( + self, + requests_per_minute: int = 60, + tokens_per_minute: int = 100000, + ): + self.requests_per_minute = requests_per_minute + self.tokens_per_minute = tokens_per_minute + + # Per-key buckets (key = API key or IP) + self._request_buckets: dict[str, RateLimitState] = defaultdict( + lambda: RateLimitState(tokens=requests_per_minute, last_update=time.time()) + ) + self._token_buckets: dict[str, RateLimitState] = defaultdict( + lambda: RateLimitState(tokens=tokens_per_minute, last_update=time.time()) + ) + + def _refill(self, state: RateLimitState, rate_per_minute: float) -> float: + """Refill bucket based on elapsed time.""" + now = time.time() + elapsed = now - state.last_update + refill = elapsed * (rate_per_minute / 60.0) + state.tokens = min(rate_per_minute, state.tokens + refill) + state.last_update = now + return state.tokens + + def check_request(self, key: str = "default") -> tuple[bool, float]: + """Check if request is allowed. Returns (allowed, wait_seconds).""" + state = self._request_buckets[key] + available = self._refill(state, self.requests_per_minute) + + if available >= 1: + state.tokens -= 1 + return True, 0 + + wait_seconds = (1 - available) * (60.0 / self.requests_per_minute) + return False, wait_seconds + + def check_tokens(self, key: str, token_count: int) -> tuple[bool, float]: + """Check if token usage is allowed.""" + state = self._token_buckets[key] + available = self._refill(state, self.tokens_per_minute) + + if available >= token_count: + state.tokens -= token_count + return True, 0 + + wait_seconds = (token_count - available) * (60.0 / self.tokens_per_minute) + return False, wait_seconds + + def stats(self) -> dict: + """Get rate limiter statistics.""" + return { + "requests_per_minute": self.requests_per_minute, + "tokens_per_minute": self.tokens_per_minute, + "active_keys": len(self._request_buckets), + } + + +# ============================================================================= +# Cost Tracking +# ============================================================================= + +class CostTracker: + """Track costs and enforce budgets.""" + + # Pricing per 1M tokens (input, output, cached_input) + PRICING = { + # Anthropic + "claude-3-5-sonnet": (3.00, 15.00, 0.30), + "claude-3-5-haiku": (0.80, 4.00, 0.08), + "claude-3-opus": (15.00, 75.00, 1.50), + "claude-sonnet-4": (3.00, 15.00, 0.30), + "claude-opus-4": (15.00, 75.00, 1.50), + # OpenAI + "gpt-4o": (2.50, 10.00, 1.25), + "gpt-4o-mini": (0.15, 0.60, 0.075), + "o1": (15.00, 60.00, 7.50), + "o1-mini": (1.10, 4.40, 0.55), + "o3-mini": (1.10, 4.40, 0.55), + "gpt-4-turbo": (10.00, 30.00, 5.00), + } + + def __init__(self, budget_limit_usd: float | None = None, budget_period: str = "daily"): + self.budget_limit_usd = budget_limit_usd + self.budget_period = budget_period + + # Cost tracking + self._costs: list[tuple[datetime, float]] = [] + self._total_cost_usd: float = 0 + self._total_savings_usd: float = 0 + + def _get_pricing(self, model: str) -> tuple[float, float, float] | None: + """Get pricing for model.""" + model_lower = model.lower() + for prefix, pricing in self.PRICING.items(): + if prefix in model_lower: + return pricing + return None + + def estimate_cost( + self, + model: str, + input_tokens: int, + output_tokens: int, + cached_tokens: int = 0, + ) -> float | None: + """Estimate cost in USD.""" + pricing = self._get_pricing(model) + if pricing is None: + return None + + input_price, output_price, cached_price = pricing + + regular_input = input_tokens - cached_tokens + cost = ( + (regular_input / 1_000_000) * input_price + + (cached_tokens / 1_000_000) * cached_price + + (output_tokens / 1_000_000) * output_price + ) + return cost + + def record_cost(self, cost_usd: float): + """Record a cost.""" + self._costs.append((datetime.now(), cost_usd)) + self._total_cost_usd += cost_usd + + def record_savings(self, savings_usd: float): + """Record savings from optimization.""" + self._total_savings_usd += savings_usd + + def get_period_cost(self) -> float: + """Get cost for current budget period.""" + now = datetime.now() + + if self.budget_period == "hourly": + cutoff = now - timedelta(hours=1) + elif self.budget_period == "daily": + cutoff = now.replace(hour=0, minute=0, second=0, microsecond=0) + else: # monthly + cutoff = now.replace(day=1, hour=0, minute=0, second=0, microsecond=0) + + return sum(cost for ts, cost in self._costs if ts >= cutoff) + + def check_budget(self) -> tuple[bool, float]: + """Check if within budget. Returns (allowed, remaining).""" + if self.budget_limit_usd is None: + return True, float('inf') + + period_cost = self.get_period_cost() + remaining = self.budget_limit_usd - period_cost + return remaining > 0, max(0, remaining) + + def stats(self) -> dict: + """Get cost statistics.""" + return { + "total_cost_usd": round(self._total_cost_usd, 4), + "total_savings_usd": round(self._total_savings_usd, 4), + "period_cost_usd": round(self.get_period_cost(), 4), + "budget_limit_usd": self.budget_limit_usd, + "budget_period": self.budget_period, + "budget_remaining_usd": round(self.check_budget()[1], 4) if self.budget_limit_usd else None, + } + + +# ============================================================================= +# Prometheus Metrics +# ============================================================================= + +class PrometheusMetrics: + """Prometheus-compatible metrics.""" + + def __init__(self): + self.requests_total = 0 + self.requests_by_provider: dict[str, int] = defaultdict(int) + self.requests_by_model: dict[str, int] = defaultdict(int) + self.requests_cached = 0 + self.requests_rate_limited = 0 + self.requests_failed = 0 + + self.tokens_input_total = 0 + self.tokens_output_total = 0 + self.tokens_saved_total = 0 + + self.latency_sum_ms = 0.0 + self.latency_count = 0 + + self.cost_total_usd = 0.0 + self.savings_total_usd = 0.0 + + def record_request( + self, + provider: str, + model: str, + input_tokens: int, + output_tokens: int, + tokens_saved: int, + latency_ms: float, + cached: bool = False, + cost_usd: float = 0, + savings_usd: float = 0, + ): + """Record metrics for a request.""" + self.requests_total += 1 + self.requests_by_provider[provider] += 1 + self.requests_by_model[model] += 1 + + if cached: + self.requests_cached += 1 + + self.tokens_input_total += input_tokens + self.tokens_output_total += output_tokens + self.tokens_saved_total += tokens_saved + + self.latency_sum_ms += latency_ms + self.latency_count += 1 + + self.cost_total_usd += cost_usd + self.savings_total_usd += savings_usd + + def record_rate_limited(self): + self.requests_rate_limited += 1 + + def record_failed(self): + self.requests_failed += 1 + + def export(self) -> str: + """Export metrics in Prometheus format.""" + lines = [ + "# HELP headroom_requests_total Total number of requests", + "# TYPE headroom_requests_total counter", + f"headroom_requests_total {self.requests_total}", + "", + "# HELP headroom_requests_cached_total Cached request count", + "# TYPE headroom_requests_cached_total counter", + f"headroom_requests_cached_total {self.requests_cached}", + "", + "# HELP headroom_requests_rate_limited_total Rate limited requests", + "# TYPE headroom_requests_rate_limited_total counter", + f"headroom_requests_rate_limited_total {self.requests_rate_limited}", + "", + "# HELP headroom_requests_failed_total Failed requests", + "# TYPE headroom_requests_failed_total counter", + f"headroom_requests_failed_total {self.requests_failed}", + "", + "# HELP headroom_tokens_input_total Total input tokens", + "# TYPE headroom_tokens_input_total counter", + f"headroom_tokens_input_total {self.tokens_input_total}", + "", + "# HELP headroom_tokens_output_total Total output tokens", + "# TYPE headroom_tokens_output_total counter", + f"headroom_tokens_output_total {self.tokens_output_total}", + "", + "# HELP headroom_tokens_saved_total Tokens saved by optimization", + "# TYPE headroom_tokens_saved_total counter", + f"headroom_tokens_saved_total {self.tokens_saved_total}", + "", + "# HELP headroom_latency_ms_sum Sum of request latencies", + "# TYPE headroom_latency_ms_sum counter", + f"headroom_latency_ms_sum {self.latency_sum_ms:.2f}", + "", + "# HELP headroom_cost_usd_total Total cost in USD", + "# TYPE headroom_cost_usd_total counter", + f"headroom_cost_usd_total {self.cost_total_usd:.6f}", + "", + "# HELP headroom_savings_usd_total Total savings in USD", + "# TYPE headroom_savings_usd_total counter", + f"headroom_savings_usd_total {self.savings_total_usd:.6f}", + ] + + # Per-provider metrics + lines.extend([ + "", + "# HELP headroom_requests_by_provider Requests by provider", + "# TYPE headroom_requests_by_provider counter", + ]) + for provider, count in self.requests_by_provider.items(): + lines.append(f'headroom_requests_by_provider{{provider="{provider}"}} {count}') + + # Per-model metrics + lines.extend([ + "", + "# HELP headroom_requests_by_model Requests by model", + "# TYPE headroom_requests_by_model counter", + ]) + for model, count in self.requests_by_model.items(): + lines.append(f'headroom_requests_by_model{{model="{model}"}} {count}') + + return "\n".join(lines) + + +# ============================================================================= +# Request Logger +# ============================================================================= + +class RequestLogger: + """Log requests to JSONL file.""" + + def __init__(self, log_file: str | None = None, log_full_messages: bool = False): + self.log_file = Path(log_file) if log_file else None + self.log_full_messages = log_full_messages + self._logs: list[RequestLog] = [] + + if self.log_file: + self.log_file.parent.mkdir(parents=True, exist_ok=True) + + def log(self, entry: RequestLog): + """Log a request.""" + self._logs.append(entry) + + if self.log_file: + with open(self.log_file, "a") as f: + log_dict = asdict(entry) + if not self.log_full_messages: + log_dict.pop("request_messages", None) + log_dict.pop("response_content", None) + f.write(json.dumps(log_dict) + "\n") + + def get_recent(self, n: int = 100) -> list[dict]: + """Get recent log entries.""" + entries = self._logs[-n:] + return [ + {k: v for k, v in asdict(e).items() + if k not in ("request_messages", "response_content")} + for e in entries + ] + + def stats(self) -> dict: + """Get logging statistics.""" + return { + "total_logged": len(self._logs), + "log_file": str(self.log_file) if self.log_file else None, + } + + +# ============================================================================= +# Main Proxy +# ============================================================================= + +class HeadroomProxy: + """Production-ready Headroom optimization proxy.""" + + ANTHROPIC_API_URL = "https://api.anthropic.com" + OPENAI_API_URL = "https://api.openai.com" + + def __init__(self, config: ProxyConfig): + self.config = config + + # Initialize providers + self.anthropic_provider = AnthropicProvider() + self.openai_provider = OpenAIProvider() + + # Initialize transforms + transforms = [ + CacheAligner(CacheAlignerConfig(enabled=True)), + SmartCrusher(SmartCrusherConfig( + enabled=True, + min_tokens_to_crush=config.min_tokens_to_crush, + max_items_after_crush=config.max_items_after_crush, + )), + RollingWindow(RollingWindowConfig( + enabled=True, + keep_system=True, + keep_last_turns=config.keep_last_turns, + )), + ] + + self.anthropic_pipeline = TransformPipeline( + transforms=transforms, + provider=self.anthropic_provider, + ) + self.openai_pipeline = TransformPipeline( + transforms=transforms, + provider=self.openai_provider, + ) + + # Initialize components + self.cache = SemanticCache( + max_entries=config.cache_max_entries, + ttl_seconds=config.cache_ttl_seconds, + ) if config.cache_enabled else None + + self.rate_limiter = TokenBucketRateLimiter( + requests_per_minute=config.rate_limit_requests_per_minute, + tokens_per_minute=config.rate_limit_tokens_per_minute, + ) if config.rate_limit_enabled else None + + self.cost_tracker = CostTracker( + budget_limit_usd=config.budget_limit_usd, + budget_period=config.budget_period, + ) if config.cost_tracking_enabled else None + + self.metrics = PrometheusMetrics() + + self.logger = RequestLogger( + log_file=config.log_file, + log_full_messages=config.log_full_messages, + ) if config.log_requests else None + + # HTTP client + self.http_client: httpx.AsyncClient | None = None + + # Request counter for IDs + self._request_counter = 0 + + async def startup(self): + """Initialize async resources.""" + self.http_client = httpx.AsyncClient( + timeout=httpx.Timeout( + connect=self.config.connect_timeout_seconds, + read=self.config.request_timeout_seconds, + write=self.config.request_timeout_seconds, + pool=self.config.connect_timeout_seconds, + ) + ) + logger.info("Headroom Proxy started") + logger.info(f"Optimization: {'ENABLED' if self.config.optimize else 'DISABLED'}") + logger.info(f"Caching: {'ENABLED' if self.config.cache_enabled else 'DISABLED'}") + logger.info(f"Rate Limiting: {'ENABLED' if self.config.rate_limit_enabled else 'DISABLED'}") + + async def shutdown(self): + """Cleanup async resources.""" + if self.http_client: + await self.http_client.aclose() + + # Print final stats + self._print_summary() + + def _print_summary(self): + """Print session summary.""" + m = self.metrics + logger.info("=" * 70) + logger.info("HEADROOM PROXY SESSION SUMMARY") + logger.info("=" * 70) + logger.info(f"Total requests: {m.requests_total}") + logger.info(f"Cached responses: {m.requests_cached}") + logger.info(f"Rate limited: {m.requests_rate_limited}") + logger.info(f"Failed: {m.requests_failed}") + logger.info(f"Input tokens: {m.tokens_input_total:,}") + logger.info(f"Output tokens: {m.tokens_output_total:,}") + logger.info(f"Tokens saved: {m.tokens_saved_total:,}") + if m.tokens_input_total > 0: + savings_pct = (m.tokens_saved_total / (m.tokens_input_total + m.tokens_saved_total)) * 100 + logger.info(f"Token savings: {savings_pct:.1f}%") + logger.info(f"Total cost: ${m.cost_total_usd:.4f}") + logger.info(f"Total savings: ${m.savings_total_usd:.4f}") + if m.latency_count > 0: + avg_latency = m.latency_sum_ms / m.latency_count + logger.info(f"Avg latency: {avg_latency:.0f}ms") + logger.info("=" * 70) + + def _next_request_id(self) -> str: + """Generate unique request ID.""" + self._request_counter += 1 + return f"hr_{int(time.time())}_{self._request_counter:06d}" + + def _extract_tags(self, headers: dict) -> dict[str, str]: + """Extract Headroom tags from headers.""" + tags = {} + for key, value in headers.items(): + if key.lower().startswith("x-headroom-"): + tag_name = key.lower().replace("x-headroom-", "") + tags[tag_name] = value + return tags + + async def _retry_request( + self, + method: str, + url: str, + headers: dict, + body: dict, + stream: bool = False, + ) -> httpx.Response: + """Make request with retry and exponential backoff.""" + last_error = None + + for attempt in range(self.config.retry_max_attempts): + try: + if stream: + # For streaming, we return early - retry happens at higher level + return await self.http_client.post(url, json=body, headers=headers) + else: + response = await self.http_client.post(url, json=body, headers=headers) + + # Don't retry client errors (4xx) + if 400 <= response.status_code < 500: + return response + + # Retry server errors (5xx) + if response.status_code >= 500: + raise httpx.HTTPStatusError( + f"Server error: {response.status_code}", + request=response.request, + response=response, + ) + + return response + + except (httpx.ConnectError, httpx.ReadTimeout, httpx.HTTPStatusError) as e: + last_error = e + + if not self.config.retry_enabled or attempt >= self.config.retry_max_attempts - 1: + raise + + # Exponential backoff with jitter + delay = min( + self.config.retry_base_delay_ms * (2 ** attempt), + self.config.retry_max_delay_ms, + ) + delay_with_jitter = delay * (0.5 + random.random()) + + logger.warning( + f"Request failed (attempt {attempt + 1}), retrying in {delay_with_jitter:.0f}ms: {e}" + ) + await asyncio.sleep(delay_with_jitter / 1000) + + raise last_error + + async def handle_anthropic_messages( + self, + request: Request, + ) -> Response | StreamingResponse: + """Handle Anthropic /v1/messages endpoint.""" + start_time = time.time() + request_id = self._next_request_id() + + # Parse request + body = await request.json() + model = body.get("model", "unknown") + messages = body.get("messages", []) + stream = body.get("stream", False) + + # Extract headers and tags + headers = {k: v for k, v in request.headers.items()} + headers.pop("host", None) + headers.pop("content-length", None) + tags = self._extract_tags(headers) + + # Rate limiting + if self.rate_limiter: + rate_key = headers.get("x-api-key", "default")[:16] + allowed, wait_seconds = self.rate_limiter.check_request(rate_key) + if not allowed: + self.metrics.record_rate_limited() + raise HTTPException( + status_code=429, + detail=f"Rate limited. Retry after {wait_seconds:.1f}s", + headers={"Retry-After": str(int(wait_seconds) + 1)}, + ) + + # Budget check + if self.cost_tracker: + allowed, remaining = self.cost_tracker.check_budget() + if not allowed: + raise HTTPException( + status_code=429, + detail=f"Budget exceeded for {self.config.budget_period} period", + ) + + # Check cache (non-streaming only) + cache_hit = False + if self.cache and not stream: + cached = self.cache.get(messages, model) + if cached: + cache_hit = True + optimization_latency = (time.time() - start_time) * 1000 + + self.metrics.record_request( + provider="anthropic", + model=model, + input_tokens=0, + output_tokens=0, + tokens_saved=cached.tokens_saved_per_hit, + latency_ms=optimization_latency, + cached=True, + ) + + return Response( + content=cached.response_body, + headers=cached.response_headers, + media_type="application/json", + ) + + # Count original tokens + tokenizer = get_tokenizer(model) + original_tokens = sum( + tokenizer.count_text(str(m.get("content", ""))) + for m in messages + ) + + # Apply optimization + transforms_applied = [] + optimized_messages = messages + optimized_tokens = original_tokens + + if self.config.optimize and messages: + try: + context_limit = self.anthropic_provider.get_context_limit(model) + result = self.anthropic_pipeline.apply( + messages=messages, + model=model, + model_limit=context_limit, + ) + + if result.messages != messages: + optimized_messages = result.messages + transforms_applied = result.transforms_applied + optimized_tokens = sum( + tokenizer.count_text(str(m.get("content", ""))) + for m in optimized_messages + ) + except Exception as e: + logger.warning(f"Optimization failed: {e}") + + tokens_saved = original_tokens - optimized_tokens + optimization_latency = (time.time() - start_time) * 1000 + + # Update body + body["messages"] = optimized_messages + + # Forward request + url = f"{self.ANTHROPIC_API_URL}/v1/messages" + + try: + if stream: + return await self._stream_response( + url, headers, body, "anthropic", model, request_id, + original_tokens, optimized_tokens, tokens_saved, + transforms_applied, tags, optimization_latency, + ) + else: + response = await self._retry_request("POST", url, headers, body) + total_latency = (time.time() - start_time) * 1000 + + # Parse response for output tokens + output_tokens = 0 + try: + resp_json = response.json() + usage = resp_json.get("usage", {}) + output_tokens = usage.get("output_tokens", 0) + except: + pass + + # Calculate cost + cost_usd = None + savings_usd = None + if self.cost_tracker: + cost_usd = self.cost_tracker.estimate_cost( + model, optimized_tokens, output_tokens + ) + original_cost = self.cost_tracker.estimate_cost( + model, original_tokens, output_tokens + ) + if cost_usd and original_cost: + savings_usd = original_cost - cost_usd + self.cost_tracker.record_cost(cost_usd) + self.cost_tracker.record_savings(savings_usd) + + # Cache response + if self.cache and response.status_code == 200: + self.cache.set( + messages, model, + response.content, + dict(response.headers), + tokens_saved=tokens_saved, + ) + + # Record metrics + self.metrics.record_request( + provider="anthropic", + model=model, + input_tokens=optimized_tokens, + output_tokens=output_tokens, + tokens_saved=tokens_saved, + latency_ms=total_latency, + cost_usd=cost_usd or 0, + savings_usd=savings_usd or 0, + ) + + # Log request + if self.logger: + self.logger.log(RequestLog( + request_id=request_id, + timestamp=datetime.now().isoformat(), + provider="anthropic", + model=model, + input_tokens_original=original_tokens, + input_tokens_optimized=optimized_tokens, + output_tokens=output_tokens, + tokens_saved=tokens_saved, + savings_percent=(tokens_saved / original_tokens * 100) if original_tokens > 0 else 0, + estimated_cost_usd=cost_usd, + estimated_savings_usd=savings_usd, + optimization_latency_ms=optimization_latency, + total_latency_ms=total_latency, + tags=tags, + cache_hit=cache_hit, + transforms_applied=transforms_applied, + request_messages=messages if self.config.log_full_messages else None, + )) + + # Log to console + if tokens_saved > 0: + logger.info( + f"[{request_id}] {model}: {original_tokens:,} → {optimized_tokens:,} " + f"(saved {tokens_saved:,} tokens, ${savings_usd:.4f})" if savings_usd else + f"[{request_id}] {model}: {original_tokens:,} → {optimized_tokens:,} " + f"(saved {tokens_saved:,} tokens)" + ) + + return Response( + content=response.content, + status_code=response.status_code, + headers=dict(response.headers), + ) + + except Exception as e: + self.metrics.record_failed() + logger.error(f"[{request_id}] Request failed: {e}") + + # Try fallback if enabled + if self.config.fallback_enabled and self.config.fallback_provider == "openai": + logger.info(f"[{request_id}] Attempting fallback to OpenAI") + # Convert to OpenAI format and retry + # (simplified - would need message format conversion) + + raise HTTPException(status_code=502, detail=str(e)) + + async def _stream_response( + self, + url: str, + headers: dict, + body: dict, + provider: str, + model: str, + request_id: str, + original_tokens: int, + optimized_tokens: int, + tokens_saved: int, + transforms_applied: list[str], + tags: dict[str, str], + optimization_latency: float, + ) -> StreamingResponse: + """Stream response with metrics tracking.""" + start_time = time.time() + + async def generate(): + output_chunks = [] + try: + async with self.http_client.stream("POST", url, json=body, headers=headers) as response: + async for chunk in response.aiter_bytes(): + output_chunks.append(chunk) + yield chunk + finally: + # Record metrics after stream completes + total_latency = (time.time() - start_time) * 1000 + + # Estimate output tokens from chunks (rough) + total_output = b"".join(output_chunks) + output_tokens = len(total_output) // 4 # Rough estimate + + self.metrics.record_request( + provider=provider, + model=model, + input_tokens=optimized_tokens, + output_tokens=output_tokens, + tokens_saved=tokens_saved, + latency_ms=total_latency, + ) + + if tokens_saved > 0: + logger.info( + f"[{request_id}] {model}: saved {tokens_saved:,} tokens (streaming)" + ) + + return StreamingResponse( + generate(), + media_type="text/event-stream", + ) + + async def handle_openai_chat( + self, + request: Request, + ) -> Response | StreamingResponse: + """Handle OpenAI /v1/chat/completions endpoint.""" + start_time = time.time() + request_id = self._next_request_id() + + body = await request.json() + model = body.get("model", "unknown") + messages = body.get("messages", []) + stream = body.get("stream", False) + + headers = {k: v for k, v in request.headers.items()} + headers.pop("host", None) + headers.pop("content-length", None) + tags = self._extract_tags(headers) + + # Rate limiting + if self.rate_limiter: + rate_key = headers.get("authorization", "default")[:20] + allowed, wait_seconds = self.rate_limiter.check_request(rate_key) + if not allowed: + self.metrics.record_rate_limited() + raise HTTPException( + status_code=429, + detail=f"Rate limited. Retry after {wait_seconds:.1f}s", + ) + + # Check cache + cache_hit = False + if self.cache and not stream: + cached = self.cache.get(messages, model) + if cached: + cache_hit = True + self.metrics.record_request( + provider="openai", model=model, + input_tokens=0, output_tokens=0, + tokens_saved=cached.tokens_saved_per_hit, + latency_ms=(time.time() - start_time) * 1000, + cached=True, + ) + return Response(content=cached.response_body, headers=cached.response_headers) + + # Token counting + tokenizer = get_tokenizer(model) + original_tokens = sum( + tokenizer.count_text(str(m.get("content", ""))) + for m in messages + ) + + # Optimization + transforms_applied = [] + optimized_messages = messages + optimized_tokens = original_tokens + + if self.config.optimize and messages: + try: + context_limit = self.openai_provider.get_context_limit(model) + result = self.openai_pipeline.apply( + messages=messages, + model=model, + model_limit=context_limit, + ) + if result.messages != messages: + optimized_messages = result.messages + transforms_applied = result.transforms_applied + optimized_tokens = sum( + tokenizer.count_text(str(m.get("content", ""))) + for m in optimized_messages + ) + except Exception as e: + logger.warning(f"Optimization failed: {e}") + + tokens_saved = original_tokens - optimized_tokens + optimization_latency = (time.time() - start_time) * 1000 + + body["messages"] = optimized_messages + url = f"{self.OPENAI_API_URL}/v1/chat/completions" + + try: + if stream: + return await self._stream_response( + url, headers, body, "openai", model, request_id, + original_tokens, optimized_tokens, tokens_saved, + transforms_applied, tags, optimization_latency, + ) + else: + response = await self._retry_request("POST", url, headers, body) + total_latency = (time.time() - start_time) * 1000 + + output_tokens = 0 + try: + resp_json = response.json() + usage = resp_json.get("usage", {}) + output_tokens = usage.get("completion_tokens", 0) + except: + pass + + # Cost tracking + cost_usd = savings_usd = None + if self.cost_tracker: + cost_usd = self.cost_tracker.estimate_cost(model, optimized_tokens, output_tokens) + original_cost = self.cost_tracker.estimate_cost(model, original_tokens, output_tokens) + if cost_usd and original_cost: + savings_usd = original_cost - cost_usd + self.cost_tracker.record_cost(cost_usd) + self.cost_tracker.record_savings(savings_usd) + + # Cache + if self.cache and response.status_code == 200: + self.cache.set(messages, model, response.content, dict(response.headers), tokens_saved) + + # Metrics + self.metrics.record_request( + provider="openai", model=model, + input_tokens=optimized_tokens, output_tokens=output_tokens, + tokens_saved=tokens_saved, latency_ms=total_latency, + cost_usd=cost_usd or 0, savings_usd=savings_usd or 0, + ) + + if tokens_saved > 0: + logger.info( + f"[{request_id}] {model}: {original_tokens:,} → {optimized_tokens:,} " + f"(saved {tokens_saved:,} tokens)" + ) + + return Response( + content=response.content, + status_code=response.status_code, + headers=dict(response.headers), + ) + except Exception as e: + self.metrics.record_failed() + raise HTTPException(status_code=502, detail=str(e)) + + async def handle_passthrough(self, request: Request, base_url: str) -> Response: + """Pass through request unchanged.""" + path = request.url.path + url = f"{base_url}{path}" + + headers = {k: v for k, v in request.headers.items()} + headers.pop("host", None) + + body = await request.body() + + response = await self.http_client.request( + method=request.method, + url=url, + headers=headers, + content=body, + ) + + return Response( + content=response.content, + status_code=response.status_code, + headers=dict(response.headers), + ) + + +# ============================================================================= +# FastAPI App +# ============================================================================= + +def create_app(config: ProxyConfig | None = None) -> FastAPI: + """Create FastAPI application.""" + if not FASTAPI_AVAILABLE: + raise ImportError("FastAPI required. Install: pip install fastapi uvicorn httpx") + + config = config or ProxyConfig() + + app = FastAPI( + title="Headroom Proxy", + description="Production-ready LLM optimization proxy", + version="1.0.0", + ) + + # CORS + app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], + ) + + proxy = HeadroomProxy(config) + + @app.on_event("startup") + async def startup(): + await proxy.startup() + + @app.on_event("shutdown") + async def shutdown(): + await proxy.shutdown() + + # Health & Metrics + @app.get("/health") + async def health(): + return { + "status": "healthy", + "version": "1.0.0", + "config": { + "optimize": config.optimize, + "cache": config.cache_enabled, + "rate_limit": config.rate_limit_enabled, + } + } + + @app.get("/stats") + async def stats(): + m = proxy.metrics + return { + "requests": { + "total": m.requests_total, + "cached": m.requests_cached, + "rate_limited": m.requests_rate_limited, + "failed": m.requests_failed, + }, + "tokens": { + "input": m.tokens_input_total, + "output": m.tokens_output_total, + "saved": m.tokens_saved_total, + "savings_percent": round( + (m.tokens_saved_total / (m.tokens_input_total + m.tokens_saved_total) * 100) + if m.tokens_input_total > 0 else 0, 2 + ), + }, + "cost": proxy.cost_tracker.stats() if proxy.cost_tracker else None, + "cache": proxy.cache.stats() if proxy.cache else None, + "rate_limiter": proxy.rate_limiter.stats() if proxy.rate_limiter else None, + "recent_requests": proxy.logger.get_recent(10) if proxy.logger else [], + } + + @app.get("/metrics") + async def metrics(): + """Prometheus metrics endpoint.""" + return PlainTextResponse( + proxy.metrics.export(), + media_type="text/plain; version=0.0.4", + ) + + @app.post("/cache/clear") + async def clear_cache(): + """Clear the response cache.""" + if proxy.cache: + proxy.cache.clear() + return {"status": "cleared"} + return {"status": "cache disabled"} + + # Anthropic endpoints + @app.post("/v1/messages") + async def anthropic_messages(request: Request): + return await proxy.handle_anthropic_messages(request) + + @app.post("/v1/messages/count_tokens") + async def anthropic_count_tokens(request: Request): + return await proxy.handle_passthrough(request, proxy.ANTHROPIC_API_URL) + + # OpenAI endpoints + @app.post("/v1/chat/completions") + async def openai_chat(request: Request): + return await proxy.handle_openai_chat(request) + + # Passthrough + @app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE"]) + async def passthrough(request: Request, path: str): + if "anthropic" in request.headers.get("user-agent", "").lower(): + base_url = proxy.ANTHROPIC_API_URL + else: + base_url = proxy.OPENAI_API_URL + return await proxy.handle_passthrough(request, base_url) + + return app + + +def run_server(config: ProxyConfig | None = None): + """Run the proxy server.""" + if not FASTAPI_AVAILABLE: + print("ERROR: FastAPI required. Install: pip install fastapi uvicorn httpx") + sys.exit(1) + + config = config or ProxyConfig() + app = create_app(config) + + print(f""" +╔══════════════════════════════════════════════════════════════════════╗ +║ HEADROOM PROXY SERVER ║ +╠══════════════════════════════════════════════════════════════════════╣ +║ Version: 1.0.0 ║ +║ Listening: http://{config.host}:{config.port:<5} ║ +╠══════════════════════════════════════════════════════════════════════╣ +║ FEATURES: ║ +║ Optimization: {'ENABLED ' if config.optimize else 'DISABLED'} ║ +║ Caching: {'ENABLED ' if config.cache_enabled else 'DISABLED'} (TTL: {config.cache_ttl_seconds}s) ║ +║ Rate Limiting: {'ENABLED ' if config.rate_limit_enabled else 'DISABLED'} ({config.rate_limit_requests_per_minute} req/min, {config.rate_limit_tokens_per_minute:,} tok/min) ║ +║ Retry: {'ENABLED ' if config.retry_enabled else 'DISABLED'} (max {config.retry_max_attempts} attempts) ║ +║ Cost Tracking: {'ENABLED ' if config.cost_tracking_enabled else 'DISABLED'} (budget: {'$' + str(config.budget_limit_usd) + '/' + config.budget_period if config.budget_limit_usd else 'unlimited'}) ║ +╠══════════════════════════════════════════════════════════════════════╣ +║ USAGE: ║ +║ Claude Code: ANTHROPIC_BASE_URL=http://{config.host}:{config.port} claude ║ +║ Cursor: Set base URL in settings ║ +╠══════════════════════════════════════════════════════════════════════╣ +║ ENDPOINTS: ║ +║ /health Health check ║ +║ /stats Detailed statistics ║ +║ /metrics Prometheus metrics ║ +║ /cache/clear Clear response cache ║ +╚══════════════════════════════════════════════════════════════════════╝ +""") + + uvicorn.run(app, host=config.host, port=config.port, log_level="warning") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Headroom Proxy Server") + + # Server + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, default=8787) + + # Optimization + parser.add_argument("--no-optimize", action="store_true", help="Disable optimization") + parser.add_argument("--min-tokens", type=int, default=500, help="Min tokens to crush") + parser.add_argument("--max-items", type=int, default=50, help="Max items after crush") + + # Caching + parser.add_argument("--no-cache", action="store_true", help="Disable caching") + parser.add_argument("--cache-ttl", type=int, default=3600, help="Cache TTL seconds") + + # Rate limiting + parser.add_argument("--no-rate-limit", action="store_true", help="Disable rate limiting") + parser.add_argument("--rpm", type=int, default=60, help="Requests per minute") + parser.add_argument("--tpm", type=int, default=100000, help="Tokens per minute") + + # Cost + parser.add_argument("--budget", type=float, help="Budget limit in USD") + parser.add_argument("--budget-period", choices=["hourly", "daily", "monthly"], default="daily") + + # Logging + parser.add_argument("--log-file", help="Log file path") + parser.add_argument("--log-messages", action="store_true", help="Log full messages") + + args = parser.parse_args() + + config = ProxyConfig( + host=args.host, + port=args.port, + optimize=not args.no_optimize, + min_tokens_to_crush=args.min_tokens, + max_items_after_crush=args.max_items, + cache_enabled=not args.no_cache, + cache_ttl_seconds=args.cache_ttl, + rate_limit_enabled=not args.no_rate_limit, + rate_limit_requests_per_minute=args.rpm, + rate_limit_tokens_per_minute=args.tpm, + budget_limit_usd=args.budget, + budget_period=args.budget_period, + log_file=args.log_file, + log_full_messages=args.log_messages, + ) + + run_server(config) diff --git a/headroom/py.typed b/headroom/py.typed new file mode 100644 index 000000000..e69de29bb diff --git a/headroom/relevance/bm25.py b/headroom/relevance/bm25.py index cd7cbea5f..7d77ed192 100644 --- a/headroom/relevance/bm25.py +++ b/headroom/relevance/bm25.py @@ -20,9 +20,8 @@ from __future__ import annotations import math import re from collections import Counter -from typing import Any -from .base import RelevanceScore, RelevanceScorer, default_batch_score +from .base import RelevanceScore, RelevanceScorer class BM25Scorer(RelevanceScorer): diff --git a/headroom/relevance/embedding.py b/headroom/relevance/embedding.py index 443291fec..48158b957 100644 --- a/headroom/relevance/embedding.py +++ b/headroom/relevance/embedding.py @@ -20,7 +20,7 @@ Limitations: from __future__ import annotations import logging -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING import numpy as np @@ -71,7 +71,7 @@ class EmbeddingScorer(RelevanceScorer): Requires sentence-transformers: pip install headroom[relevance] """ - _model_cache: dict[str, "SentenceTransformer"] = {} + _model_cache: dict[str, SentenceTransformer] = {} def __init__( self, @@ -93,7 +93,7 @@ class EmbeddingScorer(RelevanceScorer): self.model_name = model_name self.device = device self.cache_model = cache_model - self._model: "SentenceTransformer | None" = None + self._model: SentenceTransformer | None = None self._available: bool | None = None @classmethod @@ -110,7 +110,7 @@ class EmbeddingScorer(RelevanceScorer): except ImportError: return False - def _get_model(self) -> "SentenceTransformer": + def _get_model(self) -> SentenceTransformer: """Get or load the sentence transformer model. Returns: diff --git a/headroom/reporting/generator.py b/headroom/reporting/generator.py index 258a88096..1ffdb951d 100644 --- a/headroom/reporting/generator.py +++ b/headroom/reporting/generator.py @@ -11,7 +11,6 @@ from jinja2 import Template from ..storage import create_storage from ..utils import estimate_cost, format_cost - # HTML template embedded as string REPORT_TEMPLATE = """ diff --git a/headroom/storage/base.py b/headroom/storage/base.py index 30be5c772..a71ff6466 100644 --- a/headroom/storage/base.py +++ b/headroom/storage/base.py @@ -3,8 +3,9 @@ from __future__ import annotations from abc import ABC, abstractmethod +from collections.abc import Iterator from datetime import datetime -from typing import Any, Iterator +from typing import Any from ..config import RequestMetrics diff --git a/headroom/storage/jsonl.py b/headroom/storage/jsonl.py index a0a8b7b0a..b44ef2535 100644 --- a/headroom/storage/jsonl.py +++ b/headroom/storage/jsonl.py @@ -3,9 +3,10 @@ from __future__ import annotations import json +from collections.abc import Iterator from datetime import datetime from pathlib import Path -from typing import Any, Iterator +from typing import Any from ..config import RequestMetrics from ..utils import format_timestamp, parse_timestamp diff --git a/headroom/storage/sqlite.py b/headroom/storage/sqlite.py index a1081a252..dbb0343d0 100644 --- a/headroom/storage/sqlite.py +++ b/headroom/storage/sqlite.py @@ -4,9 +4,10 @@ from __future__ import annotations import json import sqlite3 +from collections.abc import Iterator from datetime import datetime from pathlib import Path -from typing import Any, Iterator +from typing import Any from ..config import RequestMetrics from ..utils import format_timestamp, parse_timestamp diff --git a/headroom/tokenizers/__init__.py b/headroom/tokenizers/__init__.py new file mode 100644 index 000000000..7522a188c --- /dev/null +++ b/headroom/tokenizers/__init__.py @@ -0,0 +1,72 @@ +"""Pluggable tokenizer system for universal LLM support. + +This module provides a registry-based tokenizer system that supports +multiple backends: + +1. tiktoken - OpenAI models (GPT-3.5, GPT-4, GPT-4o) +2. HuggingFace - Open models (Llama, Mistral, Falcon, etc.) +3. Anthropic - Claude models (via SDK or estimation) +4. Estimation - Fallback for unknown models + +Usage: + from headroom.tokenizers import TokenizerRegistry, get_tokenizer + + # Auto-detect tokenizer from model name + tokenizer = get_tokenizer("gpt-4o") + tokens = tokenizer.count_text("Hello, world!") + + # Get tokenizer for specific backend + tokenizer = get_tokenizer("llama-3-8b", backend="huggingface") + + # Register custom tokenizer + TokenizerRegistry.register("my-model", my_tokenizer) +""" + +from .base import BaseTokenizer, TokenCounter +from .estimator import CharacterCounter, EstimatingTokenCounter +from .registry import ( + TokenizerRegistry, + get_tokenizer, + list_supported_models, + register_tokenizer, +) +from .tiktoken_counter import TiktokenCounter + + +# Lazy imports for optional dependencies +def get_huggingface_tokenizer(): + """Get HuggingFaceTokenizer class (requires transformers).""" + from .huggingface import HuggingFaceTokenizer + return HuggingFaceTokenizer + + +def get_mistral_tokenizer(): + """Get MistralTokenizer class (requires mistral-common).""" + from .mistral import MistralTokenizer + return MistralTokenizer + + +def is_mistral_tokenizer_available() -> bool: + """Check if Mistral tokenizer is available.""" + from .mistral import is_mistral_available + return is_mistral_available() + + +__all__ = [ + # Registry + "TokenizerRegistry", + "get_tokenizer", + "register_tokenizer", + "list_supported_models", + # Base classes + "TokenCounter", + "BaseTokenizer", + # Implementations + "TiktokenCounter", + "EstimatingTokenCounter", + "CharacterCounter", + # Lazy loaders + "get_huggingface_tokenizer", + "get_mistral_tokenizer", + "is_mistral_tokenizer_available", +] diff --git a/headroom/tokenizers/base.py b/headroom/tokenizers/base.py new file mode 100644 index 000000000..21a95397a --- /dev/null +++ b/headroom/tokenizers/base.py @@ -0,0 +1,203 @@ +"""Base classes for tokenizer implementations. + +Defines the TokenCounter protocol and BaseTokenizer class that all +tokenizer backends must implement. +""" + +from __future__ import annotations + +import json +from abc import ABC, abstractmethod +from typing import Any, Protocol, runtime_checkable + + +@runtime_checkable +class TokenCounter(Protocol): + """Protocol for token counting implementations. + + Any class implementing this protocol can be used with Headroom + for token counting. This allows integration with various + tokenizer backends (tiktoken, HuggingFace, custom, etc.). + """ + + def count_text(self, text: str) -> int: + """Count tokens in a text string. + + Args: + text: The text to count tokens for. + + Returns: + Number of tokens in the text. + """ + ... + + def count_messages(self, messages: list[dict[str, Any]]) -> int: + """Count tokens in a list of chat messages. + + Args: + messages: List of message dicts with 'role' and 'content'. + + Returns: + Total token count including message overhead. + """ + ... + + +class BaseTokenizer(ABC): + """Abstract base class for tokenizer implementations. + + Provides common functionality for counting messages while + requiring subclasses to implement text tokenization. + """ + + # Token overhead per message (role, formatting, etc.) + # Override in subclasses for model-specific overhead + MESSAGE_OVERHEAD = 4 + REPLY_OVERHEAD = 3 # Assistant reply start tokens + + @abstractmethod + def count_text(self, text: str) -> int: + """Count tokens in a text string. Must be implemented by subclasses.""" + pass + + def count_messages(self, messages: list[dict[str, Any]]) -> int: + """Count tokens in a list of chat messages. + + Uses OpenAI-style message counting as the baseline, which + works well for most models. + + Args: + messages: List of message dicts. + + Returns: + Total token count. + """ + total = 0 + + for message in messages: + # Base message overhead + total += self.MESSAGE_OVERHEAD + + # Count role + role = message.get("role", "") + total += self.count_text(role) + + # Count content + content = message.get("content") + if content is not None: + if isinstance(content, str): + total += self.count_text(content) + elif isinstance(content, list): + # Multi-part content (images, tool results, etc.) + total += self._count_content_parts(content) + + # Count tool calls + tool_calls = message.get("tool_calls") + if tool_calls: + total += self._count_tool_calls(tool_calls) + + # Count function call (legacy) + function_call = message.get("function_call") + if function_call: + total += self._count_function_call(function_call) + + # Count name field + name = message.get("name") + if name: + total += self.count_text(name) + total += 1 # Name field overhead + + # Reply start overhead + total += self.REPLY_OVERHEAD + + return total + + def _count_content_parts(self, parts: list[Any]) -> int: + """Count tokens in multi-part content.""" + total = 0 + for part in parts: + if isinstance(part, dict): + part_type = part.get("type", "") + + if part_type == "text": + total += self.count_text(part.get("text", "")) + elif part_type == "image_url": + # Images have fixed token cost (varies by model) + total += 85 # Base image token count + elif part_type == "tool_result": + content = part.get("content", "") + if isinstance(content, str): + total += self.count_text(content) + else: + total += self.count_text(json.dumps(content)) + elif part_type == "tool_use": + total += self.count_text(part.get("name", "")) + total += self.count_text(json.dumps(part.get("input", {}))) + else: + # Unknown type - estimate from JSON + total += self.count_text(json.dumps(part)) + elif isinstance(part, str): + total += self.count_text(part) + + return total + + def _count_tool_calls(self, tool_calls: list[dict[str, Any]]) -> int: + """Count tokens in tool calls.""" + total = 0 + for call in tool_calls: + total += 4 # Tool call overhead + + if "function" in call: + func = call["function"] + total += self.count_text(func.get("name", "")) + total += self.count_text(func.get("arguments", "")) + + if "id" in call: + total += self.count_text(call["id"]) + + return total + + def _count_function_call(self, function_call: dict[str, Any]) -> int: + """Count tokens in legacy function call.""" + total = 4 # Function call overhead + total += self.count_text(function_call.get("name", "")) + total += self.count_text(function_call.get("arguments", "")) + return total + + def encode(self, text: str) -> list[int]: + """Encode text to token IDs. + + Optional method - not all backends support encoding. + Default implementation raises NotImplementedError. + + Args: + text: Text to encode. + + Returns: + List of token IDs. + + Raises: + NotImplementedError: If encoding is not supported. + """ + raise NotImplementedError( + f"{self.__class__.__name__} does not support encoding" + ) + + def decode(self, tokens: list[int]) -> str: + """Decode token IDs to text. + + Optional method - not all backends support decoding. + Default implementation raises NotImplementedError. + + Args: + tokens: List of token IDs. + + Returns: + Decoded text. + + Raises: + NotImplementedError: If decoding is not supported. + """ + raise NotImplementedError( + f"{self.__class__.__name__} does not support decoding" + ) diff --git a/headroom/tokenizers/estimator.py b/headroom/tokenizers/estimator.py new file mode 100644 index 000000000..d2537c019 --- /dev/null +++ b/headroom/tokenizers/estimator.py @@ -0,0 +1,199 @@ +"""Estimation-based token counter for fallback scenarios. + +When no exact tokenizer is available (e.g., unknown models, missing +dependencies), this provides a reasonable approximation based on +character/word heuristics calibrated against real tokenizers. +""" + +from __future__ import annotations + +import json +import re +from typing import Any + +from .base import BaseTokenizer + + +class EstimatingTokenCounter(BaseTokenizer): + """Token counter using estimation heuristics. + + This is the fallback tokenizer used when: + - Model is unknown/unsupported + - Required tokenizer library not installed + - Speed is prioritized over accuracy + + The estimation is calibrated against tiktoken cl100k_base and + provides ~90% accuracy for typical text. It tends to slightly + overestimate, which is safer for context window management. + + Estimation Strategy: + - Base: ~4 characters per token (calibrated against GPT-4) + - Adjustments for code, URLs, numbers, whitespace + - Special handling for JSON structure + + Example: + counter = EstimatingTokenCounter() + tokens = counter.count_text("Hello, world!") + print(f"Estimated tokens: {tokens}") + """ + + # Calibration constants (derived from tiktoken analysis) + CHARS_PER_TOKEN = 4.0 # Average for English text + CHARS_PER_TOKEN_CODE = 3.5 # Code is denser + CHARS_PER_TOKEN_JSON = 3.2 # JSON has more structure + + # Patterns for content type detection + CODE_PATTERN = re.compile( + r'(?:def |class |function |const |let |var |import |from |' + r'if \(|for \(|while \(|switch \(|try \{|catch \(|' + r'=>|->|\{\{|\}\}|;$)', + re.MULTILINE + ) + JSON_PATTERN = re.compile(r'^\s*[\[\{]') + URL_PATTERN = re.compile(r'https?://\S+') + UUID_PATTERN = re.compile( + r'[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}', + re.IGNORECASE + ) + + def __init__(self, chars_per_token: float | None = None): + """Initialize estimating counter. + + Args: + chars_per_token: Override default chars per token ratio. + If None, auto-detects based on content type. + """ + self._fixed_ratio = chars_per_token + + def count_text(self, text: str) -> int: + """Estimate token count for text. + + Args: + text: Text to count tokens for. + + Returns: + Estimated number of tokens. + """ + if not text: + return 0 + + # Use fixed ratio if provided + if self._fixed_ratio is not None: + return max(1, int(len(text) / self._fixed_ratio + 0.5)) + + # Auto-detect content type and adjust ratio + ratio = self._detect_ratio(text) + + # Apply ratio with minimum of 1 token + base_count = int(len(text) / ratio + 0.5) + + # Add overhead for special patterns + overhead = self._count_special_overhead(text) + + return max(1, base_count + overhead) + + def _detect_ratio(self, text: str) -> float: + """Detect optimal chars-per-token ratio based on content. + + Args: + text: Text to analyze. + + Returns: + Chars per token ratio. + """ + # Check for JSON + if self.JSON_PATTERN.match(text): + try: + json.loads(text) + return self.CHARS_PER_TOKEN_JSON + except (json.JSONDecodeError, ValueError): + pass + + # Check for code + code_matches = len(self.CODE_PATTERN.findall(text)) + if code_matches > len(text) / 500: # ~2 matches per KB + return self.CHARS_PER_TOKEN_CODE + + return self.CHARS_PER_TOKEN + + def _count_special_overhead(self, text: str) -> int: + """Count additional tokens for special patterns. + + URLs and UUIDs often tokenize into more tokens than + character count would suggest. + + Args: + text: Text to analyze. + + Returns: + Additional token overhead. + """ + overhead = 0 + + # URLs typically tokenize to more tokens + urls = self.URL_PATTERN.findall(text) + for url in urls: + # Each URL component adds overhead + overhead += url.count('/') + url.count('?') + url.count('&') + + # UUIDs are typically 8-10 tokens despite being 36 chars + uuids = self.UUID_PATTERN.findall(text) + overhead += len(uuids) * 2 # Each UUID adds ~2 extra tokens + + return overhead + + def count_messages(self, messages: list[dict[str, Any]]) -> int: + """Estimate tokens in chat messages. + + Uses the base class implementation with estimation-based + text counting. + + Args: + messages: List of chat messages. + + Returns: + Estimated total token count. + """ + # Use base class implementation + return super().count_messages(messages) + + def __repr__(self) -> str: + if self._fixed_ratio: + return f"EstimatingTokenCounter(chars_per_token={self._fixed_ratio})" + return "EstimatingTokenCounter(auto)" + + +class CharacterCounter(BaseTokenizer): + """Simple character-based counter. + + Uses a fixed character-to-token ratio. Useful for: + - Quick approximations + - Testing + - Models with unknown tokenization + + This is less accurate than EstimatingTokenCounter but faster. + """ + + def __init__(self, chars_per_token: float = 4.0): + """Initialize character counter. + + Args: + chars_per_token: Characters per token ratio. + """ + self.chars_per_token = chars_per_token + + def count_text(self, text: str) -> int: + """Count tokens based on character count. + + Args: + text: Text to count. + + Returns: + Estimated token count. + """ + if not text: + return 0 + return max(1, int(len(text) / self.chars_per_token + 0.5)) + + def __repr__(self) -> str: + return f"CharacterCounter(chars_per_token={self.chars_per_token})" diff --git a/headroom/tokenizers/huggingface.py b/headroom/tokenizers/huggingface.py new file mode 100644 index 000000000..3cac028f8 --- /dev/null +++ b/headroom/tokenizers/huggingface.py @@ -0,0 +1,316 @@ +"""HuggingFace tokenizer wrapper for open models. + +Supports Llama, Mistral, Falcon, and other models with HuggingFace +tokenizers. Requires the `transformers` library. +""" + +from __future__ import annotations + +import logging +from functools import lru_cache +from typing import Any + +from .base import BaseTokenizer + +logger = logging.getLogger(__name__) + + +# Model name to HuggingFace tokenizer mapping +# Maps common model names to their HuggingFace tokenizer identifiers +MODEL_TO_TOKENIZER: dict[str, str] = { + # Llama 3 family + "llama-3": "meta-llama/Meta-Llama-3-8B", + "llama-3-8b": "meta-llama/Meta-Llama-3-8B", + "llama-3-70b": "meta-llama/Meta-Llama-3-70B", + "llama-3.1-8b": "meta-llama/Llama-3.1-8B", + "llama-3.1-70b": "meta-llama/Llama-3.1-70B", + "llama-3.1-405b": "meta-llama/Llama-3.1-405B", + "llama-3.2-1b": "meta-llama/Llama-3.2-1B", + "llama-3.2-3b": "meta-llama/Llama-3.2-3B", + "llama-3.3-70b": "meta-llama/Llama-3.3-70B-Instruct", + # Llama 2 family + "llama-2": "meta-llama/Llama-2-7b-hf", + "llama-2-7b": "meta-llama/Llama-2-7b-hf", + "llama-2-13b": "meta-llama/Llama-2-13b-hf", + "llama-2-70b": "meta-llama/Llama-2-70b-hf", + # CodeLlama + "codellama": "codellama/CodeLlama-7b-hf", + "codellama-7b": "codellama/CodeLlama-7b-hf", + "codellama-13b": "codellama/CodeLlama-13b-hf", + "codellama-34b": "codellama/CodeLlama-34b-hf", + # Mistral family + "mistral": "mistralai/Mistral-7B-v0.1", + "mistral-7b": "mistralai/Mistral-7B-v0.1", + "mistral-7b-v0.2": "mistralai/Mistral-7B-Instruct-v0.2", + "mistral-7b-v0.3": "mistralai/Mistral-7B-Instruct-v0.3", + "mistral-nemo": "mistralai/Mistral-Nemo-Base-2407", + "mistral-small": "mistralai/Mistral-Small-Instruct-2409", + "mistral-large": "mistralai/Mistral-Large-Instruct-2407", + # Mixtral + "mixtral": "mistralai/Mixtral-8x7B-v0.1", + "mixtral-8x7b": "mistralai/Mixtral-8x7B-v0.1", + "mixtral-8x22b": "mistralai/Mixtral-8x22B-v0.1", + # Qwen family + "qwen": "Qwen/Qwen-7B", + "qwen-7b": "Qwen/Qwen-7B", + "qwen-14b": "Qwen/Qwen-14B", + "qwen-72b": "Qwen/Qwen-72B", + "qwen2": "Qwen/Qwen2-7B", + "qwen2-7b": "Qwen/Qwen2-7B", + "qwen2-72b": "Qwen/Qwen2-72B", + "qwen2.5": "Qwen/Qwen2.5-7B", + "qwen2.5-7b": "Qwen/Qwen2.5-7B", + "qwen2.5-72b": "Qwen/Qwen2.5-72B", + # DeepSeek + "deepseek": "deepseek-ai/deepseek-llm-7b-base", + "deepseek-7b": "deepseek-ai/deepseek-llm-7b-base", + "deepseek-67b": "deepseek-ai/deepseek-llm-67b-base", + "deepseek-coder": "deepseek-ai/deepseek-coder-6.7b-base", + "deepseek-v2": "deepseek-ai/DeepSeek-V2", + "deepseek-v3": "deepseek-ai/DeepSeek-V3", + # Yi family + "yi": "01-ai/Yi-6B", + "yi-6b": "01-ai/Yi-6B", + "yi-34b": "01-ai/Yi-34B", + "yi-1.5": "01-ai/Yi-1.5-6B", + # Phi family + "phi-2": "microsoft/phi-2", + "phi-3": "microsoft/Phi-3-mini-4k-instruct", + "phi-3-mini": "microsoft/Phi-3-mini-4k-instruct", + "phi-3-small": "microsoft/Phi-3-small-8k-instruct", + "phi-3-medium": "microsoft/Phi-3-medium-4k-instruct", + # Falcon + "falcon": "tiiuae/falcon-7b", + "falcon-7b": "tiiuae/falcon-7b", + "falcon-40b": "tiiuae/falcon-40b", + "falcon-180b": "tiiuae/falcon-180B", + # StarCoder + "starcoder": "bigcode/starcoder", + "starcoder2": "bigcode/starcoder2-15b", + "starcoder2-3b": "bigcode/starcoder2-3b", + "starcoder2-7b": "bigcode/starcoder2-7b", + "starcoder2-15b": "bigcode/starcoder2-15b", + # MPT + "mpt-7b": "mosaicml/mpt-7b", + "mpt-30b": "mosaicml/mpt-30b", + # Gemma + "gemma": "google/gemma-7b", + "gemma-2b": "google/gemma-2b", + "gemma-7b": "google/gemma-7b", + "gemma-2": "google/gemma-2-9b", + "gemma-2-9b": "google/gemma-2-9b", + "gemma-2-27b": "google/gemma-2-27b", +} + + +@lru_cache(maxsize=16) +def _load_tokenizer(tokenizer_name: str): + """Load and cache HuggingFace tokenizer. + + Args: + tokenizer_name: HuggingFace model/tokenizer name. + + Returns: + Loaded tokenizer, or None if unavailable. + """ + from transformers import AutoTokenizer + + try: + return AutoTokenizer.from_pretrained( + tokenizer_name, + trust_remote_code=True, + ) + except Exception as e: + logger.warning(f"Failed to load tokenizer {tokenizer_name}: {e}") + return None + + +def get_tokenizer_name(model: str) -> str: + """Get HuggingFace tokenizer name for a model. + + Args: + model: Model name. + + Returns: + HuggingFace tokenizer identifier. + """ + model_lower = model.lower() + + # Direct lookup + if model_lower in MODEL_TO_TOKENIZER: + return MODEL_TO_TOKENIZER[model_lower] + + # Try prefix matching + for key, value in MODEL_TO_TOKENIZER.items(): + if model_lower.startswith(key): + return value + + # Assume model name is the tokenizer name + return model + + +class HuggingFaceTokenizer(BaseTokenizer): + """Token counter using HuggingFace tokenizers. + + Supports any model with a HuggingFace tokenizer, including: + - Llama family (Llama 2, Llama 3, CodeLlama) + - Mistral family (Mistral, Mixtral) + - Qwen family + - DeepSeek family + - Phi family + - Falcon, StarCoder, MPT, Gemma, etc. + + Requires the `transformers` library: + pip install transformers + + Some models may require authentication: + huggingface-cli login + + Example: + counter = HuggingFaceTokenizer("llama-3-8b") + tokens = counter.count_text("Hello, world!") + """ + + # Overhead per message (varies by model, this is a reasonable default) + MESSAGE_OVERHEAD = 4 + REPLY_OVERHEAD = 3 + + def __init__(self, model: str): + """Initialize HuggingFace tokenizer. + + Args: + model: Model name (e.g., 'llama-3-8b', 'mistral-7b'). + """ + self.model = model + self.tokenizer_name = get_tokenizer_name(model) + self._tokenizer = None # Lazy load + + @property + def tokenizer(self): + """Lazy-load the tokenizer.""" + if self._tokenizer is None: + loaded = _load_tokenizer(self.tokenizer_name) + if loaded is not None: + self._tokenizer = loaded + else: + # Mark as unavailable + self._tokenizer = False + return self._tokenizer if self._tokenizer is not False else None + + def _use_fallback(self) -> bool: + """Check if we need to use fallback estimation.""" + return self.tokenizer is None + + def count_text(self, text: str) -> int: + """Count tokens in text. + + Falls back to estimation if tokenizer unavailable. + + Args: + text: Text to tokenize. + + Returns: + Number of tokens. + """ + if not text: + return 0 + if self._use_fallback(): + # Fall back to ~4 chars per token estimation + return max(1, int(len(text) / 4 + 0.5)) + tokens = self.tokenizer.encode(text, add_special_tokens=False) + return len(tokens) + + def count_messages(self, messages: list[dict[str, Any]]) -> int: + """Count tokens in chat messages. + + Uses the model's chat template if available, otherwise + falls back to base class implementation. + + Args: + messages: List of chat messages. + + Returns: + Total token count. + """ + if self._use_fallback(): + # Use base class implementation with estimation + return super().count_messages(messages) + + # Try to use chat template for accurate counting + if hasattr(self.tokenizer, "apply_chat_template"): + try: + # Apply chat template and count + formatted = self.tokenizer.apply_chat_template( + messages, + tokenize=True, + add_generation_prompt=True, + ) + return len(formatted) + except Exception: + # Fall back to base implementation + pass + + return super().count_messages(messages) + + def encode(self, text: str) -> list[int]: + """Encode text to token IDs. + + Args: + text: Text to encode. + + Returns: + List of token IDs. + + Raises: + NotImplementedError: If tokenizer not available. + """ + if self._use_fallback(): + raise NotImplementedError( + f"Encoding not available for {self.model} - " + f"tokenizer {self.tokenizer_name} could not be loaded" + ) + return self.tokenizer.encode(text, add_special_tokens=False) + + def decode(self, tokens: list[int]) -> str: + """Decode token IDs to text. + + Args: + tokens: List of token IDs. + + Returns: + Decoded text. + + Raises: + NotImplementedError: If tokenizer not available. + """ + if self._use_fallback(): + raise NotImplementedError( + f"Decoding not available for {self.model} - " + f"tokenizer {self.tokenizer_name} could not be loaded" + ) + return self.tokenizer.decode(tokens) + + @classmethod + def is_available(cls) -> bool: + """Check if HuggingFace tokenizers are available. + + Returns: + True if transformers is installed. + """ + try: + import transformers + return True + except ImportError: + return False + + @classmethod + def list_supported_models(cls) -> list[str]: + """List models with known tokenizer mappings. + + Returns: + List of supported model names. + """ + return list(MODEL_TO_TOKENIZER.keys()) + + def __repr__(self) -> str: + return f"HuggingFaceTokenizer(model={self.model!r}, tokenizer={self.tokenizer_name!r})" diff --git a/headroom/tokenizers/mistral.py b/headroom/tokenizers/mistral.py new file mode 100644 index 000000000..12317b84b --- /dev/null +++ b/headroom/tokenizers/mistral.py @@ -0,0 +1,244 @@ +"""Mistral tokenizer using the official mistral-common package. + +Mistral AI released their tokenizer publicly, making accurate +token counting possible without API calls. + +Requires: pip install mistral-common +""" + +from __future__ import annotations + +import logging +from functools import lru_cache +from typing import Any + +from .base import BaseTokenizer + +logger = logging.getLogger(__name__) + +# Check if mistral-common is available +try: + from mistral_common.protocol.instruct.messages import ( + AssistantMessage, + SystemMessage, + UserMessage, + ) + from mistral_common.protocol.instruct.request import ChatCompletionRequest + from mistral_common.tokens.tokenizers.mistral import MistralTokenizer as _MistralTokenizer + MISTRAL_AVAILABLE = True +except ImportError: + MISTRAL_AVAILABLE = False + _MistralTokenizer = None + + +def is_mistral_available() -> bool: + """Check if mistral-common is installed.""" + return MISTRAL_AVAILABLE + + +# Model to tokenizer version mapping +MODEL_TO_VERSION = { + # Mistral models use v3 tokenizer (tekken) + "mistral-large": "v3", + "mistral-large-latest": "v3", + "mistral-small": "v3", + "mistral-small-latest": "v3", + "ministral-8b": "v3", + "ministral-3b": "v3", + "mistral-nemo": "v3", + "pixtral-12b": "v3", + "codestral": "v3", + "codestral-latest": "v3", + # Mixtral uses v1 + "mixtral-8x7b": "v1", + "mixtral-8x22b": "v1", + "open-mixtral-8x7b": "v1", + "open-mixtral-8x22b": "v1", + # Mistral 7B uses v1 + "mistral-7b": "v1", + "open-mistral-7b": "v1", + "mistral-7b-instruct": "v1", +} + + +@lru_cache(maxsize=4) +def _get_tokenizer(version: str): + """Get and cache Mistral tokenizer by version.""" + if not MISTRAL_AVAILABLE: + raise RuntimeError( + "mistral-common is required for MistralTokenizer. " + "Install with: pip install mistral-common" + ) + + if version == "v3": + return _MistralTokenizer.v3(is_tekken=True) + elif version == "v2": + return _MistralTokenizer.v2() + else: # v1 + return _MistralTokenizer.v1() + + +def get_tokenizer_version(model: str) -> str: + """Get tokenizer version for a model.""" + model_lower = model.lower() + + # Direct lookup + if model_lower in MODEL_TO_VERSION: + return MODEL_TO_VERSION[model_lower] + + # Prefix matching + for prefix, version in [ + ("mistral-large", "v3"), + ("mistral-small", "v3"), + ("ministral", "v3"), + ("codestral", "v3"), + ("pixtral", "v3"), + ("mistral-nemo", "v3"), + ("mixtral", "v1"), + ("mistral-7b", "v1"), + ("open-mistral", "v1"), + ]: + if model_lower.startswith(prefix): + return version + + # Default to v3 for newer models + return "v3" + + +class MistralTokenizer(BaseTokenizer): + """Token counter using Mistral's official tokenizer. + + Uses mistral-common package for accurate token counting. + + Requires: pip install mistral-common + + Example: + counter = MistralTokenizer("mistral-large") + tokens = counter.count_text("Hello, world!") + """ + + MESSAGE_OVERHEAD = 4 + REPLY_OVERHEAD = 3 + + def __init__(self, model: str = "mistral-large"): + """Initialize Mistral tokenizer. + + Args: + model: Mistral model name. + """ + if not MISTRAL_AVAILABLE: + raise RuntimeError( + "mistral-common is required for MistralTokenizer. " + "Install with: pip install mistral-common" + ) + + self.model = model + self.version = get_tokenizer_version(model) + self._tokenizer = None # Lazy load + + @property + def tokenizer(self): + """Lazy-load the tokenizer (MistralTokenizer object).""" + if self._tokenizer is None: + self._tokenizer = _get_tokenizer(self.version) + return self._tokenizer + + @property + def _text_tokenizer(self): + """Get the underlying text tokenizer for encode/decode.""" + return self.tokenizer.instruct_tokenizer.tokenizer + + def count_text(self, text: str) -> int: + """Count tokens in text. + + Args: + text: Text to tokenize. + + Returns: + Number of tokens. + """ + if not text: + return 0 + tokens = self._text_tokenizer.encode(text, bos=False, eos=False) + return len(tokens) + + def count_messages(self, messages: list[dict[str, Any]]) -> int: + """Count tokens in chat messages. + + Uses Mistral's chat template for accurate counting. + + Args: + messages: List of chat messages. + + Returns: + Total token count. + """ + if not messages: + return 0 + + try: + # Convert to Mistral message format + mistral_messages = [] + for msg in messages: + role = msg.get("role", "user") + content = msg.get("content", "") + + if isinstance(content, list): + # Multi-part content - extract text + text_parts = [] + for part in content: + if isinstance(part, dict) and part.get("type") == "text": + text_parts.append(part.get("text", "")) + elif isinstance(part, str): + text_parts.append(part) + content = "\n".join(text_parts) + + if role == "user": + mistral_messages.append(UserMessage(content=content)) + elif role == "assistant": + mistral_messages.append(AssistantMessage(content=content)) + elif role == "system": + mistral_messages.append(SystemMessage(content=content)) + else: + # Tool messages etc - treat as user + mistral_messages.append(UserMessage(content=content)) + + # Encode with chat template + request = ChatCompletionRequest(messages=mistral_messages) + tokenized = self.tokenizer.encode_chat_completion(request) + return len(tokenized.tokens) + + except Exception as e: + logger.debug(f"Mistral chat encoding failed: {e}, falling back to text counting") + # Fallback to base implementation + return super().count_messages(messages) + + def encode(self, text: str) -> list[int]: + """Encode text to token IDs. + + Args: + text: Text to encode. + + Returns: + List of token IDs. + """ + return self._text_tokenizer.encode(text, bos=False, eos=False) + + def decode(self, tokens: list[int]) -> str: + """Decode token IDs to text. + + Args: + tokens: List of token IDs. + + Returns: + Decoded text. + """ + return self._text_tokenizer.decode(tokens) + + @classmethod + def is_available(cls) -> bool: + """Check if Mistral tokenizer is available.""" + return MISTRAL_AVAILABLE + + def __repr__(self) -> str: + return f"MistralTokenizer(model={self.model!r}, version={self.version!r})" diff --git a/headroom/tokenizers/registry.py b/headroom/tokenizers/registry.py new file mode 100644 index 000000000..4934868da --- /dev/null +++ b/headroom/tokenizers/registry.py @@ -0,0 +1,398 @@ +"""Tokenizer registry for universal model support. + +Provides automatic tokenizer selection based on model name with +support for multiple backends and custom tokenizers. +""" + +from __future__ import annotations + +import logging +import re +from collections.abc import Callable +from typing import TYPE_CHECKING + +from .base import TokenCounter +from .estimator import EstimatingTokenCounter + +if TYPE_CHECKING: + pass + +logger = logging.getLogger(__name__) + + +# Model pattern matching for tokenizer selection +# Order matters - more specific patterns first +MODEL_PATTERNS: list[tuple[str, str]] = [ + # OpenAI models -> tiktoken + (r"^gpt-4o", "tiktoken"), + (r"^gpt-4", "tiktoken"), + (r"^gpt-3\.5", "tiktoken"), + (r"^o1", "tiktoken"), + (r"^o3", "tiktoken"), + (r"^text-embedding", "tiktoken"), + (r"^text-davinci", "tiktoken"), + (r"^code-", "tiktoken"), + (r"^davinci", "tiktoken"), + (r"^curie", "tiktoken"), + (r"^babbage", "tiktoken"), + (r"^ada", "tiktoken"), + # Anthropic models -> estimation (Claude uses custom tokenizer) + (r"^claude-", "anthropic"), + # Llama family -> huggingface (when available) + (r"^llama", "huggingface"), + (r"^meta-llama", "huggingface"), + (r"^codellama", "huggingface"), + # Mistral family -> official mistral tokenizer + (r"^mistral", "mistral"), + (r"^mixtral", "mistral"), + (r"^codestral", "mistral"), + (r"^ministral", "mistral"), + (r"^pixtral", "mistral"), + # Google models -> estimation (Gemini uses SentencePiece) + (r"^gemini", "google"), + (r"^palm", "google"), + # Cohere models -> estimation + (r"^command", "cohere"), + # Open models commonly served via OpenAI-compatible APIs + (r"^phi-", "huggingface"), + (r"^qwen", "huggingface"), + (r"^deepseek", "huggingface"), + (r"^yi-", "huggingface"), + (r"^falcon", "huggingface"), + (r"^mpt-", "huggingface"), + (r"^starcoder", "huggingface"), + (r"^codegen", "huggingface"), +] + + +class TokenizerRegistry: + """Registry for tokenizer instances and factories. + + Supports: + - Automatic tokenizer selection based on model name + - Custom tokenizer registration + - Multiple backends (tiktoken, huggingface, estimation) + - Lazy loading of tokenizer dependencies + + Example: + # Auto-detect tokenizer + tokenizer = TokenizerRegistry.get("gpt-4o") + + # Register custom tokenizer + TokenizerRegistry.register("my-model", my_tokenizer) + + # Use specific backend + tokenizer = TokenizerRegistry.get("llama-3", backend="huggingface") + """ + + # Singleton registry instance + _instance: TokenizerRegistry | None = None + + # Registered tokenizers (model -> tokenizer instance) + _tokenizers: dict[str, TokenCounter] = {} + + # Registered factories (backend -> factory function) + _factories: dict[str, Callable[[str], TokenCounter]] = {} + + # Cache for auto-detected tokenizers + _cache: dict[str, TokenCounter] = {} + + def __new__(cls) -> TokenizerRegistry: + """Singleton pattern.""" + if cls._instance is None: + cls._instance = super().__new__(cls) + cls._instance._init_factories() + return cls._instance + + def _init_factories(self) -> None: + """Initialize default tokenizer factories.""" + self._factories = { + "tiktoken": self._create_tiktoken, + "huggingface": self._create_huggingface, + "anthropic": self._create_anthropic, + "google": self._create_google, + "cohere": self._create_cohere, + "mistral": self._create_mistral, + "estimation": self._create_estimation, + } + + @classmethod + def get( + cls, + model: str, + backend: str | None = None, + fallback: bool = True, + ) -> TokenCounter: + """Get tokenizer for a model. + + Args: + model: Model name (e.g., 'gpt-4o', 'claude-3-sonnet'). + backend: Force specific backend ('tiktoken', 'huggingface', etc.). + If None, auto-detects based on model name. + fallback: If True, fall back to estimation on errors. + + Returns: + TokenCounter instance for the model. + + Raises: + ValueError: If backend not found and fallback=False. + """ + registry = cls() + model_lower = model.lower() + + # Check for explicitly registered tokenizer + if model_lower in registry._tokenizers: + return registry._tokenizers[model_lower] + + # Check cache + cache_key = f"{model_lower}:{backend or 'auto'}" + if cache_key in registry._cache: + return registry._cache[cache_key] + + # Create tokenizer + try: + tokenizer = registry._create_tokenizer(model, backend) + registry._cache[cache_key] = tokenizer + return tokenizer + except Exception as e: + if fallback: + logger.warning( + f"Failed to create tokenizer for {model}: {e}. " + "Falling back to estimation." + ) + tokenizer = EstimatingTokenCounter() + registry._cache[cache_key] = tokenizer + return tokenizer + raise ValueError(f"No tokenizer available for {model}: {e}") from e + + @classmethod + def register( + cls, + model: str, + tokenizer: TokenCounter | None = None, + factory: Callable[[str], TokenCounter] | None = None, + ) -> None: + """Register a tokenizer or factory for a model. + + Args: + model: Model name to register. + tokenizer: Pre-instantiated tokenizer instance. + factory: Factory function that creates tokenizer for model. + + Raises: + ValueError: If neither tokenizer nor factory provided. + """ + registry = cls() + model_lower = model.lower() + + if tokenizer is not None: + registry._tokenizers[model_lower] = tokenizer + elif factory is not None: + registry._factories[model_lower] = factory + else: + raise ValueError("Must provide either tokenizer or factory") + + # Clear cache for this model + keys_to_remove = [k for k in registry._cache if k.startswith(model_lower)] + for key in keys_to_remove: + del registry._cache[key] + + @classmethod + def register_backend( + cls, + backend: str, + factory: Callable[[str], TokenCounter], + ) -> None: + """Register a backend factory. + + Args: + backend: Backend name. + factory: Factory function (model: str) -> TokenCounter. + """ + registry = cls() + registry._factories[backend] = factory + + @classmethod + def list_backends(cls) -> list[str]: + """List available backends.""" + registry = cls() + return list(registry._factories.keys()) + + @classmethod + def list_registered(cls) -> list[str]: + """List explicitly registered models.""" + registry = cls() + return list(registry._tokenizers.keys()) + + @classmethod + def clear_cache(cls) -> None: + """Clear the tokenizer cache.""" + registry = cls() + registry._cache.clear() + + def _create_tokenizer( + self, + model: str, + backend: str | None, + ) -> TokenCounter: + """Create tokenizer for model. + + Args: + model: Model name. + backend: Backend to use (or None for auto-detect). + + Returns: + TokenCounter instance. + """ + if backend is None: + backend = self._detect_backend(model) + + factory = self._factories.get(backend) + if factory is None: + raise ValueError(f"Unknown backend: {backend}") + + return factory(model) + + def _create_mistral(self, model: str) -> TokenCounter: + """Create Mistral tokenizer using official mistral-common.""" + try: + from .mistral import MistralTokenizer, is_mistral_available + if is_mistral_available(): + return MistralTokenizer(model) + except ImportError: + pass + + logger.warning( + "mistral-common not installed for Mistral tokenizer. " + "Install with: pip install mistral-common" + ) + return EstimatingTokenCounter() + + def _detect_backend(self, model: str) -> str: + """Detect best backend for model. + + Args: + model: Model name. + + Returns: + Backend name. + """ + model_lower = model.lower() + + for pattern, backend in MODEL_PATTERNS: + if re.match(pattern, model_lower): + return backend + + # Default to estimation for unknown models + return "estimation" + + def _create_tiktoken(self, model: str) -> TokenCounter: + """Create tiktoken-based tokenizer.""" + try: + from .tiktoken_counter import TiktokenCounter + return TiktokenCounter(model) + except ImportError: + logger.warning( + "tiktoken not installed. Install with: pip install tiktoken" + ) + return EstimatingTokenCounter() + + def _create_huggingface(self, model: str) -> TokenCounter: + """Create HuggingFace-based tokenizer.""" + try: + from .huggingface import HuggingFaceTokenizer + return HuggingFaceTokenizer(model) + except ImportError: + logger.warning( + "transformers not installed for HuggingFace tokenizer. " + "Install with: pip install transformers" + ) + return EstimatingTokenCounter() + except Exception as e: + logger.warning(f"Failed to load HuggingFace tokenizer for {model}: {e}") + return EstimatingTokenCounter() + + def _create_anthropic(self, model: str) -> TokenCounter: + """Create Anthropic tokenizer. + + Anthropic uses a custom tokenizer that's not publicly available. + We use estimation calibrated for Claude models. + """ + # Claude models use ~3.5 chars per token on average + return EstimatingTokenCounter(chars_per_token=3.5) + + def _create_google(self, model: str) -> TokenCounter: + """Create Google tokenizer. + + Gemini uses SentencePiece which isn't easily accessible. + We use estimation calibrated for Gemini models. + """ + # Gemini models use ~4 chars per token + return EstimatingTokenCounter(chars_per_token=4.0) + + def _create_cohere(self, model: str) -> TokenCounter: + """Create Cohere tokenizer. + + Cohere has its own tokenizer, we use estimation. + """ + return EstimatingTokenCounter(chars_per_token=4.0) + + def _create_estimation(self, model: str) -> TokenCounter: + """Create estimation-based tokenizer.""" + return EstimatingTokenCounter() + + +# Convenience functions +def get_tokenizer( + model: str, + backend: str | None = None, + fallback: bool = True, +) -> TokenCounter: + """Get tokenizer for a model. + + This is the main entry point for getting tokenizers. + + Args: + model: Model name (e.g., 'gpt-4o', 'claude-3-sonnet'). + backend: Force specific backend ('tiktoken', 'huggingface', etc.). + fallback: If True, fall back to estimation on errors. + + Returns: + TokenCounter instance. + + Example: + tokenizer = get_tokenizer("gpt-4o") + tokens = tokenizer.count_text("Hello, world!") + """ + return TokenizerRegistry.get(model, backend, fallback) + + +def register_tokenizer( + model: str, + tokenizer: TokenCounter | None = None, + factory: Callable[[str], TokenCounter] | None = None, +) -> None: + """Register a custom tokenizer for a model. + + Args: + model: Model name. + tokenizer: Tokenizer instance. + factory: Factory function. + + Example: + # Register instance + register_tokenizer("my-model", MyTokenizer()) + + # Register factory + register_tokenizer("my-model", factory=lambda m: MyTokenizer(m)) + """ + TokenizerRegistry.register(model, tokenizer, factory) + + +def list_supported_models() -> dict[str, str]: + """List models with known tokenizer mappings. + + Returns: + Dict mapping model pattern to backend. + """ + return {pattern: backend for pattern, backend in MODEL_PATTERNS} diff --git a/headroom/tokenizers/tiktoken_counter.py b/headroom/tokenizers/tiktoken_counter.py new file mode 100644 index 000000000..498e84ac4 --- /dev/null +++ b/headroom/tokenizers/tiktoken_counter.py @@ -0,0 +1,247 @@ +"""Tiktoken-based token counter for OpenAI models. + +Tiktoken is OpenAI's fast BPE tokenizer used by GPT models. +It supports multiple encodings: +- cl100k_base: GPT-4, GPT-3.5-turbo, text-embedding-ada-002 +- o200k_base: GPT-4o, GPT-4o-mini +- p50k_base: Codex models, text-davinci-002/003 +- r50k_base: GPT-3 models (davinci, curie, etc.) +""" + +from __future__ import annotations + +from functools import lru_cache +from typing import Any + +from .base import BaseTokenizer + +# Model to encoding mapping +MODEL_TO_ENCODING = { + # GPT-4o family (o200k_base) + "gpt-4o": "o200k_base", + "gpt-4o-mini": "o200k_base", + "gpt-4o-2024-05-13": "o200k_base", + "gpt-4o-2024-08-06": "o200k_base", + "gpt-4o-2024-11-20": "o200k_base", + "gpt-4o-mini-2024-07-18": "o200k_base", + # o1 reasoning models (o200k_base) + "o1": "o200k_base", + "o1-mini": "o200k_base", + "o1-preview": "o200k_base", + "o3-mini": "o200k_base", + # GPT-4 family (cl100k_base) + "gpt-4": "cl100k_base", + "gpt-4-turbo": "cl100k_base", + "gpt-4-turbo-preview": "cl100k_base", + "gpt-4-0314": "cl100k_base", + "gpt-4-0613": "cl100k_base", + "gpt-4-32k": "cl100k_base", + "gpt-4-32k-0314": "cl100k_base", + "gpt-4-32k-0613": "cl100k_base", + "gpt-4-1106-preview": "cl100k_base", + "gpt-4-0125-preview": "cl100k_base", + "gpt-4-turbo-2024-04-09": "cl100k_base", + # GPT-3.5 family (cl100k_base) + "gpt-3.5-turbo": "cl100k_base", + "gpt-3.5-turbo-0301": "cl100k_base", + "gpt-3.5-turbo-0613": "cl100k_base", + "gpt-3.5-turbo-1106": "cl100k_base", + "gpt-3.5-turbo-0125": "cl100k_base", + "gpt-3.5-turbo-16k": "cl100k_base", + "gpt-3.5-turbo-16k-0613": "cl100k_base", + "gpt-3.5-turbo-instruct": "cl100k_base", + # Embeddings (cl100k_base) + "text-embedding-ada-002": "cl100k_base", + "text-embedding-3-small": "cl100k_base", + "text-embedding-3-large": "cl100k_base", + # Codex (p50k_base) + "code-davinci-002": "p50k_base", + "code-davinci-001": "p50k_base", + "code-cushman-002": "p50k_base", + "code-cushman-001": "p50k_base", + # Legacy GPT-3 (r50k_base) + "text-davinci-003": "p50k_base", + "text-davinci-002": "p50k_base", + "text-davinci-001": "r50k_base", + "text-curie-001": "r50k_base", + "text-babbage-001": "r50k_base", + "text-ada-001": "r50k_base", + "davinci": "r50k_base", + "curie": "r50k_base", + "babbage": "r50k_base", + "ada": "r50k_base", +} + +# Default encoding for unknown models +DEFAULT_ENCODING = "cl100k_base" + + +@lru_cache(maxsize=8) +def _get_encoding(encoding_name: str): + """Get tiktoken encoding, cached for performance.""" + import tiktoken + return tiktoken.get_encoding(encoding_name) + + +def get_encoding_for_model(model: str) -> str: + """Get the tiktoken encoding name for a model. + + Args: + model: Model name (e.g., 'gpt-4o', 'gpt-3.5-turbo'). + + Returns: + Encoding name (e.g., 'o200k_base', 'cl100k_base'). + """ + # Direct lookup + if model in MODEL_TO_ENCODING: + return MODEL_TO_ENCODING[model] + + # Try prefix matching for versioned models + for prefix in ["gpt-4o", "gpt-4-turbo", "gpt-4", "gpt-3.5", "o1", "o3"]: + if model.startswith(prefix): + # Find any model with this prefix + for known_model, encoding in MODEL_TO_ENCODING.items(): + if known_model.startswith(prefix): + return encoding + + return DEFAULT_ENCODING + + +class TiktokenCounter(BaseTokenizer): + """Token counter using tiktoken (OpenAI's tokenizer). + + This is the most accurate tokenizer for OpenAI models and provides + a good approximation for many other models that use similar BPE + tokenization. + + Example: + counter = TiktokenCounter("gpt-4o") + tokens = counter.count_text("Hello, world!") + print(f"Token count: {tokens}") + """ + + # OpenAI-specific message overhead + MESSAGE_OVERHEAD = 3 + REPLY_OVERHEAD = 3 + + def __init__(self, model: str = "gpt-4o"): + """Initialize tiktoken counter. + + Args: + model: Model name to determine encoding. + Defaults to 'gpt-4o' (o200k_base encoding). + """ + self.model = model + self.encoding_name = get_encoding_for_model(model) + self._encoding = None # Lazy load + + @property + def encoding(self): + """Lazy-load the encoding.""" + if self._encoding is None: + self._encoding = _get_encoding(self.encoding_name) + return self._encoding + + def count_text(self, text: str) -> int: + """Count tokens in text using tiktoken. + + Args: + text: Text to tokenize. + + Returns: + Number of tokens. + """ + if not text: + return 0 + return len(self.encoding.encode(text)) + + def count_messages(self, messages: list[dict[str, Any]]) -> int: + """Count tokens in messages using OpenAI's exact formula. + + This matches OpenAI's token counting for chat completions. + + Args: + messages: List of chat messages. + + Returns: + Total token count. + """ + total = 0 + + for message in messages: + # Every message has overhead for role and formatting + total += self.MESSAGE_OVERHEAD + + for key, value in message.items(): + if value is None: + continue + + if key == "content": + if isinstance(value, str): + total += self.count_text(value) + elif isinstance(value, list): + # Multi-part content + for part in value: + if isinstance(part, dict): + if part.get("type") == "text": + total += self.count_text(part.get("text", "")) + elif part.get("type") == "image_url": + # Image tokens vary by detail level + detail = part.get("image_url", {}).get("detail", "auto") + if detail == "low": + total += 85 + else: + total += 170 # Base for high detail + else: + total += self.count_text(str(part)) + elif isinstance(part, str): + total += self.count_text(part) + elif key == "role": + total += self.count_text(value) + elif key == "name": + total += self.count_text(value) + total += 1 # Name adds 1 token + elif key == "tool_calls": + for tool_call in value: + total += 3 # Tool call overhead + if "function" in tool_call: + func = tool_call["function"] + total += self.count_text(func.get("name", "")) + total += self.count_text(func.get("arguments", "")) + if "id" in tool_call: + total += self.count_text(tool_call["id"]) + elif key == "tool_call_id": + total += self.count_text(value) + elif key == "function_call": + total += self.count_text(value.get("name", "")) + total += self.count_text(value.get("arguments", "")) + + # Every reply is primed with assistant + total += self.REPLY_OVERHEAD + + return total + + def encode(self, text: str) -> list[int]: + """Encode text to token IDs. + + Args: + text: Text to encode. + + Returns: + List of token IDs. + """ + return self.encoding.encode(text) + + def decode(self, tokens: list[int]) -> str: + """Decode token IDs to text. + + Args: + tokens: List of token IDs. + + Returns: + Decoded text. + """ + return self.encoding.decode(tokens) + + def __repr__(self) -> str: + return f"TiktokenCounter(model={self.model!r}, encoding={self.encoding_name!r})" diff --git a/headroom/transforms/pipeline.py b/headroom/transforms/pipeline.py index 6ffaa9623..59b2b871f 100644 --- a/headroom/transforms/pipeline.py +++ b/headroom/transforms/pipeline.py @@ -2,14 +2,13 @@ from __future__ import annotations -from typing import Any, TYPE_CHECKING +from typing import TYPE_CHECKING, Any from ..config import ( CacheAlignerConfig, DiffArtifact, HeadroomConfig, RollingWindowConfig, - SmartCrusherConfig, ToolCrusherConfig, TransformDiff, TransformResult, diff --git a/headroom/transforms/smart_crusher.py b/headroom/transforms/smart_crusher.py index fda4e8e14..dcdeec7ce 100644 --- a/headroom/transforms/smart_crusher.py +++ b/headroom/transforms/smart_crusher.py @@ -34,11 +34,10 @@ import statistics from collections import Counter from dataclasses import dataclass, field from enum import Enum -from typing import TYPE_CHECKING, Any +from typing import Any from ..config import RelevanceScorerConfig, TransformResult -from ..relevance import BM25Scorer, RelevanceScorer, create_scorer - +from ..relevance import RelevanceScorer, create_scorer # Legacy patterns for backwards compatibility (extract_query_anchors) _UUID_PATTERN = re.compile( diff --git a/headroom/transforms/tool_crusher.py b/headroom/transforms/tool_crusher.py index 44c980dfa..02d16cd53 100644 --- a/headroom/transforms/tool_crusher.py +++ b/headroom/transforms/tool_crusher.py @@ -2,7 +2,6 @@ from __future__ import annotations -import json from typing import Any from ..config import ToolCrusherConfig, TransformResult diff --git a/pyproject.toml b/pyproject.toml index ca84580e9..d0ea3100a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,39 +4,67 @@ build-backend = "hatchling.build" [project] name = "headroom" -version = "0.1.0" -description = "A safe, deterministic Context Budget Controller for LLM APIs" +version = "0.2.0" +description = "The Context Optimization Layer for LLM Applications - Cut costs by 50-90%" readme = "README.md" -license = "MIT" +license = "Apache-2.0" requires-python = ">=3.10" authors = [ - { name = "Headroom Team" } + { name = "Headroom Contributors" } +] +maintainers = [ + { name = "Headroom Contributors" } ] keywords = [ "llm", "openai", + "anthropic", + "claude", + "gpt", "context", "token", "optimization", + "compression", "caching", + "proxy", + "ai", + "machine-learning", ] classifiers = [ - "Development Status :: 3 - Alpha", + "Development Status :: 4 - Beta", "Intended Audience :: Developers", - "License :: OSI Approved :: MIT License", + "License :: OSI Approved :: Apache Software License", + "Operating System :: OS Independent", "Programming Language :: Python :: 3", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", + "Topic :: Scientific/Engineering :: Artificial Intelligence", "Topic :: Software Development :: Libraries :: Python Modules", + "Typing :: Typed", ] dependencies = [ "tiktoken>=0.5.0", "pydantic>=2.0.0", - "jinja2>=3.0.0", ] [project.optional-dependencies] +# Semantic relevance scoring with embeddings +relevance = [ + "sentence-transformers>=2.2.0", + "numpy>=1.24.0", +] +# Proxy server +proxy = [ + "fastapi>=0.100.0", + "uvicorn>=0.23.0", + "httpx>=0.24.0", +] +# Report generation +reports = [ + "jinja2>=3.0.0", +] +# Development dependencies dev = [ "pytest>=7.0.0", "pytest-cov>=4.0.0", @@ -44,16 +72,36 @@ dev = [ "ruff>=0.1.0", "mypy>=1.0.0", "openai>=1.0.0", + "anthropic>=0.18.0", ] +# All optional dependencies +all = [ + "headroom[relevance,proxy,reports]", +] + +[project.scripts] +headroom = "headroom.cli:main" [project.urls] Homepage = "https://github.com/headroom-sdk/headroom" -Documentation = "https://github.com/headroom-sdk/headroom#readme" +Documentation = "https://headroom.dev/docs" Repository = "https://github.com/headroom-sdk/headroom" +Issues = "https://github.com/headroom-sdk/headroom/issues" +Changelog = "https://github.com/headroom-sdk/headroom/blob/main/CHANGELOG.md" [tool.hatch.build.targets.wheel] packages = ["headroom"] +[tool.hatch.build.targets.sdist] +include = [ + "/headroom", + "/tests", + "/LICENSE", + "/NOTICE", + "/README.md", + "/CHANGELOG.md", +] + [tool.ruff] target-version = "py310" line-length = 100 @@ -71,19 +119,43 @@ select = [ ignore = [ "E501", # line too long (handled by formatter) "B008", # do not perform function calls in argument defaults + "B905", # zip without strict parameter ] [tool.ruff.lint.isort] known-first-party = ["headroom"] +[tool.ruff.format] +quote-style = "double" +indent-style = "space" + [tool.mypy] python_version = "3.10" warn_return_any = true warn_unused_configs = true disallow_untyped_defs = true +ignore_missing_imports = true [tool.pytest.ini_options] testpaths = ["tests"] python_files = ["test_*.py"] python_functions = ["test_*"] addopts = "-v --tb=short" +asyncio_mode = "auto" + +[tool.coverage.run] +source = ["headroom"] +branch = true +omit = [ + "headroom/cli.py", + "*/tests/*", +] + +[tool.coverage.report] +exclude_lines = [ + "pragma: no cover", + "def __repr__", + "raise NotImplementedError", + "if TYPE_CHECKING:", + "if __name__ == .__main__.:", +] diff --git a/tests/test_models.py b/tests/test_models.py new file mode 100644 index 000000000..24683a9f6 --- /dev/null +++ b/tests/test_models.py @@ -0,0 +1,260 @@ +"""Tests for the model registry and capabilities database.""" + +from __future__ import annotations + +import pytest +from datetime import date + +from headroom.models import ( + ModelRegistry, + ModelInfo, + get_model_info, + list_models, + register_model, +) + + +class TestModelInfo: + """Tests for ModelInfo dataclass.""" + + def test_default_values(self): + """Test default values.""" + info = ModelInfo(name="test", provider="test-provider") + assert info.context_window == 128000 + assert info.max_output_tokens == 4096 + assert info.supports_tools is True + assert info.supports_vision is False + assert info.supports_streaming is True + + def test_custom_values(self): + """Test custom values.""" + info = ModelInfo( + name="custom-model", + provider="custom", + context_window=32000, + max_output_tokens=8192, + supports_tools=False, + supports_vision=True, + input_cost_per_1m=1.5, + output_cost_per_1m=3.0, + ) + assert info.context_window == 32000 + assert info.max_output_tokens == 8192 + assert info.supports_tools is False + assert info.supports_vision is True + assert info.input_cost_per_1m == 1.5 + + def test_frozen(self): + """Test that ModelInfo is frozen (immutable).""" + info = ModelInfo(name="test", provider="test") + with pytest.raises(AttributeError): + info.name = "changed" + + +class TestModelRegistry: + """Tests for ModelRegistry.""" + + def test_get_openai_model(self): + """Test getting OpenAI model info.""" + info = ModelRegistry.get("gpt-4o") + assert info is not None + assert info.provider == "openai" + assert info.context_window == 128000 + + def test_get_anthropic_model(self): + """Test getting Anthropic model info.""" + info = ModelRegistry.get("claude-3-5-sonnet-20241022") + assert info is not None + assert info.provider == "anthropic" + assert info.context_window == 200000 + + def test_get_google_model(self): + """Test getting Google model info.""" + info = ModelRegistry.get("gemini-1.5-pro") + assert info is not None + assert info.provider == "google" + assert info.context_window == 2000000 # 2M! + + def test_get_by_alias(self): + """Test getting model by alias.""" + info = ModelRegistry.get("gpt-4o-2024-11-20") + assert info is not None + assert info.name == "gpt-4o" + + def test_get_unknown_model(self): + """Test getting unknown model returns None.""" + info = ModelRegistry.get("unknown-model-xyz") + assert info is None + + def test_get_prefix_matching(self): + """Test prefix matching for versioned models.""" + info = ModelRegistry.get("gpt-4o-new-version") + assert info is not None + assert info.name == "gpt-4o" + + def test_register_custom_model(self): + """Test registering custom model.""" + info = ModelRegistry.register( + "my-custom-model", + provider="custom", + context_window=64000, + supports_vision=True, + ) + assert info.name == "my-custom-model" + assert info.provider == "custom" + assert info.context_window == 64000 + + # Should be retrievable + retrieved = ModelRegistry.get("my-custom-model") + assert retrieved is not None + assert retrieved.context_window == 64000 + + def test_list_models_all(self): + """Test listing all models.""" + models = ModelRegistry.list_models() + assert len(models) > 0 + + def test_list_models_by_provider(self): + """Test listing models by provider.""" + openai_models = ModelRegistry.list_models(provider="openai") + assert len(openai_models) > 0 + assert all(m.provider == "openai" for m in openai_models) + + def test_list_models_with_tools(self): + """Test listing models with tool support.""" + models = ModelRegistry.list_models(supports_tools=True) + assert len(models) > 0 + assert all(m.supports_tools for m in models) + + def test_list_models_with_vision(self): + """Test listing models with vision support.""" + models = ModelRegistry.list_models(supports_vision=True) + assert len(models) > 0 + assert all(m.supports_vision for m in models) + + def test_list_models_min_context(self): + """Test listing models with minimum context.""" + models = ModelRegistry.list_models(min_context=1000000) + assert len(models) > 0 + assert all(m.context_window >= 1000000 for m in models) + + def test_list_providers(self): + """Test listing all providers.""" + providers = ModelRegistry.list_providers() + assert "openai" in providers + assert "anthropic" in providers + assert "google" in providers + + def test_get_context_limit(self): + """Test getting context limit.""" + limit = ModelRegistry.get_context_limit("gpt-4o") + assert limit == 128000 + + def test_get_context_limit_unknown(self): + """Test getting context limit for unknown model.""" + limit = ModelRegistry.get_context_limit("unknown", default=32000) + assert limit == 32000 + + def test_estimate_cost(self): + """Test cost estimation.""" + cost = ModelRegistry.estimate_cost( + model="gpt-4o", + input_tokens=1000000, + output_tokens=500000, + ) + assert cost is not None + # GPT-4o: $2.50/1M input + $10.00/1M output * 0.5 = $2.50 + $5.00 = $7.50 + assert abs(cost - 7.50) < 0.01 + + def test_estimate_cost_with_cache(self): + """Test cost estimation with cached tokens.""" + cost = ModelRegistry.estimate_cost( + model="gpt-4o", + input_tokens=1000000, + output_tokens=0, + cached_tokens=500000, # Half cached + ) + assert cost is not None + # 500K regular at $2.50/1M + 500K cached at $1.25/1M + # = $1.25 + $0.625 = $1.875 + assert abs(cost - 1.875) < 0.01 + + def test_estimate_cost_unknown_model(self): + """Test cost estimation for unknown model.""" + cost = ModelRegistry.estimate_cost( + model="unknown-model", + input_tokens=1000, + output_tokens=500, + ) + assert cost is None + + +class TestConvenienceFunctions: + """Tests for convenience functions.""" + + def test_get_model_info(self): + """Test get_model_info function.""" + info = get_model_info("gpt-4o") + assert info is not None + assert info.name == "gpt-4o" + + def test_list_models(self): + """Test list_models function.""" + models = list_models(provider="anthropic") + assert len(models) > 0 + + def test_register_model(self): + """Test register_model function.""" + info = register_model( + "test-function-model", + provider="test", + context_window=16000, + ) + assert info.name == "test-function-model" + + +class TestBuiltInModels: + """Tests for built-in model data.""" + + def test_gpt4o_info(self): + """Test GPT-4o model info.""" + info = get_model_info("gpt-4o") + assert info.provider == "openai" + assert info.context_window == 128000 + assert info.supports_tools is True + assert info.supports_vision is True + assert info.input_cost_per_1m == 2.50 + assert info.output_cost_per_1m == 10.00 + + def test_o1_info(self): + """Test o1 model info.""" + info = get_model_info("o1") + assert info.provider == "openai" + assert info.context_window == 200000 # 200K context + assert info.max_output_tokens == 100000 # 100K output + + def test_claude_info(self): + """Test Claude model info.""" + info = get_model_info("claude-3-5-sonnet-20241022") + assert info.provider == "anthropic" + assert info.context_window == 200000 + assert info.cached_input_cost_per_1m == 0.30 # 90% cache discount + + def test_gemini_info(self): + """Test Gemini model info.""" + info = get_model_info("gemini-1.5-pro") + assert info.provider == "google" + assert info.context_window == 2000000 # 2M tokens! + + def test_llama_info(self): + """Test Llama model info.""" + info = get_model_info("llama-3.1-8b") + assert info.provider == "meta" + assert info.context_window == 128000 + assert info.tokenizer_backend == "huggingface" + + def test_mistral_info(self): + """Test Mistral model info.""" + info = get_model_info("mistral-large") + assert info.provider == "mistral" + assert info.supports_tools is True diff --git a/tests/test_providers/test_cohere.py b/tests/test_providers/test_cohere.py new file mode 100644 index 000000000..b460e962c --- /dev/null +++ b/tests/test_providers/test_cohere.py @@ -0,0 +1,124 @@ +"""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 diff --git a/tests/test_providers/test_universal.py b/tests/test_providers/test_universal.py new file mode 100644 index 000000000..9b1256ba2 --- /dev/null +++ b/tests/test_providers/test_universal.py @@ -0,0 +1,293 @@ +"""Tests for universal provider support. + +Tests OpenAICompatibleProvider, GoogleProvider, and LiteLLMProvider. +""" + +from __future__ import annotations + +import pytest + + +def _transformers_available() -> bool: + """Check if transformers is available.""" + try: + import transformers # noqa: F401 + return True + except ImportError: + return False + +from headroom.providers import ( + OpenAICompatibleProvider, + ModelCapabilities, + GoogleProvider, + create_ollama_provider, + create_together_provider, + create_groq_provider, + create_vllm_provider, + create_lmstudio_provider, + is_litellm_available, +) + + +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 diff --git a/tests/test_tokenizers.py b/tests/test_tokenizers.py new file mode 100644 index 000000000..f955632e3 --- /dev/null +++ b/tests/test_tokenizers.py @@ -0,0 +1,443 @@ +"""Tests for the pluggable tokenizer system.""" + +from __future__ import annotations + +import pytest + +from headroom.tokenizers import ( + TokenizerRegistry, + get_tokenizer, + register_tokenizer, + list_supported_models, + TiktokenCounter, + EstimatingTokenCounter, + CharacterCounter, + TokenCounter, + BaseTokenizer, + is_mistral_tokenizer_available, + get_mistral_tokenizer, +) + + +class TestTiktokenCounter: + """Tests for TiktokenCounter.""" + + def test_init_default_model(self): + """Test initialization with default model.""" + counter = TiktokenCounter() + assert counter.model == "gpt-4o" + assert counter.encoding_name == "o200k_base" + + def test_init_gpt4_model(self): + """Test initialization with GPT-4.""" + counter = TiktokenCounter("gpt-4") + assert counter.model == "gpt-4" + assert counter.encoding_name == "cl100k_base" + + def test_count_text_empty(self): + """Test counting empty text.""" + counter = TiktokenCounter() + assert counter.count_text("") == 0 + + def test_count_text_simple(self): + """Test counting simple text.""" + counter = TiktokenCounter() + count = counter.count_text("Hello, world!") + assert count > 0 + assert count < 10 # Should be a few tokens + + def test_count_text_unicode(self): + """Test counting text with unicode.""" + counter = TiktokenCounter() + count = counter.count_text("Hello, 世界!") + assert count > 0 + + def test_count_messages_single(self): + """Test counting single message.""" + counter = TiktokenCounter() + messages = [{"role": "user", "content": "Hello!"}] + count = counter.count_messages(messages) + assert count > 0 + + def test_count_messages_with_tool_calls(self): + """Test counting messages with tool calls.""" + counter = TiktokenCounter() + messages = [ + {"role": "user", "content": "Search for Python"}, + { + "role": "assistant", + "tool_calls": [{ + "id": "call_123", + "type": "function", + "function": { + "name": "search", + "arguments": '{"query": "Python"}', + }, + }], + }, + { + "role": "tool", + "tool_call_id": "call_123", + "content": "Results...", + }, + ] + count = counter.count_messages(messages) + assert count > 0 + + def test_encode_decode_roundtrip(self): + """Test encode/decode roundtrip.""" + counter = TiktokenCounter() + text = "Hello, world!" + tokens = counter.encode(text) + decoded = counter.decode(tokens) + assert decoded == text + + def test_repr(self): + """Test string representation.""" + counter = TiktokenCounter("gpt-4o") + assert "TiktokenCounter" in repr(counter) + assert "gpt-4o" in repr(counter) + + +class TestEstimatingTokenCounter: + """Tests for EstimatingTokenCounter.""" + + def test_init_default(self): + """Test initialization with defaults.""" + counter = EstimatingTokenCounter() + assert counter._fixed_ratio is None + + def test_init_fixed_ratio(self): + """Test initialization with fixed ratio.""" + counter = EstimatingTokenCounter(chars_per_token=3.5) + assert counter._fixed_ratio == 3.5 + + def test_count_text_empty(self): + """Test counting empty text.""" + counter = EstimatingTokenCounter() + assert counter.count_text("") == 0 + + def test_count_text_simple(self): + """Test counting simple text.""" + counter = EstimatingTokenCounter() + text = "Hello, world!" + count = counter.count_text(text) + assert count > 0 + # Rough estimate: 13 chars / 4 chars per token ≈ 3-4 tokens + assert 2 <= count <= 6 + + def test_count_text_fixed_ratio(self): + """Test counting with fixed ratio.""" + counter = EstimatingTokenCounter(chars_per_token=5.0) + text = "x" * 50 # 50 chars + count = counter.count_text(text) + assert count == 10 # 50 / 5 = 10 + + def test_count_text_minimum_one(self): + """Test minimum of 1 token.""" + counter = EstimatingTokenCounter() + assert counter.count_text("x") >= 1 + + def test_count_messages(self): + """Test counting messages.""" + counter = EstimatingTokenCounter() + messages = [ + {"role": "user", "content": "Hello!"}, + {"role": "assistant", "content": "Hi there!"}, + ] + count = counter.count_messages(messages) + assert count > 0 + + def test_json_detection(self): + """Test JSON content detection.""" + counter = EstimatingTokenCounter() + json_text = '{"name": "test", "value": 123}' + # Should use JSON ratio + count = counter.count_text(json_text) + assert count > 0 + + def test_code_detection(self): + """Test code content detection.""" + counter = EstimatingTokenCounter() + code_text = """ +def hello(): + return "Hello, world!" +""" + count = counter.count_text(code_text) + assert count > 0 + + def test_repr(self): + """Test string representation.""" + counter = EstimatingTokenCounter() + assert "EstimatingTokenCounter" in repr(counter) + + +class TestCharacterCounter: + """Tests for CharacterCounter.""" + + def test_init_default(self): + """Test initialization with default ratio.""" + counter = CharacterCounter() + assert counter.chars_per_token == 4.0 + + def test_init_custom_ratio(self): + """Test initialization with custom ratio.""" + counter = CharacterCounter(chars_per_token=3.5) + assert counter.chars_per_token == 3.5 + + def test_count_text(self): + """Test counting text.""" + counter = CharacterCounter(chars_per_token=4.0) + text = "x" * 40 # 40 chars + count = counter.count_text(text) + assert count == 10 # 40 / 4 = 10 + + def test_count_text_empty(self): + """Test counting empty text.""" + counter = CharacterCounter() + assert counter.count_text("") == 0 + + +class TestTokenizerRegistry: + """Tests for TokenizerRegistry.""" + + def test_get_openai_model(self): + """Test getting tokenizer for OpenAI model.""" + tokenizer = get_tokenizer("gpt-4o") + assert isinstance(tokenizer, TiktokenCounter) + + def test_get_anthropic_model(self): + """Test getting tokenizer for Anthropic model.""" + tokenizer = get_tokenizer("claude-3-sonnet") + assert isinstance(tokenizer, EstimatingTokenCounter) + + def test_get_unknown_model_fallback(self): + """Test fallback for unknown model.""" + tokenizer = get_tokenizer("unknown-model-xyz") + assert isinstance(tokenizer, EstimatingTokenCounter) + + def test_get_with_specific_backend(self): + """Test forcing specific backend.""" + tokenizer = get_tokenizer("any-model", backend="estimation") + assert isinstance(tokenizer, EstimatingTokenCounter) + + def test_register_custom_tokenizer(self): + """Test registering custom tokenizer.""" + custom = EstimatingTokenCounter(chars_per_token=3.0) + register_tokenizer("my-custom-model", tokenizer=custom) + retrieved = get_tokenizer("my-custom-model") + assert retrieved is custom + + def test_list_supported_models(self): + """Test listing supported models.""" + models = list_supported_models() + assert isinstance(models, dict) + assert "gpt-4o" in str(models) or "^gpt-4o" in str(models) + + def test_clear_cache(self): + """Test clearing tokenizer cache.""" + # Get a tokenizer to populate cache + get_tokenizer("gpt-4o") + # Clear cache + TokenizerRegistry.clear_cache() + # Should still work after clearing + tokenizer = get_tokenizer("gpt-4o") + assert tokenizer is not None + + +class TestTokenCounterProtocol: + """Tests for TokenCounter protocol.""" + + def test_tiktoken_implements_protocol(self): + """Test TiktokenCounter implements protocol.""" + counter = TiktokenCounter() + assert isinstance(counter, TokenCounter) + + def test_estimating_implements_protocol(self): + """Test EstimatingTokenCounter implements protocol.""" + counter = EstimatingTokenCounter() + assert isinstance(counter, TokenCounter) + + def test_character_implements_protocol(self): + """Test CharacterCounter implements protocol.""" + counter = CharacterCounter() + assert isinstance(counter, TokenCounter) + + +class TestBaseTokenizer: + """Tests for BaseTokenizer base class.""" + + def test_message_overhead_constant(self): + """Test message overhead constant.""" + assert BaseTokenizer.MESSAGE_OVERHEAD == 4 + + def test_reply_overhead_constant(self): + """Test reply overhead constant.""" + assert BaseTokenizer.REPLY_OVERHEAD == 3 + + +class TestMistralTokenizer: + """Tests for Mistral tokenizer using official mistral-common.""" + + def test_is_available(self): + """Test availability check.""" + result = is_mistral_tokenizer_available() + assert isinstance(result, bool) + + @pytest.mark.skipif( + not is_mistral_tokenizer_available(), + reason="mistral-common not installed", + ) + def test_get_mistral_tokenizer_class(self): + """Test getting MistralTokenizer class.""" + MistralTokenizer = get_mistral_tokenizer() + assert MistralTokenizer is not None + assert hasattr(MistralTokenizer, "count_text") + + @pytest.mark.skipif( + not is_mistral_tokenizer_available(), + reason="mistral-common not installed", + ) + def test_init_default_model(self): + """Test initialization with default model.""" + MistralTokenizer = get_mistral_tokenizer() + counter = MistralTokenizer() + assert counter.model == "mistral-large" + assert counter.version == "v3" + + @pytest.mark.skipif( + not is_mistral_tokenizer_available(), + reason="mistral-common not installed", + ) + def test_init_mixtral_model(self): + """Test initialization with Mixtral model (uses v1).""" + MistralTokenizer = get_mistral_tokenizer() + counter = MistralTokenizer("mixtral-8x7b") + assert counter.version == "v1" + + @pytest.mark.skipif( + not is_mistral_tokenizer_available(), + reason="mistral-common not installed", + ) + def test_count_text_empty(self): + """Test counting empty text.""" + MistralTokenizer = get_mistral_tokenizer() + counter = MistralTokenizer() + assert counter.count_text("") == 0 + + @pytest.mark.skipif( + not is_mistral_tokenizer_available(), + reason="mistral-common not installed", + ) + def test_count_text_simple(self): + """Test counting simple text.""" + MistralTokenizer = get_mistral_tokenizer() + counter = MistralTokenizer() + count = counter.count_text("Hello, world!") + assert count > 0 + assert count < 10 + + @pytest.mark.skipif( + not is_mistral_tokenizer_available(), + reason="mistral-common not installed", + ) + def test_count_text_unicode(self): + """Test counting text with unicode.""" + MistralTokenizer = get_mistral_tokenizer() + counter = MistralTokenizer() + count = counter.count_text("Hello, 世界!") + assert count > 0 + + @pytest.mark.skipif( + not is_mistral_tokenizer_available(), + reason="mistral-common not installed", + ) + def test_count_messages(self): + """Test counting messages.""" + MistralTokenizer = get_mistral_tokenizer() + counter = MistralTokenizer() + messages = [ + {"role": "user", "content": "Hello!"}, + {"role": "assistant", "content": "Hi there!"}, + ] + count = counter.count_messages(messages) + assert count > 0 + + @pytest.mark.skipif( + not is_mistral_tokenizer_available(), + reason="mistral-common not installed", + ) + def test_count_messages_with_system(self): + """Test counting messages with system prompt.""" + MistralTokenizer = get_mistral_tokenizer() + counter = MistralTokenizer() + messages = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Hello!"}, + ] + count = counter.count_messages(messages) + assert count > 0 + + @pytest.mark.skipif( + not is_mistral_tokenizer_available(), + reason="mistral-common not installed", + ) + def test_encode_decode_roundtrip(self): + """Test encode/decode roundtrip.""" + MistralTokenizer = get_mistral_tokenizer() + counter = MistralTokenizer() + text = "Hello, world!" + tokens = counter.encode(text) + decoded = counter.decode(tokens) + assert decoded == text + + @pytest.mark.skipif( + not is_mistral_tokenizer_available(), + reason="mistral-common not installed", + ) + def test_implements_protocol(self): + """Test MistralTokenizer implements TokenCounter protocol.""" + MistralTokenizer = get_mistral_tokenizer() + counter = MistralTokenizer() + assert isinstance(counter, TokenCounter) + + @pytest.mark.skipif( + not is_mistral_tokenizer_available(), + reason="mistral-common not installed", + ) + def test_repr(self): + """Test string representation.""" + MistralTokenizer = get_mistral_tokenizer() + counter = MistralTokenizer("mistral-large") + assert "MistralTokenizer" in repr(counter) + assert "mistral-large" in repr(counter) + + @pytest.mark.skipif( + not is_mistral_tokenizer_available(), + reason="mistral-common not installed", + ) + def test_registry_returns_mistral_for_mistral_models(self): + """Test registry returns Mistral tokenizer for Mistral models.""" + tokenizer = get_tokenizer("mistral-large") + MistralTokenizer = get_mistral_tokenizer() + assert isinstance(tokenizer, MistralTokenizer) + + @pytest.mark.skipif( + not is_mistral_tokenizer_available(), + reason="mistral-common not installed", + ) + def test_registry_returns_mistral_for_mixtral(self): + """Test registry returns Mistral tokenizer for Mixtral models.""" + tokenizer = get_tokenizer("mixtral-8x7b") + MistralTokenizer = get_mistral_tokenizer() + assert isinstance(tokenizer, MistralTokenizer) + + @pytest.mark.skipif( + not is_mistral_tokenizer_available(), + reason="mistral-common not installed", + ) + def test_registry_returns_mistral_for_codestral(self): + """Test registry returns Mistral tokenizer for Codestral models.""" + tokenizer = get_tokenizer("codestral") + MistralTokenizer = get_mistral_tokenizer() + assert isinstance(tokenizer, MistralTokenizer)