Create-knowledge-find-decode/acquisition/queries/builder.py

"""Backward-compatible façade for the query planning bounded context."""
from __future__ import annotations

from typing import Any

from acquisition.repositories.base import AcquisitionRepository
from query_planning import (
    GenerationRequest,
    GeneratorKind,
    QueryBatchWriteSpec,
    QueryBatchWriter,
    UnifiedQueryGenerationService,
    planning_store_for_repository,
)
from query_planning.cartesian import (
    DEFAULT_ACTIVE_FAMILY_KEYS,
    TREES,
    QueryBuildOptions,
    build_creation_query_batch,
)


def persist_query_batch(
    repo: AcquisitionRepository,
    generated: dict[str, Any],
    *,
    name: str,
    source_type: str = "generated",
    generation_method: str = "creation_demo_v1",
    target_platforms: list[str] | None = None,
) -> tuple[Any, int]:
    """Persist legacy family JSON through the unified planning and writer flow."""
    candidates: list[dict[str, Any]] = []
    for family in generated.get("families") or []:
        for item in family.get("items") or []:
            if not item.get("keep", True):
                continue
            candidates.append(
                {
                    "query_text": item.get("query") or "",
                    "axes": item.get("parts") or {},
                    "filter_reason": item.get("reason") or "",
                    "priority": item.get("priority") or 0,
                    "source_refs": [
                        {
                            "generator_kind": GeneratorKind.CARTESIAN.value,
                            "family_key": family.get("key"),
                            "family_name": family.get("name"),
                            "axes": item.get("parts") or {},
                        }
                    ],
                    "metadata": {
                        "family_key": family.get("key"),
                        "family_name": family.get("name"),
                        "family_axes": family.get("axes") or [],
                    },
                }
            )
    platforms = tuple(target_platforms or ["xiaohongshu", "weixin", "douyin"])
    result = UnifiedQueryGenerationService().generate(
        GenerationRequest(
            generator_kind=GeneratorKind.CARTESIAN,
            name=name,
            target_platforms=platforms,
            payload={
                "candidates": candidates,
                "generator_config": generated.get("metadata") or {},
                "input_snapshot": {
                    "active_family_keys": (generated.get("metadata") or {}).get(
                        "active_family_keys", []
                    ),
                    "family_count": len(generated.get("families") or []),
                    "candidate_count": len(candidates),
                },
            },
            metadata={"generation_method": generation_method},
        )
    )
    write_result = QueryBatchWriter(
        legacy_sink=repo,
        planning_store=planning_store_for_repository(repo),
    ).write(
        result,
        QueryBatchWriteSpec(
            name=name,
            source_type=source_type,
            generation_method=generation_method,
            target_platforms=platforms,
            metadata=generated.get("metadata") or {},
        ),
    )
    return write_result.batch, len(write_result.queries)


__all__ = [
    "DEFAULT_ACTIVE_FAMILY_KEYS",
    "QueryBuildOptions",
    "TREES",
    "build_creation_query_batch",
    "persist_query_batch",
]