541 lines
17 KiB
Python
541 lines
17 KiB
Python
import os
|
|
import json
|
|
import uuid
|
|
import threading
|
|
from typing import Dict, Optional, List, Type, Any, cast
|
|
from datetime import datetime
|
|
import redis
|
|
from SimpleLLMFunc import OpenAICompatible
|
|
|
|
from context.sketch_pad import SketchPadBackend, RedisFileSketchPadBackend
|
|
from config.config import get_config
|
|
|
|
|
|
class SketchManager:
|
|
"""
|
|
General-purpose SketchPad manager that supports different backend implementations.
|
|
|
|
Main responsibilities:
|
|
1. Manage creation and lifecycle of SketchPadBackend instances
|
|
2. Provide advanced convenience interfaces
|
|
3. Handle batch operations and cleanup tasks
|
|
4. Support pluggable backend implementations
|
|
"""
|
|
|
|
_instance = None
|
|
_lock: threading.Lock = threading.Lock()
|
|
|
|
def __new__(cls, backend_class: Type[SketchPadBackend] = RedisFileSketchPadBackend):
|
|
"""Singleton pattern implementation."""
|
|
if cls._instance is None:
|
|
with cls._lock:
|
|
if cls._instance is None:
|
|
cls._instance = super(SketchManager, cls).__new__(cls)
|
|
cls._instance.backend_class = backend_class
|
|
return cls._instance
|
|
|
|
def __init__(self, backend_class: Type[SketchPadBackend]):
|
|
"""
|
|
Initialize the SketchPad manager.
|
|
|
|
Args:
|
|
backend_class: Backend implementation class, defaulting to RedisFileSketchPadBackend
|
|
"""
|
|
# Prevent duplicate initialization.
|
|
if hasattr(self, "_initialized"):
|
|
return
|
|
|
|
self.backend_class = backend_class
|
|
self.config = get_config()
|
|
self.sketch_dir = self.config.SKETCH_DIR
|
|
self._active_sketches: Dict[str, SketchPadBackend] = {}
|
|
|
|
# Ensure the directory exists.
|
|
os.makedirs(self.sketch_dir, exist_ok=True)
|
|
|
|
self._initialized = True
|
|
|
|
def _redis_client(self) -> redis.Redis:
|
|
return redis.Redis(
|
|
host=self.config.REDIS_HOST,
|
|
port=int(self.config.REDIS_PORT),
|
|
db=int(self.config.REDIS_DB),
|
|
decode_responses=True,
|
|
)
|
|
|
|
def _list_sketch_ids_from_redis(self) -> set[str]:
|
|
sketch_ids: set[str] = set()
|
|
try:
|
|
client = self._redis_client()
|
|
raw_keys = cast(Any, client.keys("sketch_pad:*:*"))
|
|
for key in cast(List[str], raw_keys):
|
|
parts = key.split(":", 2)
|
|
if len(parts) >= 3 and parts[0] == "sketch_pad" and parts[1]:
|
|
sketch_ids.add(parts[1])
|
|
except Exception as e:
|
|
print(f"Warning: Failed to list sketch ids from Redis: {e}")
|
|
return sketch_ids
|
|
|
|
def _delete_sketch_redis_keys(self, sketch_id: str) -> bool:
|
|
try:
|
|
client = self._redis_client()
|
|
raw_keys = cast(Any, client.keys(f"sketch_pad:{sketch_id}:*"))
|
|
keys = cast(List[str], raw_keys)
|
|
if not keys:
|
|
return False
|
|
deleted = cast(Any, client.delete(*keys))
|
|
return int(deleted) > 0
|
|
except Exception as e:
|
|
print(f"Warning: Failed to delete Redis sketch keys for {sketch_id}: {e}")
|
|
return False
|
|
|
|
def create_sketch_pad(
|
|
self,
|
|
sketch_id: Optional[str] = None,
|
|
**backend_kwargs,
|
|
) -> SketchPadBackend:
|
|
"""
|
|
Create a new SketchPad object.
|
|
|
|
Args:
|
|
sketch_id: SketchPad ID; generated automatically if None
|
|
**backend_kwargs: Extra parameters passed to the backend
|
|
|
|
Returns:
|
|
SketchPadBackend: Created SketchPad object
|
|
"""
|
|
with self._lock:
|
|
if sketch_id is None:
|
|
sketch_id = str(uuid.uuid4())
|
|
|
|
# Check whether it already exists.
|
|
if sketch_id in self._active_sketches:
|
|
return self._active_sketches[sketch_id]
|
|
|
|
# Generate the file path if the backend needs one.
|
|
if "file_path" not in backend_kwargs:
|
|
sketch_file = os.path.join(self.sketch_dir, f"skt_{sketch_id}.json")
|
|
backend_kwargs["file_path"] = sketch_file
|
|
|
|
# Create the SketchPad object.
|
|
sketch_pad = self.backend_class(
|
|
sketch_pad_id=sketch_id,
|
|
**backend_kwargs,
|
|
)
|
|
|
|
# Add it to the active SketchPad list.
|
|
self._active_sketches[sketch_id] = sketch_pad
|
|
|
|
return sketch_pad
|
|
|
|
def get_sketch_pad(self, sketch_id: str) -> Optional[SketchPadBackend]:
|
|
"""
|
|
Get the SketchPad object with the specified ID.
|
|
|
|
Args:
|
|
sketch_id: SketchPad ID
|
|
|
|
Returns:
|
|
SketchPadBackend: SketchPad object, or None if it does not exist
|
|
"""
|
|
with self._lock:
|
|
# First check active SketchPads.
|
|
if sketch_id in self._active_sketches:
|
|
return self._active_sketches[sketch_id]
|
|
|
|
# Try to load from file if the backend supports it.
|
|
sketch_file = os.path.join(self.sketch_dir, f"skt_{sketch_id}.json")
|
|
if os.path.exists(sketch_file):
|
|
try:
|
|
sketch_pad = self.backend_class(
|
|
sketch_pad_id=sketch_id,
|
|
file_path=sketch_file,
|
|
)
|
|
self._active_sketches[sketch_id] = sketch_pad
|
|
return sketch_pad
|
|
except Exception as e:
|
|
print(f"Warning: Failed to load sketch {sketch_id}: {e}")
|
|
|
|
return None
|
|
|
|
def delete_sketch_pad(self, sketch_id: str) -> bool:
|
|
"""
|
|
Delete the SketchPad object with the specified ID.
|
|
|
|
Args:
|
|
sketch_id: SketchPad ID
|
|
|
|
Returns:
|
|
bool: Whether deletion succeeded
|
|
"""
|
|
with self._lock:
|
|
success = False
|
|
|
|
# Remove it from active SketchPads.
|
|
if sketch_id in self._active_sketches:
|
|
del self._active_sketches[sketch_id]
|
|
success = True
|
|
|
|
# Delete sketch keys from Redis.
|
|
if self._delete_sketch_redis_keys(sketch_id):
|
|
success = True
|
|
|
|
# Delete the file if it exists.
|
|
sketch_file = os.path.join(self.sketch_dir, f"skt_{sketch_id}.json")
|
|
if os.path.exists(sketch_file):
|
|
try:
|
|
os.remove(sketch_file)
|
|
success = True
|
|
except Exception as e:
|
|
print(f"Warning: Failed to delete sketch file {sketch_file}: {e}")
|
|
|
|
return success
|
|
|
|
def list_sketch_ids(self) -> List[str]:
|
|
"""List all known SketchPad IDs, including Redis and the file system."""
|
|
sketch_ids = set(self._active_sketches.keys())
|
|
sketch_ids.update(self._list_sketch_ids_from_redis())
|
|
|
|
try:
|
|
for filename in os.listdir(self.sketch_dir):
|
|
if filename.startswith("skt_") and filename.endswith(".json"):
|
|
sketch_ids.add(filename[4:-5])
|
|
except Exception as e:
|
|
print(f"Warning: Failed to scan sketch dir for ids: {e}")
|
|
|
|
return sorted(sketch_ids)
|
|
|
|
def list_sketch_pads(self) -> List[Dict[str, Any]]:
|
|
"""
|
|
List all available SketchPads.
|
|
|
|
Returns:
|
|
List[Dict]: SketchPad information list
|
|
"""
|
|
sketches = []
|
|
|
|
# Scan SketchPad files in the file system.
|
|
try:
|
|
for filename in os.listdir(self.sketch_dir):
|
|
if filename.startswith("skt_") and filename.endswith(".json"):
|
|
sketch_id = filename[4:-5] # Remove the "skt_" prefix and ".json" suffix.
|
|
|
|
sketch_info = {
|
|
"sketch_id": sketch_id,
|
|
"file_path": os.path.join(self.sketch_dir, filename),
|
|
"is_active": sketch_id in self._active_sketches,
|
|
}
|
|
|
|
# Try to read basic information.
|
|
try:
|
|
file_path = sketch_info["file_path"]
|
|
if isinstance(file_path, str):
|
|
with open(file_path, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
sketch_info.update(
|
|
{
|
|
"total_items": len(data.get("items", {})),
|
|
"last_saved": data.get(
|
|
"serialization_timestamp"
|
|
),
|
|
"sketch_pad_id": data.get("sketch_pad_id"),
|
|
}
|
|
)
|
|
except Exception:
|
|
pass # Ignore read errors.
|
|
|
|
sketches.append(sketch_info)
|
|
|
|
except Exception as e:
|
|
print(f"Warning: Failed to list sketches: {e}")
|
|
|
|
return sketches
|
|
|
|
def save_sketch_pad(self, sketch_id: str) -> bool:
|
|
"""
|
|
Manually save the specified SketchPad to file.
|
|
|
|
Args:
|
|
sketch_id: SketchPad ID
|
|
|
|
Returns:
|
|
bool: Whether saving succeeded
|
|
"""
|
|
with self._lock:
|
|
if sketch_id in self._active_sketches:
|
|
try:
|
|
sketch_pad = self._active_sketches[sketch_id]
|
|
sketch_pad.persist()
|
|
return True
|
|
except Exception as e:
|
|
print(f"Warning: Failed to save sketch {sketch_id}: {e}")
|
|
|
|
return False
|
|
|
|
async def save_all_sketch_pads(self) -> int:
|
|
"""
|
|
Save all active SketchPads to files.
|
|
|
|
Returns:
|
|
int: Number of SketchPads successfully saved
|
|
"""
|
|
saved_count = 0
|
|
with self._lock:
|
|
for sketch_id in list(self._active_sketches.keys()):
|
|
if self.save_sketch_pad(sketch_id):
|
|
saved_count += 1
|
|
|
|
return saved_count
|
|
|
|
async def cleanup_inactive_sketches(self, max_inactive_count: int = 10) -> int:
|
|
"""
|
|
Clean up inactive SketchPads based on usage frequency.
|
|
|
|
Args:
|
|
max_inactive_count: Maximum number of SketchPads to keep active
|
|
|
|
Returns:
|
|
int: Number of cleaned SketchPads
|
|
"""
|
|
cleaned_count = 0
|
|
|
|
with self._lock:
|
|
if len(self._active_sketches) <= max_inactive_count:
|
|
return 0
|
|
|
|
# Sort by access statistics and keep the most-used items.
|
|
sketches_by_usage = []
|
|
for sketch_id, sketch_pad in self._active_sketches.items():
|
|
try:
|
|
stats = sketch_pad.get_statistics()
|
|
total_accesses = stats.total_accesses
|
|
sketches_by_usage.append((sketch_id, sketch_pad, total_accesses))
|
|
except Exception:
|
|
sketches_by_usage.append((sketch_id, sketch_pad, 0))
|
|
|
|
# Sort by access count.
|
|
sketches_by_usage.sort(key=lambda x: x[2], reverse=True)
|
|
|
|
# Save and remove low-usage SketchPads.
|
|
sketches_to_remove = sketches_by_usage[max_inactive_count:]
|
|
for sketch_id, sketch_pad, _ in sketches_to_remove:
|
|
try:
|
|
# Save to file.
|
|
sketch_pad.persist()
|
|
|
|
# Remove from the active list.
|
|
del self._active_sketches[sketch_id]
|
|
cleaned_count += 1
|
|
except Exception as e:
|
|
print(f"Warning: Error cleaning sketch {sketch_id}: {e}")
|
|
|
|
return cleaned_count
|
|
|
|
# ===== Convenience interfaces =====
|
|
|
|
async def set_item(
|
|
self,
|
|
sketch_id: str,
|
|
key: str,
|
|
value: Any,
|
|
ttl: Optional[int] = None,
|
|
summary: Optional[str] = None,
|
|
tags: Optional[set] = None,
|
|
) -> Optional[str]:
|
|
"""
|
|
Convenience method for setting an item.
|
|
|
|
Args:
|
|
sketch_id: SketchPad ID
|
|
key: Key name
|
|
value: Value
|
|
ttl: Expiration time in seconds
|
|
summary: Summary
|
|
tags: Tags
|
|
|
|
Returns:
|
|
Optional[str]: Set key name, or None on failure
|
|
"""
|
|
sketch_pad = self.get_sketch_pad(sketch_id)
|
|
if not sketch_pad:
|
|
return None
|
|
|
|
try:
|
|
return await sketch_pad.set_item(key, value, ttl, summary, tags)
|
|
except Exception as e:
|
|
print(f"Warning: Failed to set item: {e}")
|
|
return None
|
|
|
|
def get_item(self, sketch_id: str, key: str) -> Optional[Any]:
|
|
"""
|
|
Get an item.
|
|
|
|
Args:
|
|
sketch_id: SketchPad ID
|
|
key: Key name
|
|
|
|
Returns:
|
|
Optional[Any]: Item value
|
|
"""
|
|
sketch_pad = self.get_sketch_pad(sketch_id)
|
|
if not sketch_pad:
|
|
return None
|
|
|
|
return sketch_pad.get_item(key)
|
|
|
|
def get_value(self, sketch_id: str, key: str) -> Optional[Any]:
|
|
"""
|
|
Get a value.
|
|
|
|
Args:
|
|
sketch_id: SketchPad ID
|
|
key: Key name
|
|
|
|
Returns:
|
|
Optional[Any]: Value
|
|
"""
|
|
sketch_pad = self.get_sketch_pad(sketch_id)
|
|
if not sketch_pad:
|
|
return None
|
|
|
|
return sketch_pad.get_value(key)
|
|
|
|
def search_by_tags(
|
|
self, sketch_id: str, tags: set, match_all: bool = False
|
|
) -> List[tuple]:
|
|
"""
|
|
Search by tags.
|
|
|
|
Args:
|
|
sketch_id: SketchPad ID
|
|
tags: Tag set
|
|
match_all: Whether to match all tags
|
|
|
|
Returns:
|
|
List[tuple]: Search results
|
|
"""
|
|
sketch_pad = self.get_sketch_pad(sketch_id)
|
|
if not sketch_pad:
|
|
return []
|
|
|
|
return sketch_pad.search_by_tags(tags, match_all)
|
|
|
|
def search_by_content(
|
|
self, sketch_id: str, query: str, limit: int = 5
|
|
) -> List[tuple]:
|
|
"""
|
|
Search by content.
|
|
|
|
Args:
|
|
sketch_id: SketchPad ID
|
|
query: Search query
|
|
limit: Result count limit
|
|
|
|
Returns:
|
|
List[tuple]: Search results
|
|
"""
|
|
sketch_pad = self.get_sketch_pad(sketch_id)
|
|
if not sketch_pad:
|
|
return []
|
|
|
|
return sketch_pad.search_by_content(query, limit)
|
|
|
|
def delete_item(self, sketch_id: str, key: str) -> bool:
|
|
"""
|
|
Delete an item.
|
|
|
|
Args:
|
|
sketch_id: SketchPad ID
|
|
key: Key name
|
|
|
|
Returns:
|
|
bool: Whether deletion succeeded
|
|
"""
|
|
sketch_pad = self.get_sketch_pad(sketch_id)
|
|
if not sketch_pad:
|
|
return False
|
|
|
|
return sketch_pad.delete(key)
|
|
|
|
def get_statistics(self, sketch_id: str) -> Optional[Any]:
|
|
"""
|
|
Get statistics.
|
|
|
|
Args:
|
|
sketch_id: SketchPad ID
|
|
|
|
Returns:
|
|
Optional[Any]: Statistics
|
|
"""
|
|
sketch_pad = self.get_sketch_pad(sketch_id)
|
|
if not sketch_pad:
|
|
return None
|
|
|
|
try:
|
|
return sketch_pad.get_statistics()
|
|
except Exception as e:
|
|
print(f"Warning: Failed to get statistics: {e}")
|
|
return None
|
|
|
|
def list_items(self, sketch_id: str, include_value: bool = False) -> List[Any]:
|
|
"""
|
|
List all items.
|
|
|
|
Args:
|
|
sketch_id: SketchPad ID
|
|
include_value: Whether to include values
|
|
|
|
Returns:
|
|
List[Any]: Item list
|
|
"""
|
|
sketch_pad = self.get_sketch_pad(sketch_id)
|
|
if not sketch_pad:
|
|
return []
|
|
|
|
try:
|
|
return sketch_pad.list_items(include_value)
|
|
except Exception as e:
|
|
print(f"Warning: Failed to list items: {e}")
|
|
return []
|
|
|
|
|
|
# Global instance.
|
|
_global_sketch_manager: Optional[SketchManager] = None
|
|
|
|
|
|
def get_sketch_manager() -> SketchManager:
|
|
"""Get the global SketchManager instance."""
|
|
global _global_sketch_manager
|
|
if _global_sketch_manager is None:
|
|
# Get Redis configuration from config.
|
|
config = get_config()
|
|
|
|
# redis config
|
|
redis_host = config.REDIS_HOST
|
|
redis_port = int(config.REDIS_PORT)
|
|
redis_db = int(config.REDIS_DB)
|
|
|
|
# Create a custom backend class with preconfigured Redis parameters.
|
|
class ConfiguredRedisFileSketchPadBackend(RedisFileSketchPadBackend):
|
|
def __init__(
|
|
self,
|
|
sketch_pad_id: str,
|
|
redis_host: str = redis_host,
|
|
redis_port: int = redis_port,
|
|
redis_db: int = redis_db,
|
|
file_path: Optional[str] = None,
|
|
):
|
|
super().__init__(
|
|
sketch_pad_id=sketch_pad_id,
|
|
redis_host=redis_host,
|
|
redis_port=redis_port,
|
|
redis_db=redis_db,
|
|
file_path=file_path,
|
|
)
|
|
|
|
# Create SketchManager using the configured backend class.
|
|
_global_sketch_manager = SketchManager(
|
|
backend_class=ConfiguredRedisFileSketchPadBackend
|
|
)
|
|
return _global_sketch_manager
|