content-find-agent-new/tests/test_progressive_screening.py

from __future__ import annotations

from pathlib import Path
from typing import Any

from content_agent.business_modules import (
    learning_review,
    progressive_screening,
    result_source_lookup,
    run_record,
    source_seed,
    walk_engine,
)
from content_agent.business_modules.run_record.validation import validate_run
from content_agent.constants import RUNTIME_RECORD_SCHEMA_VERSION
from content_agent.integrations.policy_json import JsonPolicyBundleStore
from content_agent.integrations.runtime_files import LocalRuntimeFileStore
from content_agent.record_payload import with_raw_payload
from tests.gemini_helpers import FakeGeminiVideoClient, fake_gemini_pool, fake_gemini_review
from tests.p1_helpers import real_source_payload


class FakeProgressivePlatformClient:
    def __init__(self, pages: dict[str, list[dict[str, Any]]]) -> None:
        self.pages = pages
        self.calls: list[dict[str, Any]] = []

    def search_full_page(self, query: dict[str, Any]) -> list[dict[str, Any]]:
        self.calls.append(dict(query))
        cursor = str(query.get("page_cursor") or "")
        return [dict(item) for item in self.pages.get(cursor, [])]

    def search(self, query: dict[str, Any]) -> list[dict[str, Any]]:
        return self.search_full_page(query)


class FakeHydratingProgressivePlatformClient(FakeProgressivePlatformClient):
    def __init__(self, pages: dict[str, list[dict[str, Any]]]) -> None:
        super().__init__(pages)
        self.metadata_calls: list[dict[str, Any]] = []
        self.hydrate_calls: list[str] = []

    def search_full_page_metadata(self, query: dict[str, Any]) -> list[dict[str, Any]]:
        self.metadata_calls.append(dict(query))
        cursor = str(query.get("page_cursor") or "")
        return [
            {**dict(item), "_platform_search_item": dict(item)}
            for item in self.pages.get(cursor, [])
        ]

    def hydrate_search_result_media(
        self,
        result: dict[str, Any],
        query: dict[str, Any],
    ) -> dict[str, Any]:
        content_id = str(result.get("platform_content_id") or "")
        self.hydrate_calls.append(content_id)
        hydrated = {
            key: value
            for key, value in result.items()
            if not str(key).startswith("_platform_")
        }
        hydrated["play_url"] = f"https://video.example/{content_id}.mp4"
        return hydrated


class FakeClock:
    def __init__(self) -> None:
        self.now = 0.0
        self.sleeps: list[float] = []

    def monotonic(self) -> float:
        return self.now

    def sleep(self, seconds: float) -> None:
        self.sleeps.append(seconds)
        self.now += seconds


def test_progressive_screening_two_item_page_has_no_remainder(tmp_path):
    context = _context(tmp_path)
    client = FakeProgressivePlatformClient({"": _page(["c1", "c2"], has_more=True, cursor="next")})
    gemini = FakeGeminiVideoClient(default_result=fake_gemini_pool())

    result = _run_screening(context, client, gemini)

    assert len(client.calls) == 1
    assert [item["platform_content_id"] for item in result["platform_results"]] == ["c1", "c2"]
    assert len(gemini.calls) == 2


def test_progressive_screening_archives_batch_before_gemini_judgment(tmp_path):
    context = _context(tmp_path)
    client = FakeProgressivePlatformClient({"": _page(["c1", "c2", "c3"], has_more=True, cursor="next")})
    archive_dispatcher = BatchArchiveDispatcher()
    gemini = OssRequiredGeminiVideoClient()

    result = _run_screening(context, client, gemini, archive_dispatcher=archive_dispatcher)

    assert archive_dispatcher.archived_content_ids == ["c1", "c2", "c3"]
    assert [call["media"]["oss_url"] for call in gemini.calls] == [
        "https://res.example/c1.mp4",
        "https://res.example/c2.mp4",
        "https://res.example/c3.mp4",
    ]
    assert {row["content_media_status"] for row in result["content_media_records"]} == {"oss_uploaded"}


def test_progressive_screening_uses_actual_page_size_without_hardcoded_ten(tmp_path):
    context = _context(tmp_path)
    ids = [f"c{i}" for i in range(15)]
    client = FakeProgressivePlatformClient({"": _page(ids, has_more=True, cursor="next")})
    gemini = FakeGeminiVideoClient(default_result=fake_gemini_review())

    result = _run_screening(context, client, gemini)

    assert len(client.calls) == 1
    assert [item["platform_content_id"] for item in result["platform_results"]] == ["c0", "c1", "c2"]
    assert {item["progressive_page_item_count"] for item in result["discovered_content_items"]} == {15}


