Spaces:

babaTEEpe
/

Toun

Paused

Toun / session_manager.py

Upload 17 files

513d6d1 verified 6 months ago

13.9 kB

import os
import shutil
import uuid
import time
from datetime import datetime
import logging
from typing import Optional, Dict, Any
import glob
import re
import numpy as np

logger = logging.getLogger(__name__)

class WorldState:
    """
    Structured situational memory (The sentinel Mirror).
    Tracks persistent entities, environmental conditions, and mission facts.
    """
    def __init__(self, session_id: str):
        self.session_id = session_id
        self.last_update = time.time()
        
        # ── The Entity Mirror (The core of reality) ──
        # List of Entities. Each entity is a dict:
        # {
        #   "id": int, 
        #   "type": str, 
        #   "attributes": dict, 
        #   "state": str, 
        #   "last_seen": float, 
        #   "confidence": float, 
        #   "history": List[Dict] # Temporal trail of [state, attributes, timestamp]
        # }
        self.entities = [] 
        self.event_history = [] # Global timeline of significant semantic changes
        
        # ── Environmental Context ──
        self.lighting = "unknown"
        self.noise_level = 0.0
        self.spatial_context = "unknown" # e.g., 'corridor', 'outside'
        
        # ── Mission Logic ──
        self.target_status = "searching" # searching, visible, occluded, lost
        self.active_strategies = []
        self.diagnostics = {}
        self.confidence_history = []

    def log_event(self, text: str, event_type: str = "semantic"):
        """Appends a timestamped event to the global history."""
        self.event_history.append({
            "timestamp": time.time(),
            "type": event_type,
            "text": text
        })
        # Increased capacity for long research sessions
        if len(self.event_history) > 200:
            self.event_history.pop(0)

    def get_events_near(self, target_ts: float, count: int = 5) -> list:
        """Returns the `count` closest events to a target timestamp."""
        if not self.event_history:
            return []
        # Sort by distance to target
        scored = [(abs(e["timestamp"] - target_ts), e) for e in self.event_history]
        scored.sort(key=lambda x: x[0])
        return [e for _, e in scored[:count]]

    def get_timeline(self, window_seconds: int = 60, anchor_ts: float = None) -> str:
        """Returns a formatted string of events. 
        Uses anchor_ts (last event time) instead of wall clock to survive restarts.
        """
        # Use the last event's timestamp as anchor, NOT time.time()
        if anchor_ts is None:
            if self.event_history:
                anchor_ts = self.event_history[-1]["timestamp"]
            else:
                anchor_ts = time.time()
        
        start_ts = anchor_ts - window_seconds
        
        # 1. SCENE CONTEXT
        now = time.time()
        active_entities = [e for e in self.entities if now - e.get("last_seen", 0) < 30]
        context_lines = []
        if active_entities:
            type_counts = {}
            for e in active_entities:
                t = e["type"]
                type_counts[t] = type_counts.get(t, 0) + 1
            summary_parts = [f"{count} {t}" for t, count in type_counts.items()]
            context_lines.append(f"[CURRENT SCENE CONTEXT]")
            context_lines.append(f"Currently in view: {', '.join(summary_parts)}")
            context_lines.append("")

        # 2. RECENT EVENTS (anchored to last observation, not wall clock)
        relevant = [e for e in self.event_history if e["timestamp"] >= start_ts]
        event_lines = []
        for e in relevant:
            dt = datetime.fromtimestamp(e["timestamp"]).strftime("%H:%M:%S")
            # Add human-readable time-ago
            ago = int(anchor_ts - e["timestamp"])
            ago_str = f"{ago}s ago" if ago < 60 else f"{ago // 60}m {ago % 60}s ago"
            event_lines.append(f"[{dt}] ({ago_str}) {e['text']}")
            
        if not event_lines:
            event_lines.append(f"No observations recorded in the last {window_seconds}s.")

        output = context_lines + ["[RECENT LOGGED EVENTS]"] + event_lines
        return "\n".join(output)

    def to_dict(self) -> Dict[str, Any]:
        return self.__dict__

    @classmethod
    def from_dict(cls, data: Dict[str, Any]):
        ws = cls(data.get("session_id", "unknown"))
        ws.__dict__.update(data)
        return ws

