Changeset 0.5.0 (#35)
This commit is contained in:
396
server/providers/capabilities.go
Normal file
396
server/providers/capabilities.go
Normal file
@@ -0,0 +1,396 @@
|
||||
package providers
|
||||
|
||||
import (
|
||||
"regexp"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// ModelCapabilities describes what a model can do and its limits.
|
||||
// Zero values mean "unknown / use heuristic".
|
||||
type ModelCapabilities struct {
|
||||
Streaming bool `json:"streaming"`
|
||||
ToolCalling bool `json:"tool_calling"`
|
||||
Vision bool `json:"vision"`
|
||||
Thinking bool `json:"thinking"`
|
||||
Reasoning bool `json:"reasoning"`
|
||||
CodeOptimized bool `json:"code_optimized"`
|
||||
WebSearch bool `json:"web_search"`
|
||||
MaxContext int `json:"max_context"`
|
||||
MaxOutputTokens int `json:"max_output_tokens"`
|
||||
}
|
||||
|
||||
// ── Known Model Defaults ────────────────────
|
||||
// Authoritative output limits for models where the provider API
|
||||
// doesn't report them. Keyed by exact model ID or prefix.
|
||||
//
|
||||
// Sources:
|
||||
// Anthropic: https://docs.anthropic.com/en/docs/about-claude/models
|
||||
// OpenAI: https://platform.openai.com/docs/models
|
||||
// Meta: Model cards on Hugging Face
|
||||
// Google: https://ai.google.dev/gemini-api/docs/models
|
||||
|
||||
var knownModels = map[string]ModelCapabilities{
|
||||
// ── Anthropic ────────────────────────────
|
||||
"claude-opus-4": {
|
||||
Streaming: true, ToolCalling: true, Vision: true, Thinking: true,
|
||||
MaxContext: 200000, MaxOutputTokens: 32000,
|
||||
},
|
||||
"claude-sonnet-4": {
|
||||
Streaming: true, ToolCalling: true, Vision: true, Thinking: true,
|
||||
MaxContext: 200000, MaxOutputTokens: 16000,
|
||||
},
|
||||
"claude-3-5-sonnet": {
|
||||
Streaming: true, ToolCalling: true, Vision: true,
|
||||
MaxContext: 200000, MaxOutputTokens: 8192,
|
||||
},
|
||||
"claude-3-5-haiku": {
|
||||
Streaming: true, ToolCalling: true, Vision: true,
|
||||
MaxContext: 200000, MaxOutputTokens: 8192,
|
||||
},
|
||||
"claude-3-opus": {
|
||||
Streaming: true, ToolCalling: true, Vision: true,
|
||||
MaxContext: 200000, MaxOutputTokens: 4096,
|
||||
},
|
||||
"claude-3-haiku": {
|
||||
Streaming: true, ToolCalling: true, Vision: true,
|
||||
MaxContext: 200000, MaxOutputTokens: 4096,
|
||||
},
|
||||
|
||||
// ── OpenAI ──────────────────────────────
|
||||
"gpt-4o": {
|
||||
Streaming: true, ToolCalling: true, Vision: true,
|
||||
MaxContext: 128000, MaxOutputTokens: 16384,
|
||||
},
|
||||
"gpt-4o-mini": {
|
||||
Streaming: true, ToolCalling: true, Vision: true,
|
||||
MaxContext: 128000, MaxOutputTokens: 16384,
|
||||
},
|
||||
"gpt-4-turbo": {
|
||||
Streaming: true, ToolCalling: true, Vision: true,
|
||||
MaxContext: 128000, MaxOutputTokens: 4096,
|
||||
},
|
||||
"gpt-4": {
|
||||
Streaming: true, ToolCalling: true,
|
||||
MaxContext: 8192, MaxOutputTokens: 4096,
|
||||
},
|
||||
"o1": {
|
||||
Streaming: true, Reasoning: true,
|
||||
MaxContext: 200000, MaxOutputTokens: 100000,
|
||||
},
|
||||
"o1-mini": {
|
||||
Streaming: true, Reasoning: true,
|
||||
MaxContext: 128000, MaxOutputTokens: 65536,
|
||||
},
|
||||
"o3": {
|
||||
Streaming: true, ToolCalling: true, Reasoning: true,
|
||||
MaxContext: 200000, MaxOutputTokens: 100000,
|
||||
},
|
||||
"o3-mini": {
|
||||
Streaming: true, ToolCalling: true, Reasoning: true,
|
||||
MaxContext: 200000, MaxOutputTokens: 65536,
|
||||
},
|
||||
"o4-mini": {
|
||||
Streaming: true, ToolCalling: true, Reasoning: true,
|
||||
MaxContext: 200000, MaxOutputTokens: 100000,
|
||||
},
|
||||
|
||||
// ── Meta Llama ──────────────────────────
|
||||
"llama-3.1-405b": {
|
||||
Streaming: true, ToolCalling: true,
|
||||
MaxContext: 131072, MaxOutputTokens: 4096,
|
||||
},
|
||||
"llama-3.1-70b": {
|
||||
Streaming: true, ToolCalling: true,
|
||||
MaxContext: 131072, MaxOutputTokens: 4096,
|
||||
},
|
||||
"llama-3.1-8b": {
|
||||
Streaming: true, ToolCalling: true,
|
||||
MaxContext: 131072, MaxOutputTokens: 4096,
|
||||
},
|
||||
"llama-3.3-70b": {
|
||||
Streaming: true, ToolCalling: true,
|
||||
MaxContext: 131072, MaxOutputTokens: 4096,
|
||||
},
|
||||
"llama-4-maverick": {
|
||||
Streaming: true, ToolCalling: true, Vision: true,
|
||||
MaxContext: 1048576, MaxOutputTokens: 16384,
|
||||
},
|
||||
"llama-4-scout": {
|
||||
Streaming: true, ToolCalling: true, Vision: true,
|
||||
MaxContext: 524288, MaxOutputTokens: 16384,
|
||||
},
|
||||
|
||||
// ── Google Gemini ───────────────────────
|
||||
"gemini-2.5-pro": {
|
||||
Streaming: true, ToolCalling: true, Vision: true, Thinking: true,
|
||||
MaxContext: 1048576, MaxOutputTokens: 65536,
|
||||
},
|
||||
"gemini-2.5-flash": {
|
||||
Streaming: true, ToolCalling: true, Vision: true, Thinking: true,
|
||||
MaxContext: 1048576, MaxOutputTokens: 65536,
|
||||
},
|
||||
"gemini-2.0-flash": {
|
||||
Streaming: true, ToolCalling: true, Vision: true,
|
||||
MaxContext: 1048576, MaxOutputTokens: 8192,
|
||||
},
|
||||
|
||||
// ── DeepSeek ────────────────────────────
|
||||
"deepseek-r1": {
|
||||
Streaming: true, Reasoning: true,
|
||||
MaxContext: 65536, MaxOutputTokens: 8192,
|
||||
},
|
||||
"deepseek-v3": {
|
||||
Streaming: true, ToolCalling: true,
|
||||
MaxContext: 65536, MaxOutputTokens: 8192,
|
||||
},
|
||||
"deepseek-chat": {
|
||||
Streaming: true, ToolCalling: true,
|
||||
MaxContext: 65536, MaxOutputTokens: 8192,
|
||||
},
|
||||
|
||||
// ── Mistral ─────────────────────────────
|
||||
"mistral-large": {
|
||||
Streaming: true, ToolCalling: true,
|
||||
MaxContext: 131072, MaxOutputTokens: 8192,
|
||||
},
|
||||
"mistral-small": {
|
||||
Streaming: true, ToolCalling: true,
|
||||
MaxContext: 131072, MaxOutputTokens: 8192,
|
||||
},
|
||||
"codestral": {
|
||||
Streaming: true, ToolCalling: true, CodeOptimized: true,
|
||||
MaxContext: 262144, MaxOutputTokens: 8192,
|
||||
},
|
||||
|
||||
// ── Qwen ────────────────────────────────
|
||||
"qwen-2.5-72b": {
|
||||
Streaming: true, ToolCalling: true,
|
||||
MaxContext: 131072, MaxOutputTokens: 8192,
|
||||
},
|
||||
"qwen-2.5-coder-32b": {
|
||||
Streaming: true, ToolCalling: true, CodeOptimized: true,
|
||||
MaxContext: 131072, MaxOutputTokens: 8192,
|
||||
},
|
||||
"qwq-32b": {
|
||||
Streaming: true, Reasoning: true,
|
||||
MaxContext: 131072, MaxOutputTokens: 8192,
|
||||
},
|
||||
}
|
||||
|
||||
// LookupKnownModel finds capabilities for a model by exact ID or prefix match.
|
||||
// Returns the caps and true if found, zero value and false if not.
|
||||
func LookupKnownModel(modelID string) (ModelCapabilities, bool) {
|
||||
id := strings.ToLower(modelID)
|
||||
|
||||
// Strip provider prefix (e.g. "anthropic/claude-sonnet-4-20250514" → "claude-sonnet-4-20250514")
|
||||
if idx := strings.Index(id, "/"); idx >= 0 {
|
||||
id = id[idx+1:]
|
||||
}
|
||||
|
||||
// Exact match first
|
||||
if caps, ok := knownModels[id]; ok {
|
||||
return caps, true
|
||||
}
|
||||
|
||||
// Prefix match: "claude-sonnet-4-20250514" matches "claude-sonnet-4"
|
||||
var bestKey string
|
||||
var bestCaps ModelCapabilities
|
||||
for key, caps := range knownModels {
|
||||
if strings.HasPrefix(id, key) && len(key) > len(bestKey) {
|
||||
bestKey = key
|
||||
bestCaps = caps
|
||||
}
|
||||
}
|
||||
if bestKey != "" {
|
||||
return bestCaps, true
|
||||
}
|
||||
|
||||
return ModelCapabilities{}, false
|
||||
}
|
||||
|
||||
// ── Heuristic Capability Detection ──────────
|
||||
// Ported from ai-editor's ProviderRegistry.
|
||||
// Used for Ollama, LM Studio, and other generic endpoints
|
||||
// that don't report capabilities in their /models response.
|
||||
|
||||
var (
|
||||
toolPatterns = []*regexp.Regexp{
|
||||
regexp.MustCompile(`(?i)llama[_-]?[34]`),
|
||||
regexp.MustCompile(`(?i)granite[_-]?[34]`),
|
||||
regexp.MustCompile(`(?i)hermes[_-]?[23]?`),
|
||||
regexp.MustCompile(`(?i)qwen[_-]?2`),
|
||||
regexp.MustCompile(`(?i)mistral|mixtral`),
|
||||
regexp.MustCompile(`(?i)command[_-]?r`),
|
||||
regexp.MustCompile(`(?i)phi[_-]?[34]`),
|
||||
regexp.MustCompile(`(?i)deepseek[_-]?v[23]`),
|
||||
regexp.MustCompile(`(?i)functionary`),
|
||||
regexp.MustCompile(`(?i)^gpt[_-]?[34]`),
|
||||
regexp.MustCompile(`(?i)^claude`),
|
||||
regexp.MustCompile(`(?i)gemma[_-]?2`),
|
||||
regexp.MustCompile(`(?i)glm[_-]?4`),
|
||||
regexp.MustCompile(`(?i)gemini`),
|
||||
}
|
||||
|
||||
visionPatterns = []*regexp.Regexp{
|
||||
regexp.MustCompile(`(?i)llava|bakllava`),
|
||||
regexp.MustCompile(`(?i)moondream`),
|
||||
regexp.MustCompile(`(?i)llama.*vision|vision.*llama`),
|
||||
regexp.MustCompile(`(?i)minicpm[_-]?v`),
|
||||
regexp.MustCompile(`(?i)^gpt[_-]?4[_-]?o|^gpt[_-]?4[_-]?turbo`),
|
||||
regexp.MustCompile(`(?i)^claude[_-]?3`),
|
||||
regexp.MustCompile(`(?i)gemini`),
|
||||
regexp.MustCompile(`(?i)qwen.*vl|vl.*qwen`),
|
||||
regexp.MustCompile(`(?i)phi[_-]?[34].*vision`),
|
||||
regexp.MustCompile(`(?i)internvl`),
|
||||
regexp.MustCompile(`(?i)glm[_-]?4v`),
|
||||
}
|
||||
|
||||
reasoningPatterns = []*regexp.Regexp{
|
||||
regexp.MustCompile(`(?i)deepseek[_-]?r1`),
|
||||
regexp.MustCompile(`(?i)qwq`),
|
||||
regexp.MustCompile(`(?i)^o[1234][_-]|^o[1234]$`),
|
||||
}
|
||||
|
||||
codePatterns = []*regexp.Regexp{
|
||||
regexp.MustCompile(`(?i)codellama|code[_-]?llama`),
|
||||
regexp.MustCompile(`(?i)deepseek[_-]?coder`),
|
||||
regexp.MustCompile(`(?i)starcoder`),
|
||||
regexp.MustCompile(`(?i)coder|codestral`),
|
||||
regexp.MustCompile(`(?i)wizardcoder`),
|
||||
regexp.MustCompile(`(?i)codegemma`),
|
||||
}
|
||||
)
|
||||
|
||||
func matchesAny(id string, patterns []*regexp.Regexp) bool {
|
||||
for _, p := range patterns {
|
||||
if p.MatchString(id) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// InferCapabilities guesses model capabilities from the model ID string.
|
||||
// This is the fallback when neither the known model table nor the provider
|
||||
// API provides capability data.
|
||||
func InferCapabilities(modelID string) ModelCapabilities {
|
||||
id := strings.ToLower(modelID)
|
||||
// Strip provider prefix
|
||||
if idx := strings.Index(id, "/"); idx >= 0 {
|
||||
id = id[idx+1:]
|
||||
}
|
||||
|
||||
return ModelCapabilities{
|
||||
Streaming: true, // virtually everything streams
|
||||
ToolCalling: matchesAny(id, toolPatterns),
|
||||
Vision: matchesAny(id, visionPatterns),
|
||||
Reasoning: matchesAny(id, reasoningPatterns),
|
||||
CodeOptimized: matchesAny(id, codePatterns),
|
||||
}
|
||||
}
|
||||
|
||||
// ResolveMaxOutput returns the max output tokens for a model.
|
||||
// Priority:
|
||||
// 1. Explicit value in caps (from DB / provider API)
|
||||
// 2. Known model table
|
||||
// 3. Derive from context window (context/8, clamped 2048..16384)
|
||||
// 4. 4096 as absolute last resort
|
||||
//
|
||||
// This is the ONE place the default lives. Nothing else in the
|
||||
// codebase should hardcode a max_tokens value.
|
||||
func ResolveMaxOutput(modelID string, caps ModelCapabilities) int {
|
||||
// 1. Already set (from model_configs DB or provider API)
|
||||
if caps.MaxOutputTokens > 0 {
|
||||
return caps.MaxOutputTokens
|
||||
}
|
||||
|
||||
// 2. Known model table
|
||||
if known, ok := LookupKnownModel(modelID); ok && known.MaxOutputTokens > 0 {
|
||||
return known.MaxOutputTokens
|
||||
}
|
||||
|
||||
// 3. Derive from context window
|
||||
if caps.MaxContext > 0 {
|
||||
derived := caps.MaxContext / 8
|
||||
if derived < 2048 {
|
||||
derived = 2048
|
||||
}
|
||||
if derived > 16384 {
|
||||
derived = 16384
|
||||
}
|
||||
return derived
|
||||
}
|
||||
|
||||
// 4. Last resort — the ONLY place 4096 appears as a default
|
||||
return 4096
|
||||
}
|
||||
|
||||
// MergeCapabilities takes authoritative caps (from provider API or DB) and fills
|
||||
// gaps from the known model table and heuristic detection. Provider-reported data
|
||||
// always wins; known table fills missing fields; heuristics are last resort.
|
||||
func MergeCapabilities(authoritative ModelCapabilities, modelID string) ModelCapabilities {
|
||||
merged := authoritative
|
||||
|
||||
// Fill gaps from known model table
|
||||
known, found := LookupKnownModel(modelID)
|
||||
if found {
|
||||
if !merged.ToolCalling && known.ToolCalling {
|
||||
merged.ToolCalling = true
|
||||
}
|
||||
if !merged.Vision && known.Vision {
|
||||
merged.Vision = true
|
||||
}
|
||||
if !merged.Thinking && known.Thinking {
|
||||
merged.Thinking = true
|
||||
}
|
||||
if !merged.Reasoning && known.Reasoning {
|
||||
merged.Reasoning = true
|
||||
}
|
||||
if !merged.CodeOptimized && known.CodeOptimized {
|
||||
merged.CodeOptimized = true
|
||||
}
|
||||
if !merged.WebSearch && known.WebSearch {
|
||||
merged.WebSearch = true
|
||||
}
|
||||
if merged.MaxContext == 0 && known.MaxContext > 0 {
|
||||
merged.MaxContext = known.MaxContext
|
||||
}
|
||||
if merged.MaxOutputTokens == 0 && known.MaxOutputTokens > 0 {
|
||||
merged.MaxOutputTokens = known.MaxOutputTokens
|
||||
}
|
||||
return merged
|
||||
}
|
||||
|
||||
// No known model — fill gaps from heuristics
|
||||
inferred := InferCapabilities(modelID)
|
||||
if !merged.ToolCalling && inferred.ToolCalling {
|
||||
merged.ToolCalling = true
|
||||
}
|
||||
if !merged.Vision && inferred.Vision {
|
||||
merged.Vision = true
|
||||
}
|
||||
if !merged.Thinking && inferred.Thinking {
|
||||
merged.Thinking = true
|
||||
}
|
||||
if !merged.Reasoning && inferred.Reasoning {
|
||||
merged.Reasoning = true
|
||||
}
|
||||
if !merged.CodeOptimized && inferred.CodeOptimized {
|
||||
merged.CodeOptimized = true
|
||||
}
|
||||
if merged.MaxContext == 0 && inferred.MaxContext > 0 {
|
||||
merged.MaxContext = inferred.MaxContext
|
||||
}
|
||||
if merged.MaxOutputTokens == 0 && inferred.MaxOutputTokens > 0 {
|
||||
merged.MaxOutputTokens = inferred.MaxOutputTokens
|
||||
}
|
||||
|
||||
return merged
|
||||
}
|
||||
|
||||
// HasProviderData returns true if this capability set contains any data that
|
||||
// was likely reported by a provider (not just zero values).
|
||||
func (c ModelCapabilities) HasProviderData() bool {
|
||||
return c.ToolCalling || c.Vision || c.Thinking || c.Reasoning ||
|
||||
c.CodeOptimized || c.WebSearch || c.MaxContext > 0 || c.MaxOutputTokens > 0
|
||||
}
|
||||
Reference in New Issue
Block a user