New Hire Handbook Q&A

import base64
import io
import json
from typing import List

from pypdf import PdfReader

MAX_CHARS = int([[[MAX_CHARS|6000]]])

DOCUMENT_REGISTRY: List[dict] = []


def _truncate_text(text: str, limit: int = MAX_CHARS) -> str:
    text = text or ""
    text = text.replace("\x00", " ")
    text = "\n".join(line.rstrip() for line in text.splitlines())
    text = text.strip()
    if len(text) <= limit:
        return text
    return text[:limit] + "\n\n[Truncated for indexing]"


def _extract_pdf_text(pdf_bytes: bytes) -> str:
    reader = PdfReader(io.BytesIO(pdf_bytes))
    pages = []
    for page in reader.pages:
        try:
            pages.append(page.extract_text() or "")
        except Exception:
            pages.append("")
    return "\n\n".join(pages)


async def _add_uploaded_document(name: str, mime_type: str, base64_data: str) -> str:
    raw = base64.b64decode(base64_data)
    lower_name = (name or "document").lower()
    doc_type = "text"

    if lower_name.endswith(".pdf") or mime_type == "application/pdf":
        text = _extract_pdf_text(raw)
        doc_type = "pdf"
    else:
        text = raw.decode("utf-8", errors="ignore")
        doc_type = "text"

    text = _truncate_text(text)
    if not text.strip():
        return json.dumps({"ok": False, "error": f"No readable text found in {name}."})

    chunks = await agentop_rag.add_document(name, text)
    DOCUMENT_REGISTRY.append({"name": name, "type": doc_type, "chunks": chunks})
    return json.dumps({"ok": True, "name": name, "type": doc_type, "chunks": chunks})


async def _clear_all_documents() -> str:
    await agentop_rag.clear()
    DOCUMENT_REGISTRY.clear()
    return "ok"


async def process_user_query(query: str) -> str:
    q = (query or "").strip()
    if not q:
        return "Please ask a question about your uploaded onboarding documents."
    if not DOCUMENT_REGISTRY:
        return "No documents are loaded yet. Please upload your employee handbook or related guides first."

    hits = await agentop_rag.search(q, k=6)
    if not hits:
        return (
            "Not covered in the uploaded documents.\n\n"
            "I could not find relevant language in the handbook library you uploaded. "
            "Try uploading additional documents or rephrasing your question."
        )

    usable_hits = [h for h in hits if (h.get("text") or "").strip()]
    if not usable_hits:
        return (
            "Not covered in the uploaded documents.\n\n"
            "I found no readable supporting text for this question in the uploaded files."
        )

    context_blocks = []
    for i, hit in enumerate(usable_hits[:6], start=1):
        source = hit.get("source", "Unknown document")
        text = (hit.get("text", "") or "").strip()
        score = hit.get("score", 0)
        snippet = text[:1200]
        context_blocks.append(f"[{i}] Source: {source}\nRelevance: {score}\nExcerpt: {snippet}")

    grounded_prompt = (
        "Answer the employee onboarding question using only the provided document excerpts.\n"
        "Rules:\n"
        "1. Do not use outside knowledge.\n"
        "2. If the answer is not stated or not clear in the excerpts, say exactly: Not covered in the uploaded documents.\n"
        "3. Cite supporting documents by source name in a Sources section.\n"
        "4. Be concise and factual.\n"
        "5. Do not invent policy details, numbers, dates, or eligibility rules.\n\n"
        f"Question:\n{q}\n\n"
        f"Document excerpts:\n\n{'\n\n'.join(context_blocks)}\n\n"
        "Required output format:\n"
        "Answer: <short grounded answer or 'Not covered in the uploaded documents.'>\n"
        "Sources:\n- <document name>\n- <document name>"
    )

    reply = await process_user_query_wllama(
        grounded_prompt,
        globals().get("TEMPLATE_SYSTEM_PROMPT", "")
    )
    return (reply or "").strip()