Agent/test_upload.py

#!/usr/bin/env python3
"""
测试知识上传脚本
"""
import sys
import os
import uuid
import time
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
}

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 = get_embedding(knowledge_data['task'])
content_embedding = 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"]}')