demand-agent-new/examples/demand/pg_pattern_repository.py

"""Read-only PostgreSQL Pattern V2 repository for DemandAgent.

This module is the DemandAgent boundary to PG `open_aigc.public` Pattern V2
facts. It intentionally exposes the old Demand-side query shapes while reading
only PG main-truth tables:

- pattern_mining_execution
- pattern_mining_category
- pattern_mining_element
- pattern_mining_config
- pattern_itemset
- pattern_itemset_item
- pattern_itemset_post
"""

from __future__ import annotations

import json
import os
import re
from contextlib import contextmanager
from datetime import date, datetime
from typing import Any, Iterable, Mapping, Sequence

try:
    import psycopg2
    from psycopg2.extras import RealDictCursor
except Exception:  # pragma: no cover - exercised in environments without deps
    psycopg2 = None
    RealDictCursor = None


PG_PATTERN_SOURCE_SYSTEM = "pg_pattern_v2"
TOPIC_SCOPE = "topic"
TOPIC_ELEMENT_SOURCE_TABLE = "post_decode_topic_point_element"

_READONLY_SQL_RE = re.compile(r"^\s*(select|with|show|explain)\b", re.IGNORECASE | re.DOTALL)
_BLOCKED_SQL_RE = re.compile(
    r"\b(insert|update|delete|create|alter|drop|truncate|copy|grant|revoke|merge|call)\b",
    re.IGNORECASE,
)


def _load_dotenv_once() -> None:
    try:
        from dotenv import load_dotenv
    except Exception:
        return
    load_dotenv()


def _resolve_dsn() -> str:
    _load_dotenv_once()
    dsn = os.getenv("DEMAND_PATTERN_PG_DSN") or os.getenv("PGVECTOR_DSN")
    if not dsn:
        raise RuntimeError("PG Pattern V2 DSN missing: set DEMAND_PATTERN_PG_DSN or PGVECTOR_DSN")
    return dsn


def _assert_readonly_sql(sql: str) -> None:
    if not _READONLY_SQL_RE.search(sql):
        raise RuntimeError("PG Pattern repository only allows SELECT/SHOW/EXPLAIN queries")
    if _BLOCKED_SQL_RE.search(sql):
        raise RuntimeError("PG Pattern repository blocked a non-read-only SQL keyword")


@contextmanager
def _connect_readonly():
    if psycopg2 is None:
        raise RuntimeError("psycopg2 is required for PG Pattern V2 reads")
    conn = psycopg2.connect(_resolve_dsn())
    try:
        conn.set_session(readonly=True, autocommit=True)
        yield conn
    finally:
        conn.close()


def _fetch_all(sql: str, params: Sequence[Any] | Mapping[str, Any] | None = None) -> list[dict[str, Any]]:
    _assert_readonly_sql(sql)
    with _connect_readonly() as conn:
        with conn.cursor(cursor_factory=RealDictCursor) as cursor:
            cursor.execute(sql, params)
            return [_row_to_dict(row) for row in cursor.fetchall()]


def _fetch_one(sql: str, params: Sequence[Any] | Mapping[str, Any] | None = None) -> dict[str, Any] | None:
    rows = _fetch_all(sql, params)
    return rows[0] if rows else None


def _cursor_fetch_all(cursor: Any, sql: str, params: Sequence[Any] | Mapping[str, Any] | None = None) -> list[dict[str, Any]]:
    _assert_readonly_sql(sql)
    cursor.execute(sql, params)
    return [_row_to_dict(row) for row in cursor.fetchall()]


def _cursor_fetch_one(cursor: Any, sql: str, params: Sequence[Any] | Mapping[str, Any] | None = None) -> dict[str, Any] | None:
    rows = _cursor_fetch_all(cursor, sql, params)
    return rows[0] if rows else None


def _serialize_db_value(value: Any) -> Any:
    if isinstance(value, (datetime, date)):
        return value.isoformat()
    return value


def _row_to_dict(row: Any) -> dict[str, Any]:
    if row is None:
        return {}
    mapping = getattr(row, "_mapping", row)
    return {key: _serialize_db_value(value) for key, value in dict(mapping).items()}


