53 lines
1.9 KiB
Python
53 lines
1.9 KiB
Python
|
from unittest.mock import patch
|
||
|
|
|
||
|
|
from django.test import SimpleTestCase
|
||
|
|
|
||
|
|
from rag.chat import build_chat_context
|
||
|
|
|
||
|
|
|
||
|
|
class ChatContextTests(SimpleTestCase):
|
||
|
|
@patch("rag.chat.search_with_query")
|
||
|
|
@patch("rag.chat._rank_text_chunks_by_query")
|
||
|
|
@patch("rag.chat.chunk_text")
|
||
|
|
def test_build_chat_context_combines_farm_and_kb_context(
|
||
|
|
self,
|
||
|
|
mock_chunk_text,
|
||
|
|
mock_rank_text_chunks_by_query,
|
||
|
|
mock_search_with_query,
|
||
|
|
):
|
||
|
|
mock_chunk_text.return_value = ["chunk-a", "chunk-b"]
|
||
|
|
mock_rank_text_chunks_by_query.return_value = ["chunk-b"]
|
||
|
|
mock_search_with_query.return_value = [
|
||
|
|
{"text": "kb text 1"},
|
||
|
|
{"text": "kb text 2"},
|
||
|
|
]
|
||
|
|
|
||
|
|
context = build_chat_context(
|
||
|
|
query="وضعیت مزرعه چطور است؟",
|
||
|
|
farm_uuid="farm-123",
|
||
|
|
farm_details={"sensor_payload": {"sensor-7-1": {"soil_moisture": 30}}},
|
||
|
|
)
|
||
|
|
|
||
|
|
self.assertIn("[بخشهای مرتبط بازیابیشده از اطلاعات مزرعه]", context)
|
||
|
|
self.assertIn("chunk-b", context)
|
||
|
|
self.assertIn("[اطلاعات بازیابیشده از پایگاه دانش]", context)
|
||
|
|
self.assertIn("kb text 1", context)
|
||
|
|
self.assertIn("kb text 2", context)
|
||
|
|
|
||
|
|
@patch("rag.chat.search_with_query", return_value=[])
|
||
|
|
@patch("rag.chat._rank_text_chunks_by_query", return_value=[])
|
||
|
|
@patch("rag.chat.chunk_text", return_value=["farm chunk"])
|
||
|
|
def test_build_chat_context_falls_back_to_full_farm_context(
|
||
|
|
self,
|
||
|
|
_mock_chunk_text,
|
||
|
|
_mock_rank_text_chunks_by_query,
|
||
|
|
_mock_search_with_query,
|
||
|
|
):
|
||
|
|
context = build_chat_context(
|
||
|
|
query="رطوبت چقدر است؟",
|
||
|
|
farm_uuid="farm-123",
|
||
|
|
farm_details={"sensor_payload": {"sensor-7-1": {"soil_moisture": 30}}},
|
||
|
|
)
|
||
|
|
|
||
|
|
self.assertEqual(context, "")
|