tests/test_llamacpp_gpu.py

import sys
from pathlib import Path
from unittest.mock import patch

import pytest

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from src.extraction.llamacpp_gpu import LlamaCppGPUExtractor, _compose_prompt
from src.extraction.llamacpp_vision import llamacpp_vision_enabled
from src.openbmb_client import EXTRACTION_PROMPT


def test_llamacpp_vision_enabled(monkeypatch):
    monkeypatch.delenv("LLAMACPP_VISION", raising=False)
    assert llamacpp_vision_enabled() is False

    monkeypatch.setenv("LLAMACPP_VISION", "1")
    assert llamacpp_vision_enabled() is True

    monkeypatch.setenv("LLAMACPP_VISION", "yes")
    assert llamacpp_vision_enabled() is True


def test_compose_prompt_rejects_image_only_without_vision():
    parts = [{"type": "image_url", "image_url": {"url": "data:image/png;base64,abc"}}]
    with pytest.raises(RuntimeError, match="LLAMACPP_VISION=1"):
        _compose_prompt(parts)


def test_compose_prompt_keeps_text_files():
    parts = [{"type": "text", "text": "Hemoglobin 11.2 g/dL"}]
    prompt = _compose_prompt(parts)
    assert EXTRACTION_PROMPT in prompt
    assert "Hemoglobin 11.2 g/dL" in prompt


def test_extract_uses_vision_generation_when_enabled(monkeypatch):
    monkeypatch.setenv("LLAMACPP_VISION", "1")
    monkeypatch.delenv("BTE_SPACE_HARDWARE", raising=False)
    image_parts = [{"type": "image_url", "image_url": {"url": "data:image/png;base64,abc"}}]

    with patch(
        "src.extraction.llamacpp_gpu.document_to_payload_parts",
        return_value=image_parts,
    ), patch(
        "src.extraction.llamacpp_gpu._run_llamacpp_vision_generation",
        return_value='{"patient":{},"tests":[],"notes":[]}',
    ) as vision_run, patch(
        "src.extraction.llamacpp_gpu._run_llamacpp_generation",
    ) as text_run:
        result = LlamaCppGPUExtractor().extract("/tmp/report.pdf")

    vision_run.assert_called_once()
    text_run.assert_not_called()
    assert result.request_summary["backend"] == "llamacpp-gpu-vision"
    assert result.request_summary["vision_enabled"] is True
    assert result.request_summary["mmproj"]
    assert result.request_summary["spaces_gpu"] is True


def test_cpu_basic_uses_plain_cpu_vision_runner(monkeypatch):
    monkeypatch.setenv("BTE_SPACE_HARDWARE", "cpu-basic")
    monkeypatch.setenv("LLAMACPP_VISION", "1")
    image_parts = [{"type": "image_url", "image_url": {"url": "data:image/png;base64,abc"}}]

    with patch(
        "src.extraction.llamacpp_gpu.document_to_payload_parts",
        return_value=image_parts,
    ), patch(
        "src.extraction.llamacpp_gpu._run_llamacpp_vision_generation_cpu",
        return_value='{"patient":{},"tests":[],"notes":[]}',
    ) as cpu_run, patch(
        "src.extraction.llamacpp_gpu._run_llamacpp_vision_generation",
    ) as gpu_run:
        result = LlamaCppGPUExtractor().extract("/tmp/report.pdf")

    cpu_run.assert_called_once()
    gpu_run.assert_not_called()
    assert result.request_summary["backend"] == "llamacpp-cpu-vision"
    assert result.request_summary["spaces_gpu"] is False


def test_extract_uses_text_generation_when_vision_disabled(monkeypatch):
    monkeypatch.delenv("LLAMACPP_VISION", raising=False)
    monkeypatch.delenv("BTE_SPACE_HARDWARE", raising=False)
    text_parts = [{"type": "text", "text": "WBC 6.5"}]

    with patch(
        "src.extraction.llamacpp_gpu.document_to_payload_parts",
        return_value=text_parts,
    ), patch(
        "src.extraction.llamacpp_gpu._run_llamacpp_generation",
        return_value='{"patient":{},"tests":[],"notes":[]}',
    ) as text_run, patch(
        "src.extraction.llamacpp_gpu._run_llamacpp_vision_generation",
    ) as vision_run:
        result = LlamaCppGPUExtractor().extract("/tmp/report.txt")

    text_run.assert_called_once()
    vision_run.assert_not_called()
    assert result.request_summary["backend"] == "llamacpp-gpu"
    assert result.request_summary["vision_enabled"] is False


def test_vision_extractor_requires_mmproj_file(monkeypatch):
    monkeypatch.setenv("LLAMACPP_VISION", "1")
    monkeypatch.setenv("LLAMACPP_MMPROJ_FILE", " ")
    with pytest.raises(ValueError, match="LLAMACPP_MMPROJ_FILE"):
        LlamaCppGPUExtractor()