backend/app/workflow/nodes/knowledge.py

from __future__ import annotations

import json
import logging
import re

from app.core.sanitizer import sanitize_for_llm
from app.llm.client import LLMClient
from app.repositories.knowledge_repository import KnowledgeRepository
from app.schemas.knowledge import KnowledgeExtractionResult
from app.workflow.state import WorkflowState

logger = logging.getLogger(__name__)

# Maximum characters of document text per batch.
# Ollama/local has no limits. Cloud providers vary.
# 15K chars (~4K tokens) is a good balance for all providers.
MAX_TEXT_CHARS = 15000

SYSTEM_PROMPT = """\
/no_think
You are a knowledge extraction engine. Read document sections and extract \
discrete, atomic knowledge items. Return ONLY valid JSON. No thinking, no explanation.

Types: ENTITY, CLAIM, METHOD, METRIC, OBSERVATION, DATE, DATASET

JSON format:
{"items": [{"type": "CLAIM", "title": "short name", "value": "what the document says", "confidence": 0.9, "evidence": [{"quote": "exact verbatim quote from text", "page_number": 1, "section": "section title or null"}]}]}

Rules:
- Only extract explicitly stated facts. Never invent.
- One fact per item (atomic).
- Evidence quote must be verbatim from the source text.
- Confidence: 0.9+ = explicitly stated, 0.7-0.9 = clearly implied, below 0.7 = skip.
- Do not extract boilerplate (headers, footers, copyright).
- If nothing to extract, return {"items": []}
- Output ONLY the JSON object. No other text.
"""


def knowledge(
    state: WorkflowState,
    llm_client: LLMClient,
    repository: KnowledgeRepository,
) -> WorkflowState:
    """
    Extract structured knowledge from document sections using native JSON mode.
    Processes sections in batches that fit within MAX_TEXT_CHARS, then aggregates
    all extracted items before persisting.
    """

    state.current_node = "KNOWLEDGE_EXTRACTION"

    if not state.extracted_sections:
        state.metadata["knowledge_items_created"] = 0
        return state

    from app.workflow.progress import report_progress
    report_progress(state.workflow_run_id, "KNOWLEDGE_EXTRACTION")

    # Build individual section strings
    section_strings = []
    for index, section in enumerate(state.extracted_sections):
        metadata = section.get("metadata", {})
        text = section.get("text", "")

        sanitized = sanitize_for_llm(text, context=f"section_{index + 1}")

        section_strings.append(
            f"[SECTION {index + 1}]\n"
            f"PAGE: {metadata.get('page')}\n"
            f"SECTION TITLE: {metadata.get('section_title')}\n"
            f"SECTION TYPE: {metadata.get('section_type')}\n"
            f"ELEMENT TYPE: {metadata.get('element_type')}\n\n"
            f"{sanitized.content}\n"
        )

        if sanitized.is_suspicious:
            state.metadata.setdefault("injection_warnings", []).append(
                {"section": index + 1, "patterns": sanitized.detected_patterns}
            )

    # Group sections into batches that fit within MAX_TEXT_CHARS
    batches = _batch_sections(section_strings, MAX_TEXT_CHARS)
    logger.info(
        "Knowledge extraction: %d sections grouped into %d batch(es)",
        len(section_strings), len(batches),
    )

    # Track LLM usage
    from app.core.tracing.run_tracker import get_or_create_tracker
    tracker = get_or_create_tracker(str(state.workflow_run_id))
    tracker.start_stage("knowledge_extraction")

    # Process each batch sequentially and aggregate results
    all_items = []
    for batch_idx, batch_text in enumerate(batches):
        logger.info(
            "  Batch %d/%d (%d chars)",
            batch_idx + 1, len(batches), len(batch_text),
        )
        result = _extract_with_json_mode(llm_client, batch_text, tracker)
        if result is not None:
            all_items.extend(result.items)
            logger.info("    → extracted %d items", len(result.items))
        else:
            logger.warning("    → batch %d failed, skipping", batch_idx + 1)
            

    tracker.end_stage("knowledge_extraction")

    if not all_items:
        logger.error("All knowledge extraction batches failed or returned no items")
        state.metadata["knowledge_items_created"] = 0
        state.metadata["extraction_error"] = "No items extracted from any batch"
        return state

    # Persist all aggregated knowledge items
    created_count = 0
    for item in all_items:
        repository.create_with_evidence(
            workspace_id=state.workspace_id,
            document_version_id=state.document_version_id,
            data=item,
        )
        created_count += 1

    state.metadata["knowledge_items_created"] = created_count
    return state


