edu_assistant/learning_tasks/coding_problem.py

from langchain import ConversationChain, PromptTemplate
from langchain.chains import ConversationalRetrievalChain
from langchain.chains.base import Chain
from langchain.memory import ConversationBufferMemory
from langchain.schema import BaseRetriever
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_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.
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}
---

Current conversation:
{{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.
"""

DEFAULT_FIRST_QUESTION = "请问这段代码中有什么问题吗？"


class CodingProblem(Model):
    _primary_key_field: str = "title"
    title: str = Field()
    question: str = Field()
    standard_answer: str = Field(default="")
    analysis: str = Field(default="")
    language: str = Field(default="")
    extra: str = Field(default_factory=lambda: list())

    # TODO: Add cache to expr function with pydantic 2 computed_field decorator.
    # Wait for langchain to support pydantic2.

    @staticmethod
    def enable_redis_orm():
        store = Store(name="coding_problems", redis_config=get_redis_config(), life_span_in_seconds=3600 * 24 * 30)

        store.register_model(CodingProblem)

    def expr(self, lang=""):
        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 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 += "## Extra\n\n" + escape_for_prompt("".join(self.extra)) + "\n"
        return expr

    def __str__(self):
        return self.expr()


class CodingAnswer(BaseModel):
    answer: str = Field()
    extra: list[str] = Field(default="")

    def expr(self, lang=""):
        expr = f"Answer:\n```{lang}\n{escape_for_prompt(self.answer)}\n```\n"
        expr += escape_for_prompt("".join(self.extra)) + "\n"
        return expr

    def __str__(self):
        return self.expr()


class CodingProblemAnalysis(BaseTask):
    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
        return CodingProblem(question=question, standard_answer=standard_answer, analysis=analysis, extra=extra)

    @staticmethod
    def build_coding_answer(answer: str, extra: list[str] = None):
        extra = [] if extra is None else extra
        return CodingAnswer(answer=answer, extra=extra)

    def start_analysis(self, problem: CodingProblem, answer: CodingAnswer, first_question: str = None) -> dict:
        """start analysis of a coding problem and incorrect answer.

        Args:
            problem (CodingProblem): a coding problem
            answer (CodingAnswer): a coding problem answer

        Returns:
            dict: question answer and metadata
        """
        chain = self._build_chain(problem, answer)
        session_id = self._create_session_id()
        self._session_store[session_id] = chain

        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

        return result

    def ask(self, question: str, session_id: str) -> dict:
        """further ask question on a coding problem.

        Args:
            question (str): question to llm.
            session_id (str): specify a problem and answer session.

        Returns:
            dict: question answer and metadata
        """
        assert question

        if session_id not in self._session_store:
            return {}

        chain = self._session_store[session_id]

        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:
        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=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=self._main_llm,
                memory=memory,
                retriever=self._knowledge,
                condense_question_llm=self._secondary_llm,
                return_source_documents=True,
                combine_docs_chain_kwargs={"prompt": prompt},
            )