app_kit/project.py

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Callable
import json
from pathlib import Path
import sys

from .demo_packs import DemoPack, read_text_inputs
from .embedding import extract_keywords
from .logging_utils import setup_logging
from .model_runtime import DEFAULT_MODEL_ID, generate_text_completion, load_llama, resolve_model_path
from .storage import SQLiteStore

LOGGER = setup_logging("p1_elder_paperwork", stream=sys.stderr)


@dataclass(frozen=True)
class ProjectSpec:
    key: str
    title: str
    description: str
    data_subdir: str
    search_enabled: bool
    inbox_label: str
    processor: Callable[[DemoPack, SQLiteStore, Any], dict[str, Any]]

    def run_pack(self, pack: DemoPack, store: SQLiteStore, config: Any) -> dict[str, Any]:
        return self.processor(pack, store, config)


def _base_result(
    pack: DemoPack,
    store: SQLiteStore,
    project: str,
    title: str,
    primary_text: str,
    payload: dict[str, Any],
    search_text: str | None = None,
) -> dict[str, Any]:
    record_id = store.store_record(project, pack.pack_id, title, primary_text, payload, status='ready')
    store.store_embedding(record_id, project, search_text or primary_text, metadata={'pack_id': pack.pack_id})
    return {'record_id': record_id, 'pack_id': pack.pack_id, 'project': project, **payload}


def _first_input_kind(pack: DemoPack) -> str:
    inputs = pack.manifest.get('inputs', [])
    if isinstance(inputs, list) and inputs:
        first = inputs[0]
        if isinstance(first, dict):
            kind = first.get('kind')
            if isinstance(kind, str) and kind.strip():
                return kind.strip()
    return 'text'


def _build_user_prompt(text: str, excerpt: str, keywords: tuple[str, ...], pack: DemoPack) -> str:
    questions = [
        'What is this document about?',
        'What action is requested?',
        'Is there a deadline or date mentioned?',
        'Is there an amount, phone number, or next step mentioned?',
    ]
    return (
        "You are an elder paperwork triage assistant. Answer using only the document content below. "
        "Return strict JSON with these keys: triage, summary, qa, citations, ocr_preview, safety. "
        "triage must be one of urgent, important, FYI, informational. "
        "qa must be a list of four objects, each with question, answer, and citation fields. "
        "citations must be a list of four objects, each with question and snippet fields. "
        "safety must be an object with missing_info_policy and invented_values fields. "
        "Do not invent facts; quote short source snippets for citations when possible.\n\n"
        f"PACK_ID: {pack.pack_id}\n"
        f"EXPECTED_SIGNALS: {json.dumps(pack.expected_signals, ensure_ascii=False)}\n"
        f"DOCUMENT_KIND: {_first_input_kind(pack)}\n"
        f"KEYWORDS: {', '.join(keywords) or 'none'}\n"
        f"DOCUMENT_EXCERPT:\n{excerpt}\n\n"
        f"DOCUMENT_TEXT:\n{text[:4000]}\n\n"
        "QUESTIONS:\n"
        + "\n".join(f"{idx}. {question}" for idx, question in enumerate(questions, start=1))
    )


