rework llm

1 files changed, 83 insertions, 191 deletions

diff --git a/internal/llm/models/models.go b/internal/llm/models/models.go
index 1c1b75ad0..2f75db9c8 100644
--- a/internal/llm/models/models.go
+++ b/internal/llm/models/models.go
@@ -1,230 +1,122 @@
 package models
 
-import (
-	"context"
-	"errors"
-	"log"
-
-	"github.com/cloudwego/eino-ext/components/model/claude"
-	"github.com/cloudwego/eino-ext/components/model/openai"
-	"github.com/cloudwego/eino/components/model"
-	"github.com/spf13/viper"
-)
-
 type (
 	ModelID       string
 	ModelProvider string
 )
 
 type Model struct {
-	ID           ModelID       `json:"id"`
-	Name         string        `json:"name"`
-	Provider     ModelProvider `json:"provider"`
-	APIModel     string        `json:"api_model"`
-	CostPer1MIn  float64       `json:"cost_per_1m_in"`
-	CostPer1MOut float64       `json:"cost_per_1m_out"`
+	ID                 ModelID       `json:"id"`
+	Name               string        `json:"name"`
+	Provider           ModelProvider `json:"provider"`
+	APIModel           string        `json:"api_model"`
+	CostPer1MIn        float64       `json:"cost_per_1m_in"`
+	CostPer1MOut       float64       `json:"cost_per_1m_out"`
+	CostPer1MInCached  float64       `json:"cost_per_1m_in_cached"`
+	CostPer1MOutCached float64       `json:"cost_per_1m_out_cached"`
 }
 