# Build absolute path to sessions folder relative to this file
_BACKEND_ROOT = os.path.dirname(os.path.abspath(__file__))
DEFAULT_SESSIONS_DIR = os.path.join(_BACKEND_ROOT, "sessions")

class SessionManager:
    def __init__(self, base_dir=DEFAULT_SESSIONS_DIR, max_sessions=50, ttl_seconds=86400):
        self.base_dir = base_dir
        self.max_sessions = max_sessions
        self.ttl_seconds = ttl_seconds
        self.active_session_id: Optional[str] = None  # Tracks user's chosen session
        os.makedirs(self.base_dir, exist_ok=True)
        self.cleanup()

    def set_active_session(self, session_id: str):
        """Set the user's chosen active session. Called from the frontend dropdown."""
        if os.path.exists(self.get_session_path(session_id)):
            self.active_session_id = session_id
            logger.info(f"[SESSION] Active session set to: {session_id}")
            return True
        return False

    def get_or_create_session(self, session_id: Optional[str] = None) -> str:
        """Returns existing session or creates a new one.
        Priority: explicit session_id > active_session_id > create new.
        """
        # Sanitize common string representations of null from frontend
        if isinstance(session_id, str):
            if session_id.lower().strip() in ["", "none", "null", "undefined"]:
                session_id = None

        # 1. Use explicit session_id if it exists on disk
        if session_id and os.path.exists(self.get_session_path(session_id)):
            return session_id

        # 2. Use the user's chosen active session (set via dropdown)
        if not session_id and self.active_session_id:
            if os.path.exists(self.get_session_path(self.active_session_id)):
                return self.active_session_id

        # 3. Create a new session as last resort
        new_id = self.create_session(session_id)
        self.active_session_id = new_id  # Auto-activate newly created sessions
        return new_id

    def create_session(self, session_id: Optional[str] = None) -> str:
        """Creates a new session directory and returns the session_id."""
        if not session_id:
            session_id = str(uuid.uuid4())
            
        session_path = os.path.join(self.base_dir, session_id)
        os.makedirs(session_path, exist_ok=True)
        
        # Create subdirs for media types
        os.makedirs(os.path.join(session_path, "frames"), exist_ok=True)
        os.makedirs(os.path.join(session_path, "audio"), exist_ok=True)
        
        # Initialize world state
        self.update_world_state(session_id, WorldState(session_id))
        
        # Touch metadata file to track time
        with open(os.path.join(session_path, "metadata.txt"), "w") as f:
            f.write(f"Created: {datetime.now().isoformat()}")
            
        self.cleanup() # Run cleanup on every new session creation
        return session_id

    def update_world_state(self, session_id: str, state: WorldState):
        """Persists structured state into session storage."""
        session_path = self.get_session_path(session_id)
        if not os.path.exists(session_path): return
        
        state_path = os.path.join(session_path, "world_state.json")
        import json
        with open(state_path, "w") as f:
            json.dump(state.to_dict(), f)

    def get_world_state(self, session_id: str) -> WorldState:
        """Retrieves the latest world state."""
        session_path = self.get_session_path(session_id)
        state_path = os.path.join(session_path, "world_state.json")
        if os.path.exists(state_path):
            import json
            try:
                with open(state_path, "r") as f:
                    data = json.load(f)
                    return WorldState.from_dict(data)
            except: pass
        return WorldState(session_id)

    def update_sensory_memory(self, session_id: str, new_data: Dict[str, Any]):
        """DEPRECATED: Use update_world_state. Keeps compatibility for now."""
        ws = self.get_world_state(session_id)
        ws.diagnostics.update(new_data)
        self.update_world_state(session_id, ws)

    def get_sensory_memory(self, session_id: str) -> Dict[str, Any]:
        """DEPRECATED: Use get_world_state. Returns diagnostics."""
        ws = self.get_world_state(session_id)
        return ws.diagnostics

    def get_session_path(self, session_id: str) -> str:
        return os.path.join(self.base_dir, session_id)

    def save_frame(self, session_id: str, frame_data, timestamp: float):
        """Saves a frame to the session directory (Throttled by Orchestrator)."""
        session_path = self.get_session_path(session_id)
        if not os.path.exists(session_path):
            return False
            
        frames_dir = os.path.join(session_path, "frames")
        os.makedirs(frames_dir, exist_ok=True)
        
        frame_name = f"frame_{int(timestamp*1000)}.jpg"
        frame_path = os.path.join(frames_dir, frame_name)
        
        try:
            # Assuming frame_data is a numpy array (OpenCV) or PIL image
            if hasattr(frame_data, "save"): # PIL
                frame_data.save(frame_path)
            elif isinstance(frame_data, np.ndarray):
                import cv2
                # OpenCV handles numpy arrays directly
                success = cv2.imwrite(frame_path, frame_data)
                if not success:
                    logger.warning(f"[SESSION] cv2.imwrite failed for {frame_path}")
                    return False
            else:
                logger.warning(f"[SESSION] Unsupported frame data type: {type(frame_data)}")
                return False
            return frame_path
        except Exception as e:
            logger.error(f"[SESSION] Error saving frame: {e}")
            return False

    def get_historic_frame(self, session_id: str, timestamp: float) -> Optional[str]:
        """
        Finds the frame closest to the requested timestamp.
        """
        session_path = self.get_session_path(session_id)
        frames_dir = os.path.join(session_path, "frames")
        if not os.path.exists(frames_dir): return None
        
        frames = glob.glob(os.path.join(frames_dir, "frame_*.jpg"))
        if not frames: return None
        
        # Convert filenames back to timestamps for comparison
        target_ms = int(timestamp * 1000)
        closest_frame = min(frames, key=lambda f: abs(int(re.search(r'frame_(\d+)', f).group(1)) - target_ms))
        
        # Only return if within a reasonable window (e.g. 30 seconds)
        found_ms = int(re.search(r'frame_(\d+)', closest_frame).group(1))
        if abs(found_ms - target_ms) > 30000:
            logger.warning(f"[SESSION] No frame found within 30s window of {timestamp}")
            return None
            
        return closest_frame

    def log_session_event(self, session_id: str, text: str, event_type: str = "semantic"):
        """Adds an event to a session's world state and persists it."""
        ws = self.get_world_state(session_id)
        ws.log_event(text, event_type)
        self.update_world_state(session_id, ws)

    def get_session_timeline(self, session_id: str, window_seconds: int = 60) -> str:
        """Retrieves formatted timeline from world state."""
        ws = self.get_world_state(session_id)
        return ws.get_timeline(window_seconds)

    def get_historic_audio_context(self, session_id: str, timestamp: float, window: int = 15) -> str:
        """
        [STUB] Retrieves audio transcripts or events from memory within a window.
        """
        # For now, we'll return the general mission finding context if relevant
        return ""

    def save_audio_segment(self, session_id: str, audio_data, start_time: float):
        """Saves an audio segment to the session directory."""
        session_path = self.get_session_path(session_id)
        if not os.path.exists(session_path):
            return False
            
        audio_name = f"audio_{int(start_time*1000)}.wav"
        audio_path = os.path.join(session_path, "audio", audio_name)
        # logic to save audio file would go here
        return audio_path

    def cleanup(self):
        """Deletes old sessions based on TTL and count limit."""
        try:
            now = time.time()
            sessions = []
            for d in os.listdir(self.base_dir):
                p = os.path.join(self.base_dir, d)
                if os.path.isdir(p):
                    mtime = os.path.getmtime(p)
                    sessions.append((d, p, mtime))
            
            # Sort by time (oldest first)
            sessions.sort(key=lambda x: x[2])
            
            # 1. Cleanup by TTL
            for sid, path, mtime in sessions:
                if now - mtime > self.ttl_seconds:
                    logger.info(f"Cleaning up expired session: {sid}")
                    shutil.rmtree(path)
                    sessions.remove((sid, path, mtime))
            
            # 2. Cleanup by Count Limit
            while len(sessions) > self.max_sessions:
                sid, path, mtime = sessions.pop(0)
                logger.info(f"Cleaning up session (limit reached): {sid}")
                shutil.rmtree(path)
                
        except Exception as e:
            logger.error(f"Error during SessionManager cleanup: {e}")

