webui/coding_problem.py

import json

import gradio as gr
from fastapi.encoders import jsonable_encoder
from langchain.callbacks import get_openai_callback

from edu_assistant.learning_tasks.coding_problem import (
    DEFAULT_FIRST_QUESTION,
    DEFAULT_INSTRUCTION,
    CodingProblem,
    CodingProblemAnalysis,
)
from edu_assistant.utils.langchain_utils import load_vectorstore, shrink_docs

CodingProblem.enable_redis_orm()


class CodingProblemUI:
    def __init__(
        self,
        *,
        instruction: str = DEFAULT_INSTRUCTION,
        first_question: str = DEFAULT_FIRST_QUESTION,
        knowledge_name: str = "example",
        enable_gpt4: bool = False,
    ):
        self._init_task(instruction, first_question, knowledge_name, enable_gpt4)
        self._init_ui()

    def ui_render(self):
        self.ui.render()

    def ui_reload(
        self,
        *,
        instruction: str = DEFAULT_INSTRUCTION,
        first_question: str = DEFAULT_FIRST_QUESTION,
        knowledge_name: str = "example",
        enable_gpt4: bool = False,
        refresh: bool = True,
    ):
        self._init_task(instruction, first_question, knowledge_name, enable_gpt4)

        if refresh:
            self.ui_render()

    def get_instruction(self):
        return self.instruction

    def get_first_question(self):
        return self.first_question

    def _init_task(self, instruction: str, first_question: str, knowledge_name: str, enable_gpt4: bool):
        self.instruction = instruction
        self.first_question = first_question
        self.knowledge = knowledge_name
        self.enable_gpt4 = enable_gpt4
        self.task = CodingProblemAnalysis(
            instruction=instruction,
            first_question=first_question,
            knowledge=load_vectorstore(knowledge_name).as_retriever(),
            enable_gpt4=enable_gpt4,
        )

    def _init_ui(self):
        self.ui = gr.Blocks()
        with self.ui:
            with gr.Row():
                with gr.Column(scale=6):
                    problem_selector = gr.Dropdown(choices=self._get_problems(), show_label=False, interactive=True)
                with gr.Column(scale=1):
                    refresh_btn = gr.Button(value="刷新")
            with gr.Row():
                with gr.Column(scale=6):
                    with gr.Tab(label="错误代码分析"):
                        with gr.Row():
                            with gr.Column(scale=3):
                                with gr.Row():
                                    problem_view = gr.Markdown(label="题目")
                                with gr.Row():
                                    code_view = gr.Textbox(label="代码", lines=10, interactive=True)
                            with gr.Column(scale=3):
                                with gr.Row():
                                    chat_box = gr.Chatbot(height=500, label="聊天记录")
                                with gr.Row():
                                    chat_input = gr.Textbox(show_label=False)
                            with gr.Column():
                                with gr.Row():
                                    analysis_btn = gr.Button(value="分析")
                                with gr.Row():
                                    clear = gr.ClearButton()
                                with gr.Row():
                                    session_id = gr.Textbox(label="Session", interactive=False, value="")
                                with gr.Row():
                                    status = gr.JSON(value={"tokens": 0}, label="Status")
                                with gr.Row():
                                    docs = gr.JSON(value=["docs"], label="Docs")

                    with gr.Tab(label="题库管理"):
                        with gr.Row():
                            with gr.Column(scale=6):
                                with gr.Row():
                                    title_edit = gr.Textbox(label="标题", interactive=True)
                                with gr.Row():
                                    language_edit = gr.Dropdown(
                                        choices=["python", "cpp", "java"],
                                        label="语言",
                                        interactive=True,
                                        allow_custom_value=True,
                                    )
                            with gr.Column(scale=1):
                                manage_update = gr.Button(value="更新")
                                manage_delete = gr.Button(value="删除", variant="stop")
                        with gr.Row():
                            with gr.Column():
                                problem_edit = gr.Textbox(label="题目", lines=10, max_lines=100, interactive=True)
                            with gr.Column():
                                analysis_edit = gr.Textbox(label="解析", lines=10, max_lines=100, interactive=True)
                        with gr.Row():
                            answer_edit = gr.Textbox(label="标准答案", lines=10, max_lines=100, interactive=True)
                        with gr.Row():
                            extra_edit = gr.Textbox(label="额外信息", lines=10, max_lines=100, interactive=True)

            refresh_btn.click(self._update_problems, [], [problem_selector])
            problem_selector.select(
                self._select_problem,
                [
                    problem_selector,
                ],
                [problem_view, title_edit, language_edit, problem_edit, analysis_edit, answer_edit, extra_edit],
            )
            analysis_btn.click(
                self._analysis_problem,
                [problem_selector, code_view],
                [chat_box, session_id, status, docs],
            )
            chat_input.submit(self._chat, [chat_input, chat_box, session_id], [chat_input, chat_box, status, docs])

            manage_update.click(
                self._update_problem,
                [title_edit, language_edit, problem_edit, analysis_edit, answer_edit, extra_edit],
                [],
            )
            manage_delete.click(
                self._delete_problem,
                [problem_selector],
                [
                    problem_selector,
                    problem_view,
                    title_edit,
                    language_edit,
                    problem_edit,
                    analysis_edit,
                    answer_edit,
                    extra_edit,
                ],
            )
            clear.click(
                self._clear,
                [],
                [
                    problem_selector,
                    problem_view,
                    code_view,
                    chat_box,
                    session_id,
                    status,
                    docs,
                    problem_view,
                    title_edit,
                    language_edit,
                    problem_edit,
                    analysis_edit,
                    answer_edit,
                    extra_edit,
                ],
            )

    def _get_problems(self) -> list[str]:
        data = CodingProblem.select(columns=["title"])
        if not data:
            return []
        titles = [problem_data["title"] for problem_data in data]
        return titles

    def _update_problems(self):
        titles = self._get_problems()
        gr.Info("更新题目列表成功")
        return gr.Dropdown.update(choices=titles)

    def _select_problem(self, title: str):
        problem: CodingProblem = CodingProblem.select(ids=[title])[0]
        return (
            problem.expr(),
            problem.title,
            problem.language,
            problem.question,
            problem.analysis,
            problem.standard_answer,
            json.dumps(problem.extra, ensure_ascii=False, indent=4),
        )

    def _update_problem(self, title, language, problem, analysis, answer, extra):
        # TODO: add language
        try:
            extra_data = json.loads(extra)
        except json.JSONDecodeError:
            extra_data = [extra]

        CodingProblem.update(
            title,
            data={
                "title": title,
                "language": language,
                "question": problem,
                "analysis": analysis,
                "standard_answer": answer,
                "extra": extra_data,
            },
        )
        gr.Info("更新题目成功")

    def _delete_problem(self, title):
        CodingProblem.delete(ids=[title])
        gr.Info("删除题目成功")

        return "", "", "", "", "", "", "", ""

    def _analysis_problem(self, title, code, extra: str = ""):
        problem = CodingProblem.select(ids=[title])[0]
        answer = CodingProblemAnalysis.build_coding_answer(answer=code, extra=[extra])

        with get_openai_callback() as cb:
            result = self.task.start_analysis(problem, answer)
            status = {"tokens": cb.total_tokens, "cost": f"${cb.total_cost:.4f}"}

        answer = result["answer"]
        session_id = result["session_id"]
        docs = jsonable_encoder(shrink_docs(result.get("source_documents", [])))
        return [(self.first_question, answer)], session_id, status, docs

    def _chat(self, message, chat_history, session_id):
        if not session_id:
            return "", "", {"tokens": 0}, []
        with get_openai_callback() as cb:
            result = self.task.ask(message, session_id=session_id)
            if not result:
                raise gr.Error("Session expired. Please recreate a new problem analysis session.")

            session_id = result["session_id"]
            docs = jsonable_encoder(result.get("source_documents", []))

            bot_message = result["answer"]
            chat_history.append((message, bot_message))

            status = {"tokens": cb.total_tokens, "cost": f"${cb.total_cost:.4f}"}

        return "", chat_history, status, docs

    def _clear(self):
        return "", "", "", [], "", {"tokens": 0}, ["docs"], "", "", "", "", "", "", ""