125 lines
5.3 KiB
Python
125 lines
5.3 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)
|
|
|
|
|
|
class CanonicalFarmContextTests(SimpleTestCase):
|
|
@patch("rag.chat.build_ai_farm_snapshot", return_value={"farm_uuid": "farm-123", "farm_metrics": {"resolved_metrics": {"soil_moisture": 41.0}}})
|
|
def test_format_farm_context_uses_canonical_ai_snapshot(self, mock_snapshot):
|
|
from rag.chat import _format_farm_context
|
|
|
|
context = _format_farm_context("farm-123")
|
|
|
|
self.assertIn("farm_metrics", context)
|
|
self.assertIn("soil_moisture", context)
|
|
mock_snapshot.assert_called_once_with("farm-123")
|
|
|
|
|
|
class CanonicalUserSoilTextTests(SimpleTestCase):
|
|
@patch(
|
|
"rag.user_data.build_ai_farm_snapshot",
|
|
return_value={
|
|
"farm_uuid": "farm-123",
|
|
"aggregation_policy": {
|
|
"sensor": "cluster_mean_then_block_mean_then_farm_mean",
|
|
"satellite": "cluster_mean_then_block_mean_then_farm_mean",
|
|
"weather": "center_location_latest_forecast",
|
|
},
|
|
"farm_metrics": {"resolved_metrics": {"soil_moisture": 42.0, "nitrogen": 18.0}},
|
|
"block_metrics": [
|
|
{"block_code": "block-1", "resolved_metrics": {"soil_moisture": 40.0}},
|
|
{"block_code": "block-2", "resolved_metrics": {"soil_moisture": 44.0}},
|
|
],
|
|
"sub_block_metrics": [
|
|
{"block_code": "block-1", "sub_block_code": "cluster-a", "resolved_metrics": {"soil_moisture": 39.0}},
|
|
{"block_code": "block-2", "sub_block_code": "cluster-b", "resolved_metrics": {"soil_moisture": 45.0}},
|
|
],
|
|
"source_metadata": {
|
|
"farm_metrics": {
|
|
"canonical_source": "farmer_block_aggregated_snapshot",
|
|
"aggregation_strategy": "farmer_block_mean",
|
|
}
|
|
},
|
|
},
|
|
)
|
|
def test_build_user_soil_text_uses_only_canonical_snapshot(self, mock_snapshot):
|
|
from rag.user_data import build_user_soil_text
|
|
|
|
text = build_user_soil_text("farm-123")
|
|
|
|
self.assertIn("سیاست تجمیع خاک", text)
|
|
self.assertIn("خلاصه تجمیعشده مزرعه", text)
|
|
self.assertIn("بلوک block-1", text)
|
|
self.assertIn("زیر-بلوک cluster-a", text)
|
|
self.assertIn("canonical_source: farmer_block_aggregated_snapshot", text)
|
|
self.assertNotIn("sensor_payload", text)
|
|
mock_snapshot.assert_called_once_with("farm-123")
|