Jerry
457 lines
16 KiB
Python
457 lines
16 KiB
Python
from __future__ import annotations
|
|
import uuid
|
|
import threading
|
|
import os
|
|
from typing import Dict, Optional, List, Any
|
|
from datetime import datetime
|
|
from dataclasses import dataclass
|
|
from SimpleLLMFunc import OpenAICompatible
|
|
from SimpleLLMFunc.logger import push_warning, push_error, app_log
|
|
from context.context_manager import get_context_manager, ContextManager
|
|
from context.sketch_manager import get_sketch_manager, SketchManager
|
|
from context.sketch_pad import SketchPadBackend
|
|
from context.context import ContextBackend
|
|
from config.config import get_config
|
|
|
|
|
|
# Global current conversation context variable.
|
|
_current_conversation: Optional[Conversation] = None
|
|
_conversation_context_lock = threading.RLock()
|
|
|
|
|
|
def get_current_conversation() -> Optional[Conversation]:
|
|
"""Get the Conversation in the current context."""
|
|
global _current_conversation
|
|
with _conversation_context_lock:
|
|
return _current_conversation
|
|
|
|
|
|
def get_current_context() -> Optional[ContextBackend]:
|
|
"""Get the Context in the current context."""
|
|
conversation = get_current_conversation()
|
|
return conversation.context if conversation else None
|
|
|
|
|
|
def get_current_sketch_pad() -> Optional[SketchPadBackend]:
|
|
"""Get the SketchPad in the current context."""
|
|
conversation = get_current_conversation()
|
|
return conversation.sketch_pad if conversation else None
|
|
|
|
|
|
@dataclass
|
|
class Conversation:
|
|
"""
|
|
Conversation data class representing a complete conversation session.
|
|
Contains a unique UUID, associated Context, and SketchPad.
|
|
Supports use as a context manager.
|
|
"""
|
|
|
|
uuid: str
|
|
context: ContextBackend
|
|
sketch_pad: SketchPadBackend
|
|
created_at: datetime
|
|
last_accessed: datetime
|
|
|
|
def update_access_time(self):
|
|
"""Update the last access time."""
|
|
self.last_accessed = datetime.now()
|
|
|
|
def __enter__(self):
|
|
"""Enter the context manager."""
|
|
global _current_conversation
|
|
with _conversation_context_lock:
|
|
if _current_conversation is not None:
|
|
raise RuntimeError("Cannot nest conversation contexts")
|
|
_current_conversation = self
|
|
self.update_access_time()
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
"""Exit the context manager."""
|
|
global _current_conversation
|
|
with _conversation_context_lock:
|
|
_current_conversation = None
|
|
return False
|
|
|
|
|
|
class ConversationManager:
|
|
"""
|
|
Conversation manager responsible for creating, managing, and coordinating Conversation lifecycles.
|
|
Each Conversation contains one Context and one SketchPad, and they share the same UUID.
|
|
|
|
ConversationManager is a global singleton that uses ContextManager and SketchManager
|
|
to manage the underlying Context and SketchPad objects.
|
|
"""
|
|
|
|
_instance = None
|
|
_lock = threading.Lock()
|
|
|
|
def __new__(cls):
|
|
"""Singleton pattern implementation."""
|
|
if cls._instance is None:
|
|
with cls._lock:
|
|
if cls._instance is None:
|
|
cls._instance = super(ConversationManager, cls).__new__(cls)
|
|
return cls._instance
|
|
|
|
def __init__(self):
|
|
"""Initialize ConversationManager."""
|
|
# Prevent duplicate initialization.
|
|
if hasattr(self, "_initialized"):
|
|
return
|
|
|
|
self.config = get_config()
|
|
self.context_manager: ContextManager = get_context_manager()
|
|
self.sketch_manager: SketchManager = get_sketch_manager()
|
|
self._active_conversations: Dict[str, Conversation] = {}
|
|
self._lock = threading.RLock()
|
|
|
|
# Create the conversations directory.
|
|
self.conversations_dir = os.path.join(
|
|
os.path.dirname(self.config.CONTEXT_DIR), "conversations"
|
|
)
|
|
os.makedirs(self.conversations_dir, exist_ok=True)
|
|
|
|
self._initialized = True
|
|
|
|
def create_conversation(
|
|
self,
|
|
conversation_id: Optional[str] = None,
|
|
llm_interface: Optional[OpenAICompatible] = None,
|
|
max_history_length: int = 5,
|
|
) -> Conversation:
|
|
"""
|
|
Create a new Conversation.
|
|
|
|
Args:
|
|
conversation_id: Conversation UUID; generated automatically if None
|
|
llm_interface: LLM interface used for Context
|
|
max_history_length: Maximum Context history length
|
|
|
|
Returns:
|
|
Conversation: Created Conversation object
|
|
"""
|
|
with self._lock:
|
|
if conversation_id is None:
|
|
conversation_id = str(uuid.uuid4())
|
|
|
|
# Check whether it already exists.
|
|
if conversation_id in self._active_conversations:
|
|
conversation = self._active_conversations[conversation_id]
|
|
conversation.update_access_time()
|
|
return conversation
|
|
|
|
# Create Context with the ctx prefix.
|
|
context = self.context_manager.create_context(
|
|
context_id=conversation_id,
|
|
llm_interface=llm_interface,
|
|
max_history_length=max_history_length,
|
|
)
|
|
|
|
# Create SketchPad with the skt prefix.
|
|
sketch_pad = self.sketch_manager.create_sketch_pad(
|
|
sketch_id=conversation_id
|
|
)
|
|
|
|
# Create the Conversation object.
|
|
now = datetime.now()
|
|
conversation = Conversation(
|
|
uuid=conversation_id,
|
|
context=context,
|
|
sketch_pad=sketch_pad,
|
|
created_at=now,
|
|
last_accessed=now,
|
|
)
|
|
|
|
# Add it to the active Conversation list.
|
|
self._active_conversations[conversation_id] = conversation
|
|
|
|
# Immediately persist context and sketch_pad to the file system.
|
|
try:
|
|
import asyncio
|
|
|
|
# Create an event loop to run the asynchronous task.
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
try:
|
|
loop.run_until_complete(context.persist())
|
|
# Synchronously call sketch_pad.persist().
|
|
sketch_pad.persist()
|
|
finally:
|
|
loop.close()
|
|
app_log(f"Conversation {conversation_id} was successfully persisted to the file system")
|
|
except Exception as e:
|
|
push_warning(f"Failed to persist conversation {conversation_id}: {e}")
|
|
|
|
# Create a persistence marker file.
|
|
self._create_conversation_marker(conversation_id)
|
|
|
|
return conversation
|
|
|
|
def get_conversation(self, conversation_id: str) -> Optional[Conversation]:
|
|
"""
|
|
Get the Conversation with the specified ID.
|
|
|
|
Args:
|
|
conversation_id: Conversation UUID
|
|
|
|
Returns:
|
|
Conversation: Conversation object, or None if it does not exist
|
|
"""
|
|
with self._lock:
|
|
# First check active Conversations.
|
|
if conversation_id in self._active_conversations:
|
|
conversation = self._active_conversations[conversation_id]
|
|
conversation.update_access_time()
|
|
return conversation
|
|
|
|
# Try to rebuild from the file system.
|
|
context = self.context_manager.get_context(conversation_id)
|
|
sketch_pad = self.sketch_manager.get_sketch_pad(conversation_id)
|
|
|
|
if context is not None and sketch_pad is not None:
|
|
# Rebuild the Conversation object.
|
|
now = datetime.now()
|
|
conversation = Conversation(
|
|
uuid=conversation_id,
|
|
context=context,
|
|
sketch_pad=sketch_pad,
|
|
created_at=now, # Use the current time as the rebuild time.
|
|
last_accessed=now,
|
|
)
|
|
|
|
self._active_conversations[conversation_id] = conversation
|
|
return conversation
|
|
|
|
return None
|
|
|
|
def delete_conversation(self, conversation_id: str) -> bool:
|
|
"""
|
|
Delete the Conversation with the specified ID.
|
|
|
|
Args:
|
|
conversation_id: Conversation UUID
|
|
|
|
Returns:
|
|
bool: Whether deletion succeeded
|
|
"""
|
|
with self._lock:
|
|
success = False
|
|
|
|
# Remove it from active Conversations.
|
|
if conversation_id in self._active_conversations:
|
|
del self._active_conversations[conversation_id]
|
|
success = True
|
|
|
|
# Delete the underlying Context and SketchPad.
|
|
context_deleted = self.context_manager.delete_context(conversation_id)
|
|
sketch_deleted = self.sketch_manager.delete_sketch_pad(conversation_id)
|
|
|
|
# Delete the marker file.
|
|
marker_file = os.path.join(
|
|
self.conversations_dir, f"conv_{conversation_id}.marker"
|
|
)
|
|
if os.path.exists(marker_file):
|
|
try:
|
|
os.remove(marker_file)
|
|
except Exception as e:
|
|
push_warning(f"Failed to delete marker file {marker_file}: {e}")
|
|
|
|
return success or context_deleted or sketch_deleted
|
|
|
|
def _discover_conversation_ids(self) -> List[str]:
|
|
"""Discover all known conversation ids across memory, files, and Redis-backed stores."""
|
|
conversation_ids = set(self._active_conversations.keys())
|
|
|
|
try:
|
|
for filename in os.listdir(self.conversations_dir):
|
|
if filename.startswith("conv_") and filename.endswith(".marker"):
|
|
conversation_ids.add(filename[5:-7])
|
|
except Exception as e:
|
|
push_warning(f"Failed to scan conversation markers: {e}")
|
|
|
|
try:
|
|
conversation_ids.update(self.context_manager.list_context_ids())
|
|
except Exception as e:
|
|
push_warning(f"Failed to collect context ids: {e}")
|
|
|
|
try:
|
|
conversation_ids.update(self.sketch_manager.list_sketch_ids())
|
|
except Exception as e:
|
|
push_warning(f"Failed to collect sketch ids: {e}")
|
|
|
|
return sorted(conversation_ids)
|
|
|
|
def delete_all_conversations(self) -> List[str]:
|
|
"""Delete all known conversations from memory, files, and Redis-backed stores."""
|
|
deleted_ids: List[str] = []
|
|
for conversation_id in self._discover_conversation_ids():
|
|
if self.delete_conversation(conversation_id):
|
|
deleted_ids.append(conversation_id)
|
|
return deleted_ids
|
|
|
|
def list_conversations(self) -> List[Dict[str, Any]]:
|
|
"""
|
|
List all available Conversations.
|
|
|
|
Returns:
|
|
List[Dict]: Conversation information list
|
|
"""
|
|
conversations = []
|
|
|
|
try:
|
|
for conversation_id in self._discover_conversation_ids():
|
|
marker_file = os.path.join(
|
|
self.conversations_dir, f"conv_{conversation_id}.marker"
|
|
)
|
|
conversation_info = {
|
|
"conversation_id": conversation_id,
|
|
"marker_file": marker_file if os.path.exists(marker_file) else None,
|
|
"is_active": conversation_id in self._active_conversations,
|
|
}
|
|
|
|
context = self.context_manager.get_context(conversation_id)
|
|
sketch_pad = self.sketch_manager.get_sketch_pad(conversation_id)
|
|
|
|
if context:
|
|
metadata = context.get_metadata()
|
|
conversation_info.update(
|
|
{
|
|
"context_start_time": metadata.get("start_time"),
|
|
"context_last_activity": metadata.get("last_activity"),
|
|
"context_total_messages": context.get_total_message_count(),
|
|
"context_has_summary": bool(context.get_summary()),
|
|
}
|
|
)
|
|
|
|
if sketch_pad:
|
|
stats = sketch_pad.get_statistics()
|
|
conversation_info.update(
|
|
{
|
|
"sketch_total_items": stats.total_items,
|
|
"sketch_max_items": stats.max_items,
|
|
"sketch_memory_usage": stats.memory_usage_percent,
|
|
}
|
|
)
|
|
|
|
conversations.append(conversation_info)
|
|
except Exception as e:
|
|
push_warning(f"Failed to list conversations: {e}")
|
|
|
|
return conversations
|
|
|
|
async def save_conversation(self, conversation_id: str) -> bool:
|
|
"""
|
|
Manually save the specified Conversation to files.
|
|
|
|
Args:
|
|
conversation_id: Conversation UUID
|
|
|
|
Returns:
|
|
bool: Whether saving succeeded
|
|
"""
|
|
with self._lock:
|
|
if conversation_id in self._active_conversations:
|
|
try:
|
|
conversation = self._active_conversations[conversation_id]
|
|
|
|
# Save Context.
|
|
context_saved = await conversation.context.persist()
|
|
|
|
# Save SketchPad.
|
|
conversation.sketch_pad.persist()
|
|
sketch_saved = True
|
|
|
|
return bool(context_saved and sketch_saved)
|
|
except Exception as e:
|
|
push_warning(f"Failed to save conversation {conversation_id}: {e}")
|
|
|
|
return False
|
|
|
|
async def save_all_conversations(self) -> int:
|
|
"""
|
|
Save all active Conversations to files.
|
|
|
|
Returns:
|
|
int: Number of Conversations successfully saved
|
|
"""
|
|
saved_count = 0
|
|
with self._lock:
|
|
for conversation_id in list(self._active_conversations.keys()):
|
|
if await self.save_conversation(conversation_id):
|
|
saved_count += 1
|
|
|
|
return saved_count
|
|
|
|
async def cleanup_inactive_conversations(
|
|
self, max_inactive_time: int = 3600
|
|
) -> int:
|
|
"""
|
|
Clean up Conversations that have been inactive for a long time.
|
|
|
|
Args:
|
|
max_inactive_time: Maximum inactive time in seconds
|
|
|
|
Returns:
|
|
int: Number of cleaned Conversations
|
|
"""
|
|
cleaned_count = 0
|
|
current_time = datetime.now()
|
|
|
|
with self._lock:
|
|
conversations_to_remove = []
|
|
|
|
for conversation_id, conversation in self._active_conversations.items():
|
|
try:
|
|
inactive_time = (
|
|
current_time - conversation.last_accessed
|
|
).total_seconds()
|
|
|
|
if inactive_time > max_inactive_time:
|
|
# Save the Conversation before removing it.
|
|
await self.save_conversation(conversation_id)
|
|
conversations_to_remove.append(conversation_id)
|
|
cleaned_count += 1
|
|
except Exception as e:
|
|
push_warning(
|
|
f"Error checking activity for conversation {conversation_id}: {e}"
|
|
)
|
|
|
|
# Remove inactive Conversations.
|
|
for conversation_id in conversations_to_remove:
|
|
del self._active_conversations[conversation_id]
|
|
|
|
return cleaned_count
|
|
|
|
def _create_conversation_marker(self, conversation_id: str) -> None:
|
|
"""
|
|
Create a Conversation marker file.
|
|
|
|
Args:
|
|
conversation_id: Conversation UUID
|
|
"""
|
|
try:
|
|
marker_file = os.path.join(
|
|
self.conversations_dir, f"conv_{conversation_id}.marker"
|
|
)
|
|
with open(marker_file, "w") as f:
|
|
f.write(
|
|
f"Conversation {conversation_id} created at {datetime.now().isoformat()}"
|
|
)
|
|
except Exception as e:
|
|
push_warning(
|
|
f"Failed to create marker file for conversation {conversation_id}: {e}"
|
|
)
|
|
|
|
|
|
# Global instance.
|
|
_global_conversation_manager: Optional[ConversationManager] = None
|
|
|
|
|
|
def get_conversation_manager() -> ConversationManager:
|
|
"""Get the global ConversationManager instance."""
|
|
global _global_conversation_manager
|
|
if _global_conversation_manager is None:
|
|
_global_conversation_manager = ConversationManager()
|
|
return _global_conversation_manager
|