agents/claim_extractor.py

import asyncio
import json
import logging
import os
import re

from groq import AsyncGroq
from dotenv import load_dotenv

load_dotenv()

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
)
logger = logging.getLogger(__name__)

GROQ_MODEL = "llama-3.3-70b-versatile"
MAX_TOKENS = 1024
TEMPERATURE = 0.0
MAX_RETRIES = 2

SYSTEM_PROMPT = """أنت نظام متخصص في استخراج الادعاءات الإخبارية القابلة للتحقق من النصوص العربية.

━━━ تعريف الادعاء القابل للتحقق ━━━
الادعاء القابل للتحقق هو جملة تحتوي على:
- حدث أو فعل موصوف (اجتماع، إعلان، تسجيل أرقام، فوز، وفاة، توقيع، إلخ)
- جهة فاعلة محددة (شخص، منظمة، دولة، وزارة، إلخ) أو رقم قابل للقياس

━━━ ما يجب استخراجه ━━━
- الادعاءات الواقعية: أحداث، أرقام، نتائج، تصريحات منسوبة لجهة محددة
- الادعاءات الرقمية: "سُجّلت 300 إصابة"، "ارتفع المؤشر بنسبة 5%"
- نتائج الأحداث: فوز، خسارة، توقيع، إلغاء، إطلاق سراح

━━━ ما يجب تجاهله ━━━
- الآراء والتحليلات: "يرى المحللون أن..."، "يُتوقع أن..."
- التعليقات المبهمة: "الوضع صعب ومعقد"
- الخلفية الوصفية التي لا تحتوي فعلاً محدداً
- الجمل التمهيدية والختامية الإخبارية

━━━ قاعدة الاكتفاء الذاتي (مهمة جداً) ━━━
كل ادعاء يجب أن يكون مفهوماً تماماً بمعزل عن النص الأصلي.
إذا كانت الجملة تستخدم ضميراً أو فعلاً مضمراً (وأضاف، وأضافت، وتابع، وأكد، وأشار...),
استبدل الضمير باسم الجهة الفاعلة الصريح المذكور في السياق.

مثال:
  النص: "أعلنت إدارة النادي توقيع مبابي، وأضافت أن التقديم سيكون في برنابيو"
  ❌ خطأ:  "وأضافت أن التقديم سيكون في برنابيو"
  ✓ صحيح: "أعلنت إدارة النادي أن التقديم الرسمي لمبابي سيكون في ملعب سانتياغو برنابيو"

━━━ قواعد الإخراج ━━━
1. أجب فقط بـ JSON صالح — لا نص قبله أو بعده، لا علامات markdown
2. كل ادعاء جملة عربية كاملة ومستقلة مع فاعل صريح دائماً
3. يمكنك إعادة صياغة الجملة فقط لاستبدال الضمير بالاسم الصريح — لا تضف معلومات جديدة
4. إذا لم يوجد أي ادعاء قابل للتحقق، أعد قائمة فارغة

شكل الإجابة (JSON فقط):
{
  "claims": [
    "الادعاء الأول مع فاعل صريح",
    "الادعاء الثاني مع فاعل صريح",
    "الادعاء الثالث مع فاعل صريح"
  ]
}"""


class ClaimExtractor:
    def __init__(self):
        api_key = os.environ.get("GROQ_API_KEY2")
        if not api_key:
            logger.warning(
                "GROQ_API_KEY2 not set — ClaimExtractor will fail at runtime.")
        self.client = AsyncGroq(api_key=api_key)

    async def run(self, text: str) -> list[str]:
        if not text or not text.strip():
            return []

        result = await self._call_groq(text)
        if result is not None:
            return result

        # Retry once
        logger.warning(
            "[ClaimExtractor] First attempt failed — retrying once.")
        result = await self._call_groq(text)
        if result is not None:
            return result

        logger.error(
            "[ClaimExtractor] Both attempts failed — returning empty list.")
        return []

    async def _call_groq(self, text: str) -> list[str] | None:
        try:
            response = await self.client.chat.completions.create(
                model=GROQ_MODEL,
                messages=[
                    {"role": "system", "content": SYSTEM_PROMPT},
                    {"role": "user", "content": f"استخرج الادعاءات القابلة للتحقق من النص التالي:\n\n{text.strip()}"},
                ],
                max_tokens=MAX_TOKENS,
                temperature=TEMPERATURE,
            )
        except Exception as e:
            logger.error(f"[ClaimExtractor] Groq API error: {e}")
            return None

        raw = response.choices[0].message.content or ""
        return self._parse_response(raw)

    def _parse_response(self, raw: str) -> list[str] | None:
        clean = re.sub(r"```json\s*", "", raw)
        clean = re.sub(r"```\s*", "", clean)
        clean = clean.strip()

        try:
            parsed = json.loads(clean)
            return self._validate_claims(parsed)
        except json.JSONDecodeError:
            pass

        match = re.search(r"\{.*\}", clean, re.DOTALL)
        if match:
            try:
                parsed = json.loads(match.group())
                return self._validate_claims(parsed)
            except json.JSONDecodeError:
                pass

        logger.warning(f"[ClaimExtractor] Could not parse JSON: {raw[:200]}")
        return None

    def _validate_claims(self, parsed: dict) -> list[str]:
        claims = parsed.get("claims", [])
        if not isinstance(claims, list):
            return []

        validated = []
        for claim in claims:
            if isinstance(claim, str):
                claim = claim.strip()
                if claim:
                    validated.append(claim)

        return validated