479 lines
16 KiB
Python
479 lines
16 KiB
Python
import os
|
|
import json
|
|
import uuid
|
|
import threading
|
|
from typing import Dict, Optional, List, Type, Any, Literal, cast
|
|
from datetime import datetime
|
|
import redis
|
|
from SimpleLLMFunc import OpenAICompatible
|
|
from context.schemas import Message
|
|
from context.context import ContextBackend, RedisFileContextBackend
|
|
from config.config import get_config
|
|
from SimpleLLMFunc.logger import push_warning, app_log
|
|
|
|
|
|
class ContextManager:
|
|
"""
|
|
General-purpose context manager that supports different backend implementations.
|
|
|
|
Main responsibilities:
|
|
1. Manage creation and lifecycle of ContextBackend 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[ContextBackend] = RedisFileContextBackend):
|
|
"""Singleton pattern implementation."""
|
|
if cls._instance is None:
|
|
with cls._lock:
|
|
if cls._instance is None:
|
|
cls._instance = super(ContextManager, cls).__new__(cls)
|
|
cls._instance.backend_class = backend_class
|
|
return cls._instance
|
|
|
|
def __init__(self, backend_class: Type[ContextBackend]):
|
|
"""
|
|
Initialize the context manager.
|
|
|
|
Args:
|
|
backend_class: Backend implementation class, defaulting to RedisFileBackend
|
|
"""
|
|
# Prevent duplicate initialization.
|
|
if hasattr(self, "_initialized"):
|
|
return
|
|
|
|
self.backend_class = backend_class
|
|
self.config = get_config()
|
|
self.context_dir = self.config.CONTEXT_DIR
|
|
self._active_contexts: Dict[str, ContextBackend] = {}
|
|
|
|
# Ensure the directory exists.
|
|
os.makedirs(self.context_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_context_ids_from_redis(self) -> set[str]:
|
|
context_ids: set[str] = set()
|
|
try:
|
|
client = self._redis_client()
|
|
raw_keys = cast(Any, client.keys("context:*:*"))
|
|
for key in cast(List[str], raw_keys):
|
|
parts = key.split(":", 2)
|
|
if len(parts) >= 3 and parts[0] == "context" and parts[1]:
|
|
context_ids.add(parts[1])
|
|
except Exception as e:
|
|
print(f"Warning: Failed to list context ids from Redis: {e}")
|
|
return context_ids
|
|
|
|
def _delete_context_redis_keys(self, context_id: str) -> bool:
|
|
try:
|
|
client = self._redis_client()
|
|
raw_keys = cast(Any, client.keys(f"context:{context_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 context keys for {context_id}: {e}")
|
|
return False
|
|
|
|
def create_context(
|
|
self,
|
|
context_id: Optional[str] = None,
|
|
llm_interface: Optional[
|
|
OpenAICompatible
|
|
] = get_config().CONTEXT_SUMMARY_INTERFACE,
|
|
max_history_length: int = get_config().CONTEXT_MAX_HISTORY_LENGTH,
|
|
auto_summarize_trigger: int = get_config().CONTEXT_AUTO_SUMMARIZE_TRIGGER,
|
|
**backend_kwargs,
|
|
) -> ContextBackend:
|
|
"""
|
|
Create a new context object.
|
|
|
|
Args:
|
|
context_id: Context ID; generated automatically if None
|
|
llm_interface: LLM interface
|
|
max_history_length: Maximum history length
|
|
auto_summarize_trigger: Automatic summary trigger threshold
|
|
**backend_kwargs: Extra parameters passed to the backend
|
|
|
|
Returns:
|
|
ContextBackend: Created context object
|
|
"""
|
|
with self._lock:
|
|
if context_id is None:
|
|
context_id = str(uuid.uuid4())
|
|
|
|
# Check whether it already exists.
|
|
if context_id in self._active_contexts:
|
|
app_log(
|
|
f"Context {context_id} already exists, and is in active contexts. Returning the existing context."
|
|
)
|
|
return self._active_contexts[context_id]
|
|
|
|
# Generate the file path if the backend needs one.
|
|
if "file_path" not in backend_kwargs:
|
|
context_file = os.path.join(self.context_dir, f"ctx_{context_id}.json")
|
|
backend_kwargs["file_path"] = context_file
|
|
push_warning(f"Context file path: {context_file}")
|
|
|
|
# Create the context object.
|
|
context = self.backend_class(
|
|
context_id=context_id,
|
|
llm_interface=llm_interface,
|
|
max_history_length=max_history_length,
|
|
auto_summarize_trigger=auto_summarize_trigger,
|
|
**backend_kwargs,
|
|
)
|
|
|
|
# Add it to the active context list.
|
|
self._active_contexts[context_id] = context
|
|
|
|
return context
|
|
|
|
def get_context(self, context_id: str) -> Optional[ContextBackend]:
|
|
"""
|
|
Get the context object with the specified ID.
|
|
|
|
Args:
|
|
context_id: Context ID
|
|
|
|
Returns:
|
|
ContextBackend: Context object, or None if it does not exist
|
|
"""
|
|
with self._lock:
|
|
# First check active contexts.
|
|
if context_id in self._active_contexts:
|
|
return self._active_contexts[context_id]
|
|
|
|
# Try to load from file if the backend supports it.
|
|
context_file = os.path.join(self.context_dir, f"ctx_{context_id}.json")
|
|
if os.path.exists(context_file):
|
|
try:
|
|
context = self.backend_class(
|
|
context_id=context_id,
|
|
llm_interface=self.config.CONTEXT_SUMMARY_INTERFACE, # Can be configured later.
|
|
max_history_length=self.config.CONTEXT_MAX_HISTORY_LENGTH,
|
|
auto_summarize_trigger=self.config.CONTEXT_AUTO_SUMMARIZE_TRIGGER,
|
|
file_path=context_file,
|
|
)
|
|
self._active_contexts[context_id] = context
|
|
return context
|
|
except Exception as e:
|
|
print(f"Warning: Failed to load context {context_id}: {e}")
|
|
|
|
return None
|
|
|
|
def delete_context(self, context_id: str) -> bool:
|
|
"""
|
|
Delete the context object with the specified ID.
|
|
|
|
Args:
|
|
context_id: Context ID
|
|
|
|
Returns:
|
|
bool: Whether deletion succeeded
|
|
"""
|
|
with self._lock:
|
|
success = False
|
|
|
|
# Remove it from active contexts.
|
|
if context_id in self._active_contexts:
|
|
del self._active_contexts[context_id]
|
|
success = True
|
|
|
|
# Delete context keys from Redis.
|
|
if self._delete_context_redis_keys(context_id):
|
|
success = True
|
|
|
|
# Delete the file if it exists.
|
|
context_file = os.path.join(self.context_dir, f"ctx_{context_id}.json")
|
|
if os.path.exists(context_file):
|
|
try:
|
|
os.remove(context_file)
|
|
success = True
|
|
except Exception as e:
|
|
print(f"Warning: Failed to delete context file {context_file}: {e}")
|
|
|
|
return success
|
|
|
|
def list_context_ids(self) -> List[str]:
|
|
"""List all known context IDs, including Redis and the file system."""
|
|
context_ids = set(self._active_contexts.keys())
|
|
context_ids.update(self._list_context_ids_from_redis())
|
|
|
|
try:
|
|
for filename in os.listdir(self.context_dir):
|
|
if filename.startswith("ctx_") and filename.endswith(".json"):
|
|
context_ids.add(filename[4:-5])
|
|
except Exception as e:
|
|
print(f"Warning: Failed to scan context dir for ids: {e}")
|
|
|
|
return sorted(context_ids)
|
|
|
|
def list_contexts(self) -> List[Dict[str, Any]]:
|
|
"""
|
|
List all available contexts.
|
|
|
|
Returns:
|
|
List[Dict]: Context information list
|
|
"""
|
|
contexts = []
|
|
|
|
# Scan context files in the file system.
|
|
try:
|
|
for filename in os.listdir(self.context_dir):
|
|
if filename.startswith("ctx_") and filename.endswith(".json"):
|
|
context_id = filename[4:-5] # Remove the "ctx_" prefix and ".json" suffix.
|
|
|
|
context_info = {
|
|
"context_id": context_id,
|
|
"file_path": os.path.join(self.context_dir, filename),
|
|
"is_active": context_id in self._active_contexts,
|
|
}
|
|
|
|
# Try to read basic information.
|
|
try:
|
|
file_path = context_info["file_path"]
|
|
if isinstance(file_path, str):
|
|
with open(file_path, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
metadata = data.get("metadata", {})
|
|
context_info.update(
|
|
{
|
|
"start_time": metadata.get("start_time"),
|
|
"last_activity": metadata.get("last_activity"),
|
|
"total_messages": metadata.get(
|
|
"total_messages", 0
|
|
),
|
|
}
|
|
)
|
|
except Exception:
|
|
pass # Ignore read errors.
|
|
|
|
contexts.append(context_info)
|
|
|
|
except Exception as e:
|
|
print(f"Warning: Failed to list contexts: {e}")
|
|
|
|
return contexts
|
|
|
|
async def save_context(self, context_id: str) -> bool:
|
|
"""
|
|
Manually save the specified context to file.
|
|
|
|
Args:
|
|
context_id: Context ID
|
|
|
|
Returns:
|
|
bool: Whether saving succeeded
|
|
"""
|
|
with self._lock:
|
|
context = self._active_contexts.get(context_id)
|
|
|
|
if context is None:
|
|
return False
|
|
|
|
try:
|
|
return await context.persist()
|
|
except Exception as e:
|
|
print(f"Warning: Failed to save context {context_id}: {e}")
|
|
return False
|
|
|
|
async def save_all_contexts(self) -> int:
|
|
"""
|
|
Save all active contexts to files.
|
|
|
|
Returns:
|
|
int: Number of contexts successfully saved
|
|
"""
|
|
saved_count = 0
|
|
with self._lock:
|
|
context_ids = list(self._active_contexts.keys())
|
|
|
|
for context_id in context_ids:
|
|
if await self.save_context(context_id):
|
|
saved_count += 1
|
|
|
|
return saved_count
|
|
|
|
async def cleanup_inactive_contexts(self, max_inactive_time: int = 3600) -> int:
|
|
"""
|
|
Clean up contexts that have been inactive for a long time.
|
|
|
|
Args:
|
|
max_inactive_time: Maximum inactive time in seconds
|
|
|
|
Returns:
|
|
int: Number of cleaned contexts
|
|
"""
|
|
cleaned_count = 0
|
|
current_time = datetime.now()
|
|
contexts_to_persist: list[tuple[str, ContextBackend]] = []
|
|
|
|
with self._lock:
|
|
contexts_to_remove = []
|
|
|
|
for context_id, context in self._active_contexts.items():
|
|
try:
|
|
metadata = context.get_metadata()
|
|
last_activity_str = metadata.get("last_activity")
|
|
if last_activity_str:
|
|
last_activity = datetime.fromisoformat(last_activity_str)
|
|
inactive_time = (current_time - last_activity).total_seconds()
|
|
|
|
if inactive_time > max_inactive_time:
|
|
# Persist the context before removing it.
|
|
contexts_to_persist.append((context_id, context))
|
|
contexts_to_remove.append(context_id)
|
|
cleaned_count += 1
|
|
except Exception as e:
|
|
print(
|
|
f"Warning: Error checking activity for context {context_id}: {e}"
|
|
)
|
|
|
|
for context_id, context in contexts_to_persist:
|
|
try:
|
|
await context.persist()
|
|
except Exception as e:
|
|
print(f"Warning: Failed to persist inactive context {context_id}: {e}")
|
|
|
|
with self._lock:
|
|
for context_id in contexts_to_remove:
|
|
del self._active_contexts[context_id]
|
|
|
|
return cleaned_count
|
|
|
|
# ===== Convenience interfaces =====
|
|
|
|
async def add_message(
|
|
self,
|
|
context_id: str,
|
|
message: Message,
|
|
) -> bool:
|
|
"""
|
|
Convenience method for adding a message.
|
|
|
|
Args:
|
|
context_id: Context ID
|
|
role: Message role
|
|
content: Message content
|
|
**message_kwargs: Other message parameters
|
|
|
|
Returns:
|
|
bool: Whether adding succeeded
|
|
"""
|
|
backend = self.get_context(context_id)
|
|
if not backend:
|
|
return False
|
|
|
|
try:
|
|
await backend.store_message(message)
|
|
return True
|
|
except Exception as e:
|
|
print(f"Warning: Failed to add message: {e}")
|
|
return False
|
|
|
|
def get_history(self, context_id: str, limit: Optional[int] = None) -> List:
|
|
"""
|
|
Get conversation history.
|
|
|
|
Args:
|
|
context_id: Context ID
|
|
limit: Limit on the number of returned messages
|
|
|
|
Returns:
|
|
List: Message history
|
|
"""
|
|
backend = self.get_context(context_id)
|
|
if not backend:
|
|
return []
|
|
|
|
return backend.retrieve_messages(limit)
|
|
|
|
async def summarize_context(self, context_id: str) -> Optional[str]:
|
|
"""
|
|
Summarize the context.
|
|
|
|
Args:
|
|
context_id: Context ID
|
|
|
|
Returns:
|
|
Optional[str]: Summary content
|
|
"""
|
|
backend = self.get_context(context_id)
|
|
if not backend:
|
|
return None
|
|
|
|
return await backend.auto_summarize()
|
|
|
|
|
|
# Global instance.
|
|
_global_context_manager: Optional[ContextManager] = None
|
|
|
|
|
|
def get_context_manager() -> ContextManager:
|
|
"""Get the global ContextManager instance."""
|
|
global _global_context_manager
|
|
if _global_context_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 ConfiguredRedisFileBackend(RedisFileContextBackend):
|
|
def __init__(
|
|
self,
|
|
context_id: str,
|
|
llm_interface: Optional[
|
|
OpenAICompatible
|
|
] = config.CONTEXT_SUMMARY_INTERFACE,
|
|
max_history_length: int = config.CONTEXT_MAX_HISTORY_LENGTH,
|
|
auto_summarize_trigger: int = config.CONTEXT_AUTO_SUMMARIZE_TRIGGER,
|
|
redis_host: str = redis_host,
|
|
redis_port: int = redis_port,
|
|
redis_db: int = redis_db,
|
|
file_path: str = "",
|
|
):
|
|
super().__init__(
|
|
context_id=context_id,
|
|
llm_interface=llm_interface,
|
|
max_history_length=max_history_length,
|
|
auto_summarize_trigger=auto_summarize_trigger,
|
|
redis_host=redis_host,
|
|
redis_port=redis_port,
|
|
redis_db=redis_db,
|
|
file_path=file_path,
|
|
)
|
|
self.file_path = file_path
|
|
self.context_id = context_id
|
|
self.llm_interface = llm_interface
|
|
self.max_history_length = max_history_length
|
|
self.auto_summarize_trigger = auto_summarize_trigger
|
|
self.redis_host = redis_host
|
|
self.redis_port = redis_port
|
|
self.redis_db = redis_db
|
|
|
|
# Create ContextManager using the configured backend class.
|
|
_global_context_manager = ContextManager(
|
|
backend_class=ConfiguredRedisFileBackend
|
|
)
|
|
return _global_context_manager
|