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