dp-zp-agent/test/test_time_budget.py

"""One answer has one time budget; independent model calls run concurrently."""
import sqlite3
import time

import pytest

from tools.answers import model_client as mc


def test_parallel_map_keeps_order_and_the_callers_deadline():
    token = mc.start_deadline(30)
    try:
        def work(value):
            time.sleep(0.2)
            return value, mc.remaining_seconds() is not None
        started = time.monotonic()
        assert mc.parallel_map(work, [1, 2, 3, 4]) == [(1, True), (2, True), (3, True), (4, True)]
        assert time.monotonic() - started < 0.6  # concurrent, not 0.8 s sequentially
    finally:
        mc.reset_deadline(token)


def test_parallel_map_propagates_errors():
    def work(value):
        if value == 2:
            raise ValueError("bad")
        return value
    with pytest.raises(ValueError):
        mc.parallel_map(work, [1, 2, 3])


def test_spent_budget_stops_a_model_call_before_any_network_request():
    token = mc.start_deadline(0.5)
    try:
        config = {"provider": "openai_compatible", "model": "m", "key": "k",
                  "endpoint": "https://unreachable.invalid/v1/chat"}
        with pytest.raises(mc.DeadlineExceeded):
            mc._complete_request([{"role": "user", "content": "x"}], config)
    finally:
        mc.reset_deadline(token)


def test_nested_deadline_keeps_the_outer_budget():
    outer = mc.start_deadline(50)
    try:
        assert mc.start_deadline(1) is None  # an inner caller cannot extend or shorten it
        assert mc.remaining_seconds() > 40
    finally:
        mc.reset_deadline(outer)
    assert mc.remaining_seconds() is None


def test_low_budget_keeps_the_checked_answer_without_prose_synthesis():
    from tools.answers import synthesis
    from tools.answers.question_plan import QuestionPlan, PersonRequest
    plan = QuestionPlan(question="Na čom pracoval Ján Holp?", intent="progress",
                        people=[PersonRequest(name="Ján Holp")])
    result = {"status": "ok", "accepted_for_delivery": True, "answer": "Overené",
              "evidence": [{"id": "u1", "text": "Implementovaný PageRank.", "phase": "result"}],
              "presentation_notes": [], "index_version": {"version": "v"}, "knowledge_version": "k"}
    assert synthesis.eligible(result, plan)  # without a budget it WOULD synthesize
    called = []
    token = mc.start_deadline(5)
    try:
        out = synthesis.synthesize(result, plan, complete=lambda m: called.append(m))
    finally:
        mc.reset_deadline(token)
    assert out["answer"] == "Overené" and not called


def test_identical_repeated_question_reuses_the_finished_answer(tmp_path, monkeypatch):
    from tools.answers import answer_delivery as ad
    calls = []
    def fake(question, **kwargs):
        calls.append(question)
        return {"status": "ok", "accepted_for_delivery": True, "answer": "A", "query": question}
    monkeypatch.setattr(ad, "_deliver_checked", fake)
    db = tmp_path / "zp_index.sqlite"
    first = ad.deliver_answer("Na čom pracoval Ján Holp?", db_file=db)
    second = ad.deliver_answer("Na čom pracoval Ján Holp?", db_file=db)
    assert first == second and len(calls) == 1
    ad.deliver_answer("Na čom pracoval Ján Holp?", db_file=db, engine="graphrag")
    assert len(calls) == 2  # another engine is another answer


def test_rejected_answers_are_not_cached(tmp_path, monkeypatch):
    from tools.answers import answer_delivery as ad
    calls = []
    def fake(question, **kwargs):
        calls.append(question)
        return {"status": "error", "accepted_for_delivery": False, "reason": "answer_rejected"}
    monkeypatch.setattr(ad, "_deliver_checked", fake)
    db = tmp_path / "zp_index.sqlite"
    ad.deliver_answer("Otázka?", db_file=db)
    ad.deliver_answer("Otázka?", db_file=db)
    assert len(calls) == 2