def _jsonish_to_list(value: Any) -> list[Any]:
    if value is None:
        return []
    if isinstance(value, list):
        return value
    if isinstance(value, (tuple, set)):
        return list(value)
    if isinstance(value, bytes):
        value = value.decode("utf-8")
    if isinstance(value, str):
        raw = value.strip()
        if not raw:
            return []
        try:
            parsed = json.loads(raw)
        except json.JSONDecodeError:
            return [part.strip() for part in raw.split(",") if part.strip()]
        if isinstance(parsed, list):
            return parsed
        if parsed is None:
            return []
        return [parsed]
    return [value]


def _to_post_id_list(value: Any) -> list[str]:
    post_ids: list[str] = []
    seen: set[str] = set()
    for item in _jsonish_to_list(value):
        post_id = str(item).strip() if item is not None else ""
        if post_id and post_id not in seen:
            seen.add(post_id)
            post_ids.append(post_id)
    return post_ids


def _to_int_list(values: Any) -> list[int]:
    result: list[int] = []
    seen: set[int] = set()
    for value in _jsonish_to_list(values):
        if value is None or value == "":
            continue
        try:
            int_value = int(value)
        except (TypeError, ValueError):
            continue
        if int_value not in seen:
            seen.add(int_value)
            result.append(int_value)
    return result


def _clean_names(names: Iterable[str]) -> list[str]:
    result: list[str] = []
    seen: set[str] = set()
    for raw in names:
        name = str(raw).strip() if raw is not None else ""
        if name and name not in seen:
            seen.add(name)
            result.append(name)
    return result


def query_execution_for_evidence(execution_id: int) -> dict[str, Any] | None:
    """Read PG Pattern V2 execution facts for evidence validation."""
    row = _fetch_one(
        """
        SELECT
            id,
            snapshot_date,
            is_current,
            is_retired,
            status,
            post_filter,
            post_count,
            post_ids_digest,
            classify_execution_id,
            category_count,
            element_count,
            topic_itemset_count AS itemset_count,
            topic_itemset_count,
            topic_element_itemset_count,
            cross_itemset_count,
            script_sequence_count,
            paragraph_itemset_count,
            start_time,
            end_time,
            error_message
        FROM pattern_mining_execution
        WHERE id = %s
        LIMIT 1
        """,
        (int(execution_id),),
    )
    if not row:
        return None
    row["pattern_source_system"] = PG_PATTERN_SOURCE_SYSTEM
    return row


def query_latest_success_execution_id() -> int | None:
    row = _fetch_one(
        """
        SELECT id
        FROM pattern_mining_execution
        WHERE status = 'success'
          AND COALESCE(topic_itemset_count, 0) > 0
        ORDER BY snapshot_date DESC NULLS LAST, id DESC
        LIMIT 1
        """
    )
    return int(row["id"]) if row else None


