initial agent setup

1 files changed, 196 insertions, 0 deletions

diff --git a/internal/llm/models/models.go b/internal/llm/models/models.go
new file mode 100644
index 000000000..1895e256b
--- /dev/null
+++ b/internal/llm/models/models.go
@@ -0,0 +1,196 @@
+package models
+
+import (
+	"context"
+	"errors"
+
+	"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"` // Actual value used when calling the API
+}
+
+const (
+	DefaultBigModel    = GPT4oMini
+	DefaultLittleModel = GPT4oMini
+)
+
+// 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"
+)
+
+const (
+	ProviderOpenAI    ModelProvider = "openai"
+	ProviderAnthropic ModelProvider = "anthropic"
+	ProviderGoogle    ModelProvider = "google"
+	ProviderXAI       ModelProvider = "xai"
+	ProviderDeepSeek  ModelProvider = "deepseek"
+	ProviderMeta      ModelProvider = "meta"
+)
+
+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",
+	},
+	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",
+	},
+	Claude3Haiku: {
+		ID:       Claude3Haiku,
+		Name:     "Claude 3 Haiku",
+		Provider: ProviderAnthropic,
+		APIModel: "claude-3-haiku",
+	},
+	Claude37Sonnet: {
+		ID:       Claude37Sonnet,
+		Name:     "Claude 3.7 Sonnet",
+		Provider: ProviderAnthropic,
+		APIModel: "claude-3-7-sonnet-20250219",
+	},
+	// 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",
+	},
+	// Meta
+	Llama3: {
+		ID:       Llama3,
+		Name:     "LLaMA 3",
+		Provider: ProviderMeta,
+		APIModel: "llama-3",
+	},
+	Llama270B: {
+		ID:       Llama270B,
+		Name:     "LLaMA 2 70B",
+		Provider: ProviderMeta,
+		APIModel: "llama-2-70b",
+	},
+}
+
+func GetModel(ctx context.Context, model ModelID) (model.ChatModel, error) {
+	provider := SupportedModels[model].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,
+		})
+
+	}
+	return nil, errors.New("unsupported provider")
+}