def _batch_sections(section_strings: list[str], max_chars: int) -> list[str]:
    """
    Group section strings into batches where each batch's combined text
    stays under max_chars. If a single section exceeds max_chars, it gets
    its own batch (truncated to max_chars).
    """
    batches = []
    current_batch = []
    current_size = 0

    for section in section_strings:
        section_len = len(section)

        # If adding this section would exceed the limit, flush current batch
        if current_batch and (current_size + section_len) > max_chars:
            batches.append("\n".join(current_batch))
            current_batch = []
            current_size = 0

        current_batch.append(section)
        current_size += section_len

    # Flush remaining
    if current_batch:
        batches.append("\n".join(current_batch))

    # Safety: truncate any individual batch that somehow exceeds max_chars
    return [b[:max_chars] for b in batches]


def _extract_with_json_mode(
    llm_client: LLMClient,
    document_text: str,
    tracker,
) -> KnowledgeExtractionResult | None:
    """
    Extract one knowledge batch with defensive JSON parsing.

    First tries the full batch once. If the provider returns malformed or
    truncated output, the batch is split into two halves and each half is
    attempted once. This avoids the old 15K -> 7.5K -> 3.75K retry explosion.
    """
    from langchain_core.messages import SystemMessage, HumanMessage

    model = llm_client.get_client()

    def _record_usage(response) -> None:
        usage = getattr(response, "usage_metadata", None)

        if usage:
            tracker.record_llm_usage(
                "knowledge_extraction",
                input_tokens=int(usage.get("input_tokens", 0) or 0),
                output_tokens=int(usage.get("output_tokens", 0) or 0),
            )
            logger.info(
                "Tokens captured: in=%d out=%d",
                int(usage.get("input_tokens", 0) or 0),
                int(usage.get("output_tokens", 0) or 0),
            )
            return

        meta = getattr(response, "response_metadata", None) or {}
        usage = meta.get("token_usage") or meta.get("usage") or {}

        input_tokens = (
            usage.get("prompt_tokens", 0)
            or usage.get("input_tokens", 0)
            or 0
        )
        output_tokens = (
            usage.get("completion_tokens", 0)
            or usage.get("output_tokens", 0)
            or 0
        )

        if input_tokens or output_tokens:
            tracker.record_llm_usage(
                "knowledge_extraction",
                input_tokens=int(input_tokens),
                output_tokens=int(output_tokens),
            )
            logger.info(
                "Tokens captured: in=%d out=%d",
                int(input_tokens),
                int(output_tokens),
            )
        else:
            logger.warning(
                "Token usage unavailable. Response metadata keys: %s",
                list(meta.keys()),
            )

    def _invoke(text: str) -> KnowledgeExtractionResult | None:
        user_message = (
            "Extract knowledge items from the following document sections.\n\n"
            f"{text}"
        )

        logger.info(
            "LLM request: system=%d chars, user=%d chars, estimated_input_tokens=%d",
            len(SYSTEM_PROMPT),
            len(user_message),
            (len(SYSTEM_PROMPT) + len(user_message)) // 4,
        )

        try:
            request_model = model.bind(max_tokens=8192)

            response = request_model.invoke([
                SystemMessage(content=SYSTEM_PROMPT),
                HumanMessage(content=user_message),
            ])

            _record_usage(response)

            parsed = _parse_json_response(response)

            if parsed is not None:
                logger.info(
                    "Knowledge response parsed successfully: %d items",
                    len(parsed.items),
                )
                return parsed

            logger.warning("Knowledge response could not be recovered as JSON")
            return None

        except Exception as e:
            logger.warning(
                "Knowledge extraction LLM call failed: %s",
                str(e)[:300],
            )
            return None

    # First attempt: full batch.
    result = _invoke(document_text)

    if result is not None:
        return result

    # Only split once. Do NOT recursively go 15K -> 7.5K -> 3.75K.
    if len(document_text) > 3000:
        midpoint = len(document_text) // 2

        # Prefer splitting at a section boundary.
        split_at = document_text.rfind(
            "\n[SECTION ",
            0,
            midpoint + 500,
        )

        if split_at < 2000:
            split_at = midpoint

        left = document_text[:split_at].strip()
        right = document_text[split_at:].strip()

        recovered_items = []

        logger.info(
            "Splitting failed knowledge batch: left=%d chars right=%d chars",
            len(left),
            len(right),
        )

        if left:
            left_result = _invoke(left)

            if left_result is not None:
                recovered_items.extend(left_result.items)

        if right:
            right_result = _invoke(right)

            if right_result is not None:
                recovered_items.extend(right_result.items)

        if recovered_items:
            logger.info(
                "Recovered split batch with %d total items",
                len(recovered_items),
            )

            return KnowledgeExtractionResult(
                items=recovered_items
            )

    return None


