import json
import gradio as gr
from context.db_loader import SQLiteContextLoader
import constants
from helpers import load_json_file


class NeuroRavenApp:
    def __init__(self):
        self.system_prompt = None
        self.current_theme = None
        self.client = self.load_model_client()
        self.context_loader = SQLiteContextLoader()
        self.conversation_themes = self.get_conversation_themes()
        if constants.NO_CONTEXT not in self.conversation_themes:
            self.conversation_themes.append(constants.NO_CONTEXT)

    @staticmethod
    def get_conversation_themes():
        themes = load_json_file(constants.CONVERSATION_THEMES_JSON)
        return themes

    def handle_query(self, message, history, context_theme, should_save):
        if should_save and isinstance(self.context_loader, SQLiteContextLoader):
            if context_theme == constants.NO_CONTEXT:
                return constants.ERROR_CANNOT_SAVE_MESSAGE
            self.context_loader.save_message(context_theme, message)
            self.current_theme = None
            return constants.MESSAGE_SAVED_TO_CONTEXT

        if context_theme != self.current_theme:
            context = self.context_loader.load_context(context_theme)
            if context:
                self.system_prompt = constants.ANSWER_USING_CONTEXT_PROMPT.format(context=context)
            else:
                self.system_prompt = constants.DEFAULT_SYSTEM_PROMPT
            self.current_theme = context_theme

        response_generator = self.client.get_response(
            message=message,
            history=history,
            sys_prompt=self.system_prompt
        )

        # Генерация ответов по мере их поступления
        for response in response_generator:
            yield response  # Передаем каждый ответ в Gradio

    def launch_app(self):
        gr.ChatInterface(
            fn=self.handle_query,
            chatbot=gr.Chatbot(height=500),
            textbox=gr.Textbox(
                placeholder=constants.INPUT_PLACEHOLDER,
                container=False,
                scale=7
            ),
            title="NeuroRaven",
            theme="soft",
            retry_btn=None,
            clear_btn=constants.CLEAR_BUTTON_LABEL,
            additional_inputs=[
                gr.Dropdown(
                    self.conversation_themes,
                    label=constants.SELECT_CONTEXT_THEME_LABEL,
                    value=constants.NO_CONTEXT
                ),
                gr.Checkbox(
                    label=constants.SAVE_MESSAGE_LABEL,
                    value=False
                )
            ],
            css="footer{display:none !important}"
        ).launch(server_name=constants.SERVER_NAME, server_port=constants.SERVER_PORT)

    @staticmethod
    def load_model_client():
        with open("config.json") as config_file:
            config = json.load(config_file)

        model_name = config.get("model", "mini_4o")  # Значение по умолчанию
        if model_name == "gemma":
            from models.gemma_model import ModelClient
        elif model_name == "lms":
            from models.lms_model import ModelClient
        elif model_name == "mini_4o":
            from models.mini_4o_model import ModelClient
        elif model_name == "turbo_4o":
            from models.turbo_4o_model import ModelClient
        else:
            raise ValueError(f"Unknown model name: {model_name}")

        return ModelClient()


if __name__ == "__main__":
    app = NeuroRavenApp()
    app.launch_app()

from gradio_client import Client


class ModelClient:
    def __init__(self):
        self.model = Client("gokaygokay/Gemma-2-llamacpp")

    def get_response(self, message, history, sys_prompt):
        result = self.model.predict(
            message=message,
            model="gemma-2-27b-it-Q5_K_M.gguf",
            system_message=sys_prompt,
            max_tokens=2048,
            temperature=0.7,
            top_p=0.95,
            top_k=40,
            repeat_penalty=1.1,
            api_name="/chat"
        )
        return result


Как мне сделать чтобы если модель без стрима а возвращает результат целиком это работало