def test_progressive_screening_pool_in_first_batch_processes_remainder_only(tmp_path):
    context = _context(tmp_path)
    ids = [f"c{i}" for i in range(10)]
    client = FakeProgressivePlatformClient({"": _page(ids, has_more=True, cursor="next")})
    gemini = FakeGeminiVideoClient(
        result_by_content_id={"c0": fake_gemini_pool()},
        default_result=fake_gemini_review(),
    )

    result = _run_screening(context, client, gemini)

    assert len(client.calls) == 1
    assert len(result["platform_results"]) == 10
    assert {decision["decision_action"] for decision in result["rule_decisions"]} == {
        "ADD_TO_CONTENT_POOL",
        "KEEP_CONTENT_FOR_REVIEW",
    }


def test_progressive_screening_remainder_and_page_two_gate_page_three(tmp_path):
    context = _context(tmp_path)
    pages = {
        "": _page([f"p1_{i}" for i in range(10)], has_more=True, cursor="cursor_2"),
        "cursor_2": _page([f"p2_{i}" for i in range(9)], has_more=True, cursor="cursor_3"),
        "cursor_3": _page([f"p3_{i}" for i in range(11)], has_more=True, cursor="cursor_4"),
    }
    client = FakeProgressivePlatformClient(pages)
    gemini = FakeGeminiVideoClient(
        result_by_content_id={
            "p1_0": fake_gemini_pool(),
            "p1_3": fake_gemini_pool(),
            "p2_0": fake_gemini_pool(),
            "p3_0": fake_gemini_pool(),
        },
        default_result=fake_gemini_review(),
    )
    clock = FakeClock()
    limiter = progressive_screening.SearchCallLimiter(
        now_fn=clock.monotonic,
        sleep_fn=clock.sleep,
    )

    result = _run_screening(context, client, gemini, limiter=limiter)

    assert [call.get("page_cursor", "") for call in client.calls] == ["", "cursor_2", "cursor_3"]
    assert len(result["platform_results"]) == 30
    assert clock.sleeps == [30.0, 30.0]


def test_progressive_screening_keep_for_review_does_not_pass(tmp_path):
    context = _context(tmp_path)
    client = FakeProgressivePlatformClient({"": _page([f"c{i}" for i in range(10)], has_more=True, cursor="next")})
    gemini = FakeGeminiVideoClient(default_result=fake_gemini_review())

    result = _run_screening(context, client, gemini)

    assert len(client.calls) == 1
    assert len(result["platform_results"]) == 3
    assert {decision["decision_action"] for decision in result["rule_decisions"]} == {
        "KEEP_CONTENT_FOR_REVIEW"
    }


def test_progressive_screening_hydrates_only_initial_batch_without_pass(tmp_path):
    context = _context(tmp_path)
    client = FakeHydratingProgressivePlatformClient(
        {"": _page([f"c{i}" for i in range(10)], has_more=True, cursor="next")}
    )
    gemini = FakeGeminiVideoClient(default_result=fake_gemini_review())

    result = _run_screening(context, client, gemini)

    assert client.metadata_calls == [context["search_queries"][0]]
    assert client.calls == []
    assert client.hydrate_calls == ["c0", "c1", "c2"]
    assert len(result["platform_results"]) == 3
    assert all("_platform_search_item" not in row for row in result["platform_results"])


def test_progressive_screening_hydrates_remainder_only_after_first_batch_pass(tmp_path):
    context = _context(tmp_path)
    client = FakeHydratingProgressivePlatformClient(
        {"": _page([f"c{i}" for i in range(10)], has_more=True, cursor="next")}
    )
    gemini = FakeGeminiVideoClient(
        result_by_content_id={"c0": fake_gemini_pool()},
        default_result=fake_gemini_review(),
    )

    result = _run_screening(context, client, gemini)

    assert client.hydrate_calls == [f"c{i}" for i in range(10)]
    assert len(result["platform_results"]) == 10


def test_progressive_screening_does_not_hydrate_duplicate_again(tmp_path):
    context = _context(tmp_path)
    client = FakeHydratingProgressivePlatformClient(
        {"": _page(["dup", "c1", "c2", "dup", "c4"], has_more=True, cursor="next")}
    )
    gemini = FakeGeminiVideoClient(
        result_by_content_id={"dup": fake_gemini_pool()},
        default_result=fake_gemini_review(),
    )

    result = _run_screening(context, client, gemini)

    assert client.hydrate_calls == ["dup", "c1", "c2", "c4"]
    assert [item["platform_content_id"] for item in result["platform_results"]] == [
        "dup",
        "c1",
        "c2",
        "c4",
    ]