def _extract_json_object(text: str) -> str | None:
    """Find the first complete outer JSON object in provider output."""

    text = text.strip()

    if text.startswith("```"):
        text = re.sub(
            r"^```(?:json)?\s*",
            "",
            text,
            flags=re.IGNORECASE,
        )
        text = re.sub(
            r"\s*```$",
            "",
            text,
        )

    # Remove reasoning blocks.
    text = re.sub(
        r"<think>[\s\S]*?</think>",
        "",
        text,
        flags=re.IGNORECASE,
    )

    text = re.sub(
        r"^[\s\S]*?</think>\s*",
        "",
        text,
        flags=re.IGNORECASE,
    )

    start = text.find("{")

    if start < 0:
        return None

    depth = 0
    in_string = False
    escape_next = False

    for i in range(start, len(text)):
        ch = text[i]

        if escape_next:
            escape_next = False
            continue

        if ch == "\\" and in_string:
            escape_next = True
            continue

        if ch == '"':
            in_string = not in_string
            continue

        if in_string:
            continue

        if ch == "{":
            depth += 1

        elif ch == "}":
            depth -= 1

            if depth == 0:
                return text[start:i + 1]

    return None


def _extract_complete_item_objects(text: str) -> list[dict]:
    """
    Recover complete objects from the items array.

    If the final item is truncated, previously complete items are preserved.
    """

    match = re.search(
        r'"items"\s*:\s*\[',
        text,
        flags=re.IGNORECASE,
    )

    if not match:
        return []

    pos = match.end()
    decoder = json.JSONDecoder()
    recovered = []

    while pos < len(text):

        # Skip whitespace and commas between items.
        while pos < len(text) and text[pos] in " \t\r\n,":
            pos += 1

        if pos >= len(text) or text[pos] == "]":
            break

        if text[pos] != "{":
            break

        try:
            obj, end = decoder.raw_decode(
                text,
                pos,
            )

        except json.JSONDecodeError:
            # Current item is incomplete.
            # Keep everything successfully decoded before it.
            break

        if isinstance(obj, dict):
            recovered.append(obj)

        pos = end

    return recovered


def _validate_items(
    raw_items: list[dict],
) -> KnowledgeExtractionResult:
    """
    Validate items independently.

    One malformed item must never invalidate the other valid items.
    """

    valid = []
    rejected = 0

    for raw_item in raw_items:

        try:
            single = KnowledgeExtractionResult(
                items=[raw_item]
            )

            valid.extend(single.items)

        except Exception as exc:
            rejected += 1

            logger.warning(
                "Rejected malformed knowledge item: %s",
                str(exc)[:200],
            )

    logger.info(
        "Knowledge item validation: valid=%d rejected=%d",
        len(valid),
        rejected,
    )

    return KnowledgeExtractionResult(
        items=valid
    )


def _parse_json_response(
    response,
) -> KnowledgeExtractionResult | None:
    """
    Robust provider-response parser.

    Order:

    1. Direct JSON
    2. JSON object embedded in surrounding text
    3. Individual complete items from truncated JSON
    """

    content = (
        response.content
        if hasattr(response, "content")
        else str(response)
    )

    # LangChain can return structured content blocks.
    if isinstance(content, list):

        parts = []

        for block in content:

            if isinstance(block, str):
                parts.append(block)

            elif isinstance(block, dict):

                value = (
                    block.get("text")
                    or block.get("content")
                )

                if value:
                    parts.append(str(value))

        content = "\n".join(parts)

    if not content or not str(content).strip():
        logger.warning("LLM returned an empty response")
        return None

    text = str(content).strip()

    # Remove reasoning.
    text = re.sub(
        r"<think>[\s\S]*?</think>",
        "",
        text,
        flags=re.IGNORECASE,
    )

    text = re.sub(
        r"^[\s\S]*?</think>\s*",
        "",
        text,
        flags=re.IGNORECASE,
    )

    # Remove markdown fences.
    if text.startswith("```"):

        text = re.sub(
            r"^```(?:json)?\s*",
            "",
            text,
            flags=re.IGNORECASE,
        )

        text = re.sub(
            r"\s*```$",
            "",
            text,
        )

    # ---------------------------------------------------------
    # 1. Direct JSON
    # ---------------------------------------------------------

    try:

        data = json.loads(text)

        if (
            isinstance(data, dict)
            and isinstance(data.get("items"), list)
        ):

            result = _validate_items(
                data["items"]
            )

            if result.items or not data["items"]:
                return result

    except Exception:
        pass

    # ---------------------------------------------------------
    # 2. JSON object embedded in surrounding text
    # ---------------------------------------------------------

    object_text = _extract_json_object(text)

    if object_text:

        try:

            data = json.loads(object_text)

            if (
                isinstance(data, dict)
                and isinstance(data.get("items"), list)
            ):

                result = _validate_items(
                    data["items"]
                )

                if result.items or not data["items"]:
                    return result

        except Exception:
            pass

    # ---------------------------------------------------------
    # 3. Salvage complete items from truncated JSON
    # ---------------------------------------------------------

    recovered = _extract_complete_item_objects(text)

    if recovered:

        result = _validate_items(
            recovered
        )

        if result.items:

            logger.info(
                "Recovered %d complete items from truncated/malformed JSON",
                len(result.items),
            )

            return result

    logger.warning(
        "Could not recover knowledge JSON. Response prefix: %s...",
        text[:300].replace("\n", " "),
    )

    return None