def query_itemset_evidence(execution_id: int, itemset_ids: Iterable[Any]) -> list[dict[str, Any]]:
    """Read PG topic-scope itemset facts and aggregate matched posts."""
    clean_ids = _to_int_list(itemset_ids or [])
    if not clean_ids:
        return []

    rows = _fetch_all(
        """
        SELECT
            i.id,
            i.execution_id,
            i.mining_config_id,
            cfg.execution_id AS mining_config_execution_id,
            cfg.scope AS mining_config_scope,
            cfg.name AS mining_config_name,
            cfg.algorithm AS mining_config_algorithm,
            cfg.params::text AS mining_config_params,
            cfg.params ->> 'dimension_mode' AS mining_config_dimension_mode,
            COALESCE(cfg.params ->> 'target_depth', cfg.params ->> 'depth') AS mining_config_target_depth,
            i.scope,
            i.combination_type,
            i.item_count,
            i.support,
            i.absolute_support,
            i.dimensions::text AS dimensions,
            i.is_cross_point,
            i.is_cross_layer,
            COALESCE(
                array_agg(p.post_id ORDER BY p.post_id)
                    FILTER (WHERE p.post_id IS NOT NULL),
                ARRAY[]::varchar[]
            ) AS matched_post_ids
        FROM pattern_itemset i
        JOIN pattern_mining_config cfg
          ON cfg.id = i.mining_config_id
         AND cfg.execution_id = i.execution_id
         AND cfg.scope = %s
        LEFT JOIN pattern_itemset_post p
          ON p.itemset_id = i.id
         AND p.execution_id = i.execution_id
        WHERE i.execution_id = %s
          AND i.scope = %s
          AND i.id = ANY(%s)
        GROUP BY
            i.id,
            i.execution_id,
            i.mining_config_id,
            cfg.execution_id,
            cfg.scope,
            cfg.name,
            cfg.algorithm,
            cfg.params::text,
            cfg.params ->> 'dimension_mode',
            COALESCE(cfg.params ->> 'target_depth', cfg.params ->> 'depth'),
            i.scope,
            i.combination_type,
            i.item_count,
            i.support,
            i.absolute_support,
            i.dimensions::text,
            i.is_cross_point,
            i.is_cross_layer
        ORDER BY i.id
        """,
        (TOPIC_SCOPE, int(execution_id), TOPIC_SCOPE, clean_ids),
    )

    result: list[dict[str, Any]] = []
    for row in rows:
        itemset = dict(row)
        for key in (
                "id",
                "execution_id",
                "mining_config_id",
                "mining_config_execution_id",
                "item_count",
                "absolute_support",
        ):
            if itemset.get(key) is not None:
                itemset[key] = int(itemset[key])
        itemset["support"] = float(itemset["support"]) if itemset.get("support") is not None else None
        itemset["is_cross_point"] = bool(itemset.get("is_cross_point"))
        itemset["is_cross_layer"] = bool(itemset.get("is_cross_layer"))
        itemset["dimensions"] = _jsonish_to_list(itemset.get("dimensions"))
        if isinstance(itemset.get("mining_config_params"), str):
            try:
                itemset["mining_config_params"] = json.loads(itemset["mining_config_params"])
            except json.JSONDecodeError:
                pass
        itemset["matched_post_ids"] = _to_post_id_list(itemset.get("matched_post_ids"))
        result.append(itemset)

    by_id = {itemset["id"]: itemset for itemset in result}
    return [by_id[itemset_id] for itemset_id in clean_ids if itemset_id in by_id]


def query_itemset_items_with_categories(
        execution_id: int,
        itemset_ids: Iterable[Any],
) -> list[dict[str, Any]]:
    """Read PG itemset items and their exact category tree nodes."""
    clean_ids = _to_int_list(itemset_ids or [])
    if not clean_ids:
        return []

    rows = _fetch_all(
        """
        SELECT
            ii.id AS itemset_item_id,
            ii.itemset_id,
            ii.layer,
            ii.point_type,
            ii.dimension,
            ii.category_id,
            ii.category_path,
            ii.element_name,
            ii.element_id,
            ii.post_count,
            c.id AS bound_category_id,
            c.execution_id AS category_execution_id,
            c.source_stable_id AS category_source_stable_id,
            c.source_type AS category_source_type,
            c.name AS category_name,
            c.path AS category_full_path,
            c.level AS category_level,
            c.parent_id AS category_parent_id,
            c.parent_source_stable_id AS category_parent_source_stable_id,
            c.classified_as AS category_nature
        FROM pattern_itemset_item ii
        JOIN pattern_itemset i
          ON i.id = ii.itemset_id
         AND i.execution_id = %s
         AND i.scope = %s
        LEFT JOIN pattern_mining_category c
          ON c.id = ii.category_id
         AND c.execution_id = i.execution_id
        WHERE ii.itemset_id = ANY(%s)
        ORDER BY ii.itemset_id, ii.id
        """,
        (int(execution_id), TOPIC_SCOPE, clean_ids),
    )
    result: list[dict[str, Any]] = []
    for row in rows:
        item = dict(row)
        for key in (
                "itemset_item_id",
                "itemset_id",
                "category_id",
                "element_id",
                "post_count",
                "bound_category_id",
                "category_execution_id",
                "category_source_stable_id",
                "category_level",
                "category_parent_id",
                "category_parent_source_stable_id",
        ):
            if item.get(key) is not None:
                item[key] = int(item[key])
        item["category_found"] = item.get("bound_category_id") is not None
        result.append(item)
    return result