# Global instance
session_manager = SessionManager()

	import os
	import shutil
	import uuid
	import time
	from datetime import datetime
	import logging
	from typing import Optional, Dict, Any
	import glob
	import re
	import numpy as np

	logger = logging.getLogger(__name__)

	class WorldState:
	"""
	Structured situational memory (The sentinel Mirror).
	Tracks persistent entities, environmental conditions, and mission facts.
	"""
	def __init__(self, session_id: str):
	self.session_id = session_id
	self.last_update = time.time()

	# ── The Entity Mirror (The core of reality) ──
	# List of Entities. Each entity is a dict:
	# {
	# "id": int,
	# "type": str,
	# "attributes": dict,
	# "state": str,
	# "last_seen": float,
	# "confidence": float,
	# "history": List[Dict] # Temporal trail of [state, attributes, timestamp]
	# }
	self.entities = []
	self.event_history = [] # Global timeline of significant semantic changes

	# ── Environmental Context ──
	self.lighting = "unknown"
	self.noise_level = 0.0
	self.spatial_context = "unknown" # e.g., 'corridor', 'outside'

	# ── Mission Logic ──
	self.target_status = "searching" # searching, visible, occluded, lost
	self.active_strategies = []
	self.diagnostics = {}
	self.confidence_history = []

	def log_event(self, text: str, event_type: str = "semantic"):
	"""Appends a timestamped event to the global history."""
	self.event_history.append({
	"timestamp": time.time(),
	"type": event_type,
	"text": text
	})
	# Increased capacity for long research sessions
	if len(self.event_history) > 200:
	self.event_history.pop(0)

	def get_events_near(self, target_ts: float, count: int = 5) -> list:
	"""Returns the `count` closest events to a target timestamp."""
	if not self.event_history:
	return []
	# Sort by distance to target
	scored = [(abs(e["timestamp"] - target_ts), e) for e in self.event_history]
	scored.sort(key=lambda x: x[0])
	return [e for _, e in scored[:count]]

	def get_timeline(self, window_seconds: int = 60, anchor_ts: float = None) -> str:
	"""Returns a formatted string of events.
	Uses anchor_ts (last event time) instead of wall clock to survive restarts.
	"""
	# Use the last event's timestamp as anchor, NOT time.time()
	if anchor_ts is None:
	if self.event_history:
	anchor_ts = self.event_history[-1]["timestamp"]
	else:
	anchor_ts = time.time()

	start_ts = anchor_ts - window_seconds

	# 1. SCENE CONTEXT
	now = time.time()
	active_entities = [e for e in self.entities if now - e.get("last_seen", 0) < 30]
	context_lines = []
	if active_entities:
	type_counts = {}
	for e in active_entities:
	t = e["type"]
	type_counts[t] = type_counts.get(t, 0) + 1
	summary_parts = [f"{count} {t}" for t, count in type_counts.items()]
	context_lines.append(f"[CURRENT SCENE CONTEXT]")
	context_lines.append(f"Currently in view: {', '.join(summary_parts)}")
	context_lines.append("")

	# 2. RECENT EVENTS (anchored to last observation, not wall clock)
	relevant = [e for e in self.event_history if e["timestamp"] >= start_ts]
	event_lines = []
	for e in relevant:
	dt = datetime.fromtimestamp(e["timestamp"]).strftime("%H:%M:%S")
	# Add human-readable time-ago
	ago = int(anchor_ts - e["timestamp"])
	ago_str = f"{ago}s ago" if ago < 60 else f"{ago // 60}m {ago % 60}s ago"
	event_lines.append(f"[{dt}] ({ago_str}) {e['text']}")

	if not event_lines:
	event_lines.append(f"No observations recorded in the last {window_seconds}s.")

	output = context_lines + ["[RECENT LOGGED EVENTS]"] + event_lines
	return "\n".join(output)

	def to_dict(self) -> Dict[str, Any]:
	return self.__dict__

	@classmethod
	def from_dict(cls, data: Dict[str, Any]):
	ws = cls(data.get("session_id", "unknown"))
	ws.__dict__.update(data)
	return ws

	# Build absolute path to sessions folder relative to this file
	_BACKEND_ROOT = os.path.dirname(os.path.abspath(__file__))
	DEFAULT_SESSIONS_DIR = os.path.join(_BACKEND_ROOT, "sessions")

	class SessionManager:
	def __init__(self, base_dir=DEFAULT_SESSIONS_DIR, max_sessions=50, ttl_seconds=86400):
	self.base_dir = base_dir
	self.max_sessions = max_sessions
	self.ttl_seconds = ttl_seconds
	self.active_session_id: Optional[str] = None # Tracks user's chosen session
	os.makedirs(self.base_dir, exist_ok=True)
	self.cleanup()

	def set_active_session(self, session_id: str):
	"""Set the user's chosen active session. Called from the frontend dropdown."""
	if os.path.exists(self.get_session_path(session_id)):
	self.active_session_id = session_id
	logger.info(f"[SESSION] Active session set to: {session_id}")
	return True
	return False

	def get_or_create_session(self, session_id: Optional[str] = None) -> str:
	"""Returns existing session or creates a new one.
	Priority: explicit session_id > active_session_id > create new.
	"""
	# Sanitize common string representations of null from frontend
	if isinstance(session_id, str):
	if session_id.lower().strip() in ["", "none", "null", "undefined"]:
	session_id = None

	# 1. Use explicit session_id if it exists on disk
	if session_id and os.path.exists(self.get_session_path(session_id)):
	return session_id

	# 2. Use the user's chosen active session (set via dropdown)
	if not session_id and self.active_session_id:
	if os.path.exists(self.get_session_path(self.active_session_id)):
	return self.active_session_id

	# 3. Create a new session as last resort
	new_id = self.create_session(session_id)
	self.active_session_id = new_id # Auto-activate newly created sessions
	return new_id

	def create_session(self, session_id: Optional[str] = None) -> str:
	"""Creates a new session directory and returns the session_id."""
	if not session_id:
	session_id = str(uuid.uuid4())

	session_path = os.path.join(self.base_dir, session_id)
	os.makedirs(session_path, exist_ok=True)

	# Create subdirs for media types
	os.makedirs(os.path.join(session_path, "frames"), exist_ok=True)
	os.makedirs(os.path.join(session_path, "audio"), exist_ok=True)

	# Initialize world state
	self.update_world_state(session_id, WorldState(session_id))

	# Touch metadata file to track time
	with open(os.path.join(session_path, "metadata.txt"), "w") as f:
	f.write(f"Created: {datetime.now().isoformat()}")

	self.cleanup() # Run cleanup on every new session creation
	return session_id

	def update_world_state(self, session_id: str, state: WorldState):
	"""Persists structured state into session storage."""
	session_path = self.get_session_path(session_id)
	if not os.path.exists(session_path): return

	state_path = os.path.join(session_path, "world_state.json")
	import json
	with open(state_path, "w") as f:
	json.dump(state.to_dict(), f)

	def get_world_state(self, session_id: str) -> WorldState:
	"""Retrieves the latest world state."""
	session_path = self.get_session_path(session_id)
	state_path = os.path.join(session_path, "world_state.json")
	if os.path.exists(state_path):
	import json
	try:
	with open(state_path, "r") as f:
	data = json.load(f)
	return WorldState.from_dict(data)
	except: pass
	return WorldState(session_id)

	def update_sensory_memory(self, session_id: str, new_data: Dict[str, Any]):
	"""DEPRECATED: Use update_world_state. Keeps compatibility for now."""
	ws = self.get_world_state(session_id)
	ws.diagnostics.update(new_data)
	self.update_world_state(session_id, ws)

	def get_sensory_memory(self, session_id: str) -> Dict[str, Any]:
	"""DEPRECATED: Use get_world_state. Returns diagnostics."""
	ws = self.get_world_state(session_id)
	return ws.diagnostics

	def get_session_path(self, session_id: str) -> str:
	return os.path.join(self.base_dir, session_id)

	def save_frame(self, session_id: str, frame_data, timestamp: float):
	"""Saves a frame to the session directory (Throttled by Orchestrator)."""
	session_path = self.get_session_path(session_id)
	if not os.path.exists(session_path):
	return False

	frames_dir = os.path.join(session_path, "frames")
	os.makedirs(frames_dir, exist_ok=True)

	frame_name = f"frame_{int(timestamp*1000)}.jpg"
	frame_path = os.path.join(frames_dir, frame_name)

	try:
	# Assuming frame_data is a numpy array (OpenCV) or PIL image
	if hasattr(frame_data, "save"): # PIL
	frame_data.save(frame_path)
	elif isinstance(frame_data, np.ndarray):
	import cv2
	# OpenCV handles numpy arrays directly
	success = cv2.imwrite(frame_path, frame_data)
	if not success:
	logger.warning(f"[SESSION] cv2.imwrite failed for {frame_path}")
	return False
	else:
	logger.warning(f"[SESSION] Unsupported frame data type: {type(frame_data)}")
	return False
	return frame_path
	except Exception as e:
	logger.error(f"[SESSION] Error saving frame: {e}")
	return False

	def get_historic_frame(self, session_id: str, timestamp: float) -> Optional[str]:
	"""
	Finds the frame closest to the requested timestamp.
	"""
	session_path = self.get_session_path(session_id)
	frames_dir = os.path.join(session_path, "frames")
	if not os.path.exists(frames_dir): return None

	frames = glob.glob(os.path.join(frames_dir, "frame_*.jpg"))
	if not frames: return None

	# Convert filenames back to timestamps for comparison
	target_ms = int(timestamp * 1000)
	closest_frame = min(frames, key=lambda f: abs(int(re.search(r'frame_(\d+)', f).group(1)) - target_ms))

	# Only return if within a reasonable window (e.g. 30 seconds)
	found_ms = int(re.search(r'frame_(\d+)', closest_frame).group(1))
	if abs(found_ms - target_ms) > 30000:
	logger.warning(f"[SESSION] No frame found within 30s window of {timestamp}")
	return None

	return closest_frame

	def log_session_event(self, session_id: str, text: str, event_type: str = "semantic"):
	"""Adds an event to a session's world state and persists it."""
	ws = self.get_world_state(session_id)
	ws.log_event(text, event_type)
	self.update_world_state(session_id, ws)

	def get_session_timeline(self, session_id: str, window_seconds: int = 60) -> str:
	"""Retrieves formatted timeline from world state."""
	ws = self.get_world_state(session_id)
	return ws.get_timeline(window_seconds)

	def get_historic_audio_context(self, session_id: str, timestamp: float, window: int = 15) -> str:
	"""
	[STUB] Retrieves audio transcripts or events from memory within a window.
	"""
	# For now, we'll return the general mission finding context if relevant
	return ""

	def save_audio_segment(self, session_id: str, audio_data, start_time: float):
	"""Saves an audio segment to the session directory."""
	session_path = self.get_session_path(session_id)
	if not os.path.exists(session_path):
	return False

	audio_name = f"audio_{int(start_time*1000)}.wav"
	audio_path = os.path.join(session_path, "audio", audio_name)
	# logic to save audio file would go here
	return audio_path

	def cleanup(self):
	"""Deletes old sessions based on TTL and count limit."""
	try:
	now = time.time()
	sessions = []
	for d in os.listdir(self.base_dir):
	p = os.path.join(self.base_dir, d)
	if os.path.isdir(p):
	mtime = os.path.getmtime(p)
	sessions.append((d, p, mtime))

	# Sort by time (oldest first)
	sessions.sort(key=lambda x: x[2])

	# 1. Cleanup by TTL
	for sid, path, mtime in sessions:
	if now - mtime > self.ttl_seconds:
	logger.info(f"Cleaning up expired session: {sid}")
	shutil.rmtree(path)
	sessions.remove((sid, path, mtime))

	# 2. Cleanup by Count Limit
	while len(sessions) > self.max_sessions:
	sid, path, mtime = sessions.pop(0)
	logger.info(f"Cleaning up session (limit reached): {sid}")
	shutil.rmtree(path)

	except Exception as e:
	logger.error(f"Error during SessionManager cleanup: {e}")

	# Global instance
	session_manager = SessionManager()