MapMil/centers/parsing/workers/workers/llm_extract.py

"""DeepSeek / OpenAI-compatible LLM extraction for unstructured text."""

from __future__ import annotations

import json
import logging
import os
import re
from datetime import datetime, timezone
from typing import Any

import httpx

logger = logging.getLogger("cp-worker.llm")

COORDS_RE = re.compile(r"(-?\d{1,3}\.\d+)\s*,\s*(-?\d{1,3}\.\d+)")

DEFAULT_EXTRACT_SCHEMA: dict[str, str] = {
    "title": "string — short event title",
    "locality": "string — place / settlement name",
    "event_date": "string — date as DD.MM.YYYY or YYYY-MM-DD if known",
    "description": "string — concise event summary",
    "coords": "string — latitude, longitude if present else empty",
    "topic": "string — short topic tag",
}

DEFAULT_INSTRUCTION = (
    "Extract structured military/news event fields from the text. "
    "If the text is not an event, return is_event=false. "
    "Respond with a single JSON object only."
)


def llm_enabled() -> bool:
    return bool(os.getenv("DEEPSEEK_API_KEY", "").strip())


def llm_settings() -> dict[str, str]:
    return {
        "api_key": os.getenv("DEEPSEEK_API_KEY", "").strip(),
        "base_url": os.getenv("DEEPSEEK_BASE_URL", "https://api.deepseek.com").rstrip("/"),
        "model": os.getenv("DEEPSEEK_MODEL", "deepseek-chat"),
    }


async def extract_event_fields(
    text: str,
    *,
    extract_schema: dict[str, str] | None = None,
    instruction: str | None = None,
) -> dict[str, Any]:
    """Ask DeepSeek to fill schema fields from free text. Returns dict (+ is_event)."""
    settings = llm_settings()
    if not settings["api_key"]:
        raise RuntimeError(
            "DEEPSEEK_API_KEY is not set. Add it to .env for LLM extract_mode."
        )

    schema = extract_schema or DEFAULT_EXTRACT_SCHEMA
    instr = instruction or DEFAULT_INSTRUCTION
    schema_lines = "\n".join(f"- {k}: {v}" for k, v in schema.items())
    user_prompt = (
        f"{instr}\n\n"
        f"Fields to extract:\n{schema_lines}\n\n"
        'Return JSON: {"is_event": true|false, "fields": {<field>: <string>}}\n\n'
        f"Text:\n{text[:12000]}"
    )

    payload = {
        "model": settings["model"],
        "messages": [
            {
                "role": "system",
                "content": (
                    "You extract structured event data for a geoint map. "
                    "Output valid JSON only, no markdown."
                ),
            },
            {"role": "user", "content": user_prompt},
        ],
        "temperature": 0.1,
        "response_format": {"type": "json_object"},
    }

    url = f"{settings['base_url']}/chat/completions"
    async with httpx.AsyncClient(timeout=90.0) as client:
        response = await client.post(
            url,
            headers={
                "Authorization": f"Bearer {settings['api_key']}",
                "Content-Type": "application/json",
            },
            json=payload,
        )
        response.raise_for_status()
        data = response.json()

    content = data["choices"][0]["message"]["content"]
    parsed = json.loads(content)
    fields = parsed.get("fields") if isinstance(parsed.get("fields"), dict) else parsed
    if not isinstance(fields, dict):
        fields = {}
    # Normalize to strings for known keys
    result = {key: str(fields.get(key) or "").strip() for key in schema}
    result["is_event"] = bool(parsed.get("is_event", True))
    return result


def fields_to_ingest(
    *,
    source_type: str,
    source_url: str,
    raw_text: str,
    fields: dict[str, Any],
    domain_profile: str = "llm",
    extra_metadata: dict | None = None,
) -> dict:
    lat, lng = parse_coords(str(fields.get("coords") or ""))
    description = str(fields.get("description") or raw_text)[:8000]
    locality = str(fields.get("locality") or "")
    title = str(fields.get("title") or locality or description.splitlines()[0][:120])
    topic = str(fields.get("topic") or domain_profile)
    event_date = parse_date(str(fields.get("event_date") or ""))

    meta = {
        "extract_mode": "llm",
        "extracted": {k: fields.get(k) for k in fields if k != "is_event"},
    }
    if extra_metadata:
        meta.update(extra_metadata)

    return {
        "source_type": source_type,
        "source_url": source_url,
        "raw_text": raw_text[:20000],
        "title": title[:255],
        "description": description,
        "locality": locality,
        "latitude": lat,
        "longitude": lng,
        "event_date": event_date.isoformat() if event_date else None,
        "region": locality or None,
        "topic": topic,
        "tags": [source_type, "llm", domain_profile],
        "metadata": meta,
    }


def parse_coords(raw: str) -> tuple[float | None, float | None]:
    match = COORDS_RE.search(raw or "")
    if not match:
        return None, None
    return float(match.group(1)), float(match.group(2))


def parse_date(raw: str) -> datetime | None:
    if not raw:
        return None
    raw = raw.strip()
    for fmt in ("%d.%m.%Y", "%d.%m.%y", "%d/%m/%Y", "%d/%m/%y", "%Y-%m-%d"):
        try:
            return datetime.strptime(raw, fmt).replace(tzinfo=timezone.utc)
        except ValueError:
            continue
    return None


def crawl4ai_llm_config():
    """Build Crawl4AI LLMConfig for DeepSeek (OpenAI-compatible)."""
    from crawl4ai import LLMConfig  # type: ignore

    settings = llm_settings()
    if not settings["api_key"]:
        raise RuntimeError("DEEPSEEK_API_KEY is not set")
    return LLMConfig(
        provider=f"openai/{settings['model']}",
        api_token=settings["api_key"],
        base_url=settings["base_url"],
    )