def query_element_bindings_for_items(
        execution_id: int,
        itemset_items: Sequence[Mapping[str, Any]],
        post_ids: Iterable[Any] | None = None,
        limit_per_item: int = 200,
) -> list[dict[str, Any]]:
    """Bind itemset items to real PG Pattern V2 source-post elements."""
    clean_post_ids = _to_post_id_list(post_ids)
    if not itemset_items:
        return []

    bindings: list[dict[str, Any]] = []
    with _connect_readonly() as conn:
        with conn.cursor(cursor_factory=RealDictCursor) as cursor:
            for raw_item in itemset_items:
                item = dict(raw_item)
                category_id = item.get("bound_category_id") or item.get("category_id")
                if category_id is None:
                    continue

                clauses = [
                    "execution_id = %s",
                    "category_id = %s",
                    "source_table = %s",
                ]
                params: list[Any] = [int(execution_id), int(category_id), TOPIC_ELEMENT_SOURCE_TABLE]

                if item.get("point_type"):
                    clauses.append("point_type = %s")
                    params.append(item["point_type"])
                if item.get("dimension"):
                    clauses.append("element_type = %s")
                    params.append(item["dimension"])
                if item.get("element_name"):
                    clauses.append("name = %s")
                    params.append(item["element_name"])
                if clean_post_ids:
                    clauses.append("post_id = ANY(%s)")
                    params.append(clean_post_ids)

                where_sql = " AND ".join(clauses)
                count_row = _cursor_fetch_one(
                    cursor,
                    f"""
                    SELECT
                        COUNT(*) AS matched_element_count,
                        COUNT(DISTINCT post_id) AS matched_post_count
                    FROM pattern_mining_element
                    WHERE {where_sql}
                    """,
                    tuple(params),
                ) or {}
                post_rows = _cursor_fetch_all(
                    cursor,
                    f"""
                    SELECT DISTINCT post_id
                    FROM pattern_mining_element
                    WHERE {where_sql}
                    ORDER BY post_id
                    LIMIT %s
                    """,
                    tuple(params + [int(limit_per_item)]),
                )
                sample_rows = _cursor_fetch_all(
                    cursor,
                    f"""
                    SELECT
                        id,
                        post_id,
                        source_table,
                        source_element_id,
                        point_type,
                        point_text,
                        element_type,
                        name,
                        category_id,
                        category_path,
                        topic_point_id
                    FROM pattern_mining_element
                    WHERE {where_sql}
                    ORDER BY post_id, id
                    LIMIT %s
                    """,
                    tuple(params + [min(int(limit_per_item), 20)]),
                )

                bindings.append(
                    {
                        "itemset_id": item.get("itemset_id"),
                        "itemset_item_id": item.get("itemset_item_id"),
                        "category_id": int(category_id),
                        "point_type": item.get("point_type"),
                        "dimension": item.get("dimension"),
                        "element_name": item.get("element_name"),
                        "matched_element_count": int(count_row.get("matched_element_count") or 0),
                        "matched_post_count": int(count_row.get("matched_post_count") or 0),
                        "matched_post_ids": _to_post_id_list([row["post_id"] for row in post_rows]),
                        "sample_elements": sample_rows,
                    }
                )
    return bindings


def query_video_ids_by_names(execution_id: int, names: Iterable[str]) -> list[str]:
    """Resolve names to PG source post IDs by category name/path leaf or element name."""
    clean_names = _clean_names(names)
    if not clean_names:
        return []
    rows = _fetch_all(
        """
        WITH target_categories AS (
            SELECT id
            FROM pattern_mining_category
            WHERE execution_id = %s
              AND (
                name = ANY(%s)
                OR regexp_replace(path, '^.*/', '') = ANY(%s)
              )
        ),
        matched_elements AS (
            SELECT post_id
            FROM pattern_mining_element
            WHERE execution_id = %s
              AND source_table = %s
              AND (
                name = ANY(%s)
                OR category_id IN (SELECT id FROM target_categories)
              )
        )
        SELECT DISTINCT post_id
        FROM matched_elements
        WHERE post_id IS NOT NULL
        ORDER BY post_id
        """,
        (int(execution_id), clean_names, clean_names, int(execution_id), TOPIC_ELEMENT_SOURCE_TABLE, clean_names),
    )
    return _to_post_id_list([row["post_id"] for row in rows])


