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

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:
        return {}
    try:
        return json.loads(cleaned)
    except (json.JSONDecodeError, ValueError):
        logger.warning("Invalid JSON returned by water_need_prediction LLM: %s", cleaned[:500])
        return {}


def _load_farm_or_error(farm_uuid: str) -> dict[str, Any]:
    farm_details = get_farm_details(farm_uuid)
    if farm_details is None:
        raise ValueError("farm_uuid نامعتبر است یا اطلاعات مزرعه پیدا نشد.")
    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 ValueError(
            "Water need prediction insight response is missing required fields: "
            + ", ".join(missing)
        )
    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 Exception as exc:
        logger.error("Water need prediction insight failed for %s: %s", farm_uuid, exc)
        _fail_audit_log(audit_log, str(exc))
        raise RuntimeError(f"Water need prediction insight failed for farm {farm_uuid}.") from exc

    parsed = _validate_prediction_insight(parsed)
    parsed["farm_uuid"] = farm_uuid
    parsed["raw_response"] = raw
    return parsed