demand-agent-new/examples/demand/pattern_builds/pg_pattern_service.py

"""PG Pattern V2 adapter for DemandAgent pattern tools.

The public function names mirror the old MySQL pattern service, but every query
reads PG `open_aigc.public` Pattern V2 facts and only exposes topic-scope
itemsets as exact evidence candidates.
"""

from __future__ import annotations

from collections import defaultdict
from typing import Any, Iterable

from examples.demand import pg_pattern_repository as repo


def _clean_ints(values: Iterable[Any] | None) -> list[int]:
    result: list[int] = []
    seen: set[int] = set()
    for raw in values or []:
        try:
            value = int(raw)
        except (TypeError, ValueError):
            continue
        if value not in seen:
            seen.add(value)
            result.append(value)
    return result


def _last_path_part(path: Any) -> str:
    text = str(path or "").strip()
    if not text:
        return ""
    return text.replace(">", "/").split("/")[-1].strip()


def _group_point_types(rows: list[dict[str, Any]]) -> list[str]:
    result: list[str] = []
    seen: set[str] = set()
    for row in rows:
        value = str(row.get("point_type") or "").strip()
        if value and value not in seen:
            seen.add(value)
            result.append(value)
    return result


def _format_item(item: dict[str, Any]) -> dict[str, Any]:
    return {
        "id": item.get("itemset_item_id"),
        "itemset_id": item.get("itemset_id"),
        "layer": item.get("layer"),
        "point_type": item.get("point_type"),
        "dimension": item.get("dimension"),
        "category_id": item.get("category_id"),
        "category_path": item.get("category_path") or item.get("category_full_path"),
        "category_name": item.get("category_name") or _last_path_part(item.get("category_path")),
        "element_name": item.get("element_name"),
        "post_count": item.get("post_count"),
    }


def get_category_tree_compact(execution_id: int, source_type: str | None = None) -> str:
    categories = repo.query_categories(execution_id, source_type=source_type)
    if not categories:
        return f"未找到 PG Pattern V2 分类树，execution_id={execution_id}, source_type={source_type}"

    lines = [
        f"PG Pattern V2 分类树快照 execution_id={execution_id}",
        "说明：本 DemandAgent MVP 只使用 scope=topic 的分类 Pattern；topic_element 不进入最终证据。",
    ]
    for cat in categories:
        level = int(cat.get("level") or 0)
        indent = "  " * max(level - 1, 0)
        lines.append(
            f"{indent}- [{cat.get('id')}] {cat.get('name')} "
            f"(type={cat.get('source_type')}, level={level}, elements={cat.get('element_count') or 0}, "
            f"path={cat.get('path')})"
        )
    return "\n".join(lines)


def search_top_itemsets(
        execution_id: int,
        top_n: int = 20,
        category_ids: list | None = None,
        dimension_mode: str | None = None,
        min_support: int | None = None,
        min_item_count: int | None = None,
        max_item_count: int | None = None,
        sort_by: str = "absolute_support",
        account_name=None,
        merge_leve2=None,
        platform=None,
) -> dict[str, Any]:
    del account_name
    rows = repo.query_topic_itemsets(
        execution_id=execution_id,
        category_ids=category_ids,
        dimension_mode=dimension_mode,
        min_support=min_support,
        min_item_count=min_item_count,
        max_item_count=max_item_count,
        sort_by=sort_by,
        limit=max(int(top_n or 20) * 10, int(top_n or 20)),
        merge_leve2=merge_leve2,
        platform=platform,
    )
    itemset_ids = [int(row["id"]) for row in rows]
    items_by_itemset: dict[int, list[dict[str, Any]]] = defaultdict(list)
    if itemset_ids:
        for item in repo.query_itemset_items_with_categories(execution_id, itemset_ids):
            items_by_itemset[int(item["itemset_id"])].append(_format_item(item))

    groups: dict[str, dict[str, Any]] = {}
    for row in rows:
        dimension_key = row.get("dimension_mode") or row.get("mining_config_scope") or "topic"
        target_depth = row.get("target_depth") or ""
        key = f"{dimension_key}/{target_depth}"
        group = groups.setdefault(
            key,
            {
                "dimension_mode": dimension_key,
                "target_depth": target_depth,
                "total": 0,
                "itemsets": [],
            },
        )
        if len(group["itemsets"]) >= int(top_n or 20):
            continue
        group["total"] += 1
        group["itemsets"].append(
            {
                "id": row.get("id"),
                "scope": row.get("scope"),
                "mining_config_id": row.get("mining_config_id"),
                "dimension_mode": dimension_key,
                "target_depth": target_depth,
                "combination_type": row.get("combination_type"),
                "item_count": row.get("item_count"),
                "support": row.get("support"),
                "absolute_support": row.get("absolute_support"),
                "filtered_absolute_support": row.get("filtered_absolute_support"),
                "scoped_post_count": row.get("scoped_post_count"),
                "scope_filter": row.get("scope_filter") or {},
                "dimensions": row.get("dimensions") or [],
                "items": items_by_itemset.get(int(row["id"]), []),
            }
        )

    return {
        "pattern_source_system": repo.PG_PATTERN_SOURCE_SYSTEM,
        "scope": repo.TOPIC_SCOPE,
        "scope_filter": rows[0].get("scope_filter") if rows else {},
        "total": len(rows),
        "showing": sum(len(group["itemsets"]) for group in groups.values()),
        "groups": groups,
    }


