tiphys.llm.models

  1"""
  2Model metadata and registry.
  3
  4Defines context window limits, token costs, and capabilities for various models.
  5"""
  6
  7from dataclasses import dataclass
  8
  9
 10@dataclass
 11class ModelInfo:
 12    """Metadata for an LLM model."""
 13
 14    id: str
 15    provider: str
 16    context_window: int
 17    max_output_tokens: int
 18    supports_vision: bool = False
 19    supports_tools: bool = True
 20    cost_per_1k_input: float = 0.0
 21    cost_per_1k_output: float = 0.0
 22
 23
 24# Registry of known models with their context limits
 25MODEL_REGISTRY: dict[str, ModelInfo] = {
 26    # Anthropic
 27    "anthropic/claude-3-5-sonnet-20240620": ModelInfo(
 28        id="anthropic/claude-3-5-sonnet-20240620",
 29        provider="anthropic",
 30        context_window=200000,
 31        max_output_tokens=8192,
 32        supports_vision=True,
 33        cost_per_1k_input=0.003,
 34        cost_per_1k_output=0.015,
 35    ),
 36    "anthropic/claude-3-5-sonnet": ModelInfo(
 37        id="anthropic/claude-3-5-sonnet",
 38        provider="anthropic",
 39        context_window=200000,
 40        max_output_tokens=8192,
 41        supports_vision=True,
 42        cost_per_1k_input=0.003,
 43        cost_per_1k_output=0.015,
 44    ),
 45    "anthropic/claude-3-opus-20240229": ModelInfo(
 46        id="anthropic/claude-3-opus-20240229",
 47        provider="anthropic",
 48        context_window=200000,
 49        max_output_tokens=4096,
 50        supports_vision=True,
 51        cost_per_1k_input=0.015,
 52        cost_per_1k_output=0.075,
 53    ),
 54    # OpenAI
 55    "openai/gpt-4o": ModelInfo(
 56        id="openai/gpt-4o",
 57        provider="openai",
 58        context_window=128000,
 59        max_output_tokens=4096,
 60        supports_vision=True,
 61        cost_per_1k_input=0.005,
 62        cost_per_1k_output=0.015,
 63    ),
 64    "openai/gpt-4-turbo": ModelInfo(
 65        id="openai/gpt-4-turbo",
 66        provider="openai",
 67        context_window=128000,
 68        max_output_tokens=4096,
 69        supports_vision=True,
 70        cost_per_1k_input=0.01,
 71        cost_per_1k_output=0.03,
 72    ),
 73    "openai/gpt-3.5-turbo": ModelInfo(
 74        id="openai/gpt-3.5-turbo",
 75        provider="openai",
 76        context_window=16385,
 77        max_output_tokens=4096,
 78        cost_per_1k_input=0.0005,
 79        cost_per_1k_output=0.0015,
 80    ),
 81    # Ollama (Defaults, can be overridden)
 82    "ollama/llama3": ModelInfo(
 83        id="ollama/llama3",
 84        provider="ollama",
 85        context_window=8192,
 86        max_output_tokens=4096,
 87    ),
 88    "ollama/llama3.1": ModelInfo(
 89        id="ollama/llama3.1",
 90        provider="ollama",
 91        context_window=128000,
 92        max_output_tokens=4096,
 93    ),
 94    "ollama/phi3": ModelInfo(
 95        id="ollama/phi3",
 96        provider="ollama",
 97        context_window=4096,
 98        max_output_tokens=4096,
 99    ),
100}
101
102
103def get_model_info(model_id: str) -> ModelInfo:
104    """
105    Get metadata for a model.
106
107    If the model is not in the registry, returns a default ModelInfo.
108    """
109    if model_id in MODEL_REGISTRY:
110        return MODEL_REGISTRY[model_id]
111
112    # Handle short names (e.g., "claude-3-5-sonnet")
113    for full_id, info in MODEL_REGISTRY.items():
114        if model_id == full_id or model_id == full_id.split("/")[-1]:
115            return info
116
117    # Default for unknown models
118    provider = model_id.split("/")[0] if "/" in model_id else "unknown"
119    return ModelInfo(
120        id=model_id,
121        provider=provider,
122        context_window=4096,  # Conservative default
123        max_output_tokens=4096,
124    )