shira-hermes/api/services/kb/classifier.py

"""AI metadata extraction for uploaded KB documents.

Phase 6 (Task #20). When a user uploads a PDF/DOCX/TXT through the
batch UI, the background classifier reads the first ~3 pages of text,
calls Claude (via the existing ai-gateway proxy), and returns a draft
metadata payload — kind, title, identifier, subject, labels, dates,
and a short summary. The user reviews these on the new bulk-upload
review screen and edits before committing.

Design notes (see ~/.claude/plans/hidden-tickling-ullman.md):
  • Tool-calling guarantees JSON shape — no prompt-engineering
    fragility around "please return only JSON".
  • Sonnet model (env CLAUDE_MODEL). Hebrew legal text is hard;
    Haiku gets citations wrong too often.
  • 15s timeout, semaphore-bounded to 5 concurrent calls so a 30-file
    drop doesn't trip ai-gateway rate limits.
  • Fail-soft: any error → empty dict. Caller transitions the job to
    `awaiting_review` with `ai_suggestions = {}`, and the UI shows
    "אין הצעות AI — מלא ידנית".

The classifier is an aid, not a gate. Wrong suggestions are an
annoyance; the user always confirms before chunks land in the index.
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
from typing import Optional

from openai import AsyncOpenAI

logger = logging.getLogger("shira.kb.classifier")

# Bound concurrency so a 10–30 file batch doesn't fan out 30 parallel
# Claude calls. Five at a time gives reasonable wall-clock latency
# (~6–10s for 30 files) without overwhelming ai-gateway.
_CONCURRENCY = int(os.environ.get("KB_CLASSIFIER_CONCURRENCY", "5"))
# Empirically the ai-gateway → Claude tool-call round-trip with a ~3KB
# Hebrew system prompt + 5KB user content runs 18-25s. Pad to 45s so
# we don't fail-soft on the long-tail latency.
_TIMEOUT_S = float(os.environ.get("KB_CLASSIFIER_TIMEOUT_S", "45"))
_MAX_INPUT_CHARS = 20_000   # ~3 pages of Hebrew legal text

_semaphore = asyncio.Semaphore(_CONCURRENCY)


_SYSTEM_PROMPT = """אתה מסווג מסמכים משפטיים בעברית עבור מאגר ידע משפטי.
תקבל את הטקסט של 3 העמודים הראשונים של מסמך, ועליך לחלץ ממנו מטא-דאטה.

הסיווג חייב לכבד את החוקים הבאים:

