forked from niuniu/llm-intelligence
chore: prepare repository for publishing
This commit is contained in:
phamnazage-jpg
129
internal/collectors/collector.go
Normal file
129
internal/collectors/collector.go
Normal file
@@ -0,0 +1,129 @@
|
||||
// internal/collectors/collector.go
|
||||
// Collector 接口定义：所有数据源采集器的统一抽象
|
||||
package collectors
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Result 采集结果
|
||||
type Result struct {
|
||||
Models []ModelInfo
|
||||
Meta CollectionMeta
|
||||
}
|
||||
|
||||
// CollectionMeta 采集元信息
|
||||
type CollectionMeta struct {
|
||||
Source string
|
||||
Count int
|
||||
Duration time.Duration
|
||||
Timestamp time.Time
|
||||
BatchID string
|
||||
CollectorVersion string
|
||||
}
|
||||
|
||||
// ModelInfo 标准模型信息（与 fetch_openrouter.go 兼容）
|
||||
type ModelInfo struct {
|
||||
ID string
|
||||
Name string
|
||||
Provider string
|
||||
ProviderID string
|
||||
Version string
|
||||
Modality string
|
||||
ContextLength int
|
||||
Capabilities []string
|
||||
Pricing ModelPricing
|
||||
Description string
|
||||
IsFree bool
|
||||
SourceURL string
|
||||
}
|
||||
|
||||
// ModelPricing 标准定价信息
|
||||
type ModelPricing struct {
|
||||
Input float64
|
||||
Output float64
|
||||
}
|
||||
|
||||
// Collector 采集器接口
|
||||
type Collector interface {
|
||||
// Name 返回采集器名称
|
||||
Name() string
|
||||
|
||||
// Collect 执行采集，返回标准模型列表
|
||||
Collect(ctx context.Context) (Result, error)
|
||||
|
||||
// Schedule 返回推荐调度周期（如 "0 8 * * *"）
|
||||
Schedule() string
|
||||
|
||||
// Timeout 返回单次采集超时时间
|
||||
Timeout() time.Duration
|
||||
|
||||
// RetryCount 返回最大重试次数
|
||||
RetryCount() int
|
||||
}
|
||||
|
||||
// BaseCollector 提供默认实现的嵌入类型
|
||||
type BaseCollector struct {
|
||||
name string
|
||||
schedule string
|
||||
timeout time.Duration
|
||||
retryCount int
|
||||
version string
|
||||
}
|
||||
|
||||
func (b *BaseCollector) Name() string { return b.name }
|
||||
func (b *BaseCollector) Schedule() string { return b.schedule }
|
||||
func (b *BaseCollector) Timeout() time.Duration { return b.timeout }
|
||||
func (b *BaseCollector) RetryCount() int { return b.retryCount }
|
||||
func (b *BaseCollector) Version() string { return b.version }
|
||||
|
||||
// NewBaseCollector 创建基础采集器配置
|
||||
func NewBaseCollector(name, schedule string, timeout time.Duration, retry int, version string) BaseCollector {
|
||||
return BaseCollector{
|
||||
name: name,
|
||||
schedule: schedule,
|
||||
timeout: timeout,
|
||||
retryCount: retry,
|
||||
version: version,
|
||||
}
|
||||
}
|
||||
|
||||
// CollectorRegistry 采集器注册表
|
||||
type CollectorRegistry struct {
|
||||
collectors map[string]Collector
|
||||
}
|
||||
|
||||
// NewRegistry 创建采集器注册表
|
||||
func NewRegistry() *CollectorRegistry {
|
||||
return &CollectorRegistry{collectors: make(map[string]Collector)}
|
||||
}
|
||||
|
||||
// Register 注册采集器
|
||||
func (r *CollectorRegistry) Register(c Collector) {
|
||||
r.collectors[c.Name()] = c
|
||||
}
|
||||
|
||||
// Get 获取采集器
|
||||
func (r *CollectorRegistry) Get(name string) (Collector, bool) {
|
||||
c, ok := r.collectors[name]
|
||||
return c, ok
|
||||
}
|
||||
|
||||
// All 返回所有已注册采集器
|
||||
func (r *CollectorRegistry) All() []Collector {
|
||||
cs := make([]Collector, 0, len(r.collectors))
|
||||
for _, c := range r.collectors {
|
||||
cs = append(cs, c)
|
||||
}
|
||||
return cs
|
||||
}
|
||||
|
||||
// Names 返回所有已注册采集器名称
|
||||
func (r *CollectorRegistry) Names() []string {
|
||||
names := make([]string, 0, len(r.collectors))
|
||||
for n := range r.collectors {
|
||||
names = append(names, n)
|
||||
}
|
||||
return names
|
||||
}
|
||||
127
internal/collectors/collector_test.go
Normal file
127
internal/collectors/collector_test.go
Normal file
@@ -0,0 +1,127 @@
|
||||
// internal/collectors/collector_test.go
|
||||
package collectors
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// mockCollector 用于测试的模拟采集器
|
||||
type mockCollector struct {
|
||||
BaseCollector
|
||||
collectFunc func(ctx context.Context) (Result, error)
|
||||
}
|
||||
|
||||
func (m *mockCollector) Collect(ctx context.Context) (Result, error) {
|
||||
return m.collectFunc(ctx)
|
||||
}
|
||||
|
||||
func TestCollectorInterface(t *testing.T) {
|
||||
c := &mockCollector{
|
||||
BaseCollector: NewBaseCollector("test", "0 8 * * *", 30*time.Second, 3, "v1.0"),
|
||||
collectFunc: func(ctx context.Context) (Result, error) {
|
||||
return Result{
|
||||
Models: []ModelInfo{{ID: "test/model-1", Name: "Test Model"}},
|
||||
Meta: CollectionMeta{Source: "test", Count: 1},
|
||||
}, nil
|
||||
},
|
||||
}
|
||||
|
||||
// 测试接口方法
|
||||
if c.Name() != "test" {
|
||||
t.Errorf("Name() = %q, want %q", c.Name(), "test")
|
||||
}
|
||||
if c.Schedule() != "0 8 * * *" {
|
||||
t.Errorf("Schedule() = %q, want %q", c.Schedule(), "0 8 * * *")
|
||||
}
|
||||
if c.Timeout() != 30*time.Second {
|
||||
t.Errorf("Timeout() = %v, want %v", c.Timeout(), 30*time.Second)
|
||||
}
|
||||
if c.RetryCount() != 3 {
|
||||
t.Errorf("RetryCount() = %d, want %d", c.RetryCount(), 3)
|
||||
}
|
||||
|
||||
// 测试 Collect
|
||||
ctx := context.Background()
|
||||
result, err := c.Collect(ctx)
|
||||
if err != nil {
|
||||
t.Fatalf("Collect() error = %v", err)
|
||||
}
|
||||
if len(result.Models) != 1 {
|
||||
t.Errorf("len(Models) = %d, want 1", len(result.Models))
|
||||
}
|
||||
if result.Meta.Count != 1 {
|
||||
t.Errorf("Meta.Count = %d, want 1", result.Meta.Count)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCollectorRegistry(t *testing.T) {
|
||||
reg := NewRegistry()
|
||||
|
||||
c1 := &mockCollector{
|
||||
BaseCollector: NewBaseCollector("openrouter", "0 8 * * *", 30*time.Second, 3, "v1.0"),
|
||||
collectFunc: func(ctx context.Context) (Result, error) { return Result{}, nil },
|
||||
}
|
||||
c2 := &mockCollector{
|
||||
BaseCollector: NewBaseCollector("siliconflow", "0 9 * * *", 30*time.Second, 3, "v1.0"),
|
||||
collectFunc: func(ctx context.Context) (Result, error) { return Result{}, nil },
|
||||
}
|
||||
|
||||
reg.Register(c1)
|
||||
reg.Register(c2)
|
||||
|
||||
// 测试 Get
|
||||
got, ok := reg.Get("openrouter")
|
||||
if !ok {
|
||||
t.Fatal("Get(openrouter) not found")
|
||||
}
|
||||
if got.Name() != "openrouter" {
|
||||
t.Errorf("Get() Name = %q, want %q", got.Name(), "openrouter")
|
||||
}
|
||||
|
||||
// 测试 Names
|
||||
names := reg.Names()
|
||||
if len(names) != 2 {
|
||||
t.Errorf("Names() len = %d, want 2", len(names))
|
||||
}
|
||||
|
||||
// 测试 All
|
||||
all := reg.All()
|
||||
if len(all) != 2 {
|
||||
t.Errorf("All() len = %d, want 2", len(all))
|
||||
}
|
||||
|
||||
// 测试不存在的采集器
|
||||
_, ok = reg.Get("nonexistent")
|
||||
if ok {
|
||||
t.Error("Get(nonexistent) should return false")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCollectorTimeout(t *testing.T) {
|
||||
c := &mockCollector{
|
||||
BaseCollector: NewBaseCollector("slow", "0 8 * * *", 100*time.Millisecond, 0, "v1.0"),
|
||||
collectFunc: func(ctx context.Context) (Result, error) {
|
||||
// 模拟耗时操作
|
||||
select {
|
||||
case <-time.After(200 * time.Millisecond):
|
||||
return Result{}, nil
|
||||
case <-ctx.Done():
|
||||
return Result{}, ctx.Err()
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
|
||||
defer cancel()
|
||||
|
||||
_, err := c.Collect(ctx)
|
||||
if err == nil {
|
||||
t.Error("Expected timeout error, got nil")
|
||||
}
|
||||
if !errors.Is(err, context.DeadlineExceeded) {
|
||||
t.Errorf("Expected DeadlineExceeded, got %v", err)
|
||||
}
|
||||
}
|
||||
115
internal/collectors/provider_mapper.go
Normal file
115
internal/collectors/provider_mapper.go
Normal file
@@ -0,0 +1,115 @@
|
||||
// internal/collectors/provider_mapper.go
|
||||
// ProviderMapper: 将 OpenRouter 模型 ID 映射为标准厂商/模型名称
|
||||
package collectors
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// ProviderInfo 标准厂商信息
|
||||
type ProviderInfo struct {
|
||||
ID string // 标准ID: "openai", "anthropic", "deepseek"...
|
||||
Name string // 英文名
|
||||
NameCN string // 中文名
|
||||
Country string // "US" / "CN" / "EU"
|
||||
}
|
||||
|
||||
// ModelMapping 模型映射结果
|
||||
type ModelMapping struct {
|
||||
Provider ProviderInfo
|
||||
ModelName string // 纯模型名，不含厂商前缀
|
||||
RawID string // 原始 OpenRouter ID
|
||||
IsFree bool // 是否免费版（:free 后缀）
|
||||
}
|
||||
|
||||
// providerNameMap 标准厂商名称映射表
|
||||
// key 为标准ID（也兼容 OpenRouter 原始格式作为别名）
|
||||
var providerNameMap = map[string]ProviderInfo{
|
||||
"openai": {ID: "openai", Name: "OpenAI", NameCN: "OpenAI", Country: "US"},
|
||||
"anthropic": {ID: "anthropic", Name: "Anthropic", NameCN: "Anthropic", Country: "US"},
|
||||
"google": {ID: "google", Name: "Google", NameCN: "谷歌", Country: "US"},
|
||||
"meta": {ID: "meta", Name: "Meta", NameCN: "Meta", Country: "US"},
|
||||
"xai": {ID: "xai", Name: "xAI", NameCN: "xAI", Country: "US"},
|
||||
"x-ai": {ID: "xai", Name: "xAI", NameCN: "xAI", Country: "US"}, // OpenRouter别名
|
||||
"deepseek": {ID: "deepseek", Name: "DeepSeek", NameCN: "深度求索", Country: "CN"},
|
||||
"qwen": {ID: "qwen", Name: "Qwen", NameCN: "通义千问", Country: "CN"},
|
||||
"alibaba": {ID: "alibaba", Name: "Alibaba", NameCN: "阿里巴巴", Country: "CN"},
|
||||
"moonshot": {ID: "moonshot", Name: "Moonshot AI", NameCN: "月之暗面", Country: "CN"},
|
||||
"moonshotai": {ID: "moonshot", Name: "Moonshot AI", NameCN: "月之暗面", Country: "CN"}, // OpenRouter别名
|
||||
"zhipu": {ID: "zhipu", Name: "Zhipu AI", NameCN: "智谱AI", Country: "CN"},
|
||||
"zhipuai": {ID: "zhipu", Name: "Zhipu AI", NameCN: "智谱AI", Country: "CN"}, // OpenRouter别名
|
||||
"bytedance": {ID: "bytedance", Name: "ByteDance", NameCN: "字节跳动", Country: "CN"},
|
||||
"baidu": {ID: "baidu", Name: "Baidu", NameCN: "百度", Country: "CN"},
|
||||
"tencent": {ID: "tencent", Name: "Tencent", NameCN: "腾讯", Country: "CN"},
|
||||
"mistral": {ID: "mistral", Name: "Mistral AI", NameCN: "Mistral", Country: "EU"},
|
||||
"cohere": {ID: "cohere", Name: "Cohere", NameCN: "Cohere", Country: "US"},
|
||||
"ai21": {ID: "ai21", Name: "AI21 Labs", NameCN: "AI21", Country: "US"},
|
||||
"perplexity": {ID: "perplexity", Name: "Perplexity", NameCN: "Perplexity", Country: "US"},
|
||||
"nvidia": {ID: "nvidia", Name: "NVIDIA", NameCN: "英伟达", Country: "US"},
|
||||
"microsoft": {ID: "microsoft", Name: "Microsoft", NameCN: "微软", Country: "US"},
|
||||
"openrouter": {ID: "openrouter", Name: "OpenRouter", NameCN: "OpenRouter", Country: "US"},
|
||||
}
|
||||
|
||||
// MapOpenRouterID 将 OpenRouter 模型 ID 映射为标准信息
|
||||
// OpenRouter ID 格式: "provider/model-name" 或 "provider/model-name:free"
|
||||
func MapOpenRouterID(rawID string) (ModelMapping, error) {
|
||||
if rawID == "" {
|
||||
return ModelMapping{}, fmt.Errorf("empty model ID")
|
||||
}
|
||||
|
||||
// 检测 :free 后缀
|
||||
isFree := false
|
||||
modelPart := rawID
|
||||
if strings.HasSuffix(rawID, ":free") {
|
||||
isFree = true
|
||||
modelPart = rawID[:len(rawID)-5]
|
||||
}
|
||||
|
||||
// 分割 provider / model
|
||||
parts := strings.SplitN(modelPart, "/", 2)
|
||||
if len(parts) < 2 {
|
||||
return ModelMapping{}, fmt.Errorf("invalid model ID format: %s", rawID)
|
||||
}
|
||||
|
||||
providerKey := strings.ToLower(parts[0])
|
||||
modelName := parts[1]
|
||||
|
||||
// 查找厂商信息
|
||||
provider, ok := providerNameMap[providerKey]
|
||||
if !ok {
|
||||
// 未识别厂商，返回通用信息
|
||||
provider = ProviderInfo{
|
||||
ID: providerKey,
|
||||
Name: providerKey,
|
||||
NameCN: providerKey,
|
||||
Country: "unknown",
|
||||
}
|
||||
}
|
||||
|
||||
return ModelMapping{
|
||||
Provider: provider,
|
||||
ModelName: modelName,
|
||||
RawID: rawID,
|
||||
IsFree: isFree,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// GetAllProviderNames 返回所有已注册的厂商ID列表（用于测试覆盖度检查）
|
||||
func GetAllProviderNames() []string {
|
||||
names := make([]string, 0, len(providerNameMap))
|
||||
for k := range providerNameMap {
|
||||
names = append(names, k)
|
||||
}
|
||||
return names
|
||||
}
|
||||
|
||||
// RegisterProvider 动态注册新厂商（用于扩展）
|
||||
func RegisterProvider(key string, info ProviderInfo) {
|
||||
providerNameMap[strings.ToLower(key)] = info
|
||||
}
|
||||
|
||||
// ProviderCount 返回已注册厂商数量
|
||||
func ProviderCount() int {
|
||||
return len(providerNameMap)
|
||||
}
|
||||
167
internal/collectors/provider_mapper_test.go
Normal file
167
internal/collectors/provider_mapper_test.go
Normal file
@@ -0,0 +1,167 @@
|
||||
// internal/collectors/provider_mapper_test.go
|
||||
package collectors
|
||||
|
||||
import (
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestMapOpenRouterID(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
rawID string
|
||||
wantErr bool
|
||||
wantProvID string
|
||||
wantProvCN string
|
||||
wantModel string
|
||||
wantFree bool
|
||||
wantCountry string
|
||||
}{
|
||||
{
|
||||
name: "OpenAI GPT-4o",
|
||||
rawID: "openai/gpt-4o",
|
||||
wantProvID: "openai",
|
||||
wantProvCN: "OpenAI",
|
||||
wantModel: "gpt-4o",
|
||||
wantFree: false,
|
||||
wantCountry: "US",
|
||||
},
|
||||
{
|
||||
name: "Anthropic Claude free",
|
||||
rawID: "anthropic/claude-3.5-sonnet:free",
|
||||
wantProvID: "anthropic",
|
||||
wantProvCN: "Anthropic",
|
||||
wantModel: "claude-3.5-sonnet",
|
||||
wantFree: true,
|
||||
wantCountry: "US",
|
||||
},
|
||||
{
|
||||
name: "DeepSeek V3",
|
||||
rawID: "deepseek/deepseek-v3",
|
||||
wantProvID: "deepseek",
|
||||
wantProvCN: "深度求索",
|
||||
wantModel: "deepseek-v3",
|
||||
wantFree: false,
|
||||
wantCountry: "CN",
|
||||
},
|
||||
{
|
||||
name: "Moonshot Kimi",
|
||||
rawID: "moonshotai/kimi-k2",
|
||||
wantProvID: "moonshot",
|
||||
wantProvCN: "月之暗面",
|
||||
wantModel: "kimi-k2",
|
||||
wantFree: false,
|
||||
wantCountry: "CN",
|
||||
},
|
||||
{
|
||||
name: "Unknown provider fallback",
|
||||
rawID: "some-new-ai/model-x",
|
||||
wantProvID: "some-new-ai",
|
||||
wantProvCN: "some-new-ai",
|
||||
wantModel: "model-x",
|
||||
wantFree: false,
|
||||
wantCountry: "unknown",
|
||||
},
|
||||
{
|
||||
name: "Empty ID",
|
||||
rawID: "",
|
||||
wantErr: true,
|
||||
},
|
||||
{
|
||||
name: "Invalid format no slash",
|
||||
rawID: "invalid-id",
|
||||
wantErr: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got, err := MapOpenRouterID(tt.rawID)
|
||||
if (err != nil) != tt.wantErr {
|
||||
t.Errorf("MapOpenRouterID(%q) error = %v, wantErr %v", tt.rawID, err, tt.wantErr)
|
||||
return
|
||||
}
|
||||
if tt.wantErr {
|
||||
return
|
||||
}
|
||||
if got.Provider.ID != tt.wantProvID {
|
||||
t.Errorf("Provider.ID = %q, want %q", got.Provider.ID, tt.wantProvID)
|
||||
}
|
||||
if got.Provider.NameCN != tt.wantProvCN {
|
||||
t.Errorf("Provider.NameCN = %q, want %q", got.Provider.NameCN, tt.wantProvCN)
|
||||
}
|
||||
if got.ModelName != tt.wantModel {
|
||||
t.Errorf("ModelName = %q, want %q", got.ModelName, tt.wantModel)
|
||||
}
|
||||
if got.IsFree != tt.wantFree {
|
||||
t.Errorf("IsFree = %v, want %v", got.IsFree, tt.wantFree)
|
||||
}
|
||||
if got.Provider.Country != tt.wantCountry {
|
||||
t.Errorf("Country = %q, want %q", got.Provider.Country, tt.wantCountry)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestProviderMapCompleteness(t *testing.T) {
|
||||
// 验证所有预定义的厂商映射
|
||||
requiredProviders := []string{
|
||||
"openai", "anthropic", "google", "meta", "xai",
|
||||
"deepseek", "qwen", "moonshot", "zhipu", "bytedance",
|
||||
"baidu", "tencent", "alibaba", "mistral", "cohere",
|
||||
"ai21", "perplexity", "nvidia", "microsoft", "openrouter",
|
||||
}
|
||||
|
||||
for _, id := range requiredProviders {
|
||||
_, ok := providerNameMap[id]
|
||||
if !ok {
|
||||
t.Errorf("Required provider %q not found in providerNameMap", id)
|
||||
}
|
||||
}
|
||||
|
||||
// 验证总数 >= 20
|
||||
if ProviderCount() < 20 {
|
||||
t.Errorf("ProviderCount() = %d, want >= 20", ProviderCount())
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegisterProvider(t *testing.T) {
|
||||
// 注册新厂商
|
||||
RegisterProvider("test-corp", ProviderInfo{
|
||||
ID: "test-corp",
|
||||
Name: "Test Corp",
|
||||
NameCN: "测试公司",
|
||||
Country: "CN",
|
||||
})
|
||||
|
||||
got, err := MapOpenRouterID("test-corp/model-1")
|
||||
if err != nil {
|
||||
t.Fatalf("MapOpenRouterID after RegisterProvider failed: %v", err)
|
||||
}
|
||||
if got.Provider.NameCN != "测试公司" {
|
||||
t.Errorf("After RegisterProvider, NameCN = %q, want %q", got.Provider.NameCN, "测试公司")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetAllProviderNames(t *testing.T) {
|
||||
names := GetAllProviderNames()
|
||||
if len(names) == 0 {
|
||||
t.Error("GetAllProviderNames() returned empty slice")
|
||||
}
|
||||
// 验证包含 openai
|
||||
found := false
|
||||
for _, n := range names {
|
||||
if n == "openai" {
|
||||
found = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Error("GetAllProviderNames() missing 'openai'")
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkMapOpenRouterID(b *testing.B) {
|
||||
for i := 0; i < b.N; i++ {
|
||||
_, _ = MapOpenRouterID("openai/gpt-4o")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user