update

.env

@@ -1,2 +1,3 @@
 DEEP_SEEK_API_KEY=
+OPEN_ROUTER_API_KEY=
 DEBUG=False

app.py

@@ -1,21 +1,15 @@
 import gradio as gr
-from langchain_openai import ChatOpenAI
 from langchain_core.messages import HumanMessage, AIMessage
-from llm import DeepSeekLLM
+from llm import DeepSeekLLM, OpenRouterLLM
 from config import settings
 
 
 deep_seek_llm = DeepSeekLLM(api_key=settings.deep_seek_api_key)
+open_router_llm = OpenRouterLLM(api_key=settings.open_router_api_key)
 
 
 def init_chat():
-    return ChatOpenAI(
-        model=deep_seek_llm.default_model,
-        api_key=deep_seek_llm.api_key,
-        base_url=deep_seek_llm.base_url,
-        temperature=deep_seek_llm.default_temperature,
-        max_tokens=deep_seek_llm.default_max_tokens,
-    )
+    return deep_seek_llm.get_chat_engine()
 
 
 def predict(message, history, chat):
@@ -33,14 +27,12 @@ def predict(message, history, chat):
         yield response_message
 
 
-def update_chat(_chat, _model: str, _temperature: float, _max_tokens: int):
-    _chat = ChatOpenAI(
-        model=_model,
-        api_key=deep_seek_llm.api_key,
-        base_url=deep_seek_llm.base_url,
-        temperature=_temperature,
-        max_tokens=_max_tokens,
-    )
+def update_chat(_provider: str, _chat, _model: str, _temperature: float, _max_tokens: int):
+    print('?????', _provider, _chat, _model, _temperature, _max_tokens)
+    if _provider == 'DeepSeek':
+        _chat = deep_seek_llm.get_chat_engine(model=_model, temperature=_temperature, max_tokens=_max_tokens)
+    if _provider == 'OpenRouter':
+        _chat = open_router_llm.get_chat_engine(model=_model, temperature=_temperature, max_tokens=_max_tokens)
     return _chat
 
 
@@ -59,30 +51,86 @@ with gr.Blocks() as app:
             with gr.Column(scale=1, min_width=300):
                 with gr.Accordion('Select Model', open=True):
                     with gr.Column():
-                        model = gr.Dropdown(
-                            label='模型',
-                            choices=deep_seek_llm.support_models,
-                            value=deep_seek_llm.default_model
-                        )
-                        temperature = gr.Slider(
-                            minimum=0.0,
-                            maximum=1.0,
-                            step=0.1,
-                            value=deep_seek_llm.default_temperature,
-                            label="Temperature",
-                            key="temperature",
-                        )
-                        max_tokens = gr.Number(
-                            minimum=1024,
-                            maximum=1024 * 20,
-                            step=128,
-                            value=deep_seek_llm.default_max_tokens,
-                            label="Max Tokens",
-                            key="max_tokens",
-                        )
-        model.change(fn=update_chat, inputs=[chat_engine, model, temperature, max_tokens], outputs=[chat_engine])
-        temperature.change(fn=update_chat, inputs=[chat_engine, model, temperature, max_tokens], outputs=[chat_engine])
-        max_tokens.change(fn=update_chat, inputs=[chat_engine, model, temperature, max_tokens], outputs=[chat_engine])
+                        provider = gr.Dropdown(label='Provider', choices=['DeepSeek', 'OpenRouter'], value='DeepSeek')
+
+                    @gr.render(inputs=provider)
+                    def show_model_config_panel(_provider):
+                        if _provider == 'DeepSeek':
+                            with gr.Column():
+                                model = gr.Dropdown(
+                                    label='模型',
+                                    choices=deep_seek_llm.support_models,
+                                    value=deep_seek_llm.default_model
+                                )
+                                temperature = gr.Slider(
+                                    minimum=0.0,
+                                    maximum=1.0,
+                                    step=0.1,
+                                    value=deep_seek_llm.default_temperature,
+                                    label="Temperature",
+                                    key="temperature",
+                                )
+                                max_tokens = gr.Number(
+                                    minimum=1024,
+                                    maximum=1024 * 20,
+                                    step=128,
+                                    value=deep_seek_llm.default_max_tokens,
+                                    label="Max Tokens",
+                                    key="max_tokens",
+                                )
+                            model.change(
+                                fn=update_chat,
+                                inputs=[provider, chat_engine, model, temperature, max_tokens],
+                                outputs=[chat_engine],
+                            )
+                            temperature.change(
+                                fn=update_chat,
+                                inputs=[provider, chat_engine, model, temperature, max_tokens],
+                                outputs=[chat_engine],
+                            )
+                            max_tokens.change(
+                                fn=update_chat,
+                                inputs=[provider, chat_engine, model, temperature, max_tokens],
+                                outputs=[chat_engine],
+                            )
+                        if _provider == 'OpenRouter':
+                            with gr.Column():
+                                model = gr.Dropdown(
+                                    label='模型',
+                                    choices=open_router_llm.support_models,
+                                    value=open_router_llm.default_model
+                                )
+                                temperature = gr.Slider(
+                                    minimum=0.0,
+                                    maximum=1.0,
+                                    step=0.1,
+                                    value=open_router_llm.default_temperature,
+                                    label="Temperature",
+                                    key="temperature",
+                                )
+                                max_tokens = gr.Number(
+                                    minimum=1024,
+                                    maximum=1024 * 20,
+                                    step=128,
+                                    value=open_router_llm.default_max_tokens,
+                                    label="Max Tokens",
+                                    key="max_tokens",
+                                )
+                            model.change(
+                                fn=update_chat,
+                                inputs=[provider, chat_engine, model, temperature, max_tokens],
+                                outputs=[chat_engine],
+                            )
+                            temperature.change(
+                                fn=update_chat,
+                                inputs=[provider, chat_engine, model, temperature, max_tokens],
+                                outputs=[chat_engine],
+                            )
+                            max_tokens.change(
+                                fn=update_chat,
+                                inputs=[provider, chat_engine, model, temperature, max_tokens],
+                                outputs=[chat_engine],
+                            )
 
     with gr.Tab('画图'):
         with gr.Row():

