Tool_Agent/migrate_postgres.py

from tool_agent.registry.registry import ToolRegistry
from tool_agent.tool.tool_store import PostgreSQLToolStore
from tool_agent.models import ToolStatus

def migrate():
    registry = ToolRegistry()
    store = PostgreSQLToolStore()
    
    local_tools = registry.list_all()
    remote_tools = {kt["id"]: kt for kt in store.list_all(limit=1000)}
    
    print(f"Loaded {len(local_tools)} local tools, {len(remote_tools)} remote tools")
    
    for tool in local_tools:
        tid = tool.tool_id
        
        # Merge logic: favor remote's name/intro if it exists
        remote_t = remote_tools.get(tid, {})
        
        tool_dict = {
            "id": tid,
            "name": remote_t.get("name") or tool.name,
            "version": remote_t.get("version") or tool.version,
            "introduction": remote_t.get("introduction") or tool.description,
            "input": tool.input_schema,  # overwrite with local JSON schema
            "output": tool.output_schema,
            "status": "已接入" if tool.status == ToolStatus.ACTIVE else "未接入"
        }
        store.insert_or_update(tool_dict)
        print(f"Migrated / Synced tool: {tid}")
        
    print("Migration complete!")

if __name__ == '__main__':
    migrate()