interface.py

import gradio as gr
from typing import List
import json
from enum import Enum

from analysis import ModelProvider, ModelName, PrecedentAnalysisWorkflow
from generation import GenerationProvider, GenerationModelName, generate_legal_position
from utils import extract_court_decision_text, get_links_html, get_links_html_lp
from search import search_with_ai_action


def create_gradio_interface():
    def update_generation_model_choices(provider):
        if provider == GenerationProvider.OPENAI.value:
            return gr.Dropdown(choices=[m.value for m in GenerationModelName if m.value.startswith("ft")])
        else:
            return gr.Dropdown(choices=[m.value for m in GenerationModelName if m.value.startswith("gemini")])

    def update_analysis_model_choices(provider):
        if provider == ModelProvider.OPENAI.value:
            return gr.Dropdown(choices=[m.value for m in ModelName if m.value.startswith("gpt")])
        else:
            return gr.Dropdown(choices=[m.value for m in ModelName if m.value.startswith("claude")])

    async def generate_position_action(url, provider, model_name, comment_input):
        try:
            court_decision_text = extract_court_decision_text(url)
            legal_position_json = generate_legal_position(court_decision_text, comment_input, provider, model_name)
            position_output_content = (
                f"**Короткий зміст позиції суду за введеним рішенням (модель: {model_name}):**\n"
                f"*{legal_position_json['title']}*: \n"
                f"{legal_position_json['text']} "
                f"**Категорія:** \n{legal_position_json['category']} "
                f"({legal_position_json['proceeding']})\n\n"
            )
            return position_output_content, legal_position_json
        except Exception as e:
            return f"Error during position generation: {str(e)}", None

    async def analyze_action(legal_position_json, question, nodes, provider, model_name):
        try:
            workflow = PrecedentAnalysisWorkflow(
                provider=ModelProvider(provider),
                model_name=ModelName(model_name)
            )

            query = (
                f"{legal_position_json['title']}: "
                f"{legal_position_json['text']}: "
                f"{legal_position_json['proceeding']}: "
                f"{legal_position_json['category']}"
            )

            response_text = await workflow.run(
                query=query,
                question=question,
                nodes=nodes
            )

            output = f"**Аналіз ШІ (модель: {model_name}):**\n{response_text}\n\n"
            output += "**Наявні в базі Правові Позицій Верховного Суду:**\n\n"

            analysis_lines = response_text.split('\n')
            for line in analysis_lines:
                if line.startswith('* ['):
                    index = line[3:line.index(']')]
                    node = nodes[int(index) - 1]
                    source_node = node.node

                    source_title = source_node.metadata.get('title', 'Невідомий заголовок')
                    source_text_lp = node.text
                    doc_ids = source_node.metadata.get('doc_id')
                    lp_id = source_node.metadata.get('lp_id')

                    links = get_links_html(doc_ids)
                    links_lp = get_links_html_lp(lp_id)

                    output += f"[{index}]: *{source_title}* | {source_text_lp} | {links_lp} | {links}\n\n"

            return output

        except Exception as e:
            return f"Error during analysis: {str(e)}"

    with gr.Blocks() as app:
        gr.Markdown("# Аналізатор релевантних Правових Позицій Верховного Суду для нового судового рішення")

        with gr.Row():
            comment_input = gr.Textbox(label="Коментар до формування короткого змісту судового рішення:")
            url_input = gr.Textbox(label="URL судового рішення:")
            question_input = gr.Textbox(label="Уточнююче питання для аналізу:")

        with gr.Row():
            # Провайдер для генерування
            generation_provider_dropdown = gr.Dropdown(
                choices=[p.value for p in GenerationProvider],
                value=GenerationProvider.GEMINI.value,
                label="Провайдер AI для генерування",
            )
            generation_model_dropdown = gr.Dropdown(
                choices=[m.value for m in GenerationModelName if m.value.startswith("gemini")],
                value=GenerationModelName.GEMINI_FLASH.value,
                label="Модель для генерування",
            )

        with gr.Row():
            # Провайдер для аналізу
            analysis_provider_dropdown = gr.Dropdown(
                choices=[p.value for p in ModelProvider],
                value=ModelProvider.OPENAI.value,
                label="Провайдер AI для аналізу",
            )
            analysis_model_dropdown = gr.Dropdown(
                choices=[m.value for m in ModelName if m.value.startswith("gpt")],
                value=ModelName.GPT4o_MINI.value,
                label="Модель для аналізу",
            )

        with gr.Row():
            generate_position_button = gr.Button("Генерувати короткий зміст позиції суду")
            search_with_ai_button = gr.Button("Пошук", interactive=False)
            analyze_button = gr.Button("Аналіз", interactive=False)

        position_output = gr.Markdown(label="Короткий зміст позиції суду за введеним рішенням")
        search_output = gr.Markdown(label="Результат пошуку")
        analysis_output = gr.Markdown(label="Результат аналізу")

        state_lp_json = gr.State()
        state_nodes = gr.State()

        # Підключення функцій до кнопок та подій
        generate_position_button.click(
            fn=generate_position_action,
            inputs=[url_input, generation_provider_dropdown, generation_model_dropdown, comment_input],
            outputs=[position_output, state_lp_json]
        ).then(
            fn=lambda: gr.update(interactive=True),
            inputs=None,
            outputs=search_with_ai_button
        )

        search_with_ai_button.click(
            fn=search_with_ai_action,
            inputs=state_lp_json,
            outputs=[search_output, state_nodes]
        ).then(
            fn=lambda: gr.update(interactive=True),
            inputs=None,
            outputs=analyze_button
        )

        analyze_button.click(
            fn=analyze_action,
            inputs=[state_lp_json, question_input, state_nodes, analysis_provider_dropdown, analysis_model_dropdown],
            outputs=analysis_output
        )

        # Оновлення списків моделей при зміні провайдера
        generation_provider_dropdown.change(
            fn=update_generation_model_choices,
            inputs=generation_provider_dropdown,
            outputs=generation_model_dropdown
        )

        analysis_provider_dropdown.change(
            fn=update_analysis_model_choices,
            inputs=analysis_provider_dropdown,
            outputs=analysis_model_dropdown
        )

    return app