582 lines
19 KiB
Python
582 lines
19 KiB
Python
"""
|
|
Memory - Unified management of 3 memory types.
|
|
|
|
Architecture:
|
|
- LTM (Long-Term Memory): Configuration, library, preferences - Persistent
|
|
- STM (Short-Term Memory): Conversation, current workflow - Volatile
|
|
- Episodic Memory: Search results, transient states - Very volatile
|
|
"""
|
|
|
|
import json
|
|
import logging
|
|
from dataclasses import dataclass, field
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
# =============================================================================
|
|
# LONG-TERM MEMORY (LTM) - Persistent
|
|
# =============================================================================
|
|
|
|
|
|
@dataclass
|
|
class LongTermMemory:
|
|
"""
|
|
Long-term memory - Persistent and static.
|
|
|
|
Stores:
|
|
- User configuration (folders, URLs)
|
|
- Preferences (quality, languages)
|
|
- Library (owned movies/TV shows)
|
|
- Followed shows (watchlist)
|
|
"""
|
|
|
|
# Folder and service configuration
|
|
config: dict[str, str] = field(default_factory=dict)
|
|
|
|
# User preferences
|
|
preferences: dict[str, Any] = field(
|
|
default_factory=lambda: {
|
|
"preferred_quality": "1080p",
|
|
"preferred_languages": ["en", "fr"],
|
|
"auto_organize": False,
|
|
"naming_format": "{title}.{year}.{quality}",
|
|
}
|
|
)
|
|
|
|
# Library of owned media
|
|
library: dict[str, list[dict]] = field(
|
|
default_factory=lambda: {"movies": [], "tv_shows": []}
|
|
)
|
|
|
|
# Followed shows (watchlist)
|
|
following: list[dict] = field(default_factory=list)
|
|
|
|
def get_config(self, key: str, default: Any = None) -> Any:
|
|
"""Get a configuration value."""
|
|
return self.config.get(key, default)
|
|
|
|
def set_config(self, key: str, value: Any) -> None:
|
|
"""Set a configuration value."""
|
|
self.config[key] = value
|
|
logger.debug(f"LTM: Set config {key}")
|
|
|
|
def has_config(self, key: str) -> bool:
|
|
"""Check if a configuration exists."""
|
|
return key in self.config and self.config[key] is not None
|
|
|
|
def add_to_library(self, media_type: str, media: dict) -> None:
|
|
"""Add a media item to the library."""
|
|
if media_type not in self.library:
|
|
self.library[media_type] = []
|
|
|
|
# Avoid duplicates by imdb_id
|
|
existing_ids = [m.get("imdb_id") for m in self.library[media_type]]
|
|
if media.get("imdb_id") not in existing_ids:
|
|
media["added_at"] = datetime.now().isoformat()
|
|
self.library[media_type].append(media)
|
|
logger.info(f"LTM: Added {media.get('title')} to {media_type}")
|
|
|
|
def get_library(self, media_type: str) -> list[dict]:
|
|
"""Get the library for a media type."""
|
|
return self.library.get(media_type, [])
|
|
|
|
def follow_show(self, show: dict) -> None:
|
|
"""Add a show to the watchlist."""
|
|
existing_ids = [s.get("imdb_id") for s in self.following]
|
|
if show.get("imdb_id") not in existing_ids:
|
|
show["followed_at"] = datetime.now().isoformat()
|
|
self.following.append(show)
|
|
logger.info(f"LTM: Now following {show.get('title')}")
|
|
|
|
def to_dict(self) -> dict:
|
|
"""Convert to dictionary for serialization."""
|
|
return {
|
|
"config": self.config,
|
|
"preferences": self.preferences,
|
|
"library": self.library,
|
|
"following": self.following,
|
|
}
|
|
|
|
@classmethod
|
|
def from_dict(cls, data: dict) -> "LongTermMemory":
|
|
"""Create an instance from a dictionary."""
|
|
return cls(
|
|
config=data.get("config", {}),
|
|
preferences=data.get(
|
|
"preferences",
|
|
{
|
|
"preferred_quality": "1080p",
|
|
"preferred_languages": ["en", "fr"],
|
|
"auto_organize": False,
|
|
"naming_format": "{title}.{year}.{quality}",
|
|
},
|
|
),
|
|
library=data.get("library", {"movies": [], "tv_shows": []}),
|
|
following=data.get("following", []),
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# SHORT-TERM MEMORY (STM) - Conversation
|
|
# =============================================================================
|
|
|
|
|
|
@dataclass
|
|
class ShortTermMemory:
|
|
"""
|
|
Short-term memory - Volatile and conversational.
|
|
|
|
Stores:
|
|
- Current conversation history
|
|
- Current workflow (what we're doing)
|
|
- Extracted entities from conversation
|
|
- Current discussion topic
|
|
"""
|
|
|
|
# Conversation message history
|
|
conversation_history: list[dict[str, str]] = field(default_factory=list)
|
|
|
|
# Current workflow
|
|
current_workflow: dict | None = None
|
|
|
|
# Extracted entities (title, year, requested quality, etc.)
|
|
extracted_entities: dict[str, Any] = field(default_factory=dict)
|
|
|
|
# Current conversation topic
|
|
current_topic: str | None = None
|
|
|
|
# Conversation language
|
|
language: str = "en"
|
|
|
|
# History message limit
|
|
max_history: int = 20
|
|
|
|
def add_message(self, role: str, content: str) -> None:
|
|
"""Add a message to history."""
|
|
self.conversation_history.append(
|
|
{"role": role, "content": content, "timestamp": datetime.now().isoformat()}
|
|
)
|
|
# Keep only the last N messages
|
|
if len(self.conversation_history) > self.max_history:
|
|
self.conversation_history = self.conversation_history[-self.max_history :]
|
|
logger.debug(f"STM: Added {role} message")
|
|
|
|
def get_recent_history(self, n: int = 10) -> list[dict]:
|
|
"""Get the last N messages."""
|
|
return self.conversation_history[-n:]
|
|
|
|
def start_workflow(self, workflow_type: str, target: dict) -> None:
|
|
"""Start a new workflow."""
|
|
self.current_workflow = {
|
|
"type": workflow_type,
|
|
"target": target,
|
|
"stage": "started",
|
|
"started_at": datetime.now().isoformat(),
|
|
}
|
|
logger.info(f"STM: Started workflow '{workflow_type}'")
|
|
|
|
def update_workflow_stage(self, stage: str) -> None:
|
|
"""Update the workflow stage."""
|
|
if self.current_workflow:
|
|
self.current_workflow["stage"] = stage
|
|
logger.debug(f"STM: Workflow stage -> {stage}")
|
|
|
|
def end_workflow(self) -> None:
|
|
"""End the current workflow."""
|
|
if self.current_workflow:
|
|
logger.info(f"STM: Ended workflow '{self.current_workflow.get('type')}'")
|
|
self.current_workflow = None
|
|
|
|
def set_entity(self, key: str, value: Any) -> None:
|
|
"""Store an extracted entity."""
|
|
self.extracted_entities[key] = value
|
|
logger.debug(f"STM: Set entity {key}={value}")
|
|
|
|
def get_entity(self, key: str, default: Any = None) -> Any:
|
|
"""Get an extracted entity."""
|
|
return self.extracted_entities.get(key, default)
|
|
|
|
def clear_entities(self) -> None:
|
|
"""Clear extracted entities."""
|
|
self.extracted_entities = {}
|
|
|
|
def set_topic(self, topic: str) -> None:
|
|
"""Set the current topic."""
|
|
self.current_topic = topic
|
|
logger.debug(f"STM: Topic -> {topic}")
|
|
|
|
def set_language(self, language: str) -> None:
|
|
"""Set the conversation language."""
|
|
self.language = language
|
|
logger.debug(f"STM: Language -> {language}")
|
|
|
|
def clear(self) -> None:
|
|
"""Reset short-term memory."""
|
|
self.conversation_history = []
|
|
self.current_workflow = None
|
|
self.extracted_entities = {}
|
|
self.current_topic = None
|
|
self.language = "en"
|
|
logger.info("STM: Cleared")
|
|
|
|
def to_dict(self) -> dict:
|
|
"""Convert to dictionary."""
|
|
return {
|
|
"conversation_history": self.conversation_history,
|
|
"current_workflow": self.current_workflow,
|
|
"extracted_entities": self.extracted_entities,
|
|
"current_topic": self.current_topic,
|
|
"language": self.language,
|
|
}
|
|
|
|
|
|
# =============================================================================
|
|
# EPISODIC MEMORY - Transient states
|
|
# =============================================================================
|
|
|
|
|
|
@dataclass
|
|
class EpisodicMemory:
|
|
"""
|
|
Episodic/sensory memory - Temporary and event-driven.
|
|
|
|
Stores:
|
|
- Last search results
|
|
- Active downloads
|
|
- Recent errors
|
|
- Pending questions awaiting user response
|
|
- Background events
|
|
"""
|
|
|
|
# Last search results
|
|
last_search_results: dict | None = None
|
|
|
|
# Active downloads
|
|
active_downloads: list[dict] = field(default_factory=list)
|
|
|
|
# Recent errors
|
|
recent_errors: list[dict] = field(default_factory=list)
|
|
|
|
# Pending question awaiting user response
|
|
pending_question: dict | None = None
|
|
|
|
# Background events (download complete, new files, etc.)
|
|
background_events: list[dict] = field(default_factory=list)
|
|
|
|
# Limits for errors/events kept
|
|
max_errors: int = 5
|
|
max_events: int = 10
|
|
|
|
def store_search_results(
|
|
self, query: str, results: list[dict], search_type: str = "torrent"
|
|
) -> None:
|
|
"""
|
|
Store search results with index.
|
|
|
|
Args:
|
|
query: The search query
|
|
results: List of results
|
|
search_type: Type of search (torrent, movie, tvshow)
|
|
"""
|
|
self.last_search_results = {
|
|
"query": query,
|
|
"type": search_type,
|
|
"timestamp": datetime.now().isoformat(),
|
|
"results": [{"index": i + 1, **r} for i, r in enumerate(results)],
|
|
}
|
|
logger.info(f"Episodic: Stored {len(results)} search results for '{query}'")
|
|
|
|
def get_result_by_index(self, index: int) -> dict | None:
|
|
"""
|
|
Get a result by its number (1-indexed).
|
|
|
|
Args:
|
|
index: Result number (1, 2, 3, ...)
|
|
|
|
Returns:
|
|
The result or None if not found
|
|
"""
|
|
if not self.last_search_results:
|
|
logger.warning("Episodic: No search results stored")
|
|
return None
|
|
|
|
for result in self.last_search_results.get("results", []):
|
|
if result.get("index") == index:
|
|
return result
|
|
|
|
logger.warning(f"Episodic: Result #{index} not found")
|
|
return None
|
|
|
|
def get_search_results(self) -> dict | None:
|
|
"""Get the last search results."""
|
|
return self.last_search_results
|
|
|
|
def clear_search_results(self) -> None:
|
|
"""Clear search results."""
|
|
self.last_search_results = None
|
|
|
|
def add_active_download(self, download: dict) -> None:
|
|
"""Add an active download."""
|
|
download["started_at"] = datetime.now().isoformat()
|
|
self.active_downloads.append(download)
|
|
logger.info(f"Episodic: Added download '{download.get('name')}'")
|
|
|
|
def update_download_progress(
|
|
self, task_id: str, progress: int, status: str = "downloading"
|
|
) -> None:
|
|
"""Update download progress."""
|
|
for dl in self.active_downloads:
|
|
if dl.get("task_id") == task_id:
|
|
dl["progress"] = progress
|
|
dl["status"] = status
|
|
dl["updated_at"] = datetime.now().isoformat()
|
|
break
|
|
|
|
def complete_download(self, task_id: str, file_path: str) -> dict | None:
|
|
"""Mark a download as complete and remove it."""
|
|
for i, dl in enumerate(self.active_downloads):
|
|
if dl.get("task_id") == task_id:
|
|
completed = self.active_downloads.pop(i)
|
|
completed["status"] = "completed"
|
|
completed["file_path"] = file_path
|
|
completed["completed_at"] = datetime.now().isoformat()
|
|
|
|
# Add a background event
|
|
self.add_background_event(
|
|
"download_complete",
|
|
{"name": completed.get("name"), "file_path": file_path},
|
|
)
|
|
|
|
logger.info(f"Episodic: Download completed '{completed.get('name')}'")
|
|
return completed
|
|
return None
|
|
|
|
def get_active_downloads(self) -> list[dict]:
|
|
"""Get active downloads."""
|
|
return self.active_downloads
|
|
|
|
def add_error(
|
|
self, action: str, error: str, context: dict | None = None
|
|
) -> None:
|
|
"""Record a recent error."""
|
|
self.recent_errors.append(
|
|
{
|
|
"timestamp": datetime.now().isoformat(),
|
|
"action": action,
|
|
"error": error,
|
|
"context": context or {},
|
|
}
|
|
)
|
|
# Keep only the last N errors
|
|
self.recent_errors = self.recent_errors[-self.max_errors :]
|
|
logger.warning(f"Episodic: Error in '{action}': {error}")
|
|
|
|
def get_recent_errors(self) -> list[dict]:
|
|
"""Get recent errors."""
|
|
return self.recent_errors
|
|
|
|
def set_pending_question(
|
|
self,
|
|
question: str,
|
|
options: list[dict],
|
|
context: dict,
|
|
question_type: str = "choice",
|
|
) -> None:
|
|
"""
|
|
Record a question awaiting user response.
|
|
|
|
Args:
|
|
question: The question asked
|
|
options: List of possible options
|
|
context: Question context
|
|
question_type: Type of question (choice, confirmation, input)
|
|
"""
|
|
self.pending_question = {
|
|
"type": question_type,
|
|
"question": question,
|
|
"options": options,
|
|
"context": context,
|
|
"timestamp": datetime.now().isoformat(),
|
|
}
|
|
logger.info(f"Episodic: Pending question set ({question_type})")
|
|
|
|
def get_pending_question(self) -> dict | None:
|
|
"""Get the pending question."""
|
|
return self.pending_question
|
|
|
|
def resolve_pending_question(
|
|
self, answer_index: int | None = None
|
|
) -> dict | None:
|
|
"""
|
|
Resolve the pending question and return the chosen option.
|
|
|
|
Args:
|
|
answer_index: Answer index (1-indexed) or None to cancel
|
|
|
|
Returns:
|
|
The chosen option or None
|
|
"""
|
|
if not self.pending_question:
|
|
return None
|
|
|
|
result = None
|
|
if answer_index is not None and self.pending_question.get("options"):
|
|
for opt in self.pending_question["options"]:
|
|
if opt.get("index") == answer_index:
|
|
result = opt
|
|
break
|
|
|
|
self.pending_question = None
|
|
logger.info("Episodic: Pending question resolved")
|
|
return result
|
|
|
|
def add_background_event(self, event_type: str, data: dict) -> None:
|
|
"""Add a background event."""
|
|
self.background_events.append(
|
|
{
|
|
"type": event_type,
|
|
"timestamp": datetime.now().isoformat(),
|
|
"data": data,
|
|
"read": False,
|
|
}
|
|
)
|
|
# Keep only the last N events
|
|
self.background_events = self.background_events[-self.max_events :]
|
|
logger.info(f"Episodic: Background event '{event_type}'")
|
|
|
|
def get_unread_events(self) -> list[dict]:
|
|
"""Get unread events and mark them as read."""
|
|
unread = [e for e in self.background_events if not e.get("read")]
|
|
for e in self.background_events:
|
|
e["read"] = True
|
|
return unread
|
|
|
|
def clear(self) -> None:
|
|
"""Reset episodic memory."""
|
|
self.last_search_results = None
|
|
self.active_downloads = []
|
|
self.recent_errors = []
|
|
self.pending_question = None
|
|
self.background_events = []
|
|
logger.info("Episodic: Cleared")
|
|
|
|
def to_dict(self) -> dict:
|
|
"""Convert to dictionary."""
|
|
return {
|
|
"last_search_results": self.last_search_results,
|
|
"active_downloads": self.active_downloads,
|
|
"recent_errors": self.recent_errors,
|
|
"pending_question": self.pending_question,
|
|
"background_events": self.background_events,
|
|
}
|
|
|
|
|
|
# =============================================================================
|
|
# MEMORY MANAGER - Unified manager
|
|
# =============================================================================
|
|
|
|
|
|
class Memory:
|
|
"""
|
|
Unified manager for the 3 memory types.
|
|
|
|
Usage:
|
|
memory = Memory("memory_data")
|
|
memory.ltm.set_config("download_folder", "/path")
|
|
memory.stm.add_message("user", "Hello")
|
|
memory.episodic.store_search_results("query", results)
|
|
memory.save()
|
|
"""
|
|
|
|
def __init__(self, storage_dir: str = "memory_data"):
|
|
"""
|
|
Initialize the memory.
|
|
|
|
Args:
|
|
storage_dir: Directory for persistent storage
|
|
"""
|
|
self.storage_dir = Path(storage_dir)
|
|
self.storage_dir.mkdir(exist_ok=True)
|
|
|
|
self.ltm_file = self.storage_dir / "ltm.json"
|
|
|
|
# Initialize the 3 memory types
|
|
self.ltm = self._load_ltm()
|
|
self.stm = ShortTermMemory()
|
|
self.episodic = EpisodicMemory()
|
|
|
|
logger.info(f"Memory initialized (storage: {storage_dir})")
|
|
|
|
def _load_ltm(self) -> LongTermMemory:
|
|
"""Load LTM from file."""
|
|
if self.ltm_file.exists():
|
|
try:
|
|
data = json.loads(self.ltm_file.read_text(encoding="utf-8"))
|
|
logger.info("LTM loaded from file")
|
|
return LongTermMemory.from_dict(data)
|
|
except (OSError, json.JSONDecodeError) as e:
|
|
logger.warning(f"Could not load LTM: {e}")
|
|
return LongTermMemory()
|
|
|
|
def save(self) -> None:
|
|
"""Save LTM (the only persistent memory)."""
|
|
try:
|
|
self.ltm_file.write_text(
|
|
json.dumps(self.ltm.to_dict(), indent=2, ensure_ascii=False),
|
|
encoding="utf-8",
|
|
)
|
|
logger.debug("LTM saved to file")
|
|
except OSError as e:
|
|
logger.error(f"Failed to save LTM: {e}")
|
|
raise
|
|
|
|
def get_context_for_prompt(self) -> dict:
|
|
"""
|
|
Generate context to include in the system prompt.
|
|
|
|
Returns:
|
|
Dictionary with relevant context from all 3 memories
|
|
"""
|
|
return {
|
|
"config": self.ltm.config,
|
|
"preferences": self.ltm.preferences,
|
|
"current_workflow": self.stm.current_workflow,
|
|
"current_topic": self.stm.current_topic,
|
|
"extracted_entities": self.stm.extracted_entities,
|
|
"last_search": {
|
|
"query": (
|
|
self.episodic.last_search_results.get("query")
|
|
if self.episodic.last_search_results
|
|
else None
|
|
),
|
|
"result_count": (
|
|
len(self.episodic.last_search_results.get("results", []))
|
|
if self.episodic.last_search_results
|
|
else 0
|
|
),
|
|
},
|
|
"active_downloads_count": len(self.episodic.active_downloads),
|
|
"pending_question": self.episodic.pending_question is not None,
|
|
"unread_events": len(
|
|
[e for e in self.episodic.background_events if not e.get("read")]
|
|
),
|
|
}
|
|
|
|
def get_full_state(self) -> dict:
|
|
"""Return the full state of all 3 memories (for debug)."""
|
|
return {
|
|
"ltm": self.ltm.to_dict(),
|
|
"stm": self.stm.to_dict(),
|
|
"episodic": self.episodic.to_dict(),
|
|
}
|
|
|
|
def clear_session(self) -> None:
|
|
"""Clear session memories (STM + Episodic)."""
|
|
self.stm.clear()
|
|
self.episodic.clear()
|
|
logger.info("Session memories cleared")
|