Fixed some bugs, optimized the code structure, and

  added control over LLM model selection.

Author: echoflying

Diff for: _aichat.py

@@ -9,76 +9,120 @@
 # ark: 
 # claude: 
 
-class LLM_ai:
-    _MODELS = {
-        "zhipuai": {"url": "",   # SDK use default
+# model info: model name, max tokens, llm name
+class AI_models:
+    _MODEL = {
+        "glm-4-flash": {"llm": "zhipuai",   
                     "model": "glm-4-flash",
-                    "key": "zhipuai-key",
                     "max_tokens": 8000,
                     },
-        "kimi": {"url": "https://api-sg.moonshot.ai/v1",   # use openAI SDK
+        "glm-4-5020": {"llm": "zhipuai",   
+                    "model": "glm-4-5020",
+                    "max_tokens": 8000,
+                    },
+        "kimi8k": {"llm": "kimi",   # use openAI SDK
                     "model": "moonshot-v1-8k",
-                    "key": "sk-kimi-key",
                     "max_tokens": 8000,
                     },
-        "ark": {"url": "https://ark.cn-beijing.volces.com/api/v3",    # 火山方舟大模型，抖音，豆包，扣子是一家
-                    "model": "ep-20240929221043-jsbgc",
-                    "key": "ark-key",
+        "doubao4": {"llm": "ark",    # 火山方舟大模型，抖音，豆包，扣子是一家
+                    "model": "ep-20241005223718-nl742",    # Doubao-pro-4k
                     "max_tokens": 4000,
                     },
-        "claude": {"url": "https://api.gptapi.us/v1/chat/completions",       # Legend's testing bed
+        "doubao32": {"llm": "ark",
+                    "model": "ep-20240929221043-jsbgc",    # Doubao-pro-32k
+                    "max_tokens": 4000,
+                    },
+        "claude35": {"llm": "claude",       # Legend's testing bed
                     "model": "claude-3-5-sonnet",
-                    "key": "sk-claud-key",
                     "max_tokens": 8000,
                     },
     }
+
+    _working_model = {
+        "llm" : "",
+        "model" : "",
+        "max_tokens" : 0
+    }
+
+    # nickname 1) key used in _MODEL; 2)llm.model.9999 for dedicated llm-model, bad format will raise runtime
+    def __init__(self, model_nickname: str):
+        if model_nickname in self._MODEL:
+            self._working_model = self._MODEL[model_nickname]
+        else:
+            llm, model, n = model_nickname.split('.')
+            self._working_model = {
+                "llm" : llm,
+                "model" : model,
+                "max_tokens" : int(n)
+            }
+
+    def __getattr__(self, attr:str):
+        return self._working_model[attr]
+
+
+# maintain llm-url, get api-key, issue chat
+class LLM_ai:
+
+    _LLM = {
+        "zhipuai": {"url": "",   # SDK use default
+                    },
+        "kimi": {"url": "https://api-sg.moonshot.ai/v1",   # use openAI SDK
+                    },
+        "ark": {"url": "https://ark.cn-beijing.volces.com/api/v3",    # 火山方舟大模型，抖音，豆包，扣子是一家
+                    },
+        "claude": {"url": "https://api.gptapi.us/v1/chat/completions",       # Legend's testing bed
+                    },
+    }
+
+    # working one
     _llm = None
+    _model = None
+    _max_tokens = None
     _client = None
 
 
-    def __init__(self, llm: str, model=""):
+    def __init__(self, llm: str, model="", max_tokens= 2000):
         self._llm = llm
-        self._MODELS[llm]["key"] = st.secrets.ai_keys[llm]
+        ai_key = st.secrets.ai_keys[llm]
 
         if llm == "zhipuai":
-            self._client = ZhipuAI(api_key=self._MODELS[llm]["key"])
+            self._client = ZhipuAI(api_key=ai_key)    # zhipuai SDK use default url
         elif llm == "kimi":
-            self._client = openai.OpenAI(api_key=self._MODELS[llm]["key"],
-                                         base_url = self._MODELS[llm]["url"])
+            self._client = openai.OpenAI(api_key=ai_key, base_url = self._LLM[llm]["url"])
         elif llm == "ark":
-            self._client = Ark(api_key=self._MODELS[llm]["key"],
-                               base_url=self._MODELS[llm]["url"])
+            self._client = Ark(api_key=ai_key, base_url=self._LLM[llm]["url"])
         elif llm == "claude":
-            self._client = openai.OpenAI(api_key=self._MODELS[llm]["key"], 
-                                         base_url=self._MODELS[llm]["url"])
+            self._client = openai.OpenAI(api_key=ai_key, base_url=self._LLM[llm]["url"])
         else:
             raise ValueError(f"Invalid llm: {llm}")
 
         if not model == "":
-            self._MODELS[llm]["model"] = model
+            self._model = model
+        self._max_tokens = max_tokens
 
     # use stream model if pass_chunk is not None
     def chat(self, prompt: str, t: str, pass_chunk= None):
         if pass_chunk is None:
             response = self._client.chat.completions.create(
-                model=self._MODELS[self._llm]["model"],
+                model = self._model,
                 messages = [
                     {"role": "system", "content": prompt},
                     {"role": "user", "content": t}
                 ],
                 temperature = 0.3,
+                max_tokens=self._max_tokens,
             )
             return response.choices[0].message.content
         else:
             response = self._client.chat.completions.create(
-                model=self._MODELS[self._llm]["model"],
-                messages=[
-                {"role": "system", "content": prompt},
-                {"role": "user", "content": t},
-            ],
-            temperature=0.3,
-            stream=True,
-            max_tokens=self._MODELS[self._llm]["max_tokens"],
+                model = self._model,
+                messages = [
+                    {"role": "system", "content": prompt},
+                    {"role": "user", "content": t},
+                ],
+                temperature=0.3,
+                stream=True,
+                max_tokens=self._max_tokens,
             )
 
             answer=""