246 lines
8.9 KiB
Python
246 lines
8.9 KiB
Python
"""
|
|
Prompt tracking system for linking generated images to prompts.
|
|
Tracks active prompt executions to associate them with generated images.
|
|
"""
|
|
|
|
import threading
|
|
import time
|
|
import uuid
|
|
import hashlib
|
|
from typing import Optional, Dict, Any
|
|
from datetime import datetime, timezone
|
|
|
|
|
|
class PromptTracker:
|
|
"""Thread-safe tracking of current prompt executions."""
|
|
|
|
def __init__(self, db_manager):
|
|
"""
|
|
Initialize the prompt tracker.
|
|
|
|
Args:
|
|
db_manager: Database manager instance
|
|
"""
|
|
self.db_manager = db_manager
|
|
self._local = threading.local()
|
|
self.active_prompts = {} # Global tracking for multiple threads
|
|
self.lock = threading.Lock()
|
|
self.cleanup_interval = 300 # 5 minutes
|
|
self.prompt_timeout = 120 # 2 minutes
|
|
|
|
# Start cleanup thread
|
|
self.cleanup_thread = threading.Thread(target=self._cleanup_expired_prompts, daemon=True)
|
|
self.cleanup_thread.start()
|
|
|
|
def set_current_prompt(self, prompt_text: str, additional_data: Optional[Dict[str, Any]] = None) -> str:
|
|
"""
|
|
Set the current prompt for this thread and store in database.
|
|
|
|
Args:
|
|
prompt_text: The prompt text being executed
|
|
additional_data: Additional prompt metadata
|
|
|
|
Returns:
|
|
Execution ID for this prompt
|
|
"""
|
|
execution_id = self.generate_execution_id()
|
|
|
|
# Save prompt to database first
|
|
prompt_hash = hashlib.sha256(prompt_text.encode('utf-8')).hexdigest()
|
|
|
|
try:
|
|
# Check if prompt already exists
|
|
existing_prompt = self.db_manager.get_prompt_by_hash(prompt_hash)
|
|
|
|
if existing_prompt:
|
|
prompt_id = existing_prompt['id']
|
|
print(f"[PromptManager] Using existing prompt ID: {prompt_id}")
|
|
else:
|
|
# Save new prompt
|
|
prompt_id = self.db_manager.save_prompt(
|
|
text=prompt_text,
|
|
prompt_hash=prompt_hash,
|
|
category=additional_data.get('category') if additional_data else None,
|
|
tags=additional_data.get('tags') if additional_data else None,
|
|
rating=additional_data.get('rating') if additional_data else None,
|
|
notes=additional_data.get('notes') if additional_data else None
|
|
)
|
|
print(f"[PromptManager] Created new prompt ID: {prompt_id}")
|
|
|
|
except Exception as e:
|
|
print(f"[PromptManager] Error saving prompt: {e}")
|
|
# Generate a temporary ID for tracking
|
|
prompt_id = f"temp_{int(time.time())}"
|
|
|
|
# Create execution context
|
|
execution_context = {
|
|
'id': prompt_id,
|
|
'execution_id': execution_id,
|
|
'text': prompt_text,
|
|
'timestamp': time.time(),
|
|
'thread_id': threading.current_thread().ident,
|
|
'additional_data': additional_data or {}
|
|
}
|
|
|
|
# Store in thread-local storage
|
|
self._local.current_prompt = execution_context
|
|
|
|
# Store in global tracking with thread safety
|
|
with self.lock:
|
|
self.active_prompts[execution_id] = execution_context
|
|
|
|
print(f"[PromptManager] Set current prompt: {execution_id} -> {prompt_text[:50]}...")
|
|
return execution_id
|
|
|
|
def get_current_prompt(self) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Get the current prompt context for this thread.
|
|
|
|
Returns:
|
|
Current prompt context or None
|
|
"""
|
|
current = getattr(self._local, 'current_prompt', None)
|
|
|
|
if current:
|
|
# Check if prompt hasn't expired
|
|
if time.time() - current['timestamp'] < self.prompt_timeout:
|
|
return current
|
|
else:
|
|
print(f"[PromptManager] Current prompt expired: {current['execution_id']}")
|
|
self.clear_current_prompt()
|
|
|
|
# Fallback: try to find recent prompt from global tracking
|
|
return self._find_recent_prompt()
|
|
|
|
def _find_recent_prompt(self) -> Optional[Dict[str, Any]]:
|
|
"""Find the most recent prompt that's still valid."""
|
|
with self.lock:
|
|
current_time = time.time()
|
|
recent_prompts = [
|
|
prompt for prompt in self.active_prompts.values()
|
|
if current_time - prompt['timestamp'] < self.prompt_timeout
|
|
]
|
|
|
|
if recent_prompts:
|
|
# Return the most recent one
|
|
return max(recent_prompts, key=lambda p: p['timestamp'])
|
|
|
|
return None
|
|
|
|
def clear_current_prompt(self):
|
|
"""Clear the current prompt context for this thread."""
|
|
current = getattr(self._local, 'current_prompt', None)
|
|
if current:
|
|
execution_id = current['execution_id']
|
|
print(f"[PromptManager] Clearing current prompt: {execution_id}")
|
|
|
|
# Clear from thread-local storage
|
|
self._local.current_prompt = None
|
|
|
|
# Remove from global tracking
|
|
with self.lock:
|
|
self.active_prompts.pop(execution_id, None)
|
|
|
|
def extend_prompt_timeout(self, execution_id: str, additional_seconds: int = 60):
|
|
"""
|
|
Extend the timeout for a specific prompt execution.
|
|
|
|
Args:
|
|
execution_id: The execution ID to extend
|
|
additional_seconds: Additional seconds to add to timeout
|
|
"""
|
|
with self.lock:
|
|
if execution_id in self.active_prompts:
|
|
self.active_prompts[execution_id]['timestamp'] = time.time()
|
|
print(f"[PromptManager] Extended timeout for prompt: {execution_id}")
|
|
|
|
def generate_execution_id(self) -> str:
|
|
"""Generate a unique execution ID."""
|
|
return f"exec_{uuid.uuid4().hex[:8]}_{int(time.time())}"
|
|
|
|
def _cleanup_expired_prompts(self):
|
|
"""Background thread to clean up expired prompts."""
|
|
while True:
|
|
try:
|
|
current_time = time.time()
|
|
expired_ids = []
|
|
|
|
with self.lock:
|
|
for exec_id, prompt_data in self.active_prompts.items():
|
|
if current_time - prompt_data['timestamp'] > self.prompt_timeout:
|
|
expired_ids.append(exec_id)
|
|
|
|
for exec_id in expired_ids:
|
|
self.active_prompts.pop(exec_id, None)
|
|
|
|
if expired_ids:
|
|
print(f"[PromptManager] Cleaned up {len(expired_ids)} expired prompts")
|
|
|
|
time.sleep(self.cleanup_interval)
|
|
|
|
except Exception as e:
|
|
print(f"[PromptManager] Error in cleanup thread: {e}")
|
|
time.sleep(60) # Wait a minute before retrying
|
|
|
|
def get_active_prompts(self) -> Dict[str, Dict[str, Any]]:
|
|
"""
|
|
Get all currently active prompts.
|
|
|
|
Returns:
|
|
Dictionary of active prompts
|
|
"""
|
|
with self.lock:
|
|
return self.active_prompts.copy()
|
|
|
|
def get_status(self) -> Dict[str, Any]:
|
|
"""
|
|
Get tracker status information.
|
|
|
|
Returns:
|
|
Status information dictionary
|
|
"""
|
|
with self.lock:
|
|
active_count = len(self.active_prompts)
|
|
|
|
current_prompt = self.get_current_prompt()
|
|
|
|
return {
|
|
'active_prompts_count': active_count,
|
|
'current_prompt_id': current_prompt['id'] if current_prompt else None,
|
|
'current_execution_id': current_prompt['execution_id'] if current_prompt else None,
|
|
'thread_id': threading.current_thread().ident,
|
|
'prompt_timeout': self.prompt_timeout,
|
|
'cleanup_interval': self.cleanup_interval
|
|
}
|
|
|
|
|
|
class PromptExecutionContext:
|
|
"""Context manager for prompt executions."""
|
|
|
|
def __init__(self, prompt_tracker: PromptTracker, prompt_text: str, **kwargs):
|
|
"""
|
|
Initialize execution context.
|
|
|
|
Args:
|
|
prompt_tracker: PromptTracker instance
|
|
prompt_text: The prompt text
|
|
**kwargs: Additional prompt metadata
|
|
"""
|
|
self.prompt_tracker = prompt_tracker
|
|
self.prompt_text = prompt_text
|
|
self.additional_data = kwargs
|
|
self.execution_id = None
|
|
|
|
def __enter__(self):
|
|
"""Enter the execution context."""
|
|
self.execution_id = self.prompt_tracker.set_current_prompt(
|
|
self.prompt_text,
|
|
self.additional_data
|
|
)
|
|
return self.execution_id
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
"""Exit the execution context."""
|
|
# Don't clear immediately - let the timeout handle it
|
|
# This allows images to be generated after the prompt execution completes
|
|
pass |