Agent/direct_insert.py

#!/usr/bin/env python3
"""
直接插入知识到数据库的脚本
"""
import sys
import os
import uuid
import time
import asyncio
from datetime import datetime, timezone

sys.path.insert(0, '/root/Agent')

# 设置环境变量
os.environ['KNOWHUB_DB'] = 'gp-t4n72471pkmt4b9q7o-master.gpdbmaster.singapore.rds.aliyuncs.com'
os.environ['KNOWHUB_PORT'] = '5432'
os.environ['KNOWHUB_USER'] = 'aiddit_aigc'
os.environ['KNOWHUB_PASSWORD'] = '%a&&yqNxg^V1$toJ*WOa^-b^X=QJ'
os.environ['KNOWHUB_DB_NAME'] = 'knowhub'

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

# 测试知识数据
knowledge_data = {
    'task': '接口测试验证条目',
    'content': '这是一条用于验证 upload 接口是否可用的测试知识，可以删除。',
    'types': ['experience'],
    'tags': {'source': 'api_test'},
    'score': 1
}

async def main():
    print('正在连接数据库...')
    store = PostgreSQLStore()

    # 生成 ID
    now = datetime.now(timezone.utc)
    knowledge_id = f'knowledge-{now.strftime("%Y%m%d")}-{uuid.uuid4().hex[:8]}'
    message_id = f'msg-{uuid.uuid4().hex[:12]}'

    # 生成 embedding
    print('正在生成 embedding...')
    task_embedding = await get_embedding(knowledge_data['task'])
    content_embedding = await get_embedding(knowledge_data['content'])

    # 构建知识记录
    knowledge = {
        'id': knowledge_id,
        'message_id': message_id,
        'task': knowledge_data['task'],
        'content': knowledge_data['content'],
        'types': knowledge_data['types'],
        'tags': knowledge_data.get('tags', {}),
        'tag_keys': list(knowledge_data.get('tags', {}).keys()),
        'scopes': ['org:cybertogether'],
        'owner': 'system',
        'source': {'category': 'execution'},
        'eval': {'score': knowledge_data.get('score', 3)},
        'task_embedding': task_embedding,
        'content_embedding': content_embedding,
        'created_at': int(time.time()),
        'updated_at': int(time.time()),
        'status': 'approved',
    }

    print(f'正在插入知识条目...')
    store.insert(knowledge)
    store.close()

    print(f'✅ 成功插入知识条目！')
    print(f'  ID: {knowledge_id}')
    print(f'  Task: {knowledge_data["task"]}')
    print(f'  Types: {knowledge_data["types"]}')
    
    return knowledge_id

if __name__ == '__main__':
    knowledge_id = asyncio.run(main())
    print(f'\n最终知识 ID: {knowledge_id}')