backend/api/prediction_attribute.py

"""预测归因 API"""
import gc
import time

from backend.models.model_manager import inference_lock
from backend.platform.oom import exit_if_oom
from backend.core.prediction_attributor import analyze_prediction_attribution
from backend.api.analyze import LOCK_WAIT_TIMEOUT
from backend.platform.access_log import get_client_ip, log_prediction_attribute_request
from backend.platform.source_page import ALLOWED_SOURCE_PAGES, normalize_source_page


def prediction_attribute(attribution_request):
    """
    对上下文文本的下一 token 预测做归因分析。

    Args:
        attribution_request: 须含 ``context``、``model``。归因目标二选一：
            省略 ``target_prediction`` 且省略 ``target_token_id`` 时为 top-1；
            或传非空 ``target_prediction``（字符串首 token）；
            或传 ``target_token_id``（非负整数词表 id）；二者不可同时出现。

    Returns:
        (响应字典, 状态码) 元组
    """
    context = attribution_request.get("context")
    target_prediction = attribution_request.get("target_prediction")
    target_token_id = attribution_request.get("target_token_id")
    model = attribution_request.get("model")
    source_page = attribution_request.get("source_page")
    flow_id = attribution_request.get("flow_id")
    flow_step = attribution_request.get("flow_step")

    if context is None:
        return {"success": False, "message": "Missing required field: context"}, 400
    if not isinstance(context, str):
        return {"success": False, "message": "context must be a string"}, 400
    if context == "":
        return {"success": False, "message": "Missing required field: context"}, 400

    if target_prediction is not None and not isinstance(target_prediction, str):
        return {"success": False, "message": "target_prediction must be a string"}, 400
    if target_prediction == "":
        return {"success": False, "message": "target_prediction must not be empty"}, 400
    if target_token_id is not None and not isinstance(target_token_id, int):
        return {"success": False, "message": "target_token_id must be an integer"}, 400
    if target_token_id is not None and target_token_id < 0:
        return {"success": False, "message": "target_token_id must be >= 0"}, 400
    if target_prediction is not None and target_token_id is not None:
        return {"success": False, "message": "target_prediction and target_token_id are mutually exclusive"}, 400

    if model is None:
        return {"success": False, "message": "Missing required field: model"}, 400
    if not isinstance(model, str):
        return {"success": False, "message": "model must be a string"}, 400
    if model not in ("base", "instruct"):
        return {"success": False, "message": 'model must be "base" or "instruct"'}, 400

    if source_page is None:
        return {"success": False, "message": "Missing required field: source_page"}, 400
    if not isinstance(source_page, str):
        return {"success": False, "message": "source_page must be a string"}, 400
    if source_page == "":
        return {"success": False, "message": "source_page must not be empty"}, 400
    normalized_source_page = normalize_source_page(source_page)
    if normalized_source_page is None:
        allowed = ", ".join(sorted(ALLOWED_SOURCE_PAGES))
        return {
            "success": False,
            "message": f"source_page must be one of: {allowed} (legacy *.html and gen_attribute accepted)",
        }, 400
    source_page = normalized_source_page

    if flow_id is not None and not isinstance(flow_id, str):
        return {"success": False, "message": "flow_id must be a string"}, 400
    if flow_id == "":
        return {"success": False, "message": "flow_id must not be empty"}, 400
    if flow_step is not None and not isinstance(flow_step, int):
        return {"success": False, "message": "flow_step must be an integer"}, 400
    if flow_step is not None and flow_step < 0:
        return {"success": False, "message": "flow_step must be >= 0"}, 400

    is_causal_flow = source_page == "causal_flow"
    if is_causal_flow:
        if flow_id is None:
            return {"success": False, "message": "Missing required field: flow_id for causal flow"}, 400
        if flow_step is None:
            return {"success": False, "message": "Missing required field: flow_step for causal flow"}, 400
    elif flow_id is not None or flow_step is not None:
        return {
            "success": False,
            "message": "flow_id/flow_step are only allowed when source_page is causal_flow",
        }, 400

    client_ip = get_client_ip()
    start_time = time.perf_counter()
    logged: dict = {"request_id": None}

    def log_fn():
        logged["request_id"] = log_prediction_attribute_request(
            context=context,
            target_prediction=target_prediction,
            target_token_id=target_token_id,
            model=model,
            source_page=source_page,
            flow_id=flow_id,
            flow_step=flow_step,
            client_ip=client_ip,
        )

    def local_fn():
        request_id = logged["request_id"]
        lock_acquired = inference_lock.acquire(timeout=LOCK_WAIT_TIMEOUT)
        if not lock_acquired:
            return {
                "success": False,
                "message": (
                    f"Queue wait exceeded {LOCK_WAIT_TIMEOUT} seconds; "
                    "server is busy, please try again later."
                ),
            }, 503

        try:
            result = analyze_prediction_attribution(
                context,
                target_prediction,
                model=model,
                target_token_id=target_token_id,
            )
        except ValueError as e:
            return {"success": False, "message": str(e)}, 400
        except Exception as e:
            import traceback
            traceback.print_exc()
            exit_if_oom(e, defer_seconds=1)
            return {"success": False, "message": str(e)}, 500
        finally:
            inference_lock.release()
            gc.collect()

        elapsed = time.perf_counter() - start_time
        from backend.platform.inference_response_log import log_prediction_attribute_response
        from backend.platform.model_routing import is_worker

        if is_worker():
            log_prediction_attribute_response(
                request_id=request_id,
                result=result,
                elapsed=elapsed,
                flow_id=flow_id,
                flow_step=flow_step,
                for_worker=True,
            )

        return {"success": True, **result}, 200

    from backend.core.prediction_attributor import slot_for_prediction_attr_model
    from backend.platform.inference_ingress import ingress_inference
    from backend.platform.inference_response_log import make_prediction_attribute_response_logger

    slot = slot_for_prediction_attr_model(model)
    return ingress_inference(
        slot=slot,
        api_path="/api/prediction-attribute",
        json_body=attribution_request,
        stream=False,
        timeout=60.0,
        log_fn=log_fn,
        local_fn=local_fn,
        response_log_fn=make_prediction_attribute_response_logger(
            logged, flow_id=flow_id, flow_step=flow_step
        ),
    )