feat: Recognise model text-unicorn

prantlf · Jun 4, 2024 · 8e83e7e · 8e83e7e
1 parent cbfe399
commit 8e83e7e
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 1 deletion.
diff --git a/internal/routes/generate.go b/internal/routes/generate.go
@@ -220,7 +220,8 @@ func HandleGenerate(w http.ResponseWriter, r *http.Request) int {
 	var handler generateModelHandler
 	if strings.HasPrefix(input.Model, "gemini") {
 		handler = &generateGeminiHandler{}
-	} else if strings.HasPrefix(input.Model, "text-bison") {
+	} else if strings.HasPrefix(input.Model, "text-bison") ||
+		strings.HasPrefix(input.Model, "text-unicorn") {
 		handler = &generateBisonHandler{}
 	} else if !canProxy {
 		return wrongInput(w, fmt.Sprintf("unrecognised model %q", input.Model))

diff --git a/internal/routes/tags.go b/internal/routes/tags.go
@@ -164,6 +164,19 @@ var googleModels = []modelInfo{
 		},
 		ExpiresAt: "2024-07-06T00:00:00.000Z",
 	},
+	{
+		Name:       "text-unicorn@001",
+		Model:      "text-unicorn@001",
+		ModifiedAt: "2023-11-30T00:00:00.000Z",
+		Details: modelDetails{
+			Family: "text-unicorn",
+			Families: []string{
+				"text-unicorn",
+			},
+			Purposes: []string{"generate"},
+		},
+		ExpiresAt: "2024-11-30T00:00:00.000Z",
+	},
 	{
 		Name:       "textembedding-gecko@002",
 		Model:      "textembedding-gecko@002",