def processor_p1(pack: DemoPack, store: SQLiteStore, config: Any) -> dict[str, Any]:
    text = read_text_inputs(pack).strip()
    if not text:
        raise RuntimeError("P1 requires document text for model inference; no readable text was found in the pack.")

    excerpt = text.splitlines()[0].strip() if text.splitlines() else text[:240].strip()
    keywords = tuple(extract_keywords(text))
    model = resolve_model_path(model_id=DEFAULT_MODEL_ID)
    llm = load_llama(str(model.model_path))

    def _final_text(raw: str) -> str:
        cleaned = raw.replace('</think>', '\n').replace('<think>', '\n')
        parts = [part.strip() for part in cleaned.splitlines() if part.strip()]
        return parts[-1] if parts else raw.strip()

    def _normalize_triage(raw: str) -> str:
        lowered = raw.lower()
        if 'urgent' in lowered:
            return 'urgent'
        if 'important' in lowered:
            return 'important'
        if 'fyi' in lowered:
            return 'FYI'
        return 'informational'

    questions = [
        'What is this document about?',
        'What action is requested?',
        'Is there a deadline or date mentioned?',
        'Is there an amount, phone number, or next step mentioned?',
    ]
    source_snippet = excerpt if excerpt else text[:180].strip()

    triage_prompt = (
        "Classify the document into exactly one label: urgent, important, FYI, informational.\n"
        "Rules:\n"
        "- urgent: immediate danger, same-day emergency, urgent medical action.\n"
        "- important: routine appointment notices, follow-up visits, insurance notices, medication lists, or forms needing action soon.\n"
        "- FYI: optional informational notices.\n"
        "- informational: archival or purely informational documents.\n"
        "For routine follow-up appointment notices, the correct label is important.\n"
        f"Document:\n{text[:3000]}\n\n"
        "Return exactly one label."
    )
    triage_raw, triage_meta = generate_text_completion(
        llm=llm,
        model=model,
        system_prompt='Return only the label.',
        user_prompt=triage_prompt,
        temperature=0.0,
        max_tokens=512,
    )
    triage = _normalize_triage(_final_text(triage_raw))

    summary_prompt = (
        "Document text:\n"
        f"{text[:3500]}\n\n"
        "Write one concise sentence summarizing the document."
    )
    summary_raw, summary_meta = generate_text_completion(
        llm=llm,
        model=model,
        system_prompt='Write one concise sentence only.',
        user_prompt=summary_prompt,
        temperature=0.0,
        max_tokens=96,
    )
    summary = _final_text(summary_raw)

    safety_raw, safety_meta = generate_text_completion(
        llm=llm,
        model=model,
        system_prompt='Return one short clause only.',
        user_prompt=(
            "Based on the document below, state whether more information is needed in one short clause. "
            "Use phrasing like 'missing info likely' or 'sufficient detail'.\n\n"
            f"{text[:2200]}"
        ),
        temperature=0.0,
        max_tokens=24,
    )
    safety_note = _final_text(safety_raw)

    qa_items: list[dict[str, str]] = []
    qa_stats: list[dict[str, Any]] = []
    for question in questions:
        answer_raw, answer_meta = generate_text_completion(
            llm=llm,
            model=model,
            system_prompt='Answer the question using only the document text.',
            user_prompt=(
                f"QUESTION: {question}\n\n"
                f"DOCUMENT TEXT:\n{text[:3500]}\n\n"
                "Return one short sentence only."
            ),
            temperature=0.0,
            max_tokens=96,
        )
        qa_items.append({'question': question, 'answer': _final_text(answer_raw), 'citation': source_snippet})
        qa_stats.append(answer_meta['generation_stats'])

    generation_stats = {
        'triage': triage_meta['generation_stats'],
        'summary': summary_meta['generation_stats'],
        'safety': safety_meta['generation_stats'],
        'qa': qa_stats,
    }
    inference_meta = {
        'model_id': model.model_id,
        'model_path': str(model.model_path),
        'model_source': model.source,
        'backend': model.backend,
        'generation_stats': generation_stats,
    }
    payload = {
        'triage': triage,
        'summary': summary,
        'qa': qa_items,
        'citations': [{'question': question, 'snippet': source_snippet} for question in questions],
        'ocr_preview': summary,
        'ocr_text': text,
        'safety': {'missing_info_policy': safety_note, 'invented_values': False},
        'inbox_items': [
            {
                'record_id': 'pending',
                'title': pack.pack_id,
                'triage': triage,
                'summary': summary,
                'file_type': _first_input_kind(pack),
            },
        ],
        'expected_signals': pack.expected_signals,
        'evidence': keywords,
        'inference': inference_meta,
        'model_id': inference_meta['model_id'],
        'adapter_name': inference_meta['backend'],
        'generation_stats': generation_stats,
        'source_excerpt': source_snippet,
    }
    search_text = ' '.join([text, summary, triage, ' '.join(keywords)])
    result = _base_result(pack, store, 'p1', f'P1: {pack.pack_id}', payload['summary'], payload, search_text)
    result['record_ids'] = [result['record_id']]
    result['documents'] = [payload]
    result['triage'] = triage
    result['summary'] = summary
    result['qa'] = qa_items
    result['citations'] = payload['citations']
    result['ocr_preview'] = payload['ocr_preview']
    result['ocr_text'] = payload['ocr_text']
    result['safety'] = payload['safety']
    result['inbox_items'] = payload['inbox_items']

    LOGGER.info(
        json.dumps(
            {
                'event': 'p1_model_inference',
                'pack_id': pack.pack_id,
                'model_id': inference_meta['model_id'],
                'adapter_name': inference_meta['backend'],
                'generation_stats': generation_stats,
                'triage': triage,
            },
            ensure_ascii=False,
        )
    )
    return result