tools/wikidata.py

import asyncio
from typing import Any

from SPARQLWrapper import JSON, SPARQLWrapper

WIKIDATA_ENDPOINT = "https://query.wikidata.org/sparql"

ARABIC_STOP_WORDS = {
    "في", "من", "على", "إلى", "عن", "مع", "هذا", "هذه", "أن", "إن",
    "كان", "كانت", "التي", "الذي", "وقد", "قد", "أو", "و", "ثم",
    "لكن", "بل", "حتى", "إذا", "لما", "بأن", "إنه", "أنه", "التقى",
    "أعلن", "أعلنت", "قال", "قالت", "أكد", "أكدت", "أشار", "أشارت",
}


def extract_entity_candidates(claim: str, *, limit: int = 3) -> list[str]:
    """Return likely entity tokens from an Arabic claim."""
    tokens = [token for token in claim.split() if len(token) > 3 and token not in ARABIC_STOP_WORDS]
    return tokens[:limit]


def wikidata_lookup(entity_name: str) -> list[dict[str, Any]]:
    """Synchronously search Wikidata by Arabic label, then English fallback."""
    sparql = SPARQLWrapper(WIKIDATA_ENDPOINT)
    sparql.addCustomHttpHeader(
        "User-Agent",
        "ArabicNewsAnalyzer/1.0 (graduation-project)",
    )

    query_ar = f"""
    SELECT DISTINCT ?item ?itemLabel ?itemDescription ?article WHERE {{
        ?item rdfs:label "{entity_name}"@ar .
        OPTIONAL {{
            ?article schema:about ?item ;
                     schema:inLanguage "ar" ;
                     schema:isPartOf <https://ar.wikipedia.org/> .
        }}
        SERVICE wikibase:label {{
            bd:serviceParam wikibase:language "ar,en" .
        }}
    }}
    LIMIT 3
    """
    sparql.setQuery(query_ar)
    sparql.setReturnFormat(JSON)

    try:
        raw = sparql.query().convert()
        bindings = raw.get("results", {}).get("bindings", [])
    except Exception:
        bindings = []

    if not bindings:
        query_en = f"""
        SELECT DISTINCT ?item ?itemLabel ?itemDescription WHERE {{
            ?item rdfs:label "{entity_name}"@en .
            SERVICE wikibase:label {{
                bd:serviceParam wikibase:language "ar,en" .
            }}
        }}
        LIMIT 2
        """
        sparql.setQuery(query_en)
        try:
            raw = sparql.query().convert()
            bindings = raw.get("results", {}).get("bindings", [])
        except Exception:
            bindings = []

    evidence: list[dict[str, Any]] = []
    for binding in bindings:
        label = binding.get("itemLabel", {}).get("value", entity_name)
        description = binding.get("itemDescription", {}).get("value", "")
        article_url = binding.get("article", {}).get("value", "")
        item_url = binding.get("item", {}).get("value", "")

        if not description:
            continue
        evidence.append(
            {
                "snippet": f"{label}: {description}",
                "source_url": article_url or item_url,
                "source": "wikidata",
                "rating": None,
            }
        )

    return evidence


async def search_wikidata(claim: str) -> list[dict[str, Any]]:
    """Search Wikidata for likely entities mentioned in a claim."""
    candidates = extract_entity_candidates(claim)
    if not candidates:
        return []

    loop = asyncio.get_event_loop()
    evidence: list[dict[str, Any]] = []

    for entity_name in candidates:
        try:
            results = await loop.run_in_executor(None, wikidata_lookup, entity_name)
            evidence.extend(results)
        except Exception as exc:
            print(f"[Wikidata] Error for '{entity_name}': {exc}")

    return evidence