toolkit/brain.py

# -*- coding: utf-8 -*-
"""Общий модуль генерации: единая точка доступа к «мозгам».

Мозги: deepseek (API), lmstudio (локальная GGUF через LM Studio), local (transformers+LoRA 3B).
Все остальные инструменты (chat, project, docs, refactor, translate, rag, api) используют это.
"""
from webui import (deepseek_available, deepseek_generate,
                   lmstudio_available, lmstudio_generate, _strip_fences)


def available():
    return {"deepseek": deepseek_available(), "lmstudio": lmstudio_available()}


def pick(backend=None):
    """Выбирает мозг: явный, иначе лучший доступный (lmstudio 14B → deepseek → local 3B).

    Локальная 14B (LM Studio) — главный мозг: офлайн и приватно.
    DeepSeek — фолбэк/судья для самопроверки, local 3B — крайний случай.
    """
    if backend in ("deepseek", "lmstudio", "local"):
        return backend
    a = available()
    if a["lmstudio"]:
        return "lmstudio"
    if a["deepseek"]:
        return "deepseek"
    return "local"


def generate(backend, messages, max_tokens=1024):
    """Генерация текста заданным мозгом. Возвращает очищенный от ```-обёрток текст."""
    if backend == "deepseek":
        return _strip_fences(deepseek_generate(messages, max_tokens=max_tokens))
    if backend == "lmstudio":
        return _strip_fences(lmstudio_generate(messages, max_tokens=max_tokens))
    from infer_smart import load, gen
    if not hasattr(generate, "_m"):
        generate._m = load()
    tok, model = generate._m
    return _strip_fences(gen(tok, model, messages, max_tokens=max_tokens))