feat(qa): :sparkles: normal qa (#9)

.github/workflows/test.yml

@@ -27,11 +27,11 @@ jobs:
                 poetry-version: '1.5.1'
 
             - name: Install Dependencies
-              run: poetry install
+              run: poetry install --with dev
 
             - name: Run Test
               run:
-                poetry run pytest --junitxml=pytest.xml --cov-report=term-missing:skip-covered --cov=codedog test/ | tee pytest-coverage.txt
+                poetry run pytest --junitxml=pytest.xml --cov-report=term-missing:skip-covered --cov=edu_assistant tests/ | tee pytest-coverage.txt
 
             - name: Pytest coverage comment
               id: coverageComment

README.md

@@ -3,3 +3,24 @@
 [![Checkstyle](https://github.com/Arcadia822/codedog/actions/workflows/flake8.yml/badge.svg)](https://github.com/codedog-ai/edu-assistant/actions/workflows/flake8.yml)
 [![Pytest](https://github.com/Arcadia822/codedog/actions/workflows/test.yml/badge.svg?branch=master)](https://github.com/codedog-ai/edu-assistant/actions/workflows/test.yml)
 [![Coverage](https://img.shields.io/endpoint?url=https://gist.githubusercontent.com/Arcadia822/ce38dae58995aeffef42065093fcfe84/raw/edu_assistant_master.json)](https://github.com/codedog-ai/edu-assistant/actions/workflows/test.yml)
+
+## config
+
+| Environment Variable | Necessary | Default | Description |
+|---|---|---|---|
+| CODEDOG_SERVER                | No  | 0.0.0.0        | Server address                               |
+| CODEDOG_PORT                  | No  | 32167          | Server port                                  |
+| CODEDOG_WORKER_NUM            | No  | 1              | Server worker number                         |
+| OPENAI_API_KEY                | Yes |                | Api Key for calling openai api               |
+| OPENAI_PROXY                  | No  |                | openai proxy                                 |
+| AZURE_OPENAI                  | No  |                | use azure openai if not blank                |
+| AZURE_OPENAI_API_KEY          | No  |                | azure openai api key                         |
+| AZURE_OPENAI_API_BASE         | No  |                | azure openai api base                        |
+| AZURE_OPENAI_DEPLOYMENT_ID    | No  |                | azure openai deployment id for gpt 3.5       |
+| AZURE_OPENAI_EMBEDDING_DEP_ID | No  |                | azure openai deployment id for embedding     |
+
+## setup
+
+- install python 3.10+
+- install poetry 1.5.1+
+- run: `poetry install`

edu_assistant/learning_tasks/__init__.py

@@ -0,0 +1,3 @@
+from .qa import QaTask
+
+__all__ = ["QaTask"]

edu_assistant/learning_tasks/base.py

@@ -0,0 +1,2 @@
+class BaseTask:
+    pass

edu_assistant/learning_tasks/qa.py

@@ -0,0 +1,116 @@
+import random
+import sys
+
+from langchain import PromptTemplate
+from langchain.chains import ConversationalRetrievalChain, ConversationChain
+from langchain.chains.base import Chain
+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, update_chat_memory
+
+TEMPLATE = """{instruction}
+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.
+
+Current conversation:
+{{history}}
+Human: {{input}}
+AI:"""
+
+
+class QaTask(BaseTask):
+    _qa: Chain
+    _session_store: dict
+
+    def __init__(
+        self,
+        instruction: str = "",
+        knowledge: BaseRetriever = None,
+        session_store: dict[int, ConversationBufferMemory] = None,
+    ):
+        """Create a new QaTask service.
+
+        Args:
+            instruction (str, optional): Instruction for this task ai. Defaults to "".
+            knowledge (BaseRetriever, optional): Answer question with this knowledge retriever.
+                If not set, will not use knowledge to answer question.
+                Defaults to None.
+            session_store (dict, optional): External chat history store. Defaults to None.
+                If not set, will use internal memory to store chat history. Which will be lost after restart and might
+                cost huge memory.
+        """
+        self._prompt = TEMPLATE.format(instruction=instruction)
+        self._qa = self._build_chain(knowledge)
+        self._session_store = {} if not session_store else session_store
+
+    def _build_chain(self, knowledge):
+        if not knowledge:
+            return ConversationChain(llm=load_llm(), prompt=PromptTemplate.from_template(self._prompt))
+        else:
+            return ConversationalRetrievalChain.from_llm(
+                llm=load_llm(),
+                retriever=knowledge,
+                condense_question_llm=load_llm(),
+                return_source_documents=True,
+                combine_docs_chain_kwargs={"prompt": PromptTemplate.from_template(self._prompt)},
+            )
+
+    def ask(self, question: str, session: bool = True, session_id: int = None, session_mem: list | None = None) -> dict:
+        """ask a question with chat history.
+
+        Args:
+            question (str): question to llm.
+            session (bool, optional): whether use and store chat history. Defaults to False.
+                if session_id is not set, a new session will be created.
+            session_id (int, optional): specify a history session. Defaults to None.
+            session_mem (list | None, optional): specify chat history. Defaults to None.
+                if session_id is also set, chat history will be fully replaced.
+
+        Returns:
+            dict: question answer and metadata.
+                contains answer.
+                contains session_id if session is True.
+        """
+
+        if session:
+            session_id = session_id or self._create_session_id()
+            memory = session_mem or self._get_session_mem(session_id)
+
+            result = self._ask(question, memory)
+
+            update_chat_memory(memory, question, result["response"])
+            self._update_session_mem(session_id, memory)
+        else:
+            result = self._ask(question)
+
+        if session_id:
+            result["session_id"] = session_id
+
+        return result
+
+    def _ask(self, question, memory: ConversationBufferMemory = None) -> dict:
+        if memory is None:
+            return self._qa({"input": question, "history": ""})
+        else:
+            return self._qa({"input": question, "history": memory.chat_memory})
+
+    def _get_session_mem(self, session_id):
+        if session_id not in self._session_store:
+            memory = self._init_memory(session_id)
+        else:
+            memory = self._session_store.get(session_id)
+
+        return memory
+
+    def _update_session_mem(self, session_id, memory):
+        self._session_store[session_id] = memory
+
+    def _create_session_id(self):
+        return random.randint(1, sys.maxsize)
+
+    def _init_memory(self, session_id) -> ConversationBufferMemory:
+        # TODO: redis memory store
+        return ConversationBufferMemory(memory_key="chat_history", return_messages=True, output_key="answer")

edu_assistant/utils/common_utils.py

@@ -0,0 +1,66 @@
+"""
+utility functions for codedog
+"""
+import hashlib
+import logging
+import time
+from logging.config import dictConfig
+
+# -- Logging ------------------------------------------------------------------
+
+
+def init_local_logging(level=logging.INFO):
+    """setup logging interface for local debugging"""
+    dictConfig(get_logging_config(level))
+
+
+def get_logging_config(level=logging.INFO):
+    return {
+        "version": 1,
+        "disable_existing_loggers": False,
+        "formatters": {
+            "plain": {
+                "format": "[%(asctime)s][%(levelname)s][%(pathname)s:%(lineno)d][%(name)s]%(message)s",
+                "datefmt": "%Y-%m-%d %H:%M:%S",
+            },
+        },
+        "handlers": {
+            "console_handler": {
+                "level": "DEBUG",
+                "formatter": "plain",
+                "class": "logging.StreamHandler",
+                "stream": "ext://sys.stdout",  # Default is stderr
+            },
+        },
+        "loggers": {
+            "": {
+                "level": level,
+                "handlers": ["console_handler"],
+            },
+        },
+    }
+
+
+# -- exception ----------------------------------------------------------------
+
+
+class CodedogError(Exception):
+    def __init__(self, message: str = None, code: int = -1):
+        self.message = "" if not message else str(message)
+        self.code = code
+
+
+# -- utility ------------------------------------------------------------------
+
+
+def get_ttl_hash(seconds=3600):
+    """Return the same value withing `seconds` time period"""
+    return round(time.time() / seconds)
+
+
+def get_sha256(text: str) -> str:
+    return hashlib.sha256(text.encode("utf-8", errors="ignore")).hexdigest()
+
+
+def get_sha1(text: str) -> str:
+    return hashlib.sha1(text.encode("utf-8", errors="ignore")).hexdigest()

edu_assistant/utils/langchain_utils.py

@@ -0,0 +1,48 @@
+import os
+
+from langchain.chat_models import AzureChatOpenAI, ChatOpenAI
+from langchain.chat_models.base import BaseChatModel
+from langchain.memory.chat_memory import BaseChatMemory
+
+
+def load_llm() -> BaseChatModel:
+    if os.environ.get("AZURE_OPENAI"):
+        llm = AzureChatOpenAI(
+            openai_api_type="azure",
+            openai_api_key=os.environ.get("AZURE_OPENAI_API_KEY"),
+            openai_api_base=os.environ.get("AZURE_OPENAI_API_BASE"),
+            openai_api_version="2023-05-15",
+            deployment_name=os.environ.get("AZURE_OPENAI_DEPLOYMENT_ID", "gpt-35-turbo"),
+            model="gpt-3.5-turbo",
+            temperature=0,
+        )
+    else:
+        llm = ChatOpenAI(
+            openai_api_key=os.environ.get("OPENAI_API_KEY"),
+            openai_proxy=os.environ.get("OPENAI_PROXY", ""),
+            model="gpt-3.5-turbo",
+        )
+
+    return llm
+
+
+def load_gpt4_llm() -> BaseChatModel:
+    llm = ChatOpenAI(
+        openai_api_key=os.environ.get("OPENAI_API_KEY"),
+        openai_proxy=os.environ.get("OPENAI_PROXY", ""),
+        model="gpt-4",
+    )
+
+    return llm
+
+
+def update_chat_memory(memory: BaseChatMemory, user_text: str, ai_text: str):
+    """update langchain memory
+
+    Args:
+        memory (BaseChatMemory): chat history memory
+        question (str): qa question str
+        result (str): chain answer result
+    """
+    memory.chat_memory.add_user_message(user_text)
+    memory.chat_memory.add_ai_message(ai_text)

examples/qa.py

@@ -0,0 +1,29 @@
+import json
+
+from edu_assistant.learning_tasks import QaTask
+from edu_assistant.utils.common_utils import init_local_logging
+
+init_local_logging()
+
+
+instruction = """
+Act as a c++ professional to answer student aged 5-10 questions. Answer properly and politely.
+"""
+task = QaTask(instruction=instruction)
+
+
+def qa_once():
+    result = task.ask("请问什么是指针？", session=False)
+    print(json.dumps(result, ensure_ascii=False, indent=4))
+
+
+def qa_twice():
+    task = QaTask()
+    task.ask("请问如何释放一个数组？")
+    result = task.ask("那指针呢？")
+    print(json.dumps(result, ensure_ascii=False, indent=4))
+
+
+if __name__ == "__main__":
+    # qa_once()
+    qa_twice()

poetry.lock

@@ -2517,4 +2517,4 @@ reference = "aliyun"
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.10"
-content-hash = "10e258b988673e31bc74da1f6e92f9afca6bd111985746fc441d717b37c576c4"
+content-hash = "bf7411574c4cfde57b9017ce88a88019ca51f05c24b086bb4477e2ff85201b51"

tests/unit_tests/learning_tasks/test_qa.py

@@ -0,0 +1,53 @@
+import unittest
+from unittest.mock import MagicMock, patch
+
+from langchain.memory import ConversationBufferMemory
+
+from edu_assistant.learning_tasks import QaTask
+
+
+class TestQaTask(unittest.TestCase):
+    @patch.object(QaTask, "_build_chain")
+    def test_init(self, mock_build_chain):
+        mock_build_chain.return_value = MagicMock()
+        QaTask()
+        mock_build_chain.assert_called_once()
+
+    @patch.object(QaTask, "_build_chain")
+    def test_ask_with_session(self, mock_build_chain):
+        mock_build_chain.return_value = MagicMock()
+        qa_task = QaTask()
+        qa_task._ask = MagicMock(return_value={"response": "test answer"})
+        result = qa_task.ask("test question")
+        qa_task._ask.assert_called_once()
+        self.assertEqual(result.get("response"), "test answer")
+        self.assertIsNotNone(result.get("session_id"))
+
+    @patch.object(QaTask, "_build_chain")
+    def test_ask_without_session(self, mock_build_chain):
+        mock_build_chain.return_value = MagicMock()
+        qa_task = QaTask()
+        qa_task._ask = MagicMock(return_value={"answer": "test answer"})
+        result = qa_task.ask("test question", session=False)
+        qa_task._ask.assert_called_once_with("test question")
+        self.assertEqual(result, {"answer": "test answer"})
+
+    @patch.object(QaTask, "_build_chain")
+    def test__get_session_mem(self, mock_build_chain):
+        mock_build_chain.return_value = MagicMock()
+        memory = ConversationBufferMemory()
+        qa_task = QaTask(session_store={1: memory})
+        result = qa_task._get_session_mem(1)
+        self.assertEqual(result, memory)
+
+    @patch.object(QaTask, "_build_chain")
+    def test__update_session_mem(self, mock_build_chain):
+        mock_build_chain.return_value = MagicMock()
+        memory = ConversationBufferMemory()
+        qa_task = QaTask()
+        qa_task._update_session_mem(1, memory)
+        self.assertEqual(qa_task._session_store[1], memory)
+
+
+if __name__ == "__main__":
+    unittest.main()