def get_itemset_posts(
        itemset_ids: list[int],
        execution_id: int | None = None,
        *,
        merge_leve2=None,
        platform=None,
) -> list[dict[str, Any]]:
    clean_ids = _clean_ints(itemset_ids)
    if not clean_ids:
        return []
    if execution_id is None:
        raise ValueError("PG Pattern V2 get_itemset_posts requires execution_id")

    itemsets = repo.query_itemset_evidence(
        execution_id=execution_id,
        itemset_ids=clean_ids,
        merge_leve2=merge_leve2,
        platform=platform,
    )
    itemset_items = repo.query_itemset_items_with_categories(execution_id=execution_id, itemset_ids=clean_ids)
    items_by_itemset: dict[int, list[dict[str, Any]]] = defaultdict(list)
    for item in itemset_items:
        items_by_itemset[int(item["itemset_id"])].append(_format_item(item))

    details: list[dict[str, Any]] = []
    for itemset in itemsets:
        details.append(
            {
                "id": itemset.get("id"),
                "execution_id": itemset.get("execution_id"),
                "scope": itemset.get("scope"),
                "mining_config_id": itemset.get("mining_config_id"),
                "dimension_mode": itemset.get("mining_config_dimension_mode"),
                "target_depth": itemset.get("mining_config_target_depth"),
                "combination_type": itemset.get("combination_type"),
                "item_count": itemset.get("item_count"),
                "support": itemset.get("support"),
                "absolute_support": itemset.get("absolute_support"),
                "filtered_absolute_support": itemset.get("filtered_absolute_support"),
                "scoped_post_count": itemset.get("scoped_post_count"),
                "scope_filter": itemset.get("scope_filter") or {},
                "dimensions": itemset.get("dimensions") or [],
                "items": items_by_itemset.get(int(itemset["id"]), []),
                "post_ids": itemset.get("matched_post_ids") or [],
                "matched_post_ids": itemset.get("matched_post_ids") or [],
                "global_matched_post_ids": itemset.get("global_matched_post_ids") or [],
            }
        )
    return details


def get_post_elements(execution_id: int, post_ids: list) -> dict[str, Any]:
    rows = repo.query_elements(execution_id, post_ids=post_ids, limit=5000)
    result: dict[str, Any] = {}
    grouped: dict[tuple[str, str, str], dict[str, Any]] = {}
    for row in rows:
        post_id = str(row.get("post_id") or "")
        point_type = str(row.get("point_type") or "未分点")
        point_text = str(row.get("point_text") or "")
        key = (post_id, point_type, point_text)
        point = grouped.setdefault(
            key,
            {
                "point_text": point_text,
                "elements": {"实质": [], "形式": [], "意图": []},
            },
        )
        element_type = str(row.get("element_type") or "其他")
        point["elements"].setdefault(element_type, []).append(
            {
                "id": row.get("id"),
                "source_element_id": row.get("source_element_id"),
                "name": row.get("name"),
                "description": row.get("description"),
                "category_id": row.get("category_id"),
                "category_path": row.get("category_path"),
            }
        )
    for (post_id, point_type, _), point in grouped.items():
        result.setdefault(post_id, {}).setdefault(point_type, []).append(point)
    return result


