Agent/knowhub/knowhub_db/pg_capability_store.py

"""
PostgreSQL capability 存储封装

用于存储和检索原子能力数据，支持向量检索。
表名：capability（从 atomic_capability 迁移）
"""

import os
import json
import psycopg2
from psycopg2.extras import RealDictCursor
from typing import List, Dict, Optional
from dotenv import load_dotenv
from knowhub.knowhub_db.cascade import cascade_delete

load_dotenv()

# 关联字段子查询
_REL_SUBQUERIES = """
    (SELECT COALESCE(json_agg(rc.requirement_id), '[]'::json)
     FROM requirement_capability rc WHERE rc.capability_id = capability.id) AS requirement_ids,
    (SELECT COALESCE(json_agg(ct.tool_id), '[]'::json)
     FROM capability_tool ct WHERE ct.capability_id = capability.id) AS tool_ids,
    (SELECT COALESCE(
        json_object_agg(ct2.tool_id, ct2.description), '{}'::json)
     FROM capability_tool ct2 WHERE ct2.capability_id = capability.id AND ct2.description != '') AS implements,
    (SELECT COALESCE(json_agg(ck.knowledge_id), '[]'::json)
     FROM capability_knowledge ck WHERE ck.capability_id = capability.id) AS knowledge_ids,
    (SELECT COALESCE(json_agg(json_build_object(
        'id', ck2.knowledge_id, 'relation_type', ck2.relation_type
     )), '[]'::json)
     FROM capability_knowledge ck2 WHERE ck2.capability_id = capability.id) AS knowledge_links,
    (SELECT COALESCE(json_agg(cr.resource_id), '[]'::json)
     FROM capability_resource cr WHERE cr.capability_id = capability.id) AS resource_ids
"""

_BASE_FIELDS = "id, name, criterion, description, version"

_SELECT_FIELDS = f"{_BASE_FIELDS}, {_REL_SUBQUERIES}"


def _normalize_links(data: Dict, links_key: str, ids_key: str, default_type: str):
    """两种输入格式统一：{links_key: [{id, relation_type}]} 或 {ids_key: [id]}"""
    if links_key in data and data[links_key] is not None:
        out = []
        for item in data[links_key]:
            if isinstance(item, dict):
                out.append((item['id'], item.get('relation_type', default_type)))
            else:
                out.append((item, default_type))
        return out
    if ids_key in data and data[ids_key] is not None:
        return [(i, default_type) for i in data[ids_key]]
    return None


class PostgreSQLCapabilityStore:
    def __init__(self):
        """初始化 PostgreSQL 连接"""
        self.conn = psycopg2.connect(
            host=os.getenv('KNOWHUB_DB'),
            port=int(os.getenv('KNOWHUB_PORT', 5432)),
            user=os.getenv('KNOWHUB_USER'),
            password=os.getenv('KNOWHUB_PASSWORD'),
            database=os.getenv('KNOWHUB_DB_NAME')
        )
        self.conn.autocommit = True
        print(f"[PostgreSQL Capability] 已连接到远程数据库: {os.getenv('KNOWHUB_DB')}")

    def _reconnect(self):
        self.conn = psycopg2.connect(
            host=os.getenv('KNOWHUB_DB'),
            port=int(os.getenv('KNOWHUB_PORT', 5432)),
            user=os.getenv('KNOWHUB_USER'),
            password=os.getenv('KNOWHUB_PASSWORD'),
            database=os.getenv('KNOWHUB_DB_NAME')
        )
        self.conn.autocommit = True

    def _ensure_connection(self):
        if self.conn.closed != 0:
            self._reconnect()
        else:
            try:
                c = self.conn.cursor()
                c.execute("SELECT 1")
                c.close()
            except (psycopg2.OperationalError, psycopg2.InterfaceError):
                self._reconnect()

    def _get_cursor(self):
        self._ensure_connection()
        return self.conn.cursor(cursor_factory=RealDictCursor)

    def _save_relations(self, cursor, cap_id: str, data: Dict):
        """保存 capability 的关联表数据"""
        if 'requirement_ids' in data:
            cursor.execute("DELETE FROM requirement_capability WHERE capability_id = %s", (cap_id,))
            for req_id in data['requirement_ids']:
                cursor.execute(
                    "INSERT INTO requirement_capability (requirement_id, capability_id) VALUES (%s, %s) ON CONFLICT DO NOTHING",
                    (req_id, cap_id))

        # tool_ids + implements 合并写入 capability_tool
        if 'tool_ids' in data or 'implements' in data:
            cursor.execute("DELETE FROM capability_tool WHERE capability_id = %s", (cap_id,))
            implements = data.get('implements', {})
            tool_ids = set(data.get('tool_ids', []))
            # 先写 tool_ids 列表中的（附带 implements 的 description）
            for tool_id in tool_ids:
                desc = implements.get(tool_id, '')
                cursor.execute(
                    "INSERT INTO capability_tool (capability_id, tool_id, description) VALUES (%s, %s, %s) ON CONFLICT DO NOTHING",
                    (cap_id, tool_id, desc))
            # 再写 implements 中有但 tool_ids 列表没有的
            for tool_id, desc in implements.items():
                if tool_id not in tool_ids:
                    cursor.execute(
                        "INSERT INTO capability_tool (capability_id, tool_id, description) VALUES (%s, %s, %s) ON CONFLICT DO NOTHING",
                        (cap_id, tool_id, desc))

        k_links = _normalize_links(data, 'knowledge_links', 'knowledge_ids', 'related')
        if k_links is not None:
            cursor.execute("DELETE FROM capability_knowledge WHERE capability_id = %s", (cap_id,))
            for kid, rtype in k_links:
                cursor.execute(
                    "INSERT INTO capability_knowledge (capability_id, knowledge_id, relation_type) "
                    "VALUES (%s, %s, %s) ON CONFLICT DO NOTHING",
                    (cap_id, kid, rtype))

        if 'resource_ids' in data and data['resource_ids'] is not None:
            cursor.execute("DELETE FROM capability_resource WHERE capability_id = %s", (cap_id,))
            for rid in data['resource_ids']:
                cursor.execute(
                    "INSERT INTO capability_resource (capability_id, resource_id) VALUES (%s, %s) ON CONFLICT DO NOTHING",
                    (cap_id, rid))

    def insert_or_update(self, cap: Dict):
        """插入或更新原子能力。AnalyticDB beam 表不支持 ON CONFLICT UPDATE 当含 ALTER 新增列，改用 DELETE+INSERT。"""
        cursor = self._get_cursor()
        try:
            cursor.execute("DELETE FROM capability WHERE id = %s", (cap['id'],))
            cursor.execute("""
                INSERT INTO capability (
                    id, name, criterion, description, embedding, version
                ) VALUES (%s, %s, %s, %s, %s, %s)
            """, (
                cap['id'],
                cap.get('name', ''),
                cap.get('criterion', ''),
                cap.get('description', ''),
                cap.get('embedding'),
                cap.get('version', 'v0'),
            ))
            self._save_relations(cursor, cap['id'], cap)
            self.conn.commit()
        finally:
            cursor.close()

    def get_by_id(self, cap_id: str) -> Optional[Dict]:
        """根据 ID 获取原子能力"""
        cursor = self._get_cursor()
        try:
            cursor.execute(f"""
                SELECT {_SELECT_FIELDS}
                FROM capability WHERE id = %s
            """, (cap_id,))
            result = cursor.fetchone()
            return self._format_result(result) if result else None
        finally:
            cursor.close()

    def search(self, query_embedding: List[float], limit: int = 10) -> List[Dict]:
        """向量检索原子能力"""
        cursor = self._get_cursor()
        try:
            cursor.execute(f"""
                SELECT {_SELECT_FIELDS},
                       1 - (embedding <=> %s::real[]) as score
                FROM capability
                WHERE embedding IS NOT NULL
                ORDER BY embedding <=> %s::real[]
                LIMIT %s
            """, (query_embedding, query_embedding, limit))
            results = cursor.fetchall()
            return [self._format_result(r) for r in results]
        finally:
            cursor.close()

    def list_all(self, limit: int = 100, offset: int = 0) -> List[Dict]:
        """列出原子能力"""
        cursor = self._get_cursor()
        try:
            cursor.execute(f"""
                SELECT {_SELECT_FIELDS}
                FROM capability
                ORDER BY id
                LIMIT %s OFFSET %s
            """, (limit, offset))
            results = cursor.fetchall()
            return [self._format_result(r) for r in results]
        finally:
            cursor.close()

    def update(self, cap_id: str, updates: Dict):
        """更新原子能力字段"""
        cursor = self._get_cursor()
        try:
            # 分离关联字段
            rel_fields = {}
            for key in ('requirement_ids', 'implements', 'tool_ids',
                        'knowledge_ids', 'knowledge_links', 'resource_ids'):
                if key in updates:
                    rel_fields[key] = updates.pop(key)

            if updates:
                set_parts = []
                params = []
                for key, value in updates.items():
                    set_parts.append(f"{key} = %s")
                    params.append(value)
                params.append(cap_id)
                cursor.execute(
                    f"UPDATE capability SET {', '.join(set_parts)} WHERE id = %s",
                    params
                )

            if rel_fields:
                self._save_relations(cursor, cap_id, rel_fields)

            self.conn.commit()
        finally:
            cursor.close()

    def delete(self, cap_id: str):
        """删除原子能力及其关联表记录"""
        cursor = self._get_cursor()
        try:
            cascade_delete(cursor, 'capability', cap_id)
            self.conn.commit()
        finally:
            cursor.close()

    def count(self) -> int:
        """统计原子能力总数"""
        cursor = self._get_cursor()
        try:
            cursor.execute("SELECT COUNT(*) as count FROM capability")
            return cursor.fetchone()['count']
        finally:
            cursor.close()

    def _format_result(self, row: Dict) -> Dict:
        """格式化查询结果"""
        if not row:
            return None
        result = dict(row)
        for field in ('requirement_ids', 'tool_ids', 'knowledge_ids',
                      'resource_ids', 'knowledge_links'):
            if field in result and isinstance(result[field], str):
                result[field] = json.loads(result[field])
            elif field in result and result[field] is None:
                result[field] = []
        if 'implements' in result:
            if isinstance(result['implements'], str):
                result['implements'] = json.loads(result['implements'])
            elif result['implements'] is None:
                result['implements'] = {}
        return result

    def add_knowledge(self, cap_id: str, knowledge_id: str, relation_type: str = 'related'):
        """增量挂接 capability-knowledge 边"""
        cursor = self._get_cursor()
        try:
            cursor.execute(
                "INSERT INTO capability_knowledge (capability_id, knowledge_id, relation_type) "
                "VALUES (%s, %s, %s) ON CONFLICT DO NOTHING",
                (cap_id, knowledge_id, relation_type))
            self.conn.commit()
        finally:
            cursor.close()

    def add_resource(self, cap_id: str, resource_id: str):
        """增量挂接 capability-resource 边"""
        cursor = self._get_cursor()
        try:
            cursor.execute(
                "INSERT INTO capability_resource (capability_id, resource_id) VALUES (%s, %s) ON CONFLICT DO NOTHING",
                (cap_id, resource_id))
            self.conn.commit()
        finally:
            cursor.close()

    def close(self):
        if self.conn:
            self.conn.close()