Ai/rag/chat.py

"""
چت RAG با استریم — استفاده از دیتای embed شده کاربر و Adapter API (GapGPT / Avalai)
"""
import logging
from pathlib import Path

from .config import load_rag_config, RAGConfig, get_service_config, ServiceConfig
from .api_provider import get_chat_client
from .retrieve import search_with_query
from .user_data import build_user_soil_text, build_user_weather_text

logger = logging.getLogger(__name__)


def _load_tone(config: RAGConfig | None) -> str:
    """بارگذاری فایل لحن پیش‌فرض (chat KB)."""
    cfg = config or load_rag_config()
    base = Path(__file__).resolve().parent.parent
    chat_kb = cfg.knowledge_bases.get("chat")
    if chat_kb:
        tone_path = base / chat_kb.tone_file
        logger.debug("Loading default tone from path=%s", tone_path)
        if tone_path.exists():
            logger.debug("Default tone file found: %s", tone_path)
            return tone_path.read_text(encoding="utf-8").strip()
        logger.warning("Default tone file not found: %s", tone_path)
    return ""


def _load_kb_tone(kb_name: str, config: RAGConfig | None = None) -> str:
    """بارگذاری فایل لحن مخصوص یک پایگاه دانش."""
    cfg = config or load_rag_config()
    kb_cfg = cfg.knowledge_bases.get(kb_name)
    if not kb_cfg:
        return ""
    base = Path(__file__).resolve().parent.parent
    tone_path = base / kb_cfg.tone_file
    logger.debug("Loading kb tone for kb=%s path=%s", kb_name, tone_path)
    if tone_path.exists():
        logger.debug("KB tone file found for kb=%s", kb_name)
        return tone_path.read_text(encoding="utf-8").strip()
    logger.warning("KB tone file not found for kb=%s path=%s", kb_name, tone_path)
    return ""


def _load_service_tone(service: ServiceConfig, config: RAGConfig | None = None) -> str:
    cfg = config or load_rag_config()
    if service.tone_file:
        base = Path(__file__).resolve().parent.parent
        tone_path = base / service.tone_file
        if tone_path.exists():
            return tone_path.read_text(encoding="utf-8").strip()
    return _load_kb_tone(service.knowledge_base, cfg)


def _detect_kb_intent(query: str) -> str:
    """تشخیص ساده نوع پایگاه دانش مورد نیاز از روی متن سوال."""
    q = query.lower()
    irrigation_keywords = {"آبیاری", "آب", "رطوبت", "irrigation", "water", "et0", "بارش", "خشکی"}
    fertilization_keywords = {"کود", "کودهی", "fertiliz", "npk", "ازت", "فسفر", "پتاسیم", "nitrogen", "phosphorus", "potassium"}
    if any(kw in q for kw in irrigation_keywords):
        return "irrigation"
    if any(kw in q for kw in fertilization_keywords):
        logger.info("Detected KB intent=fertilization")
        return "fertilization"
    logger.info("Detected KB intent=chat")
    return "chat"


def build_rag_context(
    query: str,
    sensor_uuid: str | None = None,
    config: RAGConfig | None = None,
    limit: int = 8,
    kb_name: str | None = None,
    service_id: str | None = None,
) -> str:
    """
    ساخت context برای LLM: دیتای فعلی خاک کاربر + متن‌های مرتبط از RAG.
    دیتای کاربر همیشه اول می‌آید تا LLM مقادیر واقعی (مثل pH) را ببیند.
    """
    logger.info(
        "Building RAG context sensor_uuid=%s kb_name=%s limit=%s query_len=%s",
        sensor_uuid,
        kb_name,
        limit,
        len(query or ""),
    )
    parts: list[str] = []
    cfg = config or load_rag_config()
    service = get_service_config(service_id, cfg) if service_id else None
    include_user_embeddings = service.use_user_embeddings if service else True
    resolved_kb_name = kb_name or (service.knowledge_base if service else None)

    if include_user_embeddings and sensor_uuid:
        user_soil = build_user_soil_text(sensor_uuid)
        if user_soil and user_soil.strip():
            parts.append("[داده‌های فعلی خاک شما]\n" + user_soil.strip())
            logger.debug("Included user soil section sensor_uuid=%s", sensor_uuid)
        else:
            logger.info("No user soil data found sensor_uuid=%s", sensor_uuid)

        weather_text = build_user_weather_text(sensor_uuid)
        if weather_text and weather_text.strip():
            parts.append("[پیش‌بینی هواشناسی]\n" + weather_text.strip())
            logger.debug("Included weather section sensor_uuid=%s", sensor_uuid)
        else:
            logger.info("No weather data found sensor_uuid=%s", sensor_uuid)

    results = search_with_query(
        query,
        sensor_uuid=sensor_uuid,
        limit=limit,
        config=cfg,
        kb_name=resolved_kb_name,
        service_id=service_id,
        use_user_embeddings=include_user_embeddings,
    )
    if results:
        logger.info("Retrieved RAG results count=%s sensor_uuid=%s", len(results), sensor_uuid)
        rag_texts = [r.get("text", "").strip() for r in results if r.get("text")]
        if rag_texts:
            parts.append("[متن‌های مرجع]\n" + "\n\n---\n\n".join(rag_texts))
            logger.debug("Included RAG reference texts count=%s", len(rag_texts))
    else:
        logger.info("No RAG results found sensor_uuid=%s kb_name=%s", sensor_uuid, kb_name)

    return "\n\n---\n\n".join(parts) if parts else ""


def chat_rag_stream(
    query: str,
    sensor_uuid: str | None = None,
    config: RAGConfig | None = None,
    limit: int = 5,
    system_override: str | None = None,
    kb_name: str | None = None,
    service_id: str | None = None,
):
    logger.info(
        "chat_rag_stream started sensor_uuid=%s kb_name=%s limit=%s query_len=%s",
        sensor_uuid,
        kb_name,
        limit,
        len(query or ""),
    )
    """
    چت RAG با استریم: دیتای embed شده را بازیابی می‌کند و با LLM جواب می‌دهد.
    فقط دیتای همان کاربر (sensor_uuid) قابل دسترسی است.

    Args:
        query: پیام کاربر
        sensor_uuid: شناسه سنسور کاربر — اجباری
        config: تنظیمات RAG
        limit: تعداد چانک‌های بازیابی‌شده
        system_override: جایگزین system prompt (اختیاری)

    Yields:
        تک‌تک deltaهای content به‌صورت رشته
    """
    cfg = config or load_rag_config()
    resolved_service_id = service_id or kb_name or _detect_kb_intent(query)
    service = get_service_config(resolved_service_id, cfg)
    service_llm_config = service.llm
    service_cfg = RAGConfig(
        embedding=cfg.embedding,
        qdrant=cfg.qdrant,
        chunking=cfg.chunking,
        llm=service_llm_config,
        knowledge_bases=cfg.knowledge_bases,
        services=cfg.services,
        chromadb=cfg.chromadb,
    )
    client = get_chat_client(service_cfg)
    model = service_llm_config.model
    logger.debug("Loaded service config service_id=%s model=%s", resolved_service_id, model)

    detected_kb = kb_name or service.knowledge_base
    logger.info("Using knowledge base=%s for service_id=%s", detected_kb, resolved_service_id)
    context = build_rag_context(
        query,
        sensor_uuid,
        config=cfg,
        limit=limit,
        kb_name=detected_kb,
        service_id=resolved_service_id,
    )
    logger.debug("Built context length=%s", len(context))

    if system_override is not None:
        system_content = system_override
    else:
        tone = _load_service_tone(service, cfg)
        if not tone:
            tone = _load_tone(cfg)
        system_parts = [tone] if tone else []
        if service.system_prompt:
            system_parts.append(service.system_prompt)
        system_parts.append(
            "با استفاده از بخش «داده‌های فعلی خاک شما» و «متن‌های مرجع» زیر به سوال کاربر پاسخ بده. "
            "برای سوالاتی درباره خاک کاربر (مثل pH، رطوبت، NPK) حتماً از داده‌های فعلی استفاده کن. "
            "اطلاعات هواشناسی در بخش «پیش‌بینی هواشناسی» آمده. "
            "پاسخ را به زبان کاربر بنویس."
        )
        if context:
            system_parts.append("\n\n" + context)
        system_content = "\n".join(system_parts)

    messages = [
        {"role": "system", "content": system_content},
        {"role": "user", "content": query},
    ]
    logger.info("Prepared messages for model=%s service_id=%s", model, resolved_service_id)

    stream = client.chat.completions.create(
        model=model,
        messages=messages,
        stream=True,
    )
    logger.info("Started streaming response from model=%s", model)

    for chunk in stream:
        delta = chunk.choices[0].delta if chunk.choices else None
        content = delta.content if delta else ""
        if content:
            logger.debug("Streaming chunk len=%s", len(content))
            yield content
    logger.info("chat_rag_stream completed sensor_uuid=%s", sensor_uuid)