config.py

@@ -3,6 +3,7 @@ from pydantic_settings import BaseSettings, SettingsConfigDict
 
 class Settings(BaseSettings):
     deep_seek_api_key: str
+    open_router_api_key: str
     debug: bool
 
     model_config = SettingsConfigDict(env_file=('.env', '.env.local'), env_file_encoding='utf-8')

llm.py

@@ -1,5 +1,6 @@
 from typing import List
 from abc import ABC
+from langchain_openai import ChatOpenAI
 
 
 class DeepSeekLLM(ABC):
@@ -37,3 +38,69 @@ class DeepSeekLLM(ABC):
     def default_max_tokens(self) -> int:
         return self._default_max_tokens
 
+    def get_chat_engine(self, *, model: str = None, temperature: float = None, max_tokens: int = None):
+        model = model or self.default_model
+        temperature = temperature or self.default_temperature
+        max_tokens = max_tokens or self.default_max_tokens
+        return ChatOpenAI(
+            model=model,
+            api_key=self.api_key,
+            base_url=self.base_url,
+            temperature=temperature,
+            max_tokens=max_tokens,
+        )
+
+
+class OpenRouterLLM(ABC):
+    _support_models = [
+        'anthropic/claude-3.5-sonnet', 'openai/gpt-4o',
+        'nvidia/nemotron-4-340b-instruct', 'deepseek/deepseek-coder',
+        'google/gemini-flash-1.5', 'deepseek/deepseek-chat',
+        'liuhaotian/llava-yi-34b', 'qwen/qwen-110b-chat',
+        'qwen/qwen-72b-chat', 'google/gemini-pro-1.5',
+        'cohere/command-r-plus', 'anthropic/claude-3-haiku',
+    ]
+    _base_url = 'https://openrouter.ai/api/v1'
+    _default_model = 'anthropic/claude-3.5-sonnet'
+    _api_key: str
+    _default_temperature: float = 0.5
+    _default_max_tokens: int = 4096
+
+    def __init__(self, *, api_key: str):
+        self._api_key = api_key
+
+    @property
+    def support_models(self) -> List[str]:
+        return self._support_models
+
+    @property
+    def default_model(self) -> str:
+        return self._default_model
+
+    @property
+    def base_url(self) -> str:
+        return self._base_url
+
+    @property
+    def api_key(self) -> str:
+        return self._api_key
+
+    @property
+    def default_temperature(self) -> float:
+        return self._default_temperature
+
+    @property
+    def default_max_tokens(self) -> int:
+        return self._default_max_tokens
+
+    def get_chat_engine(self, *, model: str = None, temperature: float = None, max_tokens: int = None):
+        model = model or self.default_model
+        temperature = temperature or self.default_temperature
+        max_tokens = max_tokens or self.default_max_tokens
+        return ChatOpenAI(
+            model=model,
+            api_key=self.api_key,
+            base_url=self.base_url,
+            temperature=temperature,
+            max_tokens=max_tokens,
+        )