first commit
This commit is contained in:
@@ -0,0 +1,540 @@
|
||||
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
|
||||
Reference in New Issue
Block a user