Agent/insert_knowledge.py

#!/usr/bin/env python3
"""
插入知识条目到 PostgreSQL 数据库
"""

import asyncio
import uuid
import time
import sys
import os

# 添加项目路径
sys.path.insert(0, '/root/Agent')

from knowhub.knowhub_db.pg_store import PostgreSQLStore
from knowhub.embeddings import get_embedding


async def main():
    # 知识数据
    task = "人物近景肖像"
    content = "生成人物近景半身或胸部以上画面、突出面部表情和情绪、背景虚化的核心工序要素包括：1.近景/半身构图 2.面部表情细节 3.背景虚化效果 4.皮肤质感真实 5.眼神互动感"
    types = ["strategy"]
    tags = {"domain": "AI 生图", "source": "多渠调研汇总", "task": "人物近景肖像"}
    score = 5
    source = {"category": "research"}
    
    # 生成唯一 ID
    knowledge_id = str(uuid.uuid4())
    message_id = str(uuid.uuid4())
    current_time = int(time.time())
    
    # 生成 embeddings
    print(f"正在生成 task embedding...")
    task_embedding = await get_embedding(task)
    print(f"正在生成 content embedding...")
    content_embedding = await get_embedding(content)
    
    # 构建知识对象
    knowledge = {
        'id': knowledge_id,
        'message_id': message_id,
        'task': task,
        'content': content,
        'types': types,
        'tags': tags,
        'tag_keys': list(tags.keys()),
        'scopes': ['org:cybertogether'],  # 默认可见范围
        'owner': 'system',  # 默认所有者
        'source': source,
        'eval': {'score': score},
        'task_embedding': task_embedding,
        'content_embedding': content_embedding,
        'created_at': current_time,
        'updated_at': current_time,
        'status': 'approved',
    }
    
    # 插入数据库
    print(f"正在插入知识条目到 PostgreSQL...")
    store = PostgreSQLStore()
    store.insert(knowledge)
    store.close()
    
    print(f"成功插入知识条目！")
    print(f"  ID: {knowledge_id}")
    print(f"  Task: {task}")
    print(f"  Types: {types}")
    print(f"  Score: {score}")


if __name__ == '__main__':
    asyncio.run(main())