Create-knowledge-find-decode/tests/test_pipeline_tracing.py

from __future__ import annotations

from contextlib import nullcontext
from datetime import datetime, timezone
from pathlib import Path
from uuid import uuid4

from acquisition.domain import CandidateItem, MediaAsset
from core.models import Post
from decode_content.models import GateResult, ReadResult
from decode_content.repositories.postgres import PostgresDecodeRepository
from decode_content.service import DecodeWorkflowOutput
from pipeline.decode_runner import run_decode_stage
from pipeline.tracing import TraceConfig, TraceContext, sanitize_payload
from scripts import backfill_pipeline_traces


def _method_source(text: str, method_name: str) -> str:
    start = text.index(f"    def {method_name}")
    next_method = text.find("\n    def ", start + 1)
    return text[start:] if next_method == -1 else text[start:next_method]


def test_formal_schema_declares_append_only_trace_tables():
    text = Path("db/migrations/001_creation_knowledge_schema.sql").read_text(encoding="utf-8")

    for table in (
        "pipeline_runs",
        "pipeline_jobs",
        "pipeline_run_events",
        "candidate_item_hits",
        "llm_call_traces",
    ):
        assert f"CREATE TABLE IF NOT EXISTS creation_knowledge.{table}" in text
    assert "ALTER TABLE creation_knowledge.candidate_items" not in text
    assert not Path("db/migrations/005_pipeline_trace_schema.sql").exists()


def test_formal_schema_uses_deduped_contract_artifacts():
    text = Path("db/migrations/001_creation_knowledge_schema.sql").read_text(encoding="utf-8")

    assert "CREATE TABLE IF NOT EXISTS creation_knowledge.contract_artifacts" in text
    assert "CREATE TABLE IF NOT EXISTS creation_knowledge.run_contract_artifacts" in text
    assert "UNIQUE (contract_name, contract_type, source_path, content_hash)" in text
    assert "CREATE TABLE IF NOT EXISTS creation_knowledge.contract_snapshots" not in text
    assert not Path("db/migrations/005_pipeline_trace_schema.sql").exists()


def test_postgres_decode_repository_writes_deduped_contract_artifacts():
    text = Path("decode_content/repositories/postgres.py").read_text(encoding="utf-8")
    method = _method_source(text, "save_contract_snapshot")

    assert "INSERT INTO contract_artifacts" in method
    assert "INSERT INTO run_contract_artifacts" in method
    assert "ON CONFLICT (contract_name, contract_type, source_path, content_hash)" in method
    assert "RETURNING id, contract_name, contract_type, version_label" in method
    assert "content_hash, source_path, snapshot, created_at, updated_at" in method
    assert "INSERT INTO contract_snapshots" not in text


class RecordingDecodeRepo(PostgresDecodeRepository):
    def __init__(self):
        super().__init__(conn=None)
        self.one_calls = []
        self.one_results = []

    def _one(self, sql, params):
        self.one_calls.append((sql, params))
        return self.one_results.pop(0)


def test_save_contract_snapshot_accepts_updated_at_from_contract_artifacts():
    now = datetime.now(timezone.utc)
    artifact_id = uuid4()
    repo = RecordingDecodeRepo()
    repo.one_results.append(
        {
            "id": artifact_id,
            "contract_name": "创作知识提取-skill",
            "contract_type": "skill",
            "version_label": None,
            "content_hash": "a" * 64,
            "source_path": "创作知识提取-skill",
            "snapshot": {},
            "created_at": now,
            "updated_at": now,
        }
    )

    snapshot = repo.save_contract_snapshot(
        contract_name="创作知识提取-skill",
        contract_type="skill",
        content_hash="a" * 64,
        source_path="创作知识提取-skill",
        snapshot={},
    )

    sql, _params = repo.one_calls[0]
    assert "INSERT INTO contract_artifacts" in sql
    assert "RETURNING *" not in sql
    assert "updated_at" in sql
    assert snapshot.id == artifact_id
    assert snapshot.updated_at == now


