|
1 | 1 | from __future__ import annotations as _annotations
|
2 | 2 |
|
| 3 | +from dataclasses import dataclass |
| 4 | + |
3 | 5 | from . import ModelProfile
|
4 | 6 |
|
5 | 7 |
|
| 8 | +@dataclass |
| 9 | +class AnthropicModelProfile(ModelProfile): |
| 10 | + """Profile for models used with AnthropicModel. |
| 11 | +
|
| 12 | + ALL FIELDS MUST BE `anthropic_` PREFIXED SO YOU CAN MERGE THEM WITH OTHER MODELS. |
| 13 | + """ |
| 14 | + |
| 15 | + anthropic_max_tokens_limit: int = 1024 |
| 16 | + """The maximum number of tokens that can be generated by this model, as documented at https://docs.anthropic.com/en/docs/about-claude/models/overview. Defaults to 1024.""" |
| 17 | + |
| 18 | + |
| 19 | +MAX_TOKENS = { |
| 20 | + 'claude-3-haiku': 4096, |
| 21 | + 'claude-3-opus': 4096, |
| 22 | + 'claude-3-5-haiku': 8192, |
| 23 | + 'claude-3-5-sonnet': 8192, |
| 24 | + 'claude-3-7-sonnet': 64000, |
| 25 | + 'claude-opus-4': 32000, |
| 26 | + 'claude-4-opus': 32000, |
| 27 | + 'claude-sonnet-4': 64000, |
| 28 | + 'claude-4-sonnet': 64000, |
| 29 | +} |
| 30 | + |
| 31 | + |
6 | 32 | def anthropic_model_profile(model_name: str) -> ModelProfile | None:
|
7 | 33 | """Get the model profile for an Anthropic model."""
|
8 |
| - return None |
| 34 | + max_tokens = 1024 |
| 35 | + for model_prefix, model_max_tokens in MAX_TOKENS.items(): |
| 36 | + if model_name.startswith(model_prefix): |
| 37 | + max_tokens = model_max_tokens |
| 38 | + break |
| 39 | + |
| 40 | + return AnthropicModelProfile(anthropic_max_tokens_limit=max_tokens) |
0 commit comments