def search_elements(
        execution_id: int,
        keyword: str,
        element_type: str | None = None,
        limit: int = 50,
        account_name=None,
        merge_leve2=None,
        platform=None,
) -> list[dict[str, Any]]:
    del account_name
    rows = repo.query_elements(
        execution_id,
        keyword=keyword,
        element_type=element_type,
        merge_leve2=merge_leve2,
        platform=platform,
        limit=5000,
    )
    grouped: dict[tuple[str, str, int | None, str | None], list[dict[str, Any]]] = defaultdict(list)
    for row in rows:
        key = (
            str(row.get("name") or ""),
            str(row.get("element_type") or ""),
            row.get("category_id"),
            row.get("category_path"),
        )
        grouped[key].append(row)

    items: list[dict[str, Any]] = []
    for (name, etype, category_id, category_path), group_rows in grouped.items():
        if not name:
            continue
        posts = {str(row.get("post_id")) for row in group_rows if row.get("post_id") is not None}
        items.append(
            {
                "name": name,
                "element_type": etype,
                "category_id": category_id,
                "category_path": category_path,
                "point_types": _group_point_types(group_rows),
                "occurrence_count": len(group_rows),
                "post_count": len(posts),
            }
        )
    items.sort(key=lambda item: (item["post_count"], item["occurrence_count"]), reverse=True)
    return items[: int(limit or 50)]


def get_element_category_chain(
        execution_id: int,
        element_name: str,
        element_type: str | None = None,
) -> list[dict[str, Any]]:
    rows = repo.query_elements(execution_id, keyword=element_name, element_type=element_type, limit=2000)
    exact_rows = [row for row in rows if str(row.get("name") or "") == str(element_name)]
    category_ids = _clean_ints(row.get("category_id") for row in exact_rows if row.get("category_id"))
    categories = {int(row["id"]): row for row in repo.query_categories(execution_id, category_ids=category_ids)}
    results: list[dict[str, Any]] = []
    seen: set[int] = set()
    for row in exact_rows:
        category_id = row.get("category_id")
        if category_id is None or int(category_id) in seen:
            continue
        seen.add(int(category_id))
        category = categories.get(int(category_id), {})
        results.append(
            {
                "category_id": category_id,
                "category_path": row.get("category_path") or category.get("path"),
                "element_type": row.get("element_type"),
                "point_types": _group_point_types([r for r in exact_rows if r.get("category_id") == category_id]),
                "ancestors": _ancestors_from_path(category.get("path") or row.get("category_path")),
            }
        )
    return results


def _ancestors_from_path(path: Any) -> list[dict[str, Any]]:
    parts = [part.strip() for part in str(path or "").replace(">", "/").split("/") if part.strip()]
    return [{"name": part, "level": index + 1, "path": "/".join(parts[: index + 1])} for index, part in enumerate(parts)]


def get_category_detail_with_context(execution_id: int, category_id: int) -> dict[str, Any] | None:
    categories = repo.query_categories(execution_id, category_ids=[category_id])
    if not categories:
        return None
    category = categories[0]
    all_categories = repo.query_categories(execution_id)
    children = [cat for cat in all_categories if cat.get("parent_id") == int(category_id)]
    siblings = [
        cat for cat in all_categories
        if cat.get("parent_id") == category.get("parent_id") and cat.get("id") != category.get("id")
    ][:50]
    elements = get_category_elements(category_id, execution_id=execution_id)
    return {
        "category": category,
        "ancestors": _ancestors_from_path(category.get("path")),
        "children": children[:100],
        "siblings": siblings,
        "elements": elements[:100],
    }


def search_categories(execution_id: int, keyword: str, source_type: str | None = None) -> list[dict[str, Any]]:
    categories = repo.query_categories(execution_id, source_type=source_type, keyword=keyword, limit=100)
    if not categories:
        return []
    category_ids = _clean_ints(cat.get("id") for cat in categories)
    element_rows = repo.query_elements(execution_id, category_ids=category_ids, limit=5000)
    elements_by_category: dict[int, list[dict[str, Any]]] = defaultdict(list)
    for row in element_rows:
        if row.get("category_id") is not None:
            elements_by_category[int(row["category_id"])].append(row)
    result: list[dict[str, Any]] = []
    for cat in categories:
        rows = elements_by_category.get(int(cat["id"]), [])
        result.append(
            {
                **cat,
                "point_types": _group_point_types(rows),
                "post_count": len({str(row.get("post_id")) for row in rows if row.get("post_id") is not None}),
            }
        )
    return result


def get_category_elements(
        category_id: int,
        execution_id: int,
        account_name=None,
        merge_leve2=None,
        platform=None,
) -> list[dict[str, Any]]:
    del account_name
    rows = repo.query_elements(
        execution_id,
        category_ids=[category_id],
        merge_leve2=merge_leve2,
        platform=platform,
        limit=5000,
    )
    grouped: dict[tuple[str, str], list[dict[str, Any]]] = defaultdict(list)
    for row in rows:
        name = str(row.get("name") or "").strip()
        element_type = str(row.get("element_type") or "").strip()
        if name:
            grouped[(name, element_type)].append(row)
    result: list[dict[str, Any]] = []
    for (name, element_type), group_rows in grouped.items():
        result.append(
            {
                "name": name,
                "element_type": element_type,
                "point_types": _group_point_types(group_rows),
                "occurrence_count": len(group_rows),
                "post_count": len({str(row.get("post_id")) for row in group_rows if row.get("post_id") is not None}),
            }
        )
    result.sort(key=lambda item: (item["post_count"], item["occurrence_count"]), reverse=True)
    return result