def test_profile_page_author_is_the_students_supervisor(tmp_path):
    from tools.knowledge.sources import profile_authors
    db = tmp_path / "zp_index.sqlite"
    with sqlite3.connect(db) as conn:
        conn.execute("CREATE TABLE documents (path TEXT, title TEXT, author TEXT, published INTEGER)")
        conn.executemany("INSERT INTO documents VALUES (?,?,?,?)", [
            ("pages/students/2022/eva/README.md", "Eva Nová", "Daniel Hladek", 1),
            ("pages/topics/nlp/README.md", "NLP", "Daniel Hladek", 1),       # not a student profile
            ("pages/students/2023/adam/README.md", "Adam Nový", "Iný Autor", 0),  # unpublished
        ])
    assert profile_authors(db) == {"pages/students/2022/eva/README.md": "Daniel Hladek"}


def _gate_order(entries, slots_held=1):
    """Admission order of queued callers: (name, deadline, owner)."""
    import threading
    gate = mc._ModelGate()
    order, threads = [], []
    held = [gate.acquire(None) for _ in range(slots_held)]  # block until all queue
    def caller(name, deadline, owner):
        got = gate.acquire(deadline, owner)
        order.append(name)
        time.sleep(0.05)
        gate.release(got)
    for name, deadline, owner in entries:
        thread = threading.Thread(target=caller, args=(name, deadline, owner))
        thread.start(); threads.append(thread)
        time.sleep(0.05)
    for owner in held:
        gate.release(owner)
    for thread in threads:
        thread.join(5)
    return order


def test_model_gate_admits_the_nearest_deadline_first(monkeypatch):
    monkeypatch.setenv("ANSWER_LLM_SLOTS", "1")
    now = time.monotonic()
    assert _gate_order([("later", now + 60, "a"), ("offline", None, None), ("sooner", now + 30, "b")]) == [
        "sooner", "later", "offline"]


def test_model_gate_shares_slots_between_answers(monkeypatch):
    # Answer A queues two calls first; answer B (later deadline) must still
    # get one of the two slots instead of waiting for both calls of A.
    monkeypatch.setenv("ANSWER_LLM_SLOTS", "2")
    now = time.monotonic()
    order = _gate_order([("a1", now + 30, "A"), ("a2", now + 30, "A"), ("b1", now + 31, "B")], slots_held=2)
    assert order[:2] == ["a1", "b1"]


def test_model_gate_never_sends_a_call_that_cannot_finish(monkeypatch):
    import threading
    monkeypatch.setenv("ANSWER_LLM_SLOTS", "1")
    gate = mc._ModelGate()
    holder = gate.acquire(None)
    failure = []
    def caller(floor):
        try:
            gate.acquire(time.monotonic() + floor + 0.3, "x", floor)
        except mc.DeadlineExceeded:
            failure.append(floor)
    for floor in (0, 8):  # MIN_START, then a caller-declared minimum
        thread = threading.Thread(target=caller, args=(floor,)); thread.start(); thread.join(10)
    assert failure == [0, 8] and gate.waiting == []
    gate.release(holder)
    assert gate.acquire(None) is not None  # abandoned entries do not block the slot


def test_busy_means_another_answer_uses_the_endpoint():
    token = mc.start_deadline(30)
    try:
        assert not mc.model_busy()
        other = mc._gate.acquire(None, "other answer")
        try:
            assert mc.model_busy()
        finally:
            mc._gate.release(other)
        assert not mc.model_busy()
    finally:
        mc.reset_deadline(token)


def test_reserved_budget_keeps_time_for_the_audit():
    token = mc.start_deadline(30)
    try:
        with mc.reserved_budget(14):
            assert 15 < mc.remaining_seconds() <= 16
            assert not mc.has_budget(17)
        assert mc.remaining_seconds() > 29
    finally:
        mc.reset_deadline(token)
    with mc.reserved_budget(14):
        assert mc.remaining_seconds() is None  # no answer budget: no-op