evyzacq
68 lines
2.6 KiB
Python
68 lines
2.6 KiB
Python
# Model router: creates LLM clients based on provider configuration.
|
|
|
|
import logging
|
|
|
|
from server.config import settings, DEEPSEEK_API_KEY, DEEPSEEK_BASE_URL, DASHSCOPE_API_KEY, DASHSCOPE_BASE_URL
|
|
from server.core.llm_provider.base import LLMClient
|
|
from server.core.llm_provider.mock_client import MockLLMClient
|
|
|
|
logger = logging.getLogger("testflow")
|
|
|
|
|
|
class ModelRouter:
|
|
"""Creates and caches LLM client instances.
|
|
|
|
Two independent mock toggles:
|
|
- USE_MOCK → controls pipeline (left panel): text + image clients
|
|
- CHAT_USE_MOCK → controls AI assistant (right panel): chat client
|
|
"""
|
|
|
|
def __init__(self):
|
|
self._clients: dict[str, LLMClient] = {}
|
|
self._log_status()
|
|
|
|
def _log_status(self):
|
|
pipe = "MOCK" if settings.USE_MOCK else "REAL"
|
|
chat = "MOCK" if settings.CHAT_USE_MOCK else "REAL"
|
|
logger.info("[ROUTER] 流水线=%s | AI助手=%s (DeepSeek %s / Qwen %s)",
|
|
pipe, chat, settings.TEXT_MODEL, settings.IMAGE_MODEL)
|
|
|
|
def get_text_client(self) -> LLMClient:
|
|
"""Pipeline text client (controlled by USE_MOCK)."""
|
|
return self._get_or_create("text", settings.USE_MOCK, settings.TEXT_MODEL,
|
|
DEEPSEEK_API_KEY, DEEPSEEK_BASE_URL)
|
|
|
|
def get_image_client(self) -> LLMClient:
|
|
"""Pipeline image client (controlled by USE_MOCK)."""
|
|
return self._get_or_create("image", settings.USE_MOCK, settings.IMAGE_MODEL,
|
|
DASHSCOPE_API_KEY, DASHSCOPE_BASE_URL)
|
|
|
|
def get_chat_client(self) -> LLMClient:
|
|
"""AI assistant chat client (controlled by CHAT_USE_MOCK)."""
|
|
return self._get_or_create("chat", settings.CHAT_USE_MOCK, settings.TEXT_MODEL,
|
|
DEEPSEEK_API_KEY, DEEPSEEK_BASE_URL)
|
|
|
|
def _get_or_create(self, key: str, use_mock: bool, model: str,
|
|
api_key: str, base_url: str) -> LLMClient:
|
|
if key not in self._clients:
|
|
if use_mock:
|
|
logger.info("[ROUTER] %s client → MOCK", key)
|
|
self._clients[key] = MockLLMClient(model_name=model)
|
|
else:
|
|
logger.info("[ROUTER] %s client → REAL (%s)", key, model)
|
|
self._clients[key] = LLMClient(
|
|
api_key=api_key, base_url=base_url, model=model,
|
|
)
|
|
return self._clients[key]
|
|
|
|
@property
|
|
def text_model(self) -> str:
|
|
return settings.TEXT_MODEL
|
|
|
|
@property
|
|
def image_model(self) -> str:
|
|
return settings.IMAGE_MODEL
|
|
|
|
|
|
router = ModelRouter()
|