def get_category_co_occurrences(
        execution_id: int,
        category_ids: list,
        top_n: int = 30,
        account_name=None,
        merge_leve2=None,
        platform=None,
) -> dict[str, Any]:
    del account_name
    clean_ids = _clean_ints(category_ids)
    if not clean_ids:
        return {"matched_post_count": 0, "co_categories": []}
    rows = repo.query_elements(
        execution_id,
        category_ids=clean_ids,
        merge_leve2=merge_leve2,
        platform=platform,
        limit=10000,
    )
    posts_by_category: dict[int, set[str]] = defaultdict(set)
    for row in rows:
        if row.get("category_id") is not None and row.get("post_id") is not None:
            posts_by_category[int(row["category_id"])].add(str(row["post_id"]))
    matched_posts: set[str] | None = None
    for category_id in clean_ids:
        current = posts_by_category.get(category_id, set())
        matched_posts = current if matched_posts is None else matched_posts & current
    matched_posts = matched_posts or set()
    if not matched_posts:
        return {"matched_post_count": 0, "matched_post_ids": [], "co_categories": []}
    co_rows = repo.query_elements(
        execution_id,
        post_ids=list(matched_posts)[:5000],
        merge_leve2=merge_leve2,
        platform=platform,
        limit=20000,
    )
    counts: dict[tuple[int, str, str], set[str]] = defaultdict(set)
    for row in co_rows:
        cat_id = row.get("category_id")
        if cat_id is None or int(cat_id) in clean_ids:
            continue
        counts[(int(cat_id), str(row.get("category_path") or ""), str(row.get("element_type") or ""))].add(
            str(row.get("post_id"))
        )
    ranked = sorted(counts.items(), key=lambda item: len(item[1]), reverse=True)
    return {
        "matched_post_count": len(matched_posts),
        "matched_post_ids": sorted(matched_posts)[:100],
        "co_categories": [
            {
                "category_id": key[0],
                "category_path": key[1],
                "element_type": key[2],
                "post_count": len(posts),
            }
            for key, posts in ranked[: int(top_n or 30)]
        ],
    }


def get_element_co_occurrences(
        execution_id: int,
        element_names: list,
        top_n: int = 30,
        account_name=None,
        merge_leve2=None,
        platform=None,
) -> dict[str, Any]:
    del account_name
    names = [str(name).strip() for name in element_names or [] if str(name).strip()]
    if not names:
        return {"matched_post_count": 0, "co_elements": []}
    posts_by_name: dict[str, set[str]] = {}
    for name in names:
        rows = repo.query_elements(
            execution_id,
            keyword=name,
            merge_leve2=merge_leve2,
            platform=platform,
            limit=10000,
        )
        posts_by_name[name] = {
            str(row.get("post_id"))
            for row in rows
            if str(row.get("name") or "") == name and row.get("post_id") is not None
        }
    matched_posts: set[str] | None = None
    for posts in posts_by_name.values():
        matched_posts = posts if matched_posts is None else matched_posts & posts
    matched_posts = matched_posts or set()
    if not matched_posts:
        return {"matched_post_count": 0, "matched_post_ids": [], "co_elements": []}
    co_rows = repo.query_elements(
        execution_id,
        post_ids=list(matched_posts)[:5000],
        merge_leve2=merge_leve2,
        platform=platform,
        limit=20000,
    )
    grouped: dict[tuple[str, str, int | None, str | None], list[dict[str, Any]]] = defaultdict(list)
    for row in co_rows:
        name = str(row.get("name") or "")
        if not name or name in names:
            continue
        grouped[(name, str(row.get("element_type") or ""), row.get("category_id"), row.get("category_path"))].append(row)
    ranked = sorted(grouped.items(), key=lambda item: len({r.get("post_id") for r in item[1]}), reverse=True)
    return {
        "matched_post_count": len(matched_posts),
        "matched_post_ids": sorted(matched_posts)[:100],
        "co_elements": [
            {
                "name": key[0],
                "element_type": key[1],
                "category_id": key[2],
                "category_path": key[3],
                "point_types": _group_point_types(rows),
                "occurrence_count": len(rows),
                "post_count": len({str(row.get("post_id")) for row in rows if row.get("post_id") is not None}),
            }
            for key, rows in ranked[: int(top_n or 30)]
        ],
    }