Ai/rag/services/water_need_prediction.py

from __future__ import annotations

import json
import logging
from typing import Any

from farm_data.services import get_farm_details
from rag.api_provider import get_chat_client
from rag.chat import (
    _complete_audit_log,
    _create_audit_log,
    _fail_audit_log,
    _load_service_tone,
    build_rag_context,
)
from rag.config import RAGConfig, get_service_config, load_rag_config
from rag.failure_contract import RAGServiceError

logger = logging.getLogger(__name__)

KB_NAME = "water_need_prediction"
SERVICE_ID = "water_need_prediction"

WATER_NEED_PROMPT = (
    "شما یک دستیار تخصصی تحليل نياز آبي کوتاه مدت مزرعه هستي. "
    "ورودي شامل محاسبات ساختاريافته نياز آبي، اطلاعات مزرعه و متن هاي بازيابي شده از پايگاه دانش است. "
    "فقط JSON معتبر با اين کليدها برگردان: "
    "summary, irrigation_outlook, recommended_action, risk_note, confidence. "
    "confidence عددي بين 0 و 1 باشد. "
    "اعداد اصلي را از داده ورودي بگير و عدد متناقض جديد نساز."
)


def _clean_json(raw: str) -> dict[str, Any]:
    cleaned = (raw or "").strip()
    if cleaned.startswith("```"):
        cleaned = cleaned.strip("`")
        if cleaned.startswith("json"):
            cleaned = cleaned[4:]
        cleaned = cleaned.strip()
    if not cleaned:
        raise RAGServiceError(
            error_code="empty_response",
            message="Water need prediction LLM response was empty.",
            source="llm",
            retriable=True,
            details={"service_id": SERVICE_ID},
            http_status=502,
        )
    try:
        parsed = json.loads(cleaned)
    except (json.JSONDecodeError, ValueError) as exc:
        logger.warning("Invalid JSON returned by water_need_prediction LLM: %s", cleaned[:500])
        raise RAGServiceError(
            error_code="invalid_json",
            message="Water need prediction LLM response was not valid JSON.",
            source="llm",
            retriable=True,
            details={"service_id": SERVICE_ID},
            http_status=502,
        ) from exc
    if not isinstance(parsed, dict):
        raise RAGServiceError(
            error_code="invalid_schema",
            message="Water need prediction LLM response root must be a JSON object.",
            source="llm",
            details={"service_id": SERVICE_ID},
            http_status=502,
        )
    return parsed


def _load_farm_or_error(farm_uuid: str) -> dict[str, Any]:
    farm_details = get_farm_details(farm_uuid)
    if farm_details is None:
        raise RAGServiceError(
            error_code="farm_not_found",
            message="farm_uuid نامعتبر است یا اطلاعات مزرعه پیدا نشد.",
            source="farm_data",
            details={"farm_uuid": farm_uuid},
            http_status=404,
        )
    return farm_details


def _build_service_client(cfg: RAGConfig):
    service = get_service_config(SERVICE_ID, cfg)
    service_cfg = RAGConfig(
        embedding=cfg.embedding,
        qdrant=cfg.qdrant,
        chunking=cfg.chunking,
        llm=service.llm,
        knowledge_bases=cfg.knowledge_bases,
        services=cfg.services,
        chromadb=cfg.chromadb,
    )
    client = get_chat_client(service_cfg)
    return service, client, service.llm.model


def _validate_prediction_insight(parsed: dict[str, Any]) -> dict[str, Any]:
    required_keys = {
        "summary",
        "irrigation_outlook",
        "recommended_action",
        "risk_note",
        "confidence",
    }
    missing = [key for key in required_keys if key not in parsed]
    if missing:
        raise RAGServiceError(
            error_code="invalid_schema",
            message="Water need prediction insight response is missing required fields: " + ", ".join(missing),
            source="llm",
            details={"missing_fields": missing, "service_id": SERVICE_ID},
            http_status=502,
        )
    return parsed


def _build_messages(
    *,
    service: Any,
    cfg: RAGConfig,
    query: str,
    rag_context: str,
    structured_context: dict[str, Any],
) -> tuple[str, list[dict[str, str]]]:
    tone = _load_service_tone(service, cfg)
    system_parts = [tone] if tone else []
    if service.system_prompt:
        system_parts.append(service.system_prompt)
    system_parts.append(WATER_NEED_PROMPT)
    system_parts.append(
        "[کانتکست ساختاريافته نياز آبي]\n"
        + json.dumps(structured_context, ensure_ascii=False, indent=2, default=str)
    )
    if rag_context:
        system_parts.append(rag_context)
    system_prompt = "\n\n".join(part for part in system_parts if part)
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": query},
    ]
    return system_prompt, messages


def get_water_need_prediction_insight(
    *,
    farm_uuid: str,
    prediction_payload: dict[str, Any],
    query: str | None = None,
) -> dict[str, Any]:
    cfg = load_rag_config()
    service, client, model = _build_service_client(cfg)
    farm_details = _load_farm_or_error(farm_uuid)
    user_query = query or "نياز آبي کوتاه مدت اين مزرعه را تفسير کن و اقدام عملياتي پيشنهاد بده."
    structured_context = {
        "farm_uuid": farm_uuid,
        "prediction_payload": prediction_payload,
    }
    rag_context = build_rag_context(
        query=user_query,
        sensor_uuid=farm_uuid,
        config=cfg,
        kb_name=KB_NAME,
        service_id=SERVICE_ID,
        farm_details=farm_details,
    )
    system_prompt, messages = _build_messages(
        service=service,
        cfg=cfg,
        query=user_query,
        rag_context=rag_context,
        structured_context=structured_context,
    )
    audit_log = _create_audit_log(
        farm_uuid=farm_uuid,
        service_id=SERVICE_ID,
        model=model,
        query=user_query,
        system_prompt=system_prompt,
        messages=messages,
    )
    try:
        response = client.chat.completions.create(model=model, messages=messages)
        raw = response.choices[0].message.content.strip()
        parsed = _clean_json(raw)
        _complete_audit_log(audit_log, raw)
    except RAGServiceError as exc:
        logger.error("Water need prediction insight failed for %s: %s", farm_uuid, exc)
        _fail_audit_log(audit_log, str(exc))
        raise
    except Exception as exc:
        logger.error("Water need prediction insight failed for %s: %s", farm_uuid, exc)
        _fail_audit_log(audit_log, str(exc))
        raise RAGServiceError(
            error_code="upstream_failure",
            message=f"Water need prediction insight failed for farm {farm_uuid}.",
            source="llm",
            retriable=True,
            details={"farm_uuid": farm_uuid, "service_id": SERVICE_ID},
            http_status=503,
        ) from exc

    parsed = _validate_prediction_insight(parsed)
    parsed["status"] = "success"
    parsed["source"] = "llm"
    parsed["farm_uuid"] = farm_uuid
    parsed["raw_response"] = raw
    return parsed