add openrouter

main.go

@@ -73,6 +73,24 @@ var (
 		"gemma-2-27b-it",
 	}
 
+	OPENROUTER_MODELS = []string{
+		"meta-llama/llama-3.1-8b-instruct",
+		"meta-llama/llama-3.1-70b-instruct",
+		"meta-llama/llama-3.1-405b-instruct",
+		"meta-llama/llama-3.1-70b-instruct:free",
+
+		"anthropic/claude-3.5-sonnet:beta",
+		"anthropic/claude-3-5-haiku-20241022:beta",
+
+		"qwen/qwen-2.5-72b-instruct",
+		"qwen/qwen-2.5-7b-instruct",
+		"qwen/qwen-2.5-coder-32b-instruct",
+
+		"eva-unit-01/eva-qwen-2.5-32b",
+		"deepseek/deepseek-chat",
+		"nousresearch/hermes-3-llama-3.1-405b",
+	}
+
 	// Default admin public key, it will be used unless --dkn-admin-public-key is given
 	DKN_ADMIN_PUBLIC_KEY = "0208ef5e65a9c656a6f92fb2c770d5d5e2ecffe02a6aade19207f75110be6ae658"
 )
@@ -154,7 +172,7 @@ func main() {
 
 	// if DKN_MODELS are still empty, pick model interactively
 	if envvars["DKN_MODELS"] == "" || *pick_model {
-		pickedModels := utils.PickModels(OPENAI_MODELS, GEMINI_MODELS, OLLAMA_MODELS)
+		pickedModels := utils.PickModels(OPENAI_MODELS, GEMINI_MODELS, OPENROUTER_MODELS, OLLAMA_MODELS)
 		if pickedModels == "" {
 			fmt.Println("No valid model picked")
 			utils.ExitWithDelay(1)
@@ -182,6 +200,15 @@ func main() {
 		envvars["GEMINI_API_KEY"] = apikey
 	}
 
+	if utils.IsOpenRouterRequired(envvars["DKN_MODELS"], &OPENROUTER_MODELS) && envvars["OPENROUTER_API_KEY"] == "" {
+		apikey := utils.GetUserInput("Enter your OpenRoute API Key", true)
+		if apikey == "" {
+			fmt.Println("Invalid input, please place your OPENROUTER_API_KEY to .env file")
+			utils.ExitWithDelay(1)
+		}
+		envvars["OPENROUTER_API_KEY"] = apikey
+	}
+
 	// check ollama environment
 	if utils.IsOllamaRequired(envvars["DKN_MODELS"], &OLLAMA_MODELS) {
 		ollamaHost, ollamaPort := utils.HandleOllamaEnv(envvars["OLLAMA_HOST"], envvars["OLLAMA_PORT"])

utils/cmd.go

@@ -139,7 +139,7 @@ func RunCommand(working_dir string, outputDest string, wait bool, timeout time.D
 //
 // Returns:
 //   - string: A comma-separated string of selected model names.
-func PickModels(openai_models, gemini_models, ollama_models []string) string {
+func PickModels(openai_models, gemini_models, or_models, ollama_models []string) string {
 
 	// column widths
 	idWidth := 4
@@ -166,8 +166,14 @@ func PickModels(openai_models, gemini_models, ollama_models []string) string {
 		provider := "Gemini"
 		fmt.Printf("| %-*d | %-*s | %-*s |\n", idWidth, modelId, providerWidth, provider, nameWidth, model)
 	}
-	for id, model := range ollama_models {
+	for id, model := range or_models {
 		modelId := len(openai_models) + len(gemini_models) + id + 1
+		provider := "OpenRouter"
+		display_name := strings.Split(model, "/")[1]
+		fmt.Printf("| %-*d | %-*s | %-*s |\n", idWidth, modelId, providerWidth, provider, nameWidth, display_name)
+	}
+	for id, model := range ollama_models {
+		modelId := len(openai_models) + len(gemini_models) + len(or_models) + id + 1
 		provider := "Ollama"
 		fmt.Printf("| %-*d | %-*s | %-*s |\n", idWidth, modelId, providerWidth, provider, nameWidth, model)
 	}
@@ -208,12 +214,19 @@ func PickModels(openai_models, gemini_models, ollama_models []string) string {
 				picked_models_map[id] = true
 				picked_models_str = fmt.Sprintf("%s,%s", picked_models_str, gemini_models[id-len(openai_models)-1])
 			}
-		} else if id > len(openai_models)+len(gemini_models) && id <= len(ollama_models)+len(gemini_models)+len(openai_models) {
+		} else if id > len(openai_models)+len(gemini_models) && id <= len(or_models)+len(gemini_models)+len(openai_models) {
+			// openrouter model picked
+			if !picked_models_map[id] {
+				// if not already picked, add it to bin
+				picked_models_map[id] = true
+				picked_models_str = fmt.Sprintf("%s,%s", picked_models_str, or_models[id-len(openai_models)-len(gemini_models)-1])
+			}
+		} else if id > len(openai_models)+len(gemini_models)+len(or_models) && id <= len(or_models)+len(ollama_models)+len(gemini_models)+len(openai_models) {
 			// ollama model picked
 			if !picked_models_map[id] {
 				// if not already picked, add it to bin
 				picked_models_map[id] = true
-				picked_models_str = fmt.Sprintf("%s,%s", picked_models_str, ollama_models[id-len(gemini_models)-len(openai_models)-1])
+				picked_models_str = fmt.Sprintf("%s,%s", picked_models_str, ollama_models[id-len(or_models)-len(gemini_models)-len(openai_models)-1])
 			}
 		} else {
 			// out of index, invalid

utils/openrouter.go

@@ -0,0 +1,24 @@
+package utils
+
+import "strings"
+
+// IsOpenRouterRequired checks if any of the picked models require OpenRouter by comparing them against a list of available OR models.
+//
+// Parameters:
+//   - picked_models: A comma-separated string of model names selected by the user.
+//   - or_models: A pointer to a slice of strings containing available OpenRouter model names.
+//
+// Returns:
+//   - bool: Returns true if any of the picked models match OR models, indicating that OpenRouter api key is required, otherwise false.
+func IsOpenRouterRequired(picked_models string, or_models *[]string) bool {
+	required := false
+	for _, model := range strings.Split(picked_models, ",") {
+		for _, or_model := range *or_models {
+			if model == or_model {
+				required = true
+				break
+			}
+		}
+	}
+	return required
+}