rag/automation_metadata.py

from __future__ import annotations
from typing import Dict, Any, List

from core.books.storage import (
    fetch_raw_docs_for_user,
    mark_raw_status,
    upsert_document_metadata,
)
from schemas.books.sources_schema import SourceForAgent, DocMetadata
from agents.books.apa_agent import run_metadata_agent


# -------- limits (VERY IMPORTANT) --------
MAX_PAGES = 5
MAX_CHARS_PER_PAGE = 5000


def _build_agent_input(raw: Dict[str, Any]) -> Dict[str, Any]:
    pages = raw.get("pages_head") or []
    # pages = pages[:MAX_PAGES]
    # pages = [p[:MAX_CHARS_PER_PAGE] for p in pages]

    return SourceForAgent(
        source_url=raw["source_url"],
        source_type=raw.get("source_type", "pdf"),
        domain=raw.get("domain", ""),
        search_title=raw.get("search_title", "") or "",
        search_snippet=raw.get("search_snippet", "") or "",
        text_pages=pages,
    ).model_dump()


def process_user_metadata(user_id: str, book_id: str | None = None) -> Dict[str, Any]:
    pending = fetch_raw_docs_for_user(
        user_id=user_id,
        book_id=book_id,
        status="pending",
    )

    total = len(pending)
    processed = 0
    failed = 0
    items: List[Dict[str, Any]] = []
    tokens = []

    if total == 0:
        return {
            "user_id": user_id,
            "book_id": book_id,
            "total": 0,
            "processed": 0,
            "failed": 0,
            "items": [],
            "mode": "sequential",
        }

    for raw in pending:
        doc_id = raw["doc_id"]
        url = raw["source_url"]

        pages_head = raw.get("pages_head") or []
        if not pages_head:
            mark_raw_status(doc_id, "failed", "no_pages_head")
            failed += 1
            items.append(
                {
                    "doc_id": doc_id,
                    "url": url,
                    "status": "failed",
                    "reason": "no_pages_head",
                }
            )
            continue

        mark_raw_status(doc_id, "processing", "")
        
        try:
            agent_in = _build_agent_input(raw)
            out, token_usage = run_metadata_agent(agent_in)
        except Exception as e:
            print(f"⚠️ LLM internal exception for doc_id {doc_id}: {e}")
            out = None

        # ---- بعد الـ try/except ----
        if not isinstance(out, dict):
            mark_raw_status(doc_id, "failed", "llm_exception")
            failed += 1
            items.append(
                {
                    "doc_id": doc_id,
                    "url": url,
                    "status": "failed",
                    "reason": "llm_exception",
                }
            )
            continue
        
        if token_usage:    
            tokens.append(token_usage)

        print(f"Metadata agent output for doc_id {doc_id}: {type(out)} - {out}")    
        print(f"is_valid: {out.get('is_valid')}")
        if not out.get("is_valid"):
            reason = out.get("reason", "invalid")
            mark_raw_status(doc_id, "failed", reason)
            failed += 1
            items.append(
                {
                    "doc_id": doc_id,
                    "url": url,
                    "status": "failed",
                    "reason": reason,
                }
            )
            continue

        md = out.get("metadata") or {}
        if "title" not in md or "authors" not in md:
            mark_raw_status(doc_id, "failed", "missing_title_or_authors")
            failed += 1
            items.append(
                {
                    "doc_id": doc_id,
                    "url": url,
                    "status": "failed",
                    "reason": "missing_title_or_authors",
                }
            )
            continue

        try:
            doc_md = DocMetadata(
                doc_id=doc_id,
                title=md["title"],
                authors=md["authors"],
                year=md.get("year"),
                publisher_or_journal=md.get("publisher_or_journal") or "",
                normalized_source_type=out.get("normalized_source_type", "pdf"),
                apa7=md.get("apa7") or "",
                metadata=out,
            )

            upsert_document_metadata(doc_md)
            mark_raw_status(doc_id, "processed", "")
            processed += 1
            items.append(
                {
                    "doc_id": doc_id,
                    "url": url,
                    "status": "processed",
                    "title": doc_md.title,
                }
            )

        except Exception:
            mark_raw_status(doc_id, "failed", "db_failed")
            failed += 1
            items.append(
                {
                    "doc_id": doc_id,
                    "url": url,
                    "status": "failed",
                    "reason": "db_failed",
                }
            )

    return {
        "user_id": user_id,
        "book_id": book_id,
        "total": total,
        "processed": processed,
        "failed": failed,
        "items": items,
        "mode": "sequential",
        "tokens": tokens,
    }