From d55b3d904886890f348e8309b19dd5ff269a2a1c Mon Sep 17 00:00:00 2001
From: dingwenjiang <dingwenjiang@bytedance.com>
Date: Sat, 30 Sep 2023 01:16:45 +0800
Subject: [PATCH] feat: add method convert str to embedding model enum

---
 embeddings.go | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/embeddings.go b/embeddings.go
index 5ba91f23..480ee0e2 100644
--- a/embeddings.go
+++ b/embeddings.go
@@ -12,9 +12,19 @@ import (
 // to generate Embedding vectors.
 type EmbeddingModel int
 
+func ConvertStr2EmbeddingModel(modelName string) EmbeddingModel {
+	if val, ok := stringToEnum[modelName]; ok {
+		return val
+	}
+	return Unknown
+}
+
 // String implements the fmt.Stringer interface.
 func (e EmbeddingModel) String() string {
-	return enumToString[e]
+	if val, ok := enumToString[e]; ok {
+		return val
+	}
+	return "Unknown"
 }
 
 // MarshalText implements the encoding.TextMarshaler interface.