def test_sanitize_payload_omits_base64_data_urls_but_keeps_metadata():
    data_url = "data:video/mp4;base64," + ("a" * 120)

    payload = sanitize_payload(
        {"video": data_url, "url": "https://cdn.test/video.mp4"},
        max_chars=50,
    )

    assert payload["url"] == "https://cdn.test/video.mp4"
    assert payload["video"]["omitted"] == "base64_data_url"
    assert payload["video"]["chars"] == len(data_url)
    assert payload["video"]["sha256"]


def test_trace_config_invalid_integer_falls_back_to_default(monkeypatch):
    monkeypatch.setenv("CK_TRACE_MAX_PAYLOAD_CHARS", "not-a-number")

    config = TraceConfig.from_env(".env.missing")

    assert config.max_payload_chars == 50_000


class CollectingTraceWriter:
    def __init__(self):
        self.events = []
        self.hits = []
        self.calls = []

    def event(self, *, context, stage, event_type, **kwargs):
        self.events.append({"context": context, "stage": stage, "event_type": event_type, **kwargs})

    def candidate_hit(self, *, context, **kwargs):
        self.hits.append({"context": context, **kwargs})

    def llm_call(self, *, context, stage, substage=None, **kwargs):
        self.calls.append({"context": context, "stage": stage, "substage": substage, **kwargs})


class CandidateRepo:
    def __init__(self):
        self.item_id = uuid4()
        self.job_id = uuid4()
        self.query_id = uuid4()

    def list_creation_candidate_items(self, *, run_id=None, limit=100):
        return [
            CandidateItem(
                id=self.item_id,
                job_id=self.job_id,
                query_id=self.query_id,
                platform="xiaohongshu",
                platform_item_id="x1",
                unique_key="xhs:x1",
                canonical_url="https://xhs.test/1",
                title="拍照姿势",
                status="candidate",
            )
        ]

    def list_media_assets_for_item(self, item_id):
        return [
            MediaAsset(
                id=uuid4(),
                item_id=item_id,
                media_type="image",
                cdn_url="https://cdn.test/1.jpg",
                position=1,
                status="done",
            )
        ]


class DecodeService:
    def decode_post(self, *, item_id, post: Post):
        assert post.image_urls == ["https://cdn.test/1.jpg"]
        return DecodeWorkflowOutput(
            item_id=item_id,
            read_result=ReadResult(text="ok"),
            gate_result=GateResult(passed=True),
            status="decoded",
        )


def test_decode_stage_emits_trace_events_without_changing_result():
    writer = CollectingTraceWriter()
    repo = CandidateRepo()
    pipeline_run_id = uuid4()

    result = run_decode_stage(
        candidate_repo=repo,
        decode_service=DecodeService(),
        trace_writer=writer,
        trace_context=TraceContext(pipeline_run_id=pipeline_run_id, stage="decode"),
    )

    assert result.decoded == 1
    event_types = [event["event_type"] for event in writer.events]
    assert event_types[0] == "decode_stage_started"
    assert "decode_item_started" in event_types
    assert "decode_item_finished" in event_types
    assert event_types[-1] == "decode_stage_finished"
    assert writer.events[1]["context"].pipeline_run_id == pipeline_run_id


def test_backfill_pipeline_traces_dry_run_does_not_write(monkeypatch):
    acquisition_run_id = uuid4()

    monkeypatch.setattr(backfill_pipeline_traces.CreationDbConfig, "from_env", lambda _env_file: object())
    monkeypatch.setattr(backfill_pipeline_traces, "transaction", lambda _config: nullcontext(object()))
    monkeypatch.setattr(
        backfill_pipeline_traces,
        "_acquisition_runs",
        lambda _conn, limit=None: [
            {
                "id": acquisition_run_id,
                "run_key": "manual-api:test",
                "status": "done",
                "job_count": 3,
                "candidate_count": 12,
            }
        ],
    )

    result = backfill_pipeline_traces.backfill(env_file=".env", apply=False)

    assert result["dry_run"] is True
    assert result["acquisition_run_count"] == 1
    assert result["candidate_hit_count"] == 12
    assert result["runs"][0]["acquisition_run_id"] == str(acquisition_run_id)