def query_category_level(execution_id: int, name: str) -> int | None:
    clean_name = str(name).strip() if name is not None else ""
    if not clean_name:
        return None
    row = _fetch_one(
        """
        SELECT level
        FROM pattern_mining_category
        WHERE execution_id = %s
          AND (name = %s OR regexp_replace(path, '^.*/', '') = %s)
        ORDER BY id DESC
        LIMIT 1
        """,
        (int(execution_id), clean_name, clean_name),
    )
    return int(row["level"]) if row and row.get("level") is not None else None


def query_case_ids_by_post_ids(post_ids: Iterable[Any]) -> list[dict[str, Any]]:
    """PG Pattern V2 has post_id semantics for MVP exact evidence.

    DemandAgent no longer reads the old MySQL decode-case table in this PG-only
    path. `case_ids` is filled from validated post IDs by the evidence builder;
    `decode_case_ids` remains an optional empty compatibility field.
    """
    return []


def query_categories(
        execution_id: int,
        *,
        source_type: str | None = None,
        keyword: str | None = None,
        category_ids: Iterable[Any] | None = None,
        limit: int | None = None,
) -> list[dict[str, Any]]:
    clauses = ["execution_id = %s"]
    params: list[Any] = [int(execution_id)]
    if source_type:
        clauses.append("source_type = %s")
        params.append(source_type)
    if keyword:
        clauses.append("(name ILIKE %s OR path ILIKE %s)")
        like = f"%{keyword}%"
        params.extend([like, like])
    clean_ids = _to_int_list(category_ids or [])
    if clean_ids:
        clauses.append("id = ANY(%s)")
        params.append(clean_ids)

    limit_sql = "LIMIT %s" if limit else ""
    if limit:
        params.append(int(limit))
    return _fetch_all(
        f"""
        SELECT
            id,
            execution_id,
            source_stable_id,
            source_type,
            name,
            description,
            path,
            level,
            parent_id,
            parent_source_stable_id,
            element_count,
            classified_as AS category_nature
        FROM pattern_mining_category
        WHERE {" AND ".join(clauses)}
        ORDER BY source_type, path, id
        {limit_sql}
        """,
        tuple(params),
    )


def query_elements(
        execution_id: int,
        *,
        keyword: str | None = None,
        element_type: str | None = None,
        post_ids: Iterable[Any] | None = None,
        category_ids: Iterable[Any] | None = None,
        limit: int | None = None,
) -> list[dict[str, Any]]:
    clauses = [
        "execution_id = %s",
        "source_table = %s",
    ]
    params: list[Any] = [int(execution_id), TOPIC_ELEMENT_SOURCE_TABLE]
    if keyword:
        clauses.append("name ILIKE %s")
        params.append(f"%{keyword}%")
    if element_type:
        clauses.append("element_type = %s")
        params.append(element_type)
    clean_posts = _to_post_id_list(post_ids)
    if clean_posts:
        clauses.append("post_id = ANY(%s)")
        params.append(clean_posts)
    clean_categories = _to_int_list(category_ids or [])
    if clean_categories:
        clauses.append("category_id = ANY(%s)")
        params.append(clean_categories)

    limit_sql = "LIMIT %s" if limit else ""
    if limit:
        params.append(int(limit))
    return _fetch_all(
        f"""
        SELECT
            id,
            execution_id,
            post_id,
            source_table,
            source_element_id,
            element_type,
            element_sub_type,
            name,
            description,
            category_id,
            category_path,
            point_type,
            point_text,
            topic_point_id
        FROM pattern_mining_element
        WHERE {" AND ".join(clauses)}
        ORDER BY post_id, id
        {limit_sql}
        """,
        tuple(params),
    )


