Ai/rag/tests/test_chat_context.py

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")