Spaces:
Runtime error
Runtime error
File size: 1,947 Bytes
e34be6c 1543ec3 e34be6c 1543ec3 e34be6c 1543ec3 e34be6c 1543ec3 e34be6c 1543ec3 e34be6c 1543ec3 e34be6c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | from unittest.mock import MagicMock, patch
from langchain import PromptTemplate
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, mocked_init_llm):
task = QaTask(instruction="test")
assert task._chat_prompt == PromptTemplate.from_template(TEMPLATE_CHAT.format(instruction="test"))
assert task._once_prompt == PromptTemplate.from_template(TEMPLATE_ONCE.format(instruction="test"))
assert task._knowledge is None
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, 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
task = QaTask(instruction="test")
with patch.object(task, "_create_session_id") as mock_create_id:
mock_create_id.return_value = 123
result = task.ask("how are you?", session=True)
mock_create_id.assert_called_once()
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, mocked_init_llm):
mocked_llm = MagicMock()
mocked_llm.run.return_value = {"result": "ok"}
mocked_build_once_chain.return_value = mocked_llm
task = QaTask(instruction="test")
result = task.ask("how are you?", session=False)
mocked_build_once_chain.assert_called_once()
assert "session_id" not in result
|