def query_topic_itemsets(
        execution_id: int,
        *,
        category_ids: Iterable[Any] | 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",
        limit: int = 100,
) -> list[dict[str, Any]]:
    clauses = [
        "i.execution_id = %s",
        "i.scope = %s",
    ]
    params: list[Any] = [int(execution_id), TOPIC_SCOPE]
    clean_category_ids = _to_int_list(category_ids or [])
    for category_id in clean_category_ids:
        clauses.append(
            """
            EXISTS (
                SELECT 1
                FROM pattern_itemset_item ii_filter
                WHERE ii_filter.itemset_id = i.id
                  AND ii_filter.category_id = %s
            )
            """
        )
        params.append(category_id)
    if dimension_mode:
        clauses.append("cfg.params ->> 'dimension_mode' = %s")
        params.append(dimension_mode)
    if min_support is not None:
        clauses.append("i.absolute_support >= %s")
        params.append(int(min_support))
    if min_item_count is not None:
        clauses.append("i.item_count >= %s")
        params.append(int(min_item_count))
    if max_item_count is not None:
        clauses.append("i.item_count <= %s")
        params.append(int(max_item_count))

    order_sql = {
        "support": "i.support DESC NULLS LAST, i.absolute_support DESC NULLS LAST",
        "item_count": "i.item_count DESC NULLS LAST, i.absolute_support DESC NULLS LAST",
        "absolute_support": "i.absolute_support DESC NULLS LAST, i.support DESC NULLS LAST",
    }.get(str(sort_by or "absolute_support"), "i.absolute_support DESC NULLS LAST, i.support DESC NULLS LAST")
    params.append(int(limit))
    return _fetch_all(
        f"""
        SELECT
            i.id,
            i.execution_id,
            i.mining_config_id,
            cfg.scope AS mining_config_scope,
            cfg.name AS mining_config_name,
            cfg.params AS mining_config_params,
            cfg.params ->> 'dimension_mode' AS dimension_mode,
            COALESCE(cfg.params ->> 'target_depth', cfg.params ->> 'depth') AS target_depth,
            i.scope,
            i.combination_type,
            i.item_count,
            i.support,
            i.absolute_support,
            i.dimensions,
            i.is_cross_point,
            i.is_cross_layer
        FROM pattern_itemset i
        JOIN pattern_mining_config cfg
          ON cfg.id = i.mining_config_id
         AND cfg.execution_id = i.execution_id
         AND cfg.scope = %s
        WHERE {" AND ".join(clauses)}
        ORDER BY {order_sql}, i.id
        LIMIT %s
        """,
        tuple([TOPIC_SCOPE, *params]),
    )


def query_weight_score_rows(
        execution_id: int,
        *,
        level: str,
        dimension: str,
        limit: int = 5000,
) -> list[dict[str, Any]]:
    """Aggregate PG topic elements into old weight-score JSON shapes."""
    if level == "元素":
        return _fetch_all(
            """
            SELECT
                name,
                element_type,
                category_id,
                category_path,
                COUNT(*) AS occurrence_count,
                COUNT(DISTINCT post_id) AS post_count
            FROM pattern_mining_element
            WHERE execution_id = %s
              AND source_table = %s
              AND element_type = %s
              AND name IS NOT NULL
              AND name <> ''
            GROUP BY name, element_type, category_id, category_path
            ORDER BY COUNT(DISTINCT post_id) DESC, COUNT(*) DESC, name
            LIMIT %s
            """,
            (int(execution_id), TOPIC_ELEMENT_SOURCE_TABLE, dimension, int(limit)),
        )
    if level == "分类":
        return _fetch_all(
            """
            SELECT
                COALESCE(c.name, regexp_replace(e.category_path, '^.*/', '')) AS category,
                e.category_id,
                e.category_path,
                e.element_type,
                COUNT(*) AS occurrence_count,
                COUNT(DISTINCT e.post_id) AS post_count,
                COALESCE(MAX(c.level), array_length(string_to_array(trim(both '/' from e.category_path), '/'), 1)) AS level
            FROM pattern_mining_element e
            LEFT JOIN pattern_mining_category c
              ON c.id = e.category_id
             AND c.execution_id = e.execution_id
            WHERE e.execution_id = %s
              AND e.source_table = %s
              AND e.element_type = %s
              AND e.category_id IS NOT NULL
              AND e.category_path IS NOT NULL
            GROUP BY e.category_id, e.category_path, e.element_type, c.name
            ORDER BY COUNT(DISTINCT e.post_id) DESC, COUNT(*) DESC, e.category_path
            LIMIT %s
            """,
            (int(execution_id), TOPIC_ELEMENT_SOURCE_TABLE, dimension, int(limit)),
        )
    raise ValueError(f"unsupported weight score level: {level}")