-const (
-	DefaultBigModel    = Claude37Sonnet
-	DefaultLittleModel = Claude37Sonnet
-)
-
 // Model IDs
 const (
-	// OpenAI
-	GPT4o     ModelID = "gpt-4o"
-	GPT4oMini ModelID = "gpt-4o-mini"
-	GPT45     ModelID = "gpt-4.5"
-	O1        ModelID = "o1"
-	O1Mini    ModelID = "o1-mini"
 	// Anthropic
 	Claude35Sonnet ModelID = "claude-3.5-sonnet"
 	Claude3Haiku   ModelID = "claude-3-haiku"
 	Claude37Sonnet ModelID = "claude-3.7-sonnet"
-	// Google
-	Gemini20Pro   ModelID = "gemini-2.0-pro"
-	Gemini15Flash ModelID = "gemini-1.5-flash"
-	Gemini20Flash ModelID = "gemini-2.0-flash"
-	// xAI
-	Grok3     ModelID = "grok-3"
-	Grok2Mini ModelID = "grok-2-mini"
-	// DeepSeek
-	DeepSeekR1    ModelID = "deepseek-r1"
-	DeepSeekCoder ModelID = "deepseek-coder"
-	// Meta
-	Llama3    ModelID = "llama-3"
-	Llama270B ModelID = "llama-2-70b"
+	// OpenAI
+	GPT4o ModelID = "gpt-4o"
+
+	// GEMINI
+	GEMINI25      ModelID = "gemini-2.5"
+	GRMINI20Flash ModelID = "gemini-2.0-flash"
+
 	// GROQ
-	GroqLlama3SpecDec ModelID = "groq-llama-3-spec-dec"
-	GroqQwen32BCoder  ModelID = "qwen-2.5-coder-32b"
+	QWENQwq ModelID = "qwen-qwq"
 )
 
 const (
 	ProviderOpenAI    ModelProvider = "openai"
 	ProviderAnthropic ModelProvider = "anthropic"
-	ProviderGoogle    ModelProvider = "google"
-	ProviderXAI       ModelProvider = "xai"
-	ProviderDeepSeek  ModelProvider = "deepseek"
-	ProviderMeta      ModelProvider = "meta"
-	ProviderGroq      ModelProvider = "groq"
+	ProviderGemini    ModelProvider = "gemini"
+	ProviderGROQ      ModelProvider = "groq"
 )
 
 var SupportedModels = map[ModelID]Model{
-	// OpenAI
-	GPT4o: {
-		ID:       GPT4o,
-		Name:     "GPT-4o",
-		Provider: ProviderOpenAI,
-		APIModel: "gpt-4o",
-	},
-	GPT4oMini: {
-		ID:           GPT4oMini,
-		Name:         "GPT-4o Mini",
-		Provider:     ProviderOpenAI,
-		APIModel:     "gpt-4o-mini",
-		CostPer1MIn:  0.150,
-		CostPer1MOut: 0.600,
-	},
-	GPT45: {
-		ID:       GPT45,
-		Name:     "GPT-4.5",
-		Provider: ProviderOpenAI,
-		APIModel: "gpt-4.5",
-	},
-	O1: {
-		ID:       O1,
-		Name:     "o1",
-		Provider: ProviderOpenAI,
-		APIModel: "o1",
-	},
-	O1Mini: {
-		ID:       O1Mini,
-		Name:     "o1 Mini",
-		Provider: ProviderOpenAI,
-		APIModel: "o1-mini",
-	},
 	// Anthropic
 	Claude35Sonnet: {
-		ID:       Claude35Sonnet,
-		Name:     "Claude 3.5 Sonnet",
-		Provider: ProviderAnthropic,
-		APIModel: "claude-3.5-sonnet",
+		ID:                 Claude35Sonnet,
+		Name:               "Claude 3.5 Sonnet",
+		Provider:           ProviderAnthropic,
+		APIModel:           "claude-3-5-sonnet-latest",
+		CostPer1MIn:        3.0,
+		CostPer1MInCached:  3.75,
+		CostPer1MOutCached: 0.30,
+		CostPer1MOut:       15.0,
 	},
 	Claude3Haiku: {
-		ID:       Claude3Haiku,
-		Name:     "Claude 3 Haiku",
-		Provider: ProviderAnthropic,
-		APIModel: "claude-3-haiku",
+		ID:                 Claude3Haiku,
+		Name:               "Claude 3 Haiku",
+		Provider:           ProviderAnthropic,
+		APIModel:           "claude-3-haiku-latest",
+		CostPer1MIn:        0.80,
+		CostPer1MInCached:  1,
+		CostPer1MOutCached: 0.08,
+		CostPer1MOut:       4,
 	},
 	Claude37Sonnet: {
-		ID:           Claude37Sonnet,
-		Name:         "Claude 3.7 Sonnet",
-		Provider:     ProviderAnthropic,
-		APIModel:     "claude-3-7-sonnet-20250219",
-		CostPer1MIn:  3.0,
-		CostPer1MOut: 15.0,
+		ID:                 Claude37Sonnet,
+		Name:               "Claude 3.7 Sonnet",
+		Provider:           ProviderAnthropic,
+		APIModel:           "claude-3-7-sonnet-latest",
+		CostPer1MIn:        3.0,
+		CostPer1MInCached:  3.75,
+		CostPer1MOutCached: 0.30,
+		CostPer1MOut:       15.0,
 	},
-	// Google
-	Gemini20Pro: {
-		ID:       Gemini20Pro,
-		Name:     "Gemini 2.0 Pro",
-		Provider: ProviderGoogle,
-		APIModel: "gemini-2.0-pro",
-	},
-	Gemini15Flash: {
-		ID:       Gemini15Flash,
-		Name:     "Gemini 1.5 Flash",
-		Provider: ProviderGoogle,
-		APIModel: "gemini-1.5-flash",
-	},
-	Gemini20Flash: {
-		ID:       Gemini20Flash,
-		Name:     "Gemini 2.0 Flash",
-		Provider: ProviderGoogle,
-		APIModel: "gemini-2.0-flash",
-	},
-	// xAI
-	Grok3: {
-		ID:       Grok3,
-		Name:     "Grok 3",
-		Provider: ProviderXAI,
-		APIModel: "grok-3",
-	},
-	Grok2Mini: {
-		ID:       Grok2Mini,
-		Name:     "Grok 2 Mini",
-		Provider: ProviderXAI,
-		APIModel: "grok-2-mini",
-	},
-	// DeepSeek
-	DeepSeekR1: {
-		ID:       DeepSeekR1,
-		Name:     "DeepSeek R1",
-		Provider: ProviderDeepSeek,
-		APIModel: "deepseek-r1",
-	},
-	DeepSeekCoder: {
-		ID:       DeepSeekCoder,
-		Name:     "DeepSeek Coder",
-		Provider: ProviderDeepSeek,
-		APIModel: "deepseek-coder",
+
+	// OpenAI
+	GPT4o: {
+		ID:                 GPT4o,
+		Name:               "GPT-4o",
+		Provider:           ProviderOpenAI,
+		APIModel:           "gpt-4o",
+		CostPer1MIn:        2.50,
+		CostPer1MInCached:  1.25,
+		CostPer1MOutCached: 0,
+		CostPer1MOut:       10.00,
 	},
-	// Meta
-	Llama3: {
-		ID:       Llama3,
-		Name:     "LLaMA 3",
-		Provider: ProviderMeta,
-		APIModel: "llama-3",
+
+	// GEMINI
+	GEMINI25: {
+		ID:                 GEMINI25,
+		Name:               "Gemini 2.5 Pro",
+		Provider:           ProviderGemini,
+		APIModel:           "gemini-2.5-pro-exp-03-25",
+		CostPer1MIn:        0,
+		CostPer1MInCached:  0,
+		CostPer1MOutCached: 0,
+		CostPer1MOut:       0,
 	},
-	Llama270B: {
-		ID:       Llama270B,
-		Name:     "LLaMA 2 70B",
-		Provider: ProviderMeta,
-		APIModel: "llama-2-70b",
+
+	GRMINI20Flash: {
+		ID:                 GRMINI20Flash,
+		Name:               "Gemini 2.0 Flash",
+		Provider:           ProviderGemini,
+		APIModel:           "gemini-2.0-flash",
+		CostPer1MIn:        0.1,
+		CostPer1MInCached:  0,
+		CostPer1MOutCached: 0.025,
+		CostPer1MOut:       0.4,
 	},
 
 	// GROQ
-	GroqLlama3SpecDec: {
-		ID:       GroqLlama3SpecDec,
-		Name:     "GROQ LLaMA 3 SpecDec",
-		Provider: ProviderGroq,
-		APIModel: "llama-3.3-70b-specdec",
-	},
-	GroqQwen32BCoder: {
-		ID:       GroqQwen32BCoder,
-		Name:     "GROQ Qwen 2.5 Coder 32B",
-		Provider: ProviderGroq,
-		APIModel: "qwen-2.5-coder-32b",
+	QWENQwq: {
+		ID:                 QWENQwq,
+		Name:               "Qwen Qwq",
+		Provider:           ProviderGROQ,
+		APIModel:           "qwen-qwq-32b",
+		CostPer1MIn:        0,
+		CostPer1MInCached:  0,
+		CostPer1MOutCached: 0,
+		CostPer1MOut:       0,
 	},
 }
-
-func GetModel(ctx context.Context, model ModelID) (model.ChatModel, error) {
-	provider := SupportedModels[model].Provider
-	log.Printf("Provider: %s", provider)
-	maxTokens := viper.GetInt("providers.common.max_tokens")
-	switch provider {
-	case ProviderOpenAI:
-		return openai.NewChatModel(ctx, &openai.ChatModelConfig{
-			APIKey:    viper.GetString("providers.openai.key"),
-			Model:     string(SupportedModels[model].APIModel),
-			MaxTokens: &maxTokens,
-		})
-	case ProviderAnthropic:
-		return claude.NewChatModel(ctx, &claude.Config{
-			APIKey:    viper.GetString("providers.anthropic.key"),
-			Model:     string(SupportedModels[model].APIModel),
-			MaxTokens: maxTokens,
-		})
-
-	case ProviderGroq:
-		return openai.NewChatModel(ctx, &openai.ChatModelConfig{
-			BaseURL:   "https://api.groq.com/openai/v1",
-			APIKey:    viper.GetString("providers.groq.key"),
-			Model:     string(SupportedModels[model].APIModel),
-			MaxTokens: &maxTokens,
-		})
-
-	}
-	return nil, errors.New("unsupported provider")
-}