72 lines
2.9 KiB
Python
72 lines
2.9 KiB
Python
|
from unittest.mock import patch
|
||
|
|
|
||
|
|
from django.test import SimpleTestCase
|
||
|
|
|
||
|
|
from rag.chat import _normalize_history_messages, build_rag_context
|
||
|
|
|
||
|
|
|
||
|
|
class ChatContextTests(SimpleTestCase):
|
||
|
|
@patch("rag.chat.search_with_texts")
|
||
|
|
@patch("rag.chat.chunk_text")
|
||
|
|
def test_build_rag_context_includes_full_farm_and_kb_results(
|
||
|
|
self,
|
||
|
|
mock_chunk_text,
|
||
|
|
mock_search_with_texts,
|
||
|
|
):
|
||
|
|
mock_chunk_text.return_value = ["farm chunk 1", "farm chunk 2"]
|
||
|
|
mock_search_with_texts.return_value = [
|
||
|
|
{"id": "kb-1", "score": 0.8, "text": "kb text 1", "metadata": {}},
|
||
|
|
{"id": "kb-2", "score": 0.7, "text": "kb text 2", "metadata": {}},
|
||
|
|
]
|
||
|
|
|
||
|
|
context = build_rag_context(
|
||
|
|
query="وضعیت مزرعه چطور است؟",
|
||
|
|
sensor_uuid="farm-123",
|
||
|
|
service_id="chat",
|
||
|
|
farm_details={"sensor_payload": {"sensor-7-1": {"soil_moisture": 30}}},
|
||
|
|
)
|
||
|
|
|
||
|
|
self.assertIn("[اطلاعات کامل مزرعه]", context)
|
||
|
|
self.assertIn("soil_moisture", context)
|
||
|
|
self.assertIn("[متنهای مرجع]", context)
|
||
|
|
self.assertIn("kb text 1", context)
|
||
|
|
self.assertIn("kb text 2", context)
|
||
|
|
mock_search_with_texts.assert_called_once()
|
||
|
|
sent_texts = mock_search_with_texts.call_args.kwargs["texts"]
|
||
|
|
self.assertEqual(sent_texts[0], "وضعیت مزرعه چطور است؟")
|
||
|
|
self.assertIn("farm chunk 1", sent_texts)
|
||
|
|
self.assertIn("farm chunk 2", sent_texts)
|
||
|
|
|
||
|
|
def test_normalize_history_messages_supports_user_images(self):
|
||
|
|
messages = _normalize_history_messages(
|
||
|
|
[
|
||
|
|
{"role": "user", "content": "این تصویر مزرعه است", "image_urls": ["https://example.com/a.jpg"]},
|
||
|
|
{"role": "assistant", "content": "تصویر دریافت شد."},
|
||
|
|
]
|
||
|
|
)
|
||
|
|
|
||
|
|
self.assertEqual(len(messages), 2)
|
||
|
|
self.assertEqual(messages[0]["role"], "user")
|
||
|
|
self.assertIsInstance(messages[0]["content"], list)
|
||
|
|
self.assertEqual(messages[0]["content"][0]["type"], "text")
|
||
|
|
self.assertEqual(messages[0]["content"][1]["type"], "image_url")
|
||
|
|
self.assertEqual(messages[1]["role"], "assistant")
|
||
|
|
self.assertEqual(messages[1]["content"], "تصویر دریافت شد.")
|
||
|
|
|
||
|
|
@patch("rag.chat.search_with_texts", return_value=[])
|
||
|
|
@patch("rag.chat.chunk_text", return_value=["farm chunk"])
|
||
|
|
def test_build_rag_context_returns_full_farm_when_kb_empty(
|
||
|
|
self,
|
||
|
|
_mock_chunk_text,
|
||
|
|
_mock_search_with_texts,
|
||
|
|
):
|
||
|
|
context = build_rag_context(
|
||
|
|
query="رطوبت چقدر است؟",
|
||
|
|
sensor_uuid="farm-123",
|
||
|
|
service_id="chat",
|
||
|
|
farm_details={"sensor_payload": {"sensor-7-1": {"soil_moisture": 30}}},
|
||
|
|
)
|
||
|
|
|
||
|
|
self.assertIn("[اطلاعات کامل مزرعه]", context)
|
||
|
|
self.assertIn("soil_moisture", context)
|