mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-02-18 20:30:51 +08:00
Implements unified model routing
Migrates the AMP module to a new unified routing system, replacing the fallback handler with a router-based approach. This change introduces a `ModelRoutingWrapper` that handles model extraction, routing decisions, and proxying based on provider availability and model mappings. It provides a more flexible and maintainable routing mechanism by centralizing routing logic. The changes include: - Introducing new `routing` package with core routing logic. - Creating characterization tests to capture existing behavior. - Implementing model extraction and rewriting. - Updating AMP module routes to utilize the new routing wrapper. - Deprecating `FallbackHandler` in favor of the new `ModelRoutingWrapper`.
This commit is contained in:
59
internal/routing/extractor.go
Normal file
59
internal/routing/extractor.go
Normal file
@@ -0,0 +1,59 @@
|
||||
package routing
|
||||
|
||||
import (
|
||||
"strings"
|
||||
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
// ModelExtractor extracts model names from request data.
|
||||
type ModelExtractor interface {
|
||||
// Extract returns the model name from the request body and gin parameters.
|
||||
// The ginParams map contains route parameters like "action" and "path".
|
||||
Extract(body []byte, ginParams map[string]string) (string, error)
|
||||
}
|
||||
|
||||
// DefaultModelExtractor is the standard implementation of ModelExtractor.
|
||||
type DefaultModelExtractor struct{}
|
||||
|
||||
// NewModelExtractor creates a new DefaultModelExtractor.
|
||||
func NewModelExtractor() *DefaultModelExtractor {
|
||||
return &DefaultModelExtractor{}
|
||||
}
|
||||
|
||||
// Extract extracts the model name from the request.
|
||||
// It checks in order:
|
||||
// 1. JSON body "model" field (OpenAI, Claude format)
|
||||
// 2. "action" parameter for Gemini standard format (e.g., "gemini-pro:generateContent")
|
||||
// 3. "path" parameter for AMP CLI Gemini format (e.g., "/publishers/google/models/gemini-3-pro:streamGenerateContent")
|
||||
func (e *DefaultModelExtractor) Extract(body []byte, ginParams map[string]string) (string, error) {
|
||||
// First try to parse from JSON body (OpenAI, Claude, etc.)
|
||||
if result := gjson.GetBytes(body, "model"); result.Exists() && result.Type == gjson.String {
|
||||
return result.String(), nil
|
||||
}
|
||||
|
||||
// For Gemini requests, model is in the URL path
|
||||
// Standard format: /models/{model}:generateContent -> :action parameter
|
||||
if action, ok := ginParams["action"]; ok && action != "" {
|
||||
// Split by colon to get model name (e.g., "gemini-pro:generateContent" -> "gemini-pro")
|
||||
parts := strings.Split(action, ":")
|
||||
if len(parts) > 0 && parts[0] != "" {
|
||||
return parts[0], nil
|
||||
}
|
||||
}
|
||||
|
||||
// AMP CLI format: /publishers/google/models/{model}:method -> *path parameter
|
||||
// Example: /publishers/google/models/gemini-3-pro-preview:streamGenerateContent
|
||||
if path, ok := ginParams["path"]; ok && path != "" {
|
||||
// Look for /models/{model}:method pattern
|
||||
if idx := strings.Index(path, "/models/"); idx >= 0 {
|
||||
modelPart := path[idx+8:] // Skip "/models/"
|
||||
// Split by colon to get model name
|
||||
if colonIdx := strings.Index(modelPart, ":"); colonIdx > 0 {
|
||||
return modelPart[:colonIdx], nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return "", nil
|
||||
}
|
||||
Reference in New Issue
Block a user