update

app.py

@@ -2,7 +2,10 @@ import gradio as gr
 from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
 from llm import DeepSeekLLM, OpenRouterLLM, TongYiLLM
 from config import settings
-from prompts import web_prompt, explain_code_template, optimize_code_template, debug_code_template, function_gen_template, translate_doc_template, backend_developer_prompt, analyst_prompt
+from prompts import (
+    web_prompt, explain_code_template, optimize_code_template, debug_code_template,
+    function_gen_template, translate_doc_template, backend_developer_prompt, analyst_prompt
+)
 from langchain_core.prompts import PromptTemplate
 from log import logging
 from utils import convert_image_to_base64
@@ -21,6 +24,12 @@ provider_model_map = dict(
     Tongyi=tongyi_llm,
 )
 
+system_prompt_map = {
+    "前端开发助手": web_prompt,
+    "后端开发助手": backend_developer_prompt,
+    "数据分析师": analyst_prompt,
+}
+
 support_vision_models = [
     'openai/gpt-4o-mini', 'anthropic/claude-3.5-sonnet', 'google/gemini-pro-1.5-exp',
     'openai/gpt-4o', 'google/gemini-flash-1.5', 'liuhaotian/llava-yi-34b', 'anthropic/claude-3-haiku',
@@ -33,29 +42,39 @@ def get_default_chat():
     return _llm.get_chat_engine()
 
 
+def get_chat_or_default(chat):
+    if chat is None:
+        chat = get_default_chat()
+    return chat
+
+
+def convert_history_to_langchain_history(history, lc_history):
+    for his_msg in history:
+        if his_msg['role'] == 'user':
+            if not hasattr(his_msg['content'], 'file'):
+                lc_history.append(HumanMessage(content=his_msg['content']))
+        if his_msg['role'] == 'assistant':
+            lc_history.append(AIMessage(content=his_msg['content']))
+    return lc_history
+
+
+def append_system_prompt(key: str, lc_history):
+    prompt = system_prompt_map[key]
+    lc_history.append(SystemMessage(content=prompt))
+    return lc_history
+
+
 def predict(message, history, _chat, _current_assistant: str):
     logger.info(f"chat predict: {message}, {history}, {_chat}, {_current_assistant}")
     files_len = len(message.files)
-    if _chat is None:
-        _chat = get_default_chat()
+    _chat = get_chat_or_default(_chat)
     if files_len > 0:
         if _chat.model_name not in support_vision_models:
             raise gr.Error("当前模型不支持图片，请更换模型。")
 
     _lc_history = []
-    assistant_prompt = web_prompt
-    if _current_assistant == '后端开发助手':
-        assistant_prompt = backend_developer_prompt
-    if _current_assistant == '数据分析师':
-        assistant_prompt = analyst_prompt
-    _lc_history.append(SystemMessage(content=assistant_prompt))
-
-    for his_msg in history:
-        if his_msg['role'] == 'user':
-            if not hasattr(his_msg['content'], 'file'):
-                _lc_history.append(HumanMessage(content=his_msg['content']))
-        if his_msg['role'] == 'assistant':
-            _lc_history.append(AIMessage(content=his_msg['content']))
+    _lc_history = append_system_prompt(_current_assistant, _lc_history)
+    _lc_history = convert_history_to_langchain_history(history, _lc_history)
 
     if files_len == 0:
         _lc_history.append(HumanMessage(content=message.text))
@@ -81,8 +100,7 @@ def update_chat(_provider: str, _model: str, _temperature: float, _max_tokens: i
 
 
 def explain_code(_code_type: str, _code: str, _chat):
-    if _chat is None:
-        _chat = get_default_chat()
+    _chat = get_chat_or_default(_chat)
     chat_messages = [
         SystemMessage(content=explain_code_template),
         HumanMessage(content=_code),
@@ -94,8 +112,7 @@ def explain_code(_code_type: str, _code: str, _chat):
 
 
 def optimize_code(_code_type: str, _code: str, _chat):
-    if _chat is None:
-        _chat = get_default_chat()
+    _chat = get_chat_or_default(_chat)
     prompt = PromptTemplate.from_template(optimize_code_template)
     prompt = prompt.format(code_type=_code_type)
     chat_messages = [
@@ -109,8 +126,7 @@ def optimize_code(_code_type: str, _code: str, _chat):
 
 
 def debug_code(_code_type: str, _code: str, _chat):
-    if _chat is None:
-        _chat = get_default_chat()
+    _chat = get_chat_or_default(_chat)
     prompt = PromptTemplate.from_template(debug_code_template)
     prompt = prompt.format(code_type=_code_type)
     chat_messages = [
@@ -124,8 +140,7 @@ def debug_code(_code_type: str, _code: str, _chat):
 
 
 def function_gen(_code_type: str, _code: str, _chat):
-    if _chat is None:
-        _chat = get_default_chat()
+    _chat = get_chat_or_default(_chat)
     prompt = PromptTemplate.from_template(function_gen_template)
     prompt = prompt.format(code_type=_code_type)
     chat_messages = [
@@ -139,8 +154,7 @@ def function_gen(_code_type: str, _code: str, _chat):
 
 
 def translate_doc(_language_input, _language_output, _doc, _chat):
-    if _chat is None:
-        _chat = get_default_chat()
+    _chat = get_chat_or_default(_chat)
     prompt = PromptTemplate.from_template(translate_doc_template)
     prompt = prompt.format(language_input=_language_input, language_output=_language_output)
     chat_messages = [

llm.py

@@ -60,8 +60,8 @@ class DeepSeekLLM(BaseLLM):
 
 class OpenRouterLLM(BaseLLM):
     _support_models = [
-        'openai/gpt-4o-mini', 'anthropic/claude-3.5-sonnet', 'google/gemini-pro-1.5-exp',
-        'mistralai/mistral-large', 'meta-llama/llama-3.1-405b-instruct', 'openai/gpt-4o',
+        'openai/gpt-4o-mini', 'anthropic/claude-3.5-sonnet', 'openai/gpt-4o-2024-08-06',
+        'google/gemini-pro-1.5-exp', 'mistralai/mistral-large', 'meta-llama/llama-3.1-405b-instruct',
         'nvidia/nemotron-4-340b-instruct', 'deepseek/deepseek-coder', 'google/gemma-2-27b-it',
         'google/gemini-flash-1.5', 'deepseek/deepseek-chat', 'qwen/qwen-2-72b-instruct',
         'liuhaotian/llava-yi-34b', 'qwen/qwen-110b-chat',