model/extractor.py

"""Ekstraksi entitas mimpi ke JSON via Qwen GGUF."""
from __future__ import annotations

import os
import time
from typing import Any

from model.loader import get_model, kunci_inferensi
from model.normalisasi import (
    entitas_memadai,
    parse_json_dari_teks,
    validasi_dan_perbaiki,
    validasi_skema_json,
)

VERSI_PROMPT = "oneiros_v1"

CONTOH_JSON_FEW_SHOT = """{
  "title": "The White Deer Forest",
  "mood": "mysterious",
  "characters": [
    {"name": "Dreamer", "role": "self", "emotion": "afraid"},
    {"name": "White Deer", "role": "creature", "emotion": "elusive"}
  ],
  "places": [
    {"name": "Night Forest", "type": "fantastical", "feeling": "chasing"}
  ],
  "symbols": [
    {"name": "Changing Face", "appearance": "friend face shifts", "significance": "identity"}
  ],
  "connections": [
    {"from": "Dreamer", "to": "White Deer", "type": "pursuit"},
    {"from": "Dreamer", "to": "Night Forest", "type": "tension"}
  ],
  "core_theme": "Chasing something familiar that stays out of reach"
}"""

SYSTEM_PROMPT = f"""You are Oneiros, a dream cartographer. Your only job is to extract structured entities from dream narratives and return valid JSON.

RULES:
- Return ONLY valid JSON, no markdown, no explanation, no preamble
- If a field is uncertain, use your best interpretation
- "connections" should capture the most emotionally significant relationships
- Keep all text short: max 5 words per field value
- Mood must be exactly one of: mysterious, joyful, anxious, surreal, peaceful

EXAMPLE OUTPUT (follow this structure exactly):
{CONTOH_JSON_FEW_SHOT}"""

EXTRACTION_PROMPT = """Extract entities from this dream and return JSON matching exactly this schema:

{{
  "title": string (3-5 word dream title),
  "mood": "mysterious|joyful|anxious|surreal|peaceful",
  "characters": [{{"name": string, "role": "self|known|stranger|creature|archetype", "emotion": string}}],
  "places": [{{"name": string, "type": "real|distorted|fantastical|transitional", "feeling": string}}],
  "symbols": [{{"name": string, "appearance": string, "significance": string}}],
  "connections": [{{"from": string, "to": string, "type": "tension|harmony|pursuit|transformation"}}],
  "core_theme": string (one sentence)
}}

DREAM:
{dream_text}

JSON:"""

PROMPT_PERBAIKAN = (
    "Your previous response was not valid JSON. "
    "Return ONLY a single valid JSON object matching the schema. No markdown."
)


def _panggil_model(pesan_sistem: str, pesan_user: str) -> str:
    model = get_model()
    with kunci_inferensi():
        respons = model.create_chat_completion(
            messages=[
                {"role": "system", "content": pesan_sistem},
                {"role": "user", "content": pesan_user},
            ],
            max_tokens=800,
            temperature=0.2,
            top_p=0.9,
        )
    isi = respons["choices"][0]["message"]["content"]
    return (isi or "").strip()


def _parse_dengan_retry(prompt: str, raw_awal: str) -> tuple[dict | None, str, str | None]:
    """Parse output; satu retry jika gagal. Kembalikan (data, raw_final, raw_pertama)."""
    raw_pertama: str | None = None
    raw = raw_awal
    data, error = parse_json_dari_teks(raw)

    if data is None:
        raw_pertama = raw
        raw = _panggil_model(
            SYSTEM_PROMPT,
            prompt + "\n\n" + PROMPT_PERBAIKAN + f"\n\nBroken output:\n{raw[:500]}",
        )
        data, error = parse_json_dari_teks(raw)

    if data is None:
        return None, raw, raw_pertama

    error_skema = validasi_skema_json(data)
    # Space CPU: default tanpa retry schema (hemat ~10–30s). Lokal: retry aktif.
    boleh_retry_skema = not os.getenv("SPACE_ID") or os.getenv("ONEIROS_ALLOW_SCHEMA_RETRY") == "1"
    if os.getenv("ONEIROS_SKIP_SCHEMA_RETRY") == "1":
        boleh_retry_skema = False

    if error_skema and boleh_retry_skema:
        if raw_pertama is None:
            raw_pertama = raw
        raw = _panggil_model(
            SYSTEM_PROMPT,
            prompt + "\n\n" + PROMPT_PERBAIKAN + f"\n\nSchema error: {error_skema}",
        )
        data2, error2 = parse_json_dari_teks(raw)
        if data2 is not None:
            data = data2
            error = error2
            error_skema = validasi_skema_json(data)
        else:
            error = error2 or error_skema

    if data is not None and error_skema:
        # Tetap normalisasi; parse_ok ditentukan nanti via entitas_memadai
        return data, raw, raw_pertama

    return data, raw, raw_pertama


def extract_entities(
    teks_mimpi: str,
) -> tuple[dict[str, Any], str, bool, str | None, str | None]:
    """
    Ekstrak entitas dari narasi mimpi.

    Returns:
        entities, raw_output, parse_ok, parse_error, raw_percobaan_pertama
    """
    teks = teks_mimpi.strip()
    # Hindari str.format pada teks mimpi (kurung kurawal dari user bisa merusak template)
    prompt = EXTRACTION_PROMPT.replace("{dream_text}", teks)
    raw_awal = _panggil_model(SYSTEM_PROMPT, prompt)
    data, raw, raw_pertama = _parse_dengan_retry(prompt, raw_awal)

    if data is None:
        kosong = validasi_dan_perbaiki(
            {
                "title": "Parse Failed",
                "mood": "mysterious",
                "characters": [],
                "places": [],
                "symbols": [],
                "connections": [],
                "core_theme": "",
            }
        )
        return kosong, raw, False, "JSON parse gagal", raw_pertama

    entities = validasi_dan_perbaiki(data)
    error_skema = validasi_skema_json(data)
    parse_ok = error_skema is None and entitas_memadai(entities)
    parse_error = None
    if error_skema:
        parse_error = error_skema
    elif not entitas_memadai(entities):
        parse_error = "Entitas kosong setelah normalisasi"

    return entities, raw, parse_ok, parse_error, raw_pertama


def extract_entities_dengan_waktu(
    teks_mimpi: str,
) -> tuple[dict[str, Any], str, bool, str | None, float, str | None]:
    """extract_entities + durasi detik + raw_percobaan_pertama."""
    mulai = time.perf_counter()
    entities, raw, ok, err, raw_pertama = extract_entities(teks_mimpi)
    return entities, raw, ok, err, time.perf_counter() - mulai, raw_pertama