Ai/rag/services/soil_anomaly.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 = "soil_anomaly"
SERVICE_ID = "soil_anomaly"

SOIL_ANOMALY_PROMPT = (
    "شما یک دستیار تخصصی تحلیل ناهنجاری داده های خاک و سنسور مزرعه هستی. "
    "ورودی شامل داده های ساختاریافته ناهنجاری، اطلاعات مزرعه، و متن های بازیابی شده از پایگاه دانش است. "
    "فقط JSON معتبر برگردان و فقط این کلیدها را تولید کن: "
    "summary, explanation, likely_cause, recommended_action, monitoring_priority, confidence. "
    "monitoring_priority فقط یکی از low, medium, high, urgent باشد. "
    "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="Soil anomaly 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 soil_anomaly LLM: %s", cleaned[:500])
        raise RAGServiceError(
            error_code="invalid_json",
            message="Soil anomaly 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="Soil anomaly LLM response root must be a JSON object.",
            source="llm",
            retriable=False,
            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_anomaly_insight(parsed: dict[str, Any]) -> dict[str, Any]:
    required_keys = {
        "summary",
        "explanation",
        "likely_cause",
        "recommended_action",
        "monitoring_priority",
        "confidence",
    }
    missing = [key for key in required_keys if key not in parsed]
    if missing:
        raise RAGServiceError(
            error_code="invalid_schema",
            message="Soil anomaly 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(SOIL_ANOMALY_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_soil_anomaly_insight(
    *,
    farm_uuid: str,
    anomaly_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,
        "anomaly_payload": anomaly_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("Soil anomaly insight failed for %s: %s", farm_uuid, exc)
        _fail_audit_log(audit_log, str(exc))
        raise
    except Exception as exc:
        logger.error("Soil anomaly insight failed for %s: %s", farm_uuid, exc)
        _fail_audit_log(audit_log, str(exc))
        raise RAGServiceError(
            error_code="upstream_failure",
            message=f"Soil anomaly 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_anomaly_insight(parsed)
    parsed["status"] = "success"
    parsed["source"] = "llm"
    parsed["farm_uuid"] = farm_uuid
    parsed["raw_response"] = raw
    return parsed