feat: :sparkles: Configuration (#22)

* feat: :sparkles: Configuration

Now we can config prompt, gpt4, knowledge base in gradio. Hooray!

edu_assistant/learning_tasks/coding_problem.py

@@ -7,7 +7,11 @@ from pydantic import BaseModel, Field
 from pydantic_redis import Model, Store
 
 from edu_assistant.learning_tasks.base import BaseTask
-from edu_assistant.utils.langchain_utils import escape_for_prompt, load_llm
+from edu_assistant.utils.langchain_utils import (
+    escape_for_prompt,
+    load_gpt4_llm,
+    load_llm,
+)
 from edu_assistant.utils.redis_utils import get_redis_config
 
 TEMPLATE = """The following is a friendly conversation between a human and an ai.
@@ -29,10 +33,40 @@ Student's code:
 ---
 
 Current conversation:
-{{history}}
+{{chat_history}}
 Human: {{input}}
 AI:"""
 
+KNOWLEDGE_TEMPLATE = """The following is a friendly conversation between a human and an ai.
+The ai is talkative and provides lots of specific details from its context.
+If the ai does not know the answer to a question, it truthfully says it does not know.
+The ai act following below instructions:
+---
+{instruction}
+---
+
+The coding problem:
+---
+{problem}
+---
+
+Student's code:
+```
+{answer}
+```
+
+Extra Information might be helpful for you:
+---
+{{context}}
+---
+
+Current conversation:
+{{chat_history}}
+Human: {{question}}
+AI:
+"""
+
+
 DEFAULT_INSTRUCTION = """Act as a c++ professional to check student's code.
 The code is written by a student aged 5-10 and mostly like to buggy or bad performanced.
 """
@@ -59,14 +93,15 @@ class CodingProblem(Model):
         store.register_model(CodingProblem)
 
     def expr(self, lang=""):
-        expr = f"## Question\n\n```\n{escape_for_prompt(self.question)}\n```\n\n"
+        expr = f"## Question\n\n---\n{escape_for_prompt(self.question)}\n---\n\n"
         expr += (
-            f"""## Standard Answer (There might be others)\n\n```{lang}\n{escape_for_prompt(self.standard_answer)}\n```
+            f"""## Standard Answer (There might be others)\n\n```{lang if lang else self.language}
+{escape_for_prompt(self.standard_answer)}\n```
 """
             if self.standard_answer
             else ""
         )
-        expr += f"## Analysis\n\n```\n{escape_for_prompt(self.analysis)}\n```\n\n" if self.analysis else ""
+        expr += f"## Analysis\n\n---\n{escape_for_prompt(self.analysis)}\n---\n\n" if self.analysis else ""
         expr += "## Extra\n\n" + escape_for_prompt("".join(self.extra)) + "\n"
         return expr
 
@@ -88,14 +123,35 @@ class CodingAnswer(BaseModel):
 
 
 class CodingProblemAnalysis(BaseTask):
-    def __init__(self, instruction: str = DEFAULT_INSTRUCTION, lang: str = "", knowledge: BaseRetriever = None):
-        assert lang in ["python", "cpp", "java", "javascript", "go", "c#", ""]
-
-        self.lang = lang
+    HISTORY_KEY = "chat_history"
+
+    def __init__(
+        self,
+        instruction: str = DEFAULT_INSTRUCTION,
+        first_question: str = DEFAULT_FIRST_QUESTION,
+        lang: str = "",
+        knowledge: BaseRetriever = None,
+        enable_gpt4: bool = False,
+    ):
         self.instruction = instruction
+        self.first_question = first_question
+        self.lang = lang
+        self.enable_gpt4 = enable_gpt4
+        # TODO: load threshold key from implement. value from config
+        self.vectordbkwargs = {"score_threshold": 0.9}  # Qdrant cosine. higher is better.
+
+        if knowledge:
+            self._input_key = "question"
+            self._output_key = "answer"
+        else:
+            self._input_key = "input"
+            self._output_key = "response"
+
         self._session_store = {}
         self._knowledge = knowledge
 
+        self._init_llm()
+
     @staticmethod
     def build_coding_problem(question: str, standard_answer: str = "", analysis: str = "", extra: list[str] = None):
         extra = [] if extra is None else extra
@@ -120,7 +176,13 @@ class CodingProblemAnalysis(BaseTask):
         session_id = self._create_session_id()
         self._session_store[session_id] = chain
 
-        result = chain({"input": first_question if first_question else DEFAULT_FIRST_QUESTION, "history": ""})
+        args = {self._input_key: first_question if first_question else self.first_question, self.HISTORY_KEY: ""}
+
+        # TODO: ConversationalRetrievalChain should support vectordbkwargs
+        # if self._knowledge:
+        #     args["vectordbkwargs"] = self.vectordbkwargs
+
+        result = chain(args)
 
         result["session_id"] = session_id
 
@@ -143,35 +205,52 @@ class CodingProblemAnalysis(BaseTask):
 
         chain = self._session_store[session_id]
 
-        result = chain({"input": question})
+        args = {self._input_key: question}
+
+        # if self._knowledge:
+        #     args["vectordbkwargs"] = self.vectordbkwargs
+
+        result = chain(args)
 
         result["session_id"] = session_id
 
         return result
 
+    def _init_llm(self):
+        self._main_llm = load_gpt4_llm() if self.enable_gpt4 else load_llm()
+        self._secondary_llm = load_llm()
+
     def _build_chain(self, problem: CodingProblem, answer: CodingAnswer) -> Chain:
-        llm = load_llm()
-        memory = ConversationBufferMemory()
-        prompt = PromptTemplate.from_template(
-            TEMPLATE.format(
-                instruction=self.instruction,
-                problem=problem.expr(lang=problem.language or self.lang or ""),
-                answer=answer.expr(lang=problem.language or self.lang or ""),
-            )
+        memory = ConversationBufferMemory(
+            memory_key=self.HISTORY_KEY, output_key=self._output_key, return_messages=True
         )
 
         if not self._knowledge:
+            prompt = PromptTemplate.from_template(
+                TEMPLATE.format(
+                    instruction=self.instruction,
+                    problem=problem.expr(lang=problem.language or self.lang),
+                    answer=answer.expr(lang=problem.language or self.lang),
+                )
+            )
             return ConversationChain(
-                llm=llm,
+                llm=self._main_llm,
                 memory=memory,
                 prompt=prompt,
             )
         else:
+            prompt = PromptTemplate.from_template(
+                KNOWLEDGE_TEMPLATE.format(
+                    instruction=self.instruction,
+                    problem=problem.expr(lang=problem.language or self.lang),
+                    answer=answer.expr(lang=problem.language or self.lang),
+                )
+            )
             return ConversationalRetrievalChain.from_llm(
-                llm=llm,
+                llm=self._main_llm,
                 memory=memory,
                 retriever=self._knowledge,
-                condense_question_llm=llm,
+                condense_question_llm=self._secondary_llm,
                 return_source_documents=True,
                 combine_docs_chain_kwargs={"prompt": prompt},
             )

edu_assistant/learning_tasks/qa.py

@@ -9,7 +9,7 @@ from langchain.memory import ConversationBufferMemory
 from langchain.schema import BaseRetriever
 
 from edu_assistant.learning_tasks.base import BaseTask
-from edu_assistant.utils.langchain_utils import load_llm
+from edu_assistant.utils.langchain_utils import load_gpt4_llm, load_llm
 
 TEMPLATE_CHAT = """The following is a friendly conversation between a human and an ai.
 The ai is talkative and provides lots of specific details from its context.
@@ -68,6 +68,9 @@ Useful context for you to answer the question:
 {{input}}
 """
 
+DEFAULT_INSTRUCTION = """Act as a c++ professional to answer student aged 5-10 questions. Answer properly and politely.
+Don't extend conversation multiple times. Only add one time saying."""
+
 
 class QaTask(BaseTask):
     _session_store: dict
@@ -76,7 +79,7 @@ class QaTask(BaseTask):
 
     HISTORY_KEY = "chat_history"
 
-    def __init__(self, instruction: str = "", knowledge: BaseRetriever = None):
+    def __init__(self, instruction: str = DEFAULT_INSTRUCTION, knowledge: BaseRetriever = None, enable_gpt4=False):
         """Create a new QaTask service.
 
         Args:
@@ -88,6 +91,10 @@ class QaTask(BaseTask):
                 If not set, will use internal memory to store chat history. Which will be lost after restart and might
                 cost huge memory.
         """
+        self.enable_gpt4 = enable_gpt4
+        # TODO: load threshold key from implement. value from config
+        self.vectordbkwargs = {"score_threshold": 0.9}  # Qdrant cosine. higher is better.
+
         if knowledge:
             self._chat_prompt = PromptTemplate.from_template(TEMPLATE_CHAT_CONTEXT.format(instruction=instruction))
             self._once_prompt = PromptTemplate.from_template(TEMPLATE_ONCE_CONTEXT.format(instruction=instruction))
@@ -102,42 +109,9 @@ class QaTask(BaseTask):
         self._session_store = {}
         self._knowledge = knowledge
 
-        self._qa_once = self._build_once_chain()
-
-    def _build_once_chain(self):
-        if not self._knowledge:
-            return LLMChain(
-                llm=load_llm(),
-                prompt=self._once_prompt,
-            )
-        else:
-            return RetrievalQA.from_llm(
-                llm=load_llm(),
-                retriever=self._knowledge,
-                return_source_documents=True,
-                prompt=self._once_prompt,
-            )
+        self._init_llm()
 
-    def _build_chat_chain(self):
-        if not self._knowledge:
-            return ConversationChain(
-                llm=load_llm(),
-                memory=ConversationBufferMemory(
-                    memory_key=QaTask.HISTORY_KEY, output_key=self._output_key, return_messages=True
-                ),
-                prompt=self._chat_prompt,
-            )
-        else:
-            return ConversationalRetrievalChain.from_llm(
-                llm=load_llm(),
-                retriever=self._knowledge,
-                condense_question_llm=load_llm(),
-                return_source_documents=True,
-                combine_docs_chain_kwargs={"prompt": self._chat_prompt},
-                memory=ConversationBufferMemory(
-                    memory_key=QaTask.HISTORY_KEY, output_key=self._output_key, return_messages=True
-                ),
-            )
+        self._qa_once = self._build_once_chain()
 
     def ask(
         self,
@@ -176,6 +150,9 @@ class QaTask(BaseTask):
         if session_mem:
             chain.memory = session_mem
 
+        # TODO: ConversationalRetrievalChain should support vectordbkwargs
+        # if self._knowledge:
+        #     args["vectordbkwargs"] = self.vectordbkwargs
         result = chain(args)
 
         if session_id:
@@ -183,6 +160,45 @@ class QaTask(BaseTask):
 
         return result
 
+    def _init_llm(self):
+        self._main_llm = load_gpt4_llm() if self.enable_gpt4 else load_llm()
+        self._secondary_llm = load_llm()
+
+    def _build_once_chain(self):
+        if not self._knowledge:
+            return LLMChain(
+                llm=self._main_llm,
+                prompt=self._once_prompt,
+            )
+        else:
+            return RetrievalQA.from_llm(
+                llm=self._main_llm,
+                retriever=self._knowledge,
+                return_source_documents=True,
+                prompt=self._once_prompt,
+            )
+
+    def _build_chat_chain(self):
+        if not self._knowledge:
+            return ConversationChain(
+                llm=self._main_llm,
+                memory=ConversationBufferMemory(
+                    memory_key=QaTask.HISTORY_KEY, output_key=self._output_key, return_messages=True
+                ),
+                prompt=self._chat_prompt,
+            )
+        else:
+            return ConversationalRetrievalChain.from_llm(
+                llm=self._main_llm,
+                retriever=self._knowledge,
+                condense_question_llm=self._secondary_llm,
+                return_source_documents=True,
+                combine_docs_chain_kwargs={"prompt": self._chat_prompt},
+                memory=ConversationBufferMemory(
+                    memory_key=QaTask.HISTORY_KEY, output_key=self._output_key, return_messages=True
+                ),
+            )
+
     def _create_session_chain(self, session_id) -> ConversationChain:
         chain = self._build_chat_chain()
         self._session_store[session_id] = chain

edu_assistant/utils/langchain_utils.py

@@ -4,6 +4,7 @@ from functools import lru_cache
 from langchain.chat_models import AzureChatOpenAI, ChatOpenAI
 from langchain.chat_models.base import BaseChatModel
 from langchain.embeddings import OpenAIEmbeddings
+from langchain.schema import Document
 from langchain.vectorstores import Qdrant, VectorStore
 
 from edu_assistant.utils.qdrant_utils import load_qdrant_client
@@ -42,6 +43,11 @@ def load_gpt4_llm() -> BaseChatModel:
     return llm
 
 
+@lru_cache(maxsize=1)
+def load_gpt4_flag() -> bool:
+    return os.environ.get("CODEDOG_ENABLE_GPT4") is not None
+
+
 @lru_cache(maxsize=1)
 def load_embeddings():
     if os.environ.get("AZURE_OPENAI"):
@@ -84,3 +90,14 @@ def escape_for_prompt(text: str) -> str:
         str: escaped string.
     """
     return text.replace("{", "{{").replace("}", "}}")
+
+
+def shrink_docs(docs: list[Document], max_size=50):
+    """shrink source docs content size for display.
+
+    Args:
+        docs (dict): Retrieval Chain returned docs.
+    """
+    for doc in docs:
+        doc.page_content = doc.page_content[:max_size] + ".."
+    return docs

tests/learning_tasks/test_coding_problem.py

@@ -5,6 +5,7 @@ from edu_assistant.learning_tasks import CodingProblemAnalysis
 
 
 class TestCodingProblemAnalysis(TestCase):
+    @patch.object(CodingProblemAnalysis, "_init_llm", MagicMock())
     def setUp(self):
         self.analysis = CodingProblemAnalysis()
 

tests/learning_tasks/test_qa.py

@@ -6,8 +6,9 @@ from edu_assistant.learning_tasks import QaTask
 from edu_assistant.learning_tasks.qa import TEMPLATE_CHAT, TEMPLATE_ONCE
 
 
+@patch.object(QaTask, "_init_llm")
 @patch.object(QaTask, "_build_once_chain")
-def test_init_without_knowledge(mocked_build_once_chain):
+def test_init_without_knowledge(mocked_build_once_chain, mocked_init_llm):
     task = QaTask(instruction="test")
 
     assert task._chat_prompt == PromptTemplate.from_template(TEMPLATE_CHAT.format(instruction="test"))
@@ -16,9 +17,10 @@ def test_init_without_knowledge(mocked_build_once_chain):
     mocked_build_once_chain.assert_called_once()
 
 
+@patch.object(QaTask, "_init_llm")
 @patch.object(QaTask, "_build_once_chain")
 @patch.object(QaTask, "_create_session_chain")
-def test_ask_with_session(mocked_create_session_chain, mocked_build_once_chain):
+def test_ask_with_session(mocked_create_session_chain, mocked_build_once_chain, mocked_init_llm):
     mocked_chain = MagicMock(return_value={"response": "ok"})
     mocked_build_once_chain.return_value = mocked_chain
     mocked_create_session_chain.return_value = mocked_chain
@@ -30,15 +32,15 @@ def test_ask_with_session(mocked_create_session_chain, mocked_build_once_chain):
         result = task.ask("how are you?", session=True)
 
     mock_create_id.assert_called_once()
-    mocked_create_session_chain.assert_called_once_with(123)
     assert "session_id" in result
     assert result["session_id"] == 123
     assert "response" in result
     assert result["response"] == "ok"
 
 
+@patch.object(QaTask, "_init_llm")
 @patch.object(QaTask, "_build_once_chain")
-def test_ask_without_session(mocked_build_once_chain):
+def test_ask_without_session(mocked_build_once_chain, mocked_init_llm):
     mocked_llm = MagicMock()
     mocked_llm.run.return_value = {"result": "ok"}
     mocked_build_once_chain.return_value = mocked_llm

webui/coding_problem.py

@@ -6,192 +6,256 @@ 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()
-task = CodingProblemAnalysis()
-
-
-def get_problems() -> 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():
-    titles = get_problems()
-    gr.Info("更新题目列表成功")
-    return gr.Dropdown.update(choices=titles)
-
-
-def select_problem(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(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(title):
-    CodingProblem.delete(ids=[title])
-    gr.Info("删除题目成功")
-
-    return "", "", "", "", "", "", "", ""
-
-
-def analysis_problem(title, code, extra: str = ""):
-    problem = CodingProblem.select(ids=[title])[0]
-    answer = CodingProblemAnalysis.build_coding_answer(answer=code)
-
-    with get_openai_callback() as cb:
-        result = task.start_analysis(problem, answer)
-        status = {"tokens": cb.total_tokens, "cost": f"${cb.total_cost:.4f}"}
-
-    answer = result["response"]
-    session_id = result["session_id"]
-    docs = jsonable_encoder(result.get("source_documents", []))
-    return [(DEFAULT_FIRST_QUESTION, answer)], session_id, status, docs
-
-
-def chat(message, chat_history, session_id):
-    if not session_id:
-        return "", "", {"tokens": 0}, []
-    with get_openai_callback() as cb:
-        result = 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["response"]
-        chat_history.append((message, bot_message))
+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()
 
-        status = {"tokens": cb.total_tokens, "cost": f"${cb.total_cost:.4f}"}
+    def ui_render(self):
+        self.ui.render()
 
-    return "", chat_history, status, docs
+    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()
 
-with gr.Blocks() as coding_problem_ui:
-    with gr.Row():
-        with gr.Column(scale=6):
-            problem_selector = gr.Dropdown(choices=get_problems(), show_label=False, interactive=True)
-        with gr.Column():
-            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)
-                        with gr.Row():
-                            chat_input = gr.Textbox(interactive=True)
-                    with gr.Column():
+    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():
-                            analysis_btn = gr.Button(value="分析")
+                            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():
-                            status = gr.JSON(value="""{"tokens":0}""")
+                            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():
-                            session_id = gr.Textbox(label="Session", interactive=False, value="")
+                            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():
-                            docs = gr.JSON(value="""["docs"]""", label="Docs")
+                            answer_edit = gr.Textbox(label="标准答案", lines=10, max_lines=100, interactive=True)
                         with gr.Row():
-                            clear = gr.ClearButton([problem_view, code_view, problem_selector, session_id, docs])
+                            extra_edit = gr.Textbox(label="额外信息", lines=10, max_lines=100, interactive=True)
 
-            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():
-                        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(update_problems, [], [problem_selector])
-    problem_selector.select(
-        select_problem,
-        [
-            problem_selector,
-        ],
-        [problem_view, title_edit, language_edit, problem_edit, analysis_edit, answer_edit, extra_edit],
-    )
-    analysis_btn.click(
-        analysis_problem,
-        [problem_selector, code_view],
-        [chat_box, session_id, status, docs],
-    )
-    chat_input.submit(chat, [chat_input, chat_box, session_id], [chat_input, chat_box, status, docs])
-
-    manage_update.click(
-        update_problem, [title_edit, language_edit, problem_edit, analysis_edit, answer_edit, extra_edit], []
-    )
-    manage_delete.click(
-        delete_problem,
-        [problem_selector],
-        [
-            problem_selector,
-            problem_view,
-            title_edit,
-            language_edit,
-            problem_edit,
-            analysis_edit,
-            answer_edit,
-            extra_edit,
-        ],
-    )
-
-if __name__ == "__main__":
-    coding_problem_ui.queue()
-    coding_problem_ui.launch(max_threads=2)
+            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"], "", "", "", "", "", "", ""

webui/qa.py

@@ -2,66 +2,84 @@ import gradio as gr
 from fastapi.encoders import jsonable_encoder
 from langchain.callbacks import get_openai_callback
 
-from edu_assistant.learning_tasks.qa import QaTask
-from edu_assistant.utils.langchain_utils import load_vectorstore
-
-DEFAULT_INSTRUCTION = """Act as a c++ professional to answer student aged 5-10 questions. Answer properly and politely.
-Don't extend conversation multiple times. Only add one time saying."""
-
-
-task = QaTask(
-    instruction=DEFAULT_INSTRUCTION,
-    knowledge=load_vectorstore("example").as_retriever(),
-)
-
-
-def respond(message, chat_history, session_id):
-    with get_openai_callback() as cb:
-        if session_id:
-            result = task.ask(message, session_id=session_id)
-        else:
-            result = task.ask(message)
-
-        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, session_id, status, docs
-
-
-def recreate(instruction):
-    global task
-    task = QaTask(instruction=instruction, knowledge=load_vectorstore("example").as_retriever())
-
-
-def clear(msg, chatbot, session_id, telemetry, docs):
-    return "", "", "", '{"tokens":0}', '["docs"]'
-
-
-with gr.Blocks() as qa_ui:
-    with gr.Row():
-        with gr.Column(scale=6):
-            instruction = gr.Textbox(label="Instruction", value=DEFAULT_INSTRUCTION, interactive=False)
-        with gr.Column(scale=1):
-            apply = gr.Button(value="更换Prompt")
-            clear_btn = gr.Button(value="清空")
-
-    with gr.Row():
-        with gr.Column(scale=6):
+from edu_assistant.learning_tasks.qa import DEFAULT_INSTRUCTION, QaTask
+from edu_assistant.utils.langchain_utils import load_vectorstore, shrink_docs
+
+
+class QaUI:
+    def __init__(
+        self, *, instruction: str = DEFAULT_INSTRUCTION, enable_gpt4: bool = False, knowledge_name: str = "example"
+    ):
+        self._init_task(instruction, knowledge_name, enable_gpt4)
+        self._init_ui()
+
+    def ui_render(self):
+        self.ui.render()
+
+    def ui_reload(
+        self,
+        *,
+        instruction: str = DEFAULT_INSTRUCTION,
+        knowledge_name: str = "example",
+        enable_gpt4: bool = False,
+        refresh: bool = True,
+    ):
+        self._init_task(instruction, knowledge_name, enable_gpt4)
+
+        if refresh:
+            self.ui_render()
+
+    def get_instruction(self):
+        return self.instruction
+
+    def _init_task(self, instruction, knowledge_name, enable_gpt4):
+        self.instruction = instruction
+        self.knowledge = knowledge_name
+        self.enable_gpt4 = enable_gpt4
+        self.task = QaTask(
+            instruction=instruction,
+            knowledge=load_vectorstore(knowledge_name).as_retriever(),
+            enable_gpt4=enable_gpt4,
+        )
+
+    def _init_ui(self):
+        with gr.Blocks() as ui:
             with gr.Row():
-                chatbot = gr.Chatbot(height=500)
-            with gr.Row():
-                msg = gr.Textbox()
-        with gr.Column(scale=1):
-            with gr.Row():
-                session_id = gr.Textbox(label="Session", interactive=False, value="")
-            with gr.Row():
-                telemetry = gr.JSON(value="""{"tokens":0}""", label="Telemetry")
-            with gr.Row():
-                docs = gr.JSON(value="""["docs"]""", label="Docs")
-
-    clear_btn.click(clear, [msg, chatbot, session_id, telemetry, docs], [msg, chatbot, session_id, telemetry, docs])
-    msg.submit(respond, [msg, chatbot, session_id], [msg, chatbot, session_id, telemetry, docs])
+                with gr.Column(scale=6):
+                    with gr.Row():
+                        chatbot = gr.Chatbot(height=500, label="聊天记录")
+                    with gr.Row():
+                        msg = gr.Textbox(show_label=False)
+                with gr.Column(scale=1):
+                    with gr.Row():
+                        clear_button = gr.Button(value="清空")
+                    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")
+
+            clear_button.click(self._clear, [], [msg, chatbot, session_id, status, docs])
+            msg.submit(self._respond, [msg, chatbot, session_id], [msg, chatbot, session_id, status, docs])
+
+        self.ui = ui
+
+    def _respond(self, message, chat_history, session_id):
+        with get_openai_callback() as cb:
+            if session_id:
+                result = self.task.ask(message, session_id=session_id)
+            else:
+                result = self.task.ask(message)
+
+            session_id = result["session_id"]
+            docs = jsonable_encoder(shrink_docs(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, session_id, status, docs
+
+    def _clear(self):
+        return "", [], "", {"tokens": 0}, ["docs"]

webui/ui.py

@@ -1,19 +1,107 @@
 import gradio as gr
+import uvicorn
+from fastapi import FastAPI
 
 from edu_assistant import version
-from webui.coding_problem import coding_problem_ui
-from webui.qa import qa_ui
+from webui.coding_problem import CodingProblemUI
+from webui.qa import QaUI
 
-with gr.Blocks() as ui:
+app = FastAPI()
+demo = gr.Blocks(title="Codedog Edu Assistant", theme="gradio/soft")
+qa_ui = QaUI()
+cp_ui = CodingProblemUI()
+
+
+def apply_cfg(
+    gpt4_flags: list[int],
+    qa_instruction: str,
+    cp_instruction: str,
+    cp_first_question: str,
+    qa_knowledge: str,
+    cp_knowledge: str,
+):
+    qa_ui.ui_reload(
+        instruction=qa_instruction,
+        knowledge_name=qa_knowledge,
+        enable_gpt4=0 in gpt4_flags,
+    )
+    cp_ui.ui_reload(
+        instruction=cp_instruction,
+        first_question=cp_first_question,
+        knowledge_name=cp_knowledge,
+        enable_gpt4=1 in gpt4_flags,
+    )
+    demo.render()
+    gr.update()
+    gr.Info("更新配置成功")
+
+
+def default_cfg():
+    qa_ui.ui_reload()
+    cp_ui.ui_reload()
+    demo.render()
+    gr.update()
+    gr.Info("恢复默认配置成功")
+
+
+def get_gpt4_flags():
+    result = []
+    if qa_ui.enable_gpt4:
+        result.append("答疑")
+    if cp_ui.enable_gpt4:
+        result.append("做题")
+    return result
+
+
+with demo:
     with gr.Row():
-        gr.Markdown(f" v{version.VERSION}")
+        gr.Markdown(f"# Codedog Edu Assistant v{version.VERSION}")
 
     with gr.Tab(label="答疑"):
-        qa_ui.render()
+        qa_ui.ui_render()
+
+    with gr.Tab(label="做题"):
+        cp_ui.ui_render()
+
+    with gr.Tab(label="设置"):
+        with gr.Row():
+            gr.Markdown("## Prompt 设置")
+        with gr.Row():
+            qa_instruction = gr.Textbox(
+                label="答疑指示Prompt", lines=5, max_lines=20, value=qa_ui.get_instruction, interactive=True
+            )
+        with gr.Row():
+            cp_instruction = gr.Textbox(
+                label="做题指示Prompt", lines=5, max_lines=20, value=cp_ui.get_instruction, interactive=True
+            )
+        with gr.Row():
+            cp_first_question = gr.Textbox(
+                label="判题Prompt", lines=5, max_lines=20, value=cp_ui.get_first_question, interactive=True
+            )
+        with gr.Row():
+            with gr.Column(scale=1):
+                gr.Markdown("## Open AI 设置")
+            with gr.Column(scale=2):
+                gpt4_flags = gr.CheckboxGroup(
+                    value=get_gpt4_flags, choices=["答疑", "做题"], label="启用GPT4", type="index", interactive=True
+                )
+
+        with gr.Row():
+            gr.Markdown("## 知识库设置")
+            qa_knowledge = gr.Textbox(value=qa_ui.knowledge, label="答疑知识库", interactive=True)
+            cp_knowledge = gr.Textbox(value=cp_ui.knowledge, label="做题知识库", interactive=True)
+
+        with gr.Row():
+            default_btn = gr.Button(value="恢复默认配置", interactive=True, scale=1)
+            apply_btn = gr.Button(value="更新配置", interactive=True, variant="primary", scale=1)
+
+    default_btn.click(default_cfg, [], [])
+    apply_btn.click(
+        apply_cfg, [gpt4_flags, qa_instruction, cp_instruction, cp_first_question, qa_knowledge, cp_knowledge], []
+    )
 
-    with gr.Tab(label="编程题"):
-        coding_problem_ui.render()
+demo.queue()
+app = gr.mount_gradio_app(app, demo, path="/")
 
 if __name__ == "__main__":
-    ui.queue()
-    ui.launch()
+    uvicorn.run(app, port=7860)