‏1. **kind** — אחד מהבאים בלבד:
   • law          — חוק (הצעת חוק, ספר חוקים, תיקון לחוק)
   • regulation   — תקנות
   • circular     — חוזר רשות / משרד / מוסד (חוזר נכות, חוזר מנכ"ל, חוזר לשכה רפואית)
   • caselaw      — פסק דין / החלטה. כותרת מתחילה לרוב ב-"בבית המשפט", "ע"א", "בג"ץ", "ב"ל",
                    "תיק", או שם תיק עם מספר.
   • tool         — כלי הערכה / מאמר / מדריך אקדמי. אינו מסמך משפטי פורמלי.
                    דוגמאות: GMFCS, MACS, CARS, Mini-MACS, מדריך הערכת תפקוד מאוניברסיטה.

‏2. **title** — כותרת המסמך כפי שהיא מופיעה בעמוד הראשון. אל תוסיף, אל תקצר.
   ללא מרכאות סוגרות.

‏3. **identifier** — זיהוי קצר ומדויק. דוגמאות לפי kind:
   • law/regulation: "חוק הביטוח הלאומי, התשנ\"ה–1995", "תקנה 37(א)"
   • circular: "חוזר נכות 1990/02", "חוזר מנכ\"ל 14/2018"
   • caselaw:  "ע\"א 1234/24", "בג\"ץ 5678/22", "ב\"ל 9876-12-23"
   • tool:     שם הגוף המוציא ("CanChild — McMaster University")
   אם אין זיהוי ברור — null. אל תמציא.

‏4. **subject** — תיאור תת-נושא קצר (≤80 תווים). למשל: "ילד נכה / אוטיזם",
   "ילד נכה / שיתוק מוחין", "נפגעי עבודה / מחלת מקצוע".

‏5. **labels** — עד 3 תוויות לסיווג מעמיק. כל תווית בפורמט "קטגוריה / תת-קטגוריה" אם ניתן.
   דוגמאות:
   • "ילד נכה / אוטיזם"
   • "ילד נכה / שיתוק מוחין"
   • "ילד נכה / תלות בזולת"
   • "ילד נכה / עיכוב התפתחותי"
   • "ילד נכה / מומים בגפיים"
   • "נפגעי עבודה / שמיעה"

‏6. **published_at / effective_at** — בפורמט YYYY-MM-DD בלבד.
   אם רק שנה ידועה — YYYY-01-01. אם לא ידוע — null.
   effective_at הוא תאריך התחילה לתוקף (אם מצוין במפורש בנפרד מתאריך הפרסום).

‏7. **summary** — משפט אחד באורך 80–200 תווים שמסכם על מה המסמך. ללא קלישאות
   ("מסמך חשוב"...). תאר את התוכן הספציפי בקצרה.

‏**אסור: לא להמציא נתונים שלא מופיעים בטקסט. עדיף null על ניחוש.**
"""


_TOOL_SCHEMA = {
    "type": "function",
    "function": {
        "name": "extract_kb_metadata",
        "description": "Return extracted metadata for the document.",
        "parameters": {
            "type": "object",
            "required": ["kind", "title", "summary"],
            "properties": {
                "kind": {
                    "type": "string",
                    "enum": ["law", "regulation", "circular", "caselaw", "tool"],
                },
                "title":        {"type": "string", "maxLength": 500},
                "identifier":   {"type": ["string", "null"], "maxLength": 200},
                "subject":      {"type": ["string", "null"], "maxLength": 80},
                "labels": {
                    "type": "array",
                    "maxItems": 3,
                    "items": {"type": "string", "maxLength": 80},
                },
                "published_at": {
                    "type": ["string", "null"],
                    "pattern": r"^\d{4}-\d{2}-\d{2}$",
                },
                "effective_at": {
                    "type": ["string", "null"],
                    "pattern": r"^\d{4}-\d{2}-\d{2}$",
                },
                "summary": {"type": "string", "minLength": 40, "maxLength": 250},
            },
            "additionalProperties": False,
        },
    },
}


async def classify(
    *,
    text: str,
    filename: str,
    topic_name: Optional[str] = None,
) -> dict:
    """Best-effort extraction. Returns {} on any error.

    `text` is the first ~3 pages of the document (caller decides — we
    truncate to _MAX_INPUT_CHARS just in case). `topic_name` lets the
    LLM bias label suggestions toward the active legal domain
    ("ביטוח לאומי" vs "דיני עבודה").
    """
    if not text or not text.strip():
        logger.info("[classifier] empty text for %s — skipping", filename)
        return {}

    base = (os.environ.get("AI_GATEWAY_URL") or "http://localhost:3000").rstrip("/")
    api_key = os.environ.get("AI_GATEWAY_API_KEY", "")
    if not api_key:
        logger.warning("[classifier] AI_GATEWAY_API_KEY not set — skipping")
        return {}

    truncated = text[:_MAX_INPUT_CHARS]
    user_msg = (
        f"topic_hint: {topic_name or 'כללי'}\n"
        f"filename: {filename}\n\n"
        f"--- מסמך ---\n{truncated}\n--- סוף ---"
    )

    client = AsyncOpenAI(base_url=f"{base}/v1", api_key=api_key)
    async with _semaphore:
        try:
            resp = await asyncio.wait_for(
                client.chat.completions.create(
                    model=os.environ.get("CLAUDE_MODEL", "sonnet"),
                    messages=[
                        {"role": "system", "content": _SYSTEM_PROMPT},
                        {"role": "user", "content": user_msg},
                    ],
                    tools=[_TOOL_SCHEMA],
                    tool_choice={
                        "type": "function",
                        "function": {"name": "extract_kb_metadata"},
                    },
                    max_tokens=600,
                    temperature=0.1,
                ),
                timeout=_TIMEOUT_S,
            )
        except asyncio.TimeoutError:
            logger.warning("[classifier] timeout on %s after %ss", filename, _TIMEOUT_S)
            return {}
        except Exception as e:
            logger.warning("[classifier] error on %s: %s", filename, e)
            return {}

    # Extract the tool call. We forced the model to call exactly one tool;
    # if it didn't, we got nothing useful.
    try:
        choice = resp.choices[0]
        tool_calls = choice.message.tool_calls or []
        if not tool_calls:
            logger.warning("[classifier] no tool call returned for %s", filename)
            return {}
        args_str = tool_calls[0].function.arguments
        parsed = json.loads(args_str)
    except (KeyError, IndexError, AttributeError, json.JSONDecodeError) as e:
        logger.warning("[classifier] failed to parse tool call for %s: %s", filename, e)
        return {}

    return _normalize(parsed)


def _normalize(d: dict) -> dict:
    """Coerce types and drop empty values from the LLM output."""
    out: dict = {}
    if isinstance(d.get("kind"), str) and d["kind"] in (
        "law", "regulation", "circular", "caselaw", "tool"
    ):
        out["kind"] = d["kind"]
    for key in ("title", "identifier", "subject", "summary"):
        v = d.get(key)
        if isinstance(v, str) and v.strip():
            out[key] = v.strip()
    for key in ("published_at", "effective_at"):
        v = d.get(key)
        if isinstance(v, str) and len(v) == 10 and v[4] == "-" and v[7] == "-":
            out[key] = v
    labels = d.get("labels") or []
    if isinstance(labels, list):
        clean: list[str] = []
        for raw in labels[:3]:
            if isinstance(raw, str) and raw.strip():
                clean.append(raw.strip())
        if clean:
            out["labels"] = clean
    return out