From 81c7b6afc5e5c2745f6c416c838067fe856cc7b9 Mon Sep 17 00:00:00 2001
From: Kevin Hu <kevinhu.sh@gmail.com>
Date: Wed, 20 Nov 2024 12:16:36 +0800
Subject: [PATCH] Make spark model robuster to model name (#3514)

### What problem does this PR solve?


### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)
---
 conf/llm_factories.json | 8 ++++----
 rag/llm/chat_model.py   | 6 +++++-
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/conf/llm_factories.json b/conf/llm_factories.json
index fced65e52..69997b493 100644
--- a/conf/llm_factories.json
+++ b/conf/llm_factories.json
@@ -3,7 +3,7 @@
         {
             "name": "OpenAI",
             "logo": "",
-            "tags": "LLM,TEXT EMBEDDING,SPEECH2TEXT,MODERATION",
+            "tags": "LLM,TEXT EMBEDDING,TTS,TEXT RE-RANK,SPEECH2TEXT,MODERATION",
             "status": "1",
             "llm": [
                 {
@@ -89,7 +89,7 @@
         {
             "name": "Tongyi-Qianwen",
             "logo": "",
-            "tags": "LLM,TEXT EMBEDDING,TEXT RE-RANK,SPEECH2TEXT,MODERATION",
+            "tags": "LLM,TEXT EMBEDDING,TEXT RE-RANK,TTS,SPEECH2TEXT,MODERATION",
             "status": "1",
             "llm": [
                 {
@@ -352,7 +352,7 @@
         {
             "name": "Xinference",
             "logo": "",
-            "tags": "LLM,TEXT EMBEDDING,SPEECH2TEXT,MODERATION,TEXT RE-RANK",
+            "tags": "LLM,TEXT EMBEDDING,TTS,SPEECH2TEXT,MODERATION,TEXT RE-RANK",
             "status": "1",
             "llm": []
         },
@@ -2303,7 +2303,7 @@
         {
             "name": "XunFei Spark",
             "logo": "",
-            "tags": "LLM",
+            "tags": "LLM,TTS",
             "status": "1",
             "llm": []
         },
diff --git a/rag/llm/chat_model.py b/rag/llm/chat_model.py
index 716d020ae..084bbc55e 100644
--- a/rag/llm/chat_model.py
+++ b/rag/llm/chat_model.py
@@ -1164,7 +1164,11 @@ class SparkChat(Base):
             "Spark-Pro-128K": "pro-128k",
             "Spark-4.0-Ultra": "4.0Ultra",
         }
-        model_version = model2version[model_name]
+        version2model = {v: k for k, v in model2version.items()}
+        assert model_name in model2version or model_name in version2model, f"The given model name is not supported yet. Support: {list(model2version.keys())}"
+        if model_name in model2version:
+            model_version = model2version[model_name]
+        else: model_version = model_name
         super().__init__(key, model_version, base_url)