513 lines
17 KiB
Python
513 lines
17 KiB
Python
"""
|
|
Performance optimizations for MaxClaw Agent
|
|
Enhanced with connection pooling, lazy loading, and advanced caching
|
|
"""
|
|
import os
|
|
import sys
|
|
import time
|
|
import functools
|
|
import threading
|
|
from typing import Dict, Any, Optional, Callable, List
|
|
from pathlib import Path
|
|
import json
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
|
|
# Add current directory to path
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
|
|
|
|
class ConnectionPool:
|
|
"""Connection pool for managing API connections"""
|
|
|
|
def __init__(self, max_connections: int = 10):
|
|
"""
|
|
Initialize connection pool
|
|
|
|
Args:
|
|
max_connections: Maximum number of connections in the pool
|
|
"""
|
|
self.max_connections = max_connections
|
|
self.connections = {}
|
|
self.lock = threading.Lock()
|
|
|
|
def get_connection(self, provider: str, **kwargs):
|
|
"""
|
|
Get a connection from the pool
|
|
|
|
Args:
|
|
provider: Provider name
|
|
**kwargs: Connection parameters
|
|
|
|
Returns:
|
|
Connection object
|
|
"""
|
|
with self.lock:
|
|
key = f"{provider}:{str(sorted(kwargs.items()))}"
|
|
|
|
if key not in self.connections:
|
|
# Create new connection
|
|
if provider == "ollama":
|
|
import requests
|
|
base_url = kwargs.get("base_url", "http://localhost:11434")
|
|
session = requests.Session()
|
|
session.headers.update({"Content-Type": "application/json"})
|
|
self.connections[key] = {
|
|
"session": session,
|
|
"base_url": base_url,
|
|
"last_used": time.time()
|
|
}
|
|
elif provider in ["openai", "openrouter", "anthropic", "google"]:
|
|
import requests
|
|
session = requests.Session()
|
|
api_key = kwargs.get("api_key", "")
|
|
if api_key:
|
|
session.headers.update({"Authorization": f"Bearer {api_key}"})
|
|
self.connections[key] = {
|
|
"session": session,
|
|
"last_used": time.time()
|
|
}
|
|
|
|
# Update last used time
|
|
if key in self.connections:
|
|
self.connections[key]["last_used"] = time.time()
|
|
|
|
return self.connections.get(key)
|
|
|
|
def cleanup_idle_connections(self, max_idle_time: int = 300):
|
|
"""
|
|
Clean up idle connections
|
|
|
|
Args:
|
|
max_idle_time: Maximum idle time in seconds
|
|
"""
|
|
with self.lock:
|
|
current_time = time.time()
|
|
keys_to_remove = []
|
|
|
|
for key, conn in self.connections.items():
|
|
if current_time - conn["last_used"] > max_idle_time:
|
|
keys_to_remove.append(key)
|
|
|
|
for key in keys_to_remove:
|
|
del self.connections[key]
|
|
|
|
def get_stats(self) -> Dict[str, Any]:
|
|
"""Get connection pool statistics"""
|
|
with self.lock:
|
|
return {
|
|
"total_connections": len(self.connections),
|
|
"max_connections": self.max_connections,
|
|
"connections": list(self.connections.keys())
|
|
}
|
|
|
|
|
|
class LazyLoader:
|
|
"""Lazy loading for heavy resources"""
|
|
|
|
def __init__(self):
|
|
self.loaded_modules = {}
|
|
self.loaded_tools = {}
|
|
self.lock = threading.Lock()
|
|
|
|
def load_module(self, module_name: str, import_path: str):
|
|
"""
|
|
Lazily load a module
|
|
|
|
Args:
|
|
module_name: Name to store the module under
|
|
import_path: Import path for the module
|
|
|
|
Returns:
|
|
Loaded module
|
|
"""
|
|
with self.lock:
|
|
if module_name not in self.loaded_modules:
|
|
try:
|
|
module = __import__(import_path, fromlist=[''])
|
|
self.loaded_modules[module_name] = module
|
|
except ImportError as e:
|
|
print(f"Warning: Could not load module {module_name}: {e}")
|
|
self.loaded_modules[module_name] = None
|
|
|
|
return self.loaded_modules.get(module_name)
|
|
|
|
def load_tool(self, tool_name: str, loader_func: Callable):
|
|
"""
|
|
Lazily load a tool
|
|
|
|
Args:
|
|
tool_name: Name to store the tool under
|
|
loader_func: Function to load the tool
|
|
|
|
Returns:
|
|
Loaded tool
|
|
"""
|
|
with self.lock:
|
|
if tool_name not in self.loaded_tools:
|
|
try:
|
|
tool = loader_func()
|
|
self.loaded_tools[tool_name] = tool
|
|
except Exception as e:
|
|
print(f"Warning: Could not load tool {tool_name}: {e}")
|
|
self.loaded_tools[tool_name] = None
|
|
|
|
return self.loaded_tools.get(tool_name)
|
|
|
|
def clear(self):
|
|
"""Clear all loaded resources"""
|
|
with self.lock:
|
|
self.loaded_modules.clear()
|
|
self.loaded_tools.clear()
|
|
|
|
|
|
class BatchProcessor:
|
|
"""Batch processing for multiple operations"""
|
|
|
|
def __init__(self, max_workers: int = 5):
|
|
"""
|
|
Initialize batch processor
|
|
|
|
Args:
|
|
max_workers: Maximum number of worker threads
|
|
"""
|
|
self.max_workers = max_workers
|
|
self.executor = ThreadPoolExecutor(max_workers=max_workers)
|
|
|
|
def process_batch(self, tasks: List[Callable], timeout: int = 30) -> List[Any]:
|
|
"""
|
|
Process a batch of tasks in parallel
|
|
|
|
Args:
|
|
tasks: List of callable tasks
|
|
timeout: Timeout in seconds
|
|
|
|
Returns:
|
|
List of results
|
|
"""
|
|
futures = [self.executor.submit(task) for task in tasks]
|
|
results = []
|
|
|
|
try:
|
|
for future in as_completed(futures, timeout=timeout):
|
|
try:
|
|
result = future.result()
|
|
results.append(result)
|
|
except Exception as e:
|
|
results.append({"error": str(e)})
|
|
except Exception as e:
|
|
print(f"Batch processing error: {e}")
|
|
|
|
return results
|
|
|
|
def shutdown(self):
|
|
"""Shutdown the executor"""
|
|
self.executor.shutdown(wait=True)
|
|
|
|
|
|
class PerformanceMonitor:
|
|
"""Monitor and optimize performance"""
|
|
|
|
def __init__(self):
|
|
self.metrics = {}
|
|
self.cache = {}
|
|
self.cache_ttl = 300 # 5 minutes default TTL
|
|
self.connection_pool = ConnectionPool()
|
|
self.lazy_loader = LazyLoader()
|
|
self.batch_processor = BatchProcessor()
|
|
|
|
def timing(self, func: Callable) -> Callable:
|
|
"""Decorator to measure function execution time"""
|
|
@functools.wraps(func)
|
|
def wrapper(*args, **kwargs):
|
|
start_time = time.time()
|
|
result = func(*args, **kwargs)
|
|
end_time = time.time()
|
|
|
|
# Store timing metric
|
|
func_name = func.__name__
|
|
if func_name not in self.metrics:
|
|
self.metrics[func_name] = []
|
|
|
|
self.metrics[func_name].append(end_time - start_time)
|
|
|
|
# Keep only last 100 measurements
|
|
if len(self.metrics[func_name]) > 100:
|
|
self.metrics[func_name] = self.metrics[func_name][-100:]
|
|
|
|
return result
|
|
return wrapper
|
|
|
|
def cache_result(self, ttl: Optional[int] = None):
|
|
"""Decorator to cache function results"""
|
|
def decorator(func: Callable) -> Callable:
|
|
@functools.wraps(func)
|
|
def wrapper(*args, **kwargs):
|
|
# Create cache key
|
|
cache_key = f"{func.__name__}:{str(args)}:{str(kwargs)}"
|
|
|
|
# Check cache
|
|
if cache_key in self.cache:
|
|
cached_time, cached_result = self.cache[cache_key]
|
|
if time.time() - cached_time < (ttl or self.cache_ttl):
|
|
return cached_result
|
|
|
|
# Execute function
|
|
result = func(*args, **kwargs)
|
|
|
|
# Cache result
|
|
self.cache[cache_key] = (time.time(), result)
|
|
|
|
return result
|
|
return wrapper
|
|
return decorator
|
|
|
|
def get_metrics(self) -> Dict[str, Any]:
|
|
"""Get performance metrics"""
|
|
metrics_summary = {}
|
|
|
|
for func_name, timings in self.metrics.items():
|
|
if timings:
|
|
metrics_summary[func_name] = {
|
|
"count": len(timings),
|
|
"avg_time": sum(timings) / len(timings),
|
|
"min_time": min(timings),
|
|
"max_time": max(timings),
|
|
"total_time": sum(timings)
|
|
}
|
|
|
|
return metrics_summary
|
|
|
|
def clear_cache(self):
|
|
"""Clear the cache"""
|
|
self.cache.clear()
|
|
|
|
def clear_metrics(self):
|
|
"""Clear performance metrics"""
|
|
self.metrics.clear()
|
|
|
|
def get_connection_pool_stats(self) -> Dict[str, Any]:
|
|
"""Get connection pool statistics"""
|
|
return self.connection_pool.get_stats()
|
|
|
|
def cleanup_resources(self):
|
|
"""Clean up resources"""
|
|
self.connection_pool.cleanup_idle_connections()
|
|
self.clear_cache()
|
|
|
|
def shutdown(self):
|
|
"""Shutdown all resources"""
|
|
self.batch_processor.shutdown()
|
|
self.cleanup_resources()
|
|
|
|
|
|
# Global performance monitor instance
|
|
performance_monitor = PerformanceMonitor()
|
|
|
|
|
|
def optimize_config_manager():
|
|
"""Optimize configuration manager with caching"""
|
|
from config_manager import config_manager
|
|
|
|
# Cache configuration loading
|
|
original_load_config = config_manager._load_config
|
|
|
|
@functools.wraps(original_load_config)
|
|
def optimized_load_config():
|
|
# Check if we have a cached version
|
|
cache_key = "config_manager:load_config"
|
|
if cache_key in performance_monitor.cache:
|
|
cached_time, cached_result = performance_monitor.cache[cache_key]
|
|
if time.time() - cached_time < 60: # Cache for 1 minute
|
|
return cached_result
|
|
|
|
# Load configuration
|
|
result = original_load_config()
|
|
|
|
# Cache result
|
|
performance_monitor.cache[cache_key] = (time.time(), result)
|
|
|
|
return result
|
|
|
|
# Replace method
|
|
config_manager._load_config = optimized_load_config
|
|
|
|
# Cache configuration retrieval
|
|
original_get_configurations = config_manager.get_configurations
|
|
|
|
@functools.wraps(original_get_configurations)
|
|
def optimized_get_configurations():
|
|
cache_key = "config_manager:get_configurations"
|
|
if cache_key in performance_monitor.cache:
|
|
cached_time, cached_result = performance_monitor.cache[cache_key]
|
|
if time.time() - cached_time < 30: # Cache for 30 seconds
|
|
return cached_result
|
|
|
|
result = original_get_configurations()
|
|
performance_monitor.cache[cache_key] = (time.time(), result)
|
|
return result
|
|
|
|
config_manager.get_configurations = optimized_get_configurations
|
|
|
|
print("✅ Configuration manager optimized with caching")
|
|
|
|
|
|
def optimize_provider_tools():
|
|
"""Optimize provider tools with connection pooling"""
|
|
from tools.provider_tools import get_provider_tools
|
|
|
|
# Cache provider tools
|
|
original_get_provider_tools = get_provider_tools
|
|
|
|
@functools.wraps(original_get_provider_tools)
|
|
def optimized_get_provider_tools(provider: str, **kwargs):
|
|
cache_key = f"provider_tools:{provider}:{str(kwargs)}"
|
|
if cache_key in performance_monitor.cache:
|
|
cached_time, cached_result = performance_monitor.cache[cache_key]
|
|
if time.time() - cached_time < 300: # Cache for 5 minutes
|
|
return cached_result
|
|
|
|
result = original_get_provider_tools(provider, **kwargs)
|
|
performance_monitor.cache[cache_key] = (time.time(), result)
|
|
return result
|
|
|
|
# Replace function
|
|
import tools.provider_tools
|
|
tools.provider_tools.get_provider_tools = optimized_get_provider_tools
|
|
|
|
print("✅ Provider tools optimized with caching")
|
|
|
|
|
|
def optimize_agent_creation():
|
|
"""Optimize agent creation with caching"""
|
|
from main import create_maxclaw_agent
|
|
|
|
# Cache agent creation
|
|
original_create_maxclaw_agent = create_maxclaw_agent
|
|
|
|
@functools.wraps(original_create_maxclaw_agent)
|
|
def optimized_create_maxclaw_agent(system_prompt=None):
|
|
cache_key = f"agent:{system_prompt}"
|
|
if cache_key in performance_monitor.cache:
|
|
cached_time, cached_result = performance_monitor.cache[cache_key]
|
|
if time.time() - cached_time < 600: # Cache for 10 minutes
|
|
return cached_result
|
|
|
|
result = original_create_maxclaw_agent(system_prompt)
|
|
performance_monitor.cache[cache_key] = (time.time(), result)
|
|
return result
|
|
|
|
# Replace function
|
|
import main
|
|
main.create_maxclaw_agent = optimized_create_maxclaw_agent
|
|
|
|
print("✅ Agent creation optimized with caching")
|
|
|
|
|
|
def optimize_api_server():
|
|
"""Optimize API server with response caching"""
|
|
try:
|
|
from api_server import app, list_configurations, get_active_configuration
|
|
|
|
# Cache configuration endpoints
|
|
original_list_configurations = list_configurations
|
|
original_get_active_configuration = get_active_configuration
|
|
|
|
@functools.wraps(original_list_configurations)
|
|
async def optimized_list_configurations():
|
|
cache_key = "api:list_configurations"
|
|
if cache_key in performance_monitor.cache:
|
|
cached_time, cached_result = performance_monitor.cache[cache_key]
|
|
if time.time() - cached_time < 10: # Cache for 10 seconds
|
|
return cached_result
|
|
|
|
result = await original_list_configurations()
|
|
performance_monitor.cache[cache_key] = (time.time(), result)
|
|
return result
|
|
|
|
@functools.wraps(original_get_active_configuration)
|
|
async def optimized_get_active_configuration():
|
|
cache_key = "api:get_active_configuration"
|
|
if cache_key in performance_monitor.cache:
|
|
cached_time, cached_result = performance_monitor.cache[cache_key]
|
|
if time.time() - cached_time < 10: # Cache for 10 seconds
|
|
return cached_result
|
|
|
|
result = await original_get_active_configuration()
|
|
performance_monitor.cache[cache_key] = (time.time(), result)
|
|
return result
|
|
|
|
# Replace functions
|
|
import api_server
|
|
api_server.list_configurations = optimized_list_configurations
|
|
api_server.get_active_configuration = optimized_get_active_configuration
|
|
|
|
print("✅ API server optimized with response caching")
|
|
except ImportError:
|
|
print("⚠️ API server not available for optimization")
|
|
|
|
|
|
def apply_all_optimizations():
|
|
"""Apply all performance optimizations"""
|
|
print("🚀 Applying performance optimizations...")
|
|
|
|
try:
|
|
optimize_config_manager()
|
|
except Exception as e:
|
|
print(f"⚠️ Could not optimize config manager: {e}")
|
|
|
|
try:
|
|
optimize_provider_tools()
|
|
except Exception as e:
|
|
print(f"⚠️ Could not optimize provider tools: {e}")
|
|
|
|
try:
|
|
optimize_agent_creation()
|
|
except Exception as e:
|
|
print(f"⚠️ Could not optimize agent creation: {e}")
|
|
|
|
try:
|
|
optimize_api_server()
|
|
except Exception as e:
|
|
print(f"⚠️ Could not optimize API server: {e}")
|
|
|
|
print("✅ All optimizations applied")
|
|
|
|
|
|
def get_performance_report() -> str:
|
|
"""Get a performance report"""
|
|
metrics = performance_monitor.get_metrics()
|
|
|
|
if not metrics:
|
|
return "No performance metrics collected yet."
|
|
|
|
report = "Performance Report:\n"
|
|
report += "=" * 50 + "\n"
|
|
|
|
for func_name, stats in metrics.items():
|
|
report += f"\n{func_name}:\n"
|
|
report += f" Calls: {stats['count']}\n"
|
|
report += f" Avg Time: {stats['avg_time']:.4f}s\n"
|
|
report += f" Min Time: {stats['min_time']:.4f}s\n"
|
|
report += f" Max Time: {stats['max_time']:.4f}s\n"
|
|
report += f" Total Time: {stats['total_time']:.4f}s\n"
|
|
|
|
# Cache statistics
|
|
cache_size = len(performance_monitor.cache)
|
|
report += f"\nCache Statistics:\n"
|
|
report += f" Cached Items: {cache_size}\n"
|
|
|
|
# Connection pool statistics
|
|
pool_stats = performance_monitor.get_connection_pool_stats()
|
|
report += f"\nConnection Pool:\n"
|
|
report += f" Total Connections: {pool_stats['total_connections']}\n"
|
|
report += f" Max Connections: {pool_stats['max_connections']}\n"
|
|
|
|
return report
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# Apply optimizations when run directly
|
|
apply_all_optimizations()
|
|
|
|
# Print performance report
|
|
print("\n" + get_performance_report()) |