Gradio_UI.py

import gradio as gr
import logging
from pathlib import Path
from typing import Generator
from datetime import datetime

logger = logging.getLogger(__name__)

class GradioUI:
    def __init__(self, agent):
        """Ініціалізація інтерфейсу"""
        self.agent = agent
        logger.info("GradioUI успішно ініціалізовано")
        
    def get_model_choices(self):
        """Отримання списку доступних моделей"""
        return [
            f"{key} - {desc}" 
            for key, desc in self.agent.model_initializer.get_available_models()
        ]
        
    def get_default_model(self):
        """Отримання моделі за замовчуванням"""
        default_model = self.agent.model_initializer.config['default_model']
        model_choices = dict(self.agent.model_initializer.get_available_models())
        return f"{default_model} - {model_choices[default_model]}"

    def _load_documentation(self):
        """Завантаження документації з markdown файлу"""
        try:
            docs_path = Path(__file__).parent / "docs" / "agent_documentation.md"
            with open(docs_path, 'r', encoding='utf-8') as f:
                return f.read()
        except Exception as e:
            logger.error(f"Помилка завантаження документації: {e}")
            return "# Помилка завантаження документації\nБудь ласка, перевірте наявність файлу docs/agent_documentation.md"
        
    def build_interface(self):
        """Побудова інтерфейсу Gradio"""
        with gr.Blocks(theme=gr.themes.Soft()) as interface:
            with gr.Tabs():
                with gr.Tab("Основний інтерфейс"):
                    with gr.Row():
                        model_dropdown = gr.Dropdown(
                            choices=self.get_model_choices(),
                            value=self.get_default_model(),
                            label="Оберіть модель",
                            info="Оберіть модель для обробки запитів"
                        )

                    # Контейнер для чату та логу з фіксованою висотою
                    with gr.Row(equal_height=True):
                        # Ліва колонка (чат)
                        with gr.Column(scale=1):
                            chatbot = gr.Chatbot(
                                label="Research Agent",
                                height=500,
                                show_copy_button=True,
                                container=True,
                            )
                        
                        # Права колонка (лог)
                        with gr.Column(scale=1):
                            steps_log = gr.Markdown(
                                value="",
                                label="Process Steps",
                                elem_id="process-steps",
                                show_label=True,
                            )
                            # CSS для фіксованої висоти та скролу
                            gr.HTML("""
                                <style>
                                    #process-steps {
                                        height: 500px;
                                        overflow-y: auto;
                                        padding: 1rem;
                                        border: 1px solid #e5e7eb;
                                        border-radius: 0.5rem;
                                        background: white;
                                    }
                                </style>
                            """)
                    
                    with gr.Row():
                        text_input = gr.Textbox(
                            label="Введіть ваш запит",
                            placeholder="Введіть ваш запит тут...",
                            lines=2
                        )
                        
                    with gr.Row():
                        submit_btn = gr.Button("Надіслати", variant="primary")
                        clear_btn = gr.Button("Очистити чат")
                        clear_context_btn = gr.Button("Очистити контекст", variant="secondary")
                        save_btn = gr.Button("Зберегти")

                    response_text = gr.Textbox(visible=False)
                    file_output = gr.File(interactive=False)

                with gr.Tab("Документація"):
                    gr.Markdown(self._load_documentation())

            # Зберігання стану
            state = gr.State([])
            current_model = gr.State(self.agent.model_initializer.config['default_model'])

            def change_model(choice):
                """Обробник зміни моделі"""
                model_key = choice.split(" - ")[0]
                self.agent.model_initializer.initialize_model(model_key)
                return model_key

            def process_steps_generator(message, chat_history) -> Generator:
                """Генератор для поступового оновлення Process Steps"""
                chat_history = list(chat_history)
                chat_history.append((message, None))
                current_steps = []
                
                # Початкове повідомлення
                current_steps.append("🚀 Початок обробки запиту...\n")
                yield chat_history, "\n".join(current_steps), None

                try:
                    def step_callback(step_data):
                        step_num = len(current_steps)
                        step_text = f"\n### Крок {step_num}: {step_data['info']}"
                        
                        if step_data['result']:
                            if isinstance(step_data['result'], dict):
                                # Форматуємо словник для кращого відображення
                                result_text = "\n".join(f"- {k}: {v}" for k, v in step_data['result'].items())
                                step_text += f"\n```\n{result_text}\n```"
                            else:
                                step_text += f"\n```\n{step_data['result']}\n```"
                        
                        current_steps.append(step_text)
                        return "\n".join(current_steps)
                    
                    self.agent.step_callback = step_callback
                    response = self.agent.process_query(message)
                    chat_history[-1] = (message, response)
                    
                    # Додаємо інформацію про контекст
                    context_info = (
                        f"\n### Інформація про контекст\n"
                        f"Кількість попередніх запитів: {len(self.agent.conversation_context['previous_queries'])}\n"
                        f"Поточна тема: {self.agent.conversation_context['last_topic']}"
                    )
                    current_steps.append(context_info)
                    current_steps.append("\n✅ Обробку завершено успішно!")
                    
                    yield chat_history, "\n".join(current_steps), response
                    
                except Exception as e:
                    logger.error(f"Помилка відповіді: {e}")
                    error_msg = f"Помилка: {str(e)}"
                    chat_history[-1] = (message, error_msg)
                    current_steps.append(f"\n❌ Помилка: {error_msg}")
                    yield chat_history, "\n".join(current_steps), None

            def save_response(response):
                """Збереження відповіді у файл"""
                if response:
                    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                    filename = f"research_report_{timestamp}.md"
                    file_path = Path("saved_reports") / filename
                    file_path.parent.mkdir(exist_ok=True)
                    with open(file_path, "w", encoding="utf-8") as f:
                        f.write(str(response))
                    return str(file_path)
                return None

            def clear_chat():
                """Очищення чату"""
                return [], "", "", None

            def clear_chat_and_context():
                """Очищення чату та контексту"""
                self.agent.clear_context()
                return [], "", "", None

            # Налаштування обробників подій
            submit_btn.click(
                fn=process_steps_generator,
                inputs=[text_input, chatbot],
                outputs=[chatbot, steps_log, response_text],
                queue=True
            )

            model_dropdown.change(
                fn=change_model,
                inputs=[model_dropdown],
                outputs=[current_model]
            )

            save_btn.click(
                fn=save_response,
                inputs=[response_text],
                outputs=[file_output]
            )
            
            clear_btn.click(
                fn=clear_chat,
                outputs=[chatbot, steps_log, response_text, file_output],
                queue=False
            )

            clear_context_btn.click(
                fn=clear_chat_and_context,
                outputs=[chatbot, steps_log, response_text, file_output],
                queue=False
            )

        return interface

    def launch(self, **kwargs):
        """Запуск Gradio інтерфейсу"""
        interface = self.build_interface()
        interface.queue()
        interface.launch(**kwargs)