def test_progressive_screening_deduplicates_across_batches(tmp_path):
    context = _context(tmp_path)
    client = FakeProgressivePlatformClient(
        {
            "": _page(["dup", "c1", "c2", "dup", "c4"], has_more=True, cursor="next"),
        }
    )
    gemini = FakeGeminiVideoClient(
        result_by_content_id={"dup": fake_gemini_pool()},
        default_result=fake_gemini_review(),
    )

    result = _run_screening(context, client, gemini)

    assert [item["platform_content_id"] for item in result["platform_results"]] == [
        "dup",
        "c1",
        "c2",
        "c4",
    ]
    assert len(result["rule_decisions"]) == 4


def test_progressive_screening_runtime_validates_and_does_not_duplicate_old_paging(tmp_path):
    context = _context(tmp_path)
    pages = {
        "": _page([f"p1_{i}" for i in range(10)], has_more=True, cursor="cursor_2"),
        "cursor_2": _page([f"p2_{i}" for i in range(10)], has_more=True, cursor="cursor_3"),
    }
    client = FakeProgressivePlatformClient(pages)
    gemini = FakeGeminiVideoClient(
        result_by_content_id={"p1_0": fake_gemini_pool(), "p1_3": fake_gemini_pool()},
        default_result=fake_gemini_review(),
    )
    clock = FakeClock()
    limiter = progressive_screening.SearchCallLimiter(
        now_fn=clock.monotonic,
        sleep_fn=clock.sleep,
    )

    result = _run_screening(context, client, gemini, limiter=limiter)
    walk_result = walk_engine.run_bounded_walk(
        run_id=context["run_id"],
        policy_run_id=context["policy_run_id"],
        pattern_seed_pack=context["pattern_seed_pack"],
        source_context=context["source_context"],
        search_queries=context["search_queries"],
        discovered_content_items=result["discovered_content_items"],
        content_media_records=result["content_media_records"],
        evidence_bundles=result["evidence_bundles"],
        rule_decisions=result["rule_decisions"],
        policy_bundle=context["policy_bundle"],
        platform_client=client,
        runtime=context["runtime"],
        gemini_video_client=gemini,
    )
    record = run_record.run(
        context["run_id"],
        context["policy_run_id"],
        context["search_queries"],
        walk_result["discovered_content_items"],
        walk_result["rule_decisions"],
        walk_result["source_path_record_basis"],
        context["policy_bundle"],
        context["runtime"],
        walk_actions=walk_result["walk_actions"],
        query_failures=result["query_failures"],
    )
    result_source_lookup.run(
        context["run_id"],
        context["policy_run_id"],
        context["policy_bundle"],
        walk_result["discovered_content_items"],
        walk_result["content_media_records"],
        walk_result["rule_decisions"],
        record["source_path_records"],
        record["search_clues"],
        context["runtime"],
    )
    learning_review.run(context["run_id"], context["policy_run_id"], context["runtime"])

    validation = validate_run(context["run_id"], context["runtime"])
    page_calls = [call for call in client.calls if call.get("page_cursor") == "cursor_2"]
    assert len(page_calls) == 1
    assert validation["status"] == "pass", validation["findings"]


def test_progressive_screening_keeps_cursor_when_next_page_was_not_consumed(tmp_path):
    context = _context(tmp_path)
    client = FakeProgressivePlatformClient({"": _page([f"c{i}" for i in range(10)], has_more=True, cursor="next")})
    gemini = FakeGeminiVideoClient(
        result_by_content_id={"c0": fake_gemini_pool()},
        default_result=fake_gemini_review(),
    )

    result = _run_screening(context, client, gemini)

    assert len(client.calls) == 1
    assert {item["next_cursor"] for item in result["discovered_content_items"]} == {"next"}
    assert {item["has_more"] for item in result["discovered_content_items"]} == {True}


def _context(tmp_path: Path) -> dict[str, Any]:
    run_id = "run_progressive"
    policy_run_id = "policy_progressive"
    runtime = LocalRuntimeFileStore(tmp_path / "runtime")
    runtime.prepare_run(run_id)
    seed = source_seed.run(run_id, policy_run_id, real_source_payload(), runtime)
    query = with_raw_payload(
        {
            "record_schema_version": RUNTIME_RECORD_SCHEMA_VERSION,
            "run_id": run_id,
            "policy_run_id": policy_run_id,
            "search_query_id": "q_001",
            "search_query": "早上好",
            "search_query_generation_method": "item_single",
            "discovery_start_source": "pattern_itemset",
            "previous_discovery_step": "search_query_generated",
            "pattern_seed_ref": {"seed_term": "早上好"},
        }
    )
    runtime.append_jsonl(run_id, "search_queries.jsonl", [query])
    return {
        "run_id": run_id,
        "policy_run_id": policy_run_id,
        "search_queries": [query],
        "source_context": seed["source_context"],
        "pattern_seed_pack": seed["pattern_seed_pack"],
        "policy_bundle": JsonPolicyBundleStore(Path(".")).load_policy_bundle("V4"),
        "runtime": runtime,
    }


def _run_screening(
    context: dict[str, Any],
    client: FakeProgressivePlatformClient,
    gemini: FakeGeminiVideoClient,
    *,
    limiter: progressive_screening.SearchCallLimiter | None = None,
    archive_dispatcher: Any | None = None,
) -> dict[str, Any]:
    kwargs = {
        "run_id": context["run_id"],
        "policy_run_id": context["policy_run_id"],
        "search_queries": context["search_queries"],
        "source_context": context["source_context"],
        "policy_bundle": context["policy_bundle"],
        "platform_client": client,
        "runtime": context["runtime"],
        "gemini_video_client": gemini,
    }
    if limiter is not None:
        kwargs["limiter"] = limiter
    if archive_dispatcher is not None:
        kwargs["archive_dispatcher"] = archive_dispatcher
    return progressive_screening.run(**kwargs)


class BatchArchiveDispatcher:
    def __init__(self) -> None:
        self.archived_content_ids: list[str] = []
        self.closed = False

    def archive_records(self, records: list[dict[str, Any]]) -> list[dict[str, Any]]:
        archived = []
        for row in records:
            content_id = row["platform_content_id"]
            self.archived_content_ids.append(content_id)
            raw_payload = {
                **(row.get("raw_payload") or {}),
                "content_media_status": "oss_uploaded",
                "oss_url": f"https://res.example/{content_id}.mp4",
                "oss_archive_status": "uploaded",
            }
            archived.append(
                {
                    **row,
                    "content_media_status": "oss_uploaded",
                    "oss_url": f"https://res.example/{content_id}.mp4",
                    "raw_payload": raw_payload,
                }
            )
        return archived

    def drain_completed(self) -> list[dict[str, Any]]:
        return []

    def enable_runtime_writes(self) -> None:
        return None

    def shutdown(self, *, wait: bool = False) -> None:
        self.closed = True


class OssRequiredGeminiVideoClient(FakeGeminiVideoClient):
    def analyze(
        self,
        content: dict[str, Any],
        media: dict[str, Any],
        source_context: dict[str, Any],
    ) -> dict[str, Any]:
        assert media["content_media_status"] == "oss_uploaded"
        assert media["oss_url"] == f"https://res.example/{content['platform_content_id']}.mp4"
        return super().analyze(content, media, source_context)


def _page(ids: list[str], *, has_more: bool, cursor: str) -> list[dict[str, Any]]:
    return [_item(content_id, has_more=has_more, cursor=cursor) for content_id in ids]


def _item(content_id: str, *, has_more: bool, cursor: str) -> dict[str, Any]:
    return {
        "content_discovery_id": f"q_001_content_{content_id}",
        "search_query_id": "q_001",
        "platform": "douyin",
        "platform_content_id": content_id,
        "platform_content_format": "video",
        "description": f"内容 {content_id}",
        "platform_author_id": f"author_{content_id}",
        "author_display_name": "作者",
        "statistics": {
            "digg_count": 1_000_000,
            "comment_count": 50_000,
            "share_count": 20_000,
            "collect_count": 100_000,
            "play_count": 0,
        },
        "tags": [],
        "play_url": f"https://video.test/{content_id}.mp4",
        "has_more": has_more,
        "next_cursor": cursor,
        "score": 72,
        "risk_level": "unknown",
        "discovery_relation": "derived_from_pattern_demand",
        "discovery_start_source": "pattern_itemset",
        "previous_discovery_step": "search_query_direct",
        "content_metadata_source": "fake_progressive_search",
        "platform_auth_mode": "no_bearer",
        "platform_raw_payload": {"channel_content_id": content_id},
    }