Agent/knowhub/knowhub_db/pg_store.py

"""
PostgreSQL 存储封装（替代 Milvus）

使用远程 PostgreSQL + pgvector/fastann 存储知识数据
"""

import os
import json
import psycopg2
from psycopg2.extras import RealDictCursor, execute_batch
from typing import List, Dict, Optional
from dotenv import load_dotenv

load_dotenv()


class PostgreSQLStore:
    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 = False
        print(f"[PostgreSQL] 已连接到远程数据库: {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 = False

    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 insert(self, knowledge: Dict):
        """插入单条知识"""
        cursor = self._get_cursor()
        try:
            cursor.execute("""
                INSERT INTO knowledge (
                    id, task_embedding, content_embedding, message_id, task, content, types, tags,
                    tag_keys, scopes, owner, resource_ids, source, eval,
                    created_at, updated_at, status, relationships,
                    support_capability, tools
                ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
            """, (
                knowledge['id'],
                knowledge.get('task_embedding') or knowledge.get('embedding'),
                knowledge.get('content_embedding'),
                knowledge['message_id'],
                knowledge['task'],
                knowledge['content'],
                knowledge.get('types', []),
                json.dumps(knowledge.get('tags', {})),
                knowledge.get('tag_keys', []),
                knowledge.get('scopes', []),
                knowledge['owner'],
                knowledge.get('resource_ids', []),
                json.dumps(knowledge.get('source', {})),
                json.dumps(knowledge.get('eval', {})),
                knowledge['created_at'],
                knowledge['updated_at'],
                knowledge.get('status', 'approved'),
                json.dumps(knowledge.get('relationships', [])),
                json.dumps(knowledge.get('support_capability', [])),
                json.dumps(knowledge.get('tools', [])),
            ))
            self.conn.commit()
        finally:
            cursor.close()

    def search(self, query_embedding: List[float], filters: Optional[str] = None, limit: int = 10) -> List[Dict]:
        """向量检索（使用余弦相似度）"""
        cursor = self._get_cursor()
        try:
            where_clause = self._build_where_clause(filters) if filters else ""
            sql = f"""
                SELECT id, message_id, task, content, types, tags, tag_keys,
                       scopes, owner, resource_ids, source, eval, created_at,
                       updated_at, status, relationships, support_capability, tools,
                       1 - (task_embedding <=> %s::real[]) as score
                FROM knowledge
                {where_clause}
                ORDER BY task_embedding <=> %s::real[]
                LIMIT %s
            """
            cursor.execute(sql, (query_embedding, query_embedding, limit))
            results = cursor.fetchall()
            return [self._format_result(r) for r in results]
        finally:
            cursor.close()

    def query(self, filters: str, limit: int = 100) -> List[Dict]:
        """纯标量查询"""
        cursor = self._get_cursor()
        try:
            where_clause = self._build_where_clause(filters)
            sql = f"""
                SELECT id, message_id, task, content, types, tags, tag_keys,
                       scopes, owner, resource_ids, source, eval, created_at,
                       updated_at, status, relationships, support_capability, tools
                FROM knowledge
                {where_clause}
                LIMIT %s
            """
            cursor.execute(sql, (limit,))
            results = cursor.fetchall()
            return [self._format_result(r) for r in results]
        finally:
            cursor.close()

    def get_by_id(self, knowledge_id: str, include_embedding: bool = False) -> Optional[Dict]:
        """根据ID获取知识（默认不返回embedding以提升性能）"""
        cursor = self._get_cursor()
        try:
            # 默认不返回embedding（1536维向量太大，详情页不需要）
            if include_embedding:
                fields = "id, task_embedding, content_embedding, message_id, task, content, types, tags, tag_keys, scopes, owner, resource_ids, source, eval, created_at, updated_at, status, relationships, support_capability, tools"
            else:
                fields = "id, message_id, task, content, types, tags, tag_keys, scopes, owner, resource_ids, source, eval, created_at, updated_at, status, relationships, support_capability, tools"

            cursor.execute(f"""
                SELECT {fields}
                FROM knowledge WHERE id = %s
            """, (knowledge_id,))
            result = cursor.fetchone()
            return self._format_result(result) if result else None
        finally:
            cursor.close()

    def update(self, knowledge_id: str, updates: Dict):
        """更新知识"""
        cursor = self._get_cursor()
        try:
            set_parts = []
            params = []
            for key, value in updates.items():
                if key in ('tags', 'source', 'eval', 'support_capability', 'tools'):
                    set_parts.append(f"{key} = %s")
                    params.append(json.dumps(value))
                elif key == 'relationships':
                    set_parts.append(f"{key} = %s")
                    params.append(json.dumps(value) if isinstance(value, list) else value)
                else:
                    set_parts.append(f"{key} = %s")
                    params.append(value)
            params.append(knowledge_id)
            sql = f"UPDATE knowledge SET {', '.join(set_parts)} WHERE id = %s"
            cursor.execute(sql, params)
            self.conn.commit()
        finally:
            cursor.close()

    def delete(self, knowledge_id: str):
        """删除知识"""
        cursor = self._get_cursor()
        try:
            cursor.execute("DELETE FROM knowledge WHERE id = %s", (knowledge_id,))
            self.conn.commit()
        finally:
            cursor.close()

    def count(self) -> int:
        """返回知识总数"""
        cursor = self._get_cursor()
        try:
            cursor.execute("SELECT COUNT(*) as count FROM knowledge")
            return cursor.fetchone()['count']
        finally:
            cursor.close()

    def _build_where_clause(self, filters: str) -> str:
        """将Milvus风格的过滤表达式转换为PostgreSQL WHERE子句"""
        if not filters:
            return ""

        where = filters

        import re

        # 替换操作符
        where = where.replace(' == ', ' = ')
        where = where.replace(' or ', ' OR ')
        where = where.replace(' and ', ' AND ')

        # 处理数组包含操作
        where = re.sub(r'array_contains\((\w+),\s*"([^"]+)"\)', r"\1 @> ARRAY['\2']", where)

        # 处理 eval["score"] 语法
        where = where.replace('eval["score"]', "(eval->>'score')::int")

        # 把所有剩余的双引号字符串值替换为单引号（PostgreSQL标准）
        where = re.sub(r'"([^"]*)"', r"'\1'", where)

        return f"WHERE {where}"

    def _format_result(self, row: Dict) -> Dict:
        """格式化查询结果"""
        if not row:
            return None
        result = dict(row)
        if 'tags' in result and isinstance(result['tags'], str):
            result['tags'] = json.loads(result['tags'])
        if 'source' in result and isinstance(result['source'], str):
            result['source'] = json.loads(result['source'])
        if 'eval' in result and isinstance(result['eval'], str):
            result['eval'] = json.loads(result['eval'])
        if 'relationships' in result and isinstance(result['relationships'], str):
            result['relationships'] = json.loads(result['relationships'])
        if 'support_capability' in result and isinstance(result['support_capability'], str):
            result['support_capability'] = json.loads(result['support_capability'])
        if 'tools' in result and isinstance(result['tools'], str):
            result['tools'] = json.loads(result['tools'])
        if 'created_at' in result and result['created_at']:
            result['created_at'] = result['created_at'] * 1000
        if 'updated_at' in result and result['updated_at']:
            result['updated_at'] = result['updated_at'] * 1000
        return result

    def close(self):
        """关闭连接"""
        if self.conn:
            self.conn.close()

    def insert_batch(self, knowledge_list: List[Dict]):
        """批量插入知识"""
        if not knowledge_list:
            return

        cursor = self._get_cursor()
        try:
            data = []
            for k in knowledge_list:
                data.append((
                    k['id'], k.get('task_embedding') or k.get('embedding'),
                    k.get('content_embedding'),
                    k['message_id'], k['task'],
                    k['content'], k.get('types', []),
                    json.dumps(k.get('tags', {})), k.get('tag_keys', []),
                    k.get('scopes', []), k['owner'], k.get('resource_ids', []),
                    json.dumps(k.get('source', {})), json.dumps(k.get('eval', {})),
                    k['created_at'], k['updated_at'], k.get('status', 'approved'),
                    json.dumps(k.get('relationships', [])),
                    json.dumps(k.get('support_capability', [])),
                    json.dumps(k.get('tools', [])),
                ))

            execute_batch(cursor, """
                INSERT INTO knowledge (
                    id, task_embedding, content_embedding, message_id, task, content, types, tags,
                    tag_keys, scopes, owner, resource_ids, source, eval,
                    created_at, updated_at, status, relationships,
                    support_capability, tools
                ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
            """, data)
            self.conn.commit()
        finally:
            cursor.close()