qwen_agent/scripts/cache_manager.py

#!/usr/bin/env python3
"""
Agent cache management utility.
Used to manage and maintain the DiskAgent cache.
"""
import sys
import os
import argparse
import json

# Add the project root directory to the Python path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from agent.agent_cache_manager import get_cache_manager
from agent.agent_memory_cache import get_memory_cache_manager


def show_stats():
    """Show cache statistics."""
    print("\n=== Agent Cache Statistics ===")

    print("\n--- Memory Cache ---")
    mem_cache = get_memory_cache_manager()
    mem_stats = mem_cache.get_stats()
    print(f"Cache Type: {mem_stats.get('type', 'memory')}")
    print(f"Total Items: {mem_stats.get('total_items', 0)}")
    print(f"Max Size: {mem_stats.get('max_size', 0)}")
    print(f"Default TTL: {mem_stats.get('default_ttl', 0)} seconds")
    print(f"Auto Renew: {mem_stats.get('auto_renew', False)}")
    print(f"Hits: {mem_stats.get('hits', 0)}")
    print(f"Misses: {mem_stats.get('misses', 0)}")
    print(f"Hit Rate: {mem_stats.get('hit_rate_percent', 0)}%")
    print(f"Memory Usage: {mem_stats.get('memory_usage_mb', 0)} MB")
    print(f"Evictions: {mem_stats.get('evictions', 0)}")

    print("\n--- Disk Cache (if used) ---")
    try:
        disk_cache = get_cache_manager()
        disk_stats = disk_cache.get_stats()
        print(f"Cache Directory: {disk_stats.get('cache_dir', 'N/A')}")
        print(f"Total Items: {disk_stats.get('total_items', 0)}")
        print(f"Current Size: {disk_stats.get('size_mb', 0)} MB")
        print(f"Size Limit: {disk_stats.get('size_limit_mb', 0)} MB")
        print(f"Usage: {disk_stats.get('size_mb', 0) / disk_stats.get('size_limit_mb', 1) * 100:.1f}%")
    except:
        print("No disk cache configured")

    print("=" * 30)


def clear_all():
    """Clear all caches."""
    print("\nWarning: This will clear ALL cached agents!")
    confirm = input("Are you sure? (yes/no): ")

    if confirm.lower() == 'yes':
        # Clear the in-memory cache
        mem_cache = get_memory_cache_manager()
        mem_success = mem_cache.clear_all()
        print(f"Memory cache cleared: {'✓' if mem_success else '✗'}")

        # Clear the disk cache if present
        try:
            disk_cache = get_cache_manager()
            disk_success = disk_cache.clear_all()
            print(f"Disk cache cleared: {'✓' if disk_success else '✗'}")
        except:
            print("No disk cache to clear")
    else:
        print("Operation cancelled")


def clear_expired():
    """Clear expired cache entries."""
    cache_manager = get_cache_manager()

    count = cache_manager.clear_expired()
    if count > 0:
        print(f"✓ Cleared {count} expired cache entries")
    else:
        print("No expired entries found")


def cleanup_old():
    """Clean up old cache entries."""
    parser = argparse.ArgumentParser(description='Clean up old cache entries')
    parser.add_argument('--max-age', type=int, default=3600,
                       help='Maximum age in seconds (default: 3600 = 1 hour)')
    args = parser.parse_args(sys.argv[2:])

    cache_manager = get_cache_manager()

    count = cache_manager.cleanup_old_entries(args.max_age)
    if count > 0:
        print(f"✓ Cleaned up {count} old cache entries (older than {args.max_age} seconds)")
    else:
        print("No old entries found")


def list_keys():
    """List all cache keys (first 10)."""
    print("\n=== Cache Keys (showing first 10) ===")

    print("\n--- Memory Cache Keys ---")
    mem_cache = get_memory_cache_manager()
    mem_keys = mem_cache.get_keys()

    for i, key in enumerate(mem_keys[:10]):
        print(f"  - {key}")

    if len(mem_keys) > 10:
        print(f"  ... and {len(mem_keys) - 10} more in memory cache")

    print(f"\nTotal memory cache keys: {len(mem_keys)}")

    print("\n--- Disk Cache Keys ---")
    try:
        disk_cache = get_cache_manager()
        disk_keys = list(disk_cache.cache.iterkeys())

        for i, key in enumerate(disk_keys[:10]):
            print(f"  - {key}")

        if len(disk_keys) > 10:
            print(f"  ... and {len(disk_keys) - 10} more in disk cache")

        print(f"\nTotal disk cache keys: {len(disk_keys)}")
    except:
        print("No disk cache configured")

    print("=" * 30)


def delete_key():
    """Delete a specific cache key."""
    if len(sys.argv) < 3:
        print("Usage: python cache_manager.py delete <cache_key>")
        return

    cache_key = sys.argv[2]
    cache_manager = get_cache_manager()

    if cache_manager.delete(cache_key):
        print(f"✓ Deleted cache key: {cache_key}")
    else:
        print(f"✗ Cache key not found: {cache_key}")


def main():
    """Main function."""
    if len(sys.argv) < 2:
        print("\nAgent Cache Manager")
        print("\nUsage:")
        print("  python cache_manager.py <command> [options]")
        print("\nCommands:")
        print("  stats      - Show cache statistics")
        print("  clear      - Clear ALL cache entries")
        print("  expired    - Clear expired entries")
        print("  cleanup    - Clean up old entries (use --max-age)")
        print("  list       - List cache keys")
        print("  delete     - Delete specific cache key")
        print("\nExamples:")
        print("  python cache_manager.py stats")
        print("  python cache_manager.py cleanup --max-age 7200")
        print("  python cache_manager.py delete my-cache-key")
        return

    command = sys.argv[1]

    try:
        if command == 'stats':
            show_stats()
        elif command == 'clear':
            clear_all()
        elif command == 'expired':
            clear_expired()
        elif command == 'cleanup':
            cleanup_old()
        elif command == 'list':
            list_keys()
        elif command == 'delete':
            delete_key()
        else:
            print(f"Unknown command: {command}")
            print("Use 'python cache_manager.py' to see available commands")
    except Exception as e:
        print(f"Error: {e}")


if __name__ == "__main__":
    main()