Ai/rag/chat.py

"""
چت RAG با استریم — استفاده از دیتای embed شده کاربر و Avalai API
"""
import os
from pathlib import Path

from openai import OpenAI

from .config import load_rag_config, RAGConfig
from .retrieve import search_with_query
from .user_data import build_user_soil_text


def _get_chat_client(config: RAGConfig | None) -> OpenAI:
    """ساخت کلاینت OpenAI برای Avalai Chat API."""
    cfg = config or load_rag_config()
    llm = cfg.llm
    env_var = llm.api_key_env or "AVALAI_API_KEY"
    api_key = os.environ.get(env_var)
    base_url = llm.base_url or os.environ.get(
        "AVALAI_BASE_URL", "https://api.avalai.ir/v1"
    )
    return OpenAI(api_key=api_key, base_url=base_url)


def _load_tone(config: RAGConfig | None) -> str:
    """بارگذاری فایل لحن."""
    cfg = config or load_rag_config()
    base = Path(__file__).resolve().parent.parent
    tone_path = base / cfg.tone_file
    if tone_path.exists():
        return tone_path.read_text(encoding="utf-8").strip()
    return ""


def build_rag_context(
    query: str,
    sensor_uuid: str,
    config: RAGConfig | None = None,
    limit: int = 8,
) -> str:
    """
    ساخت context برای LLM: دیتای فعلی خاک کاربر + متن‌های مرتبط از RAG.
    دیتای کاربر همیشه اول می‌آید تا LLM مقادیر واقعی (مثل pH) را ببیند.
    """
    parts: list[str] = []

    # ۱. دیتای فعلی خاک کاربر از DB — همیشه اول (برای سوالاتی مثل «pH خاک من چند است»)
    user_soil = build_user_soil_text(sensor_uuid)
    if user_soil and user_soil.strip():
        parts.append("[داده‌های فعلی خاک شما]\n" + user_soil.strip())

    # ۲. متن‌های مرتبط از RAG
    results = search_with_query(
        query, sensor_uuid=sensor_uuid, limit=limit, config=config
    )
    if results:
        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))

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


def chat_rag_stream(
    query: str,
    sensor_uuid: str,
    config: RAGConfig | None = None,
    limit: int = 5,
    system_override: str | None = None,
):
    """
    چت 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()
    client = _get_chat_client(cfg)
    model = cfg.llm.model

    context = build_rag_context(query, sensor_uuid, config=cfg, limit=limit)

    if system_override is not None:
        system_content = system_override
    else:
        tone = _load_tone(cfg)
        system_parts = [tone] if tone else []
        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},
    ]

    stream = client.chat.completions.create(
        model=model,
        messages=messages,
        stream=True,
    )

    for chunk in stream:
        delta = chunk.choices[0].delta if chunk.choices else None
        content = delta.content if delta else ""
        if content:
            yield content