mirror of
https://github.com/Nighthawk42/MiraiAssist.git
synced 2026-08-30 08:32:26 +00:00
379 lines
20 KiB
Python
379 lines
20 KiB
Python
# modules/context_manager.py
|
|
|
|
import datetime
|
|
import json
|
|
import logging
|
|
from pathlib import Path
|
|
import time
|
|
import shutil
|
|
from typing import List, Dict, Any, Optional, Tuple
|
|
|
|
# --- Dependency Imports with Checks ---
|
|
try:
|
|
from sentence_transformers import SentenceTransformer
|
|
SENTENCE_TRANSFORMERS_AVAILABLE = True
|
|
except ImportError:
|
|
SENTENCE_TRANSFORMERS_AVAILABLE = False
|
|
class SentenceTransformer: # Dummy for type hints
|
|
def __init__(self, *args, **kwargs): pass
|
|
def encode(self, *args, **kwargs): return []
|
|
|
|
try:
|
|
import chromadb
|
|
from chromadb.config import Settings as ChromaSettings
|
|
CHROMA_AVAILABLE = True
|
|
except ImportError:
|
|
CHROMA_AVAILABLE = False
|
|
class chromadb: # Dummy for type hints
|
|
@staticmethod
|
|
def PersistentClient(*args, **kwargs): return ChromaClientDummy()
|
|
class ChromaClientDummy:
|
|
def get_or_create_collection(self, *args, **kwargs): return ChromaCollectionDummy()
|
|
class ChromaCollectionDummy:
|
|
def count(self): return 0
|
|
def get(self, *args, **kwargs): return {"ids": []}
|
|
def add(self, *args, **kwargs): pass
|
|
def query(self, *args, **kwargs): return {"ids": [[]], "documents": [[]], "metadatas": [[]], "distances": [[]]}
|
|
def delete(self, *args, **kwargs): pass
|
|
class ChromaSettings: pass
|
|
|
|
|
|
# Local Imports
|
|
from .config_manager import ConfigManager
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class ContextManagerError(Exception):
|
|
"""Custom exception for ContextManager specific errors."""
|
|
pass
|
|
|
|
class ContextManager:
|
|
"""
|
|
Manages conversation history using Retrieval-Augmented Generation (RAG).
|
|
- Stores the full, raw conversation history chronologically in a JSON file.
|
|
- Indexes each message into a persistent ChromaDB vector store.
|
|
- Provides methods to add messages and retrieve relevant past messages.
|
|
"""
|
|
DEFAULT_STORAGE_PATH = "data/conversation_state.json"
|
|
DEFAULT_VECTOR_DB_PATH = "data/chroma_db"
|
|
DEFAULT_EMBEDDING_MODEL = "all-MiniLM-L6-v2"
|
|
DEFAULT_COLLECTION_NAME = "mirei_chat_history"
|
|
DEFAULT_RAG_N_RESULTS_FALLBACK = 3 # Fallback if n_results not provided to retrieve_relevant_context
|
|
|
|
def __init__(self, config: ConfigManager):
|
|
logger.info("Initializing ContextManager (RAG)...")
|
|
|
|
if not SENTENCE_TRANSFORMERS_AVAILABLE:
|
|
raise ContextManagerError(
|
|
"Required library 'sentence-transformers' not installed. Run: uv add sentence-transformers"
|
|
)
|
|
if not CHROMA_AVAILABLE:
|
|
raise ContextManagerError("Required library 'chromadb' not installed. Run: uv add chromadb")
|
|
|
|
cfg_section = config.get("context_manager", default={})
|
|
self.storage_path: Path = Path(cfg_section.get("storage_path", self.DEFAULT_STORAGE_PATH)).resolve()
|
|
self.vector_db_path: Path = Path(cfg_section.get("vector_db_path", self.DEFAULT_VECTOR_DB_PATH)).resolve()
|
|
self.embedding_model_name: str = cfg_section.get("embedding_model_name", self.DEFAULT_EMBEDDING_MODEL)
|
|
self.collection_name: str = cfg_section.get("collection_name", self.DEFAULT_COLLECTION_NAME)
|
|
# n_retrieval_results and n_include_recent from context_manager config are now primarily
|
|
# used as fallbacks if methods are called without specific counts, or for direct use if any.
|
|
# MemoryManager will use its own config for prompt construction.
|
|
self.n_retrieval_results_fallback: int = int(cfg_section.get("retrieval_results", self.DEFAULT_RAG_N_RESULTS_FALLBACK))
|
|
|
|
|
|
self.messages: List[Dict[str, Any]] = []
|
|
self.embedding_model: Optional[SentenceTransformer] = None
|
|
self.chroma_client: Optional[chromadb.API] = None
|
|
self.collection: Optional[chromadb.Collection] = None # Type hint for Chroma's Collection
|
|
|
|
self._load_embedding_model()
|
|
self._initialize_vector_db()
|
|
self._load_full_history() # This now ensures 'original_index' is present
|
|
self._synchronize_index()
|
|
|
|
logger.info("ContextManager (RAG) initialized successfully.")
|
|
|
|
def _load_embedding_model(self) -> None:
|
|
logger.info(f"Loading embedding model: '{self.embedding_model_name}'...")
|
|
start_time = time.time()
|
|
try:
|
|
self.embedding_model = SentenceTransformer(self.embedding_model_name)
|
|
_ = self.embedding_model.encode(["test warm-up"], show_progress_bar=False) # Warm-up/check
|
|
load_time = time.time() - start_time
|
|
logger.info(f"Embedding model '{self.embedding_model_name}' loaded in {load_time:.2f}s.")
|
|
except Exception as e:
|
|
logger.critical(f"Failed to load SentenceTransformer model '{self.embedding_model_name}': {e}", exc_info=True)
|
|
raise ContextManagerError(f"Embedding model load failed: {e}") from e
|
|
|
|
def _initialize_vector_db(self) -> None:
|
|
logger.info(f"Initializing ChromaDB client at: {self.vector_db_path}")
|
|
try:
|
|
self.vector_db_path.mkdir(parents=True, exist_ok=True)
|
|
self.chroma_client = chromadb.PersistentClient(
|
|
path=str(self.vector_db_path),
|
|
settings=ChromaSettings(anonymized_telemetry=False)
|
|
)
|
|
self.collection = self.chroma_client.get_or_create_collection(
|
|
name=self.collection_name,
|
|
# metadata={"hnsw:space": "cosine"} # Optional: Explicitly set distance metric if needed
|
|
)
|
|
logger.info(f"ChromaDB collection '{self.collection_name}' ready. Initial item count: {self.collection.count()}")
|
|
except Exception as e:
|
|
logger.critical(f"Failed to initialize ChromaDB: {e}", exc_info=True)
|
|
raise ContextManagerError(f"ChromaDB initialization failed: {e}") from e
|
|
|
|
def _load_full_history(self) -> None:
|
|
if self.storage_path.exists() and self.storage_path.is_file():
|
|
try:
|
|
logger.info(f"Loading full conversation history from {self.storage_path}")
|
|
with self.storage_path.open("r", encoding="utf-8") as f:
|
|
content = f.read()
|
|
if not content.strip():
|
|
logger.warning(f"History file '{self.storage_path}' is empty.")
|
|
self.messages = []; return
|
|
|
|
loaded_data = json.loads(content)
|
|
if isinstance(loaded_data, list):
|
|
valid_messages = []
|
|
for i, msg_dict in enumerate(loaded_data):
|
|
if (isinstance(msg_dict, dict) and
|
|
"role" in msg_dict and isinstance(msg_dict["role"], str) and
|
|
"content" in msg_dict and # Allow empty content for system messages potentially
|
|
msg_dict["role"] in ["user", "assistant", "system"]): # Allow system role
|
|
# Ensure 'original_index' is present and correct
|
|
msg_copy = msg_dict.copy()
|
|
msg_copy['original_index'] = i # The index in the loaded list is its original_index
|
|
valid_messages.append(msg_copy)
|
|
else:
|
|
logger.warning(f"Skipping invalid message format at index {i} in history: {msg_dict}")
|
|
self.messages = valid_messages
|
|
logger.info(f"Loaded {len(self.messages)} valid messages from history.")
|
|
else:
|
|
logger.warning(f"History file '{self.storage_path}' not a list. Starting fresh.")
|
|
self.messages = []
|
|
except (json.JSONDecodeError, IOError) as e:
|
|
logger.error(f"Failed to load/parse history '{self.storage_path}' ({type(e).__name__}): {e}. Backing up.", exc_info=False)
|
|
self._backup_corrupted_file(self.storage_path); self.messages = []
|
|
except Exception as e:
|
|
logger.error(f"Unexpected error loading history '{self.storage_path}': {e}. Backing up.", exc_info=True)
|
|
self._backup_corrupted_file(self.storage_path); self.messages = []
|
|
else:
|
|
logger.info(f"History file not found at '{self.storage_path}'. Starting empty history.")
|
|
self.messages = []
|
|
|
|
def _backup_corrupted_file(self, file_path: Path) -> None:
|
|
try:
|
|
backup_dir = file_path.parent / "backups"; backup_dir.mkdir(exist_ok=True)
|
|
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
backup_path = backup_dir / f"{file_path.stem}_corrupted_{timestamp}{file_path.suffix}"
|
|
shutil.move(str(file_path), str(backup_path))
|
|
logger.info(f"Backed up corrupted file to: {backup_path}")
|
|
except Exception as backup_e:
|
|
logger.error(f"Could not back up file '{file_path}': {backup_e}", exc_info=True)
|
|
|
|
def _synchronize_index(self) -> None:
|
|
if not self.collection or not self.embedding_model:
|
|
logger.error("Cannot sync index: DB or model not ready."); return
|
|
logger.info("Synchronizing vector index with loaded history...")
|
|
start_time = time.time()
|
|
messages_to_index: List[Dict[str, Any]] = [] # List of message dicts
|
|
try:
|
|
existing_db_ids_result = self.collection.get(include=[])
|
|
existing_db_ids = set(existing_db_ids_result.get('ids', []))
|
|
logger.debug(f"Found {len(existing_db_ids)} existing IDs in Chroma.")
|
|
|
|
for msg in self.messages:
|
|
# 'original_index' should have been set during _load_full_history or add_message
|
|
msg_original_idx = msg.get('original_index')
|
|
if msg_original_idx is None:
|
|
logger.error(f"Message found without original_index during sync: {str(msg)[:100]}. Skipping.")
|
|
continue
|
|
msg_id = f"msg_{msg_original_idx}"
|
|
if msg_id not in existing_db_ids:
|
|
messages_to_index.append(msg) # msg already contains 'original_index'
|
|
|
|
if not messages_to_index: logger.info("Vector index is synchronized."); return
|
|
logger.info(f"Found {len(messages_to_index)} messages from history to index...")
|
|
|
|
ids_to_add, docs_to_add, metas_to_add = [], [], []
|
|
for message in messages_to_index:
|
|
content = message.get("content", "")
|
|
role = message.get("role", "unknown")
|
|
msg_original_idx = message['original_index'] # Should exist
|
|
if content: # Only index messages with content
|
|
msg_id = f"msg_{msg_original_idx}"
|
|
ids_to_add.append(msg_id)
|
|
docs_to_add.append(content)
|
|
metas_to_add.append({"role": role, "original_index": msg_original_idx})
|
|
|
|
if ids_to_add:
|
|
logger.debug(f"Encoding {len(docs_to_add)} documents for batch indexing...")
|
|
embeddings = self.embedding_model.encode(docs_to_add, show_progress_bar=False).tolist()
|
|
logger.debug(f"Adding {len(ids_to_add)} items to ChromaDB...")
|
|
self.collection.add(ids=ids_to_add, embeddings=embeddings, documents=docs_to_add, metadatas=metas_to_add)
|
|
logger.info(f"Successfully indexed {len(ids_to_add)} messages from history.")
|
|
else:
|
|
logger.info("No valid messages found to index after filtering.")
|
|
except Exception as e:
|
|
logger.error(f"Error during index synchronization: {e}", exc_info=True)
|
|
finally:
|
|
sync_time = time.time() - start_time
|
|
logger.info(f"Index sync check completed in {sync_time:.2f}s.")
|
|
|
|
def _index_message(self, msg_original_index: int, message: Dict[str, Any]) -> None:
|
|
if not self.collection or not self.embedding_model:
|
|
logger.error("Cannot index: DB or model not ready."); return
|
|
msg_id = f"msg_{msg_original_index}"
|
|
content = message.get("content", "")
|
|
role = message.get("role", "unknown")
|
|
if not content: logger.debug(f"Skipping indexing for msg_{msg_original_index} (no content)."); return
|
|
try:
|
|
embedding = self.embedding_model.encode([content], show_progress_bar=False)[0].tolist()
|
|
self.collection.add(
|
|
ids=[msg_id], embeddings=[embedding], documents=[content],
|
|
metadatas=[{"role": role, "original_index": msg_original_index}]
|
|
)
|
|
logger.debug(f"Indexed message: {msg_id} (Role: {role}, OrigIdx: {msg_original_index})")
|
|
except Exception as e:
|
|
logger.error(f"Failed to index msg {msg_id}: {e}", exc_info=True)
|
|
|
|
def add_message(self, role: str, content: str) -> None:
|
|
if role not in ("user", "assistant", "system"):
|
|
raise ValueError(f"Invalid role: '{role}'. Must be 'user', 'assistant', or 'system'.")
|
|
if not isinstance(content, str):
|
|
content = str(content)
|
|
|
|
new_message = {"role": role, "content": content}
|
|
# Assign 'original_index' based on its future position in self.messages
|
|
new_message['original_index'] = len(self.messages)
|
|
self.messages.append(new_message)
|
|
logger.debug(
|
|
f"Added message to memory (OrigIdx: {new_message['original_index']}): Role={role}, Content='{content[:50]}...'"
|
|
)
|
|
self._index_message(new_message['original_index'], new_message)
|
|
|
|
def retrieve_relevant_context(self, query: str, n_results: Optional[int] = None) -> List[Dict[str, Any]]:
|
|
if not self.collection or not self.embedding_model:
|
|
logger.error("Cannot retrieve: DB or model not ready."); return []
|
|
if not query: logger.warning("Cannot retrieve context for empty query."); return []
|
|
|
|
collection_count = self.collection.count()
|
|
if collection_count == 0: logger.debug("Skipping retrieval: Vector store empty."); return []
|
|
|
|
num_to_retrieve = n_results if n_results is not None else self.n_retrieval_results_fallback
|
|
num_to_retrieve = min(num_to_retrieve, collection_count)
|
|
if num_to_retrieve <= 0: logger.debug("No results to retrieve."); return []
|
|
|
|
try:
|
|
logger.info(f"Retrieving up to {num_to_retrieve} messages for query: '{query[:60]}...'")
|
|
start_time = time.time()
|
|
query_embedding = self.embedding_model.encode([query], show_progress_bar=False)[0].tolist()
|
|
results = self.collection.query(
|
|
query_embeddings=[query_embedding], n_results=num_to_retrieve,
|
|
include=["documents", "metadatas", "distances"]
|
|
)
|
|
retrieval_time = time.time() - start_time
|
|
logger.debug(f"ChromaDB query in {retrieval_time:.3f}s.")
|
|
|
|
retrieved_messages = []
|
|
if results and results.get("ids") and results["ids"][0]:
|
|
for i, doc_id in enumerate(results["ids"][0]):
|
|
metadata = results["metadatas"][0][i] if results["metadatas"] and results["metadatas"][0] else {}
|
|
original_idx_val = metadata.get("original_index", -1)
|
|
try: original_idx = int(original_idx_val)
|
|
except (ValueError, TypeError): original_idx = -1; logger.warning(f"Invalid original_index {original_idx_val}")
|
|
|
|
retrieved_messages.append({
|
|
"role": metadata.get("role", "unknown"),
|
|
"content": results["documents"][0][i] if results["documents"] and results["documents"][0] else "",
|
|
"metadata": {"original_index": original_idx, "distance": results["distances"][0][i] if results["distances"] and results["distances"][0] else float('inf')}
|
|
})
|
|
logger.info(f"Retrieved {len(retrieved_messages)} relevant messages.")
|
|
return retrieved_messages
|
|
except Exception as e:
|
|
logger.error(f"Error during context retrieval: {e}", exc_info=True); return []
|
|
|
|
def get_recent_messages(self, num_messages_to_fetch: int) -> List[Dict[str, Any]]:
|
|
"""Gets the last N messages. Ensures 'original_index' is present."""
|
|
if num_messages_to_fetch <= 0: return []
|
|
|
|
start_idx_slice = max(0, len(self.messages) - num_messages_to_fetch)
|
|
recent_slice = self.messages[start_idx_slice:]
|
|
|
|
# Ensure all messages in the slice have 'original_index'.
|
|
# This primarily safeguards against older data formats if any were loaded
|
|
# without 'original_index' (though _load_full_history attempts to add it).
|
|
processed_recent: List[Dict[str, Any]] = []
|
|
for i, msg_dict in enumerate(recent_slice):
|
|
msg_copy = msg_dict.copy() # Work with a copy
|
|
if 'original_index' not in msg_copy or not isinstance(msg_copy['original_index'], int):
|
|
# Fallback: if somehow original_index is missing or invalid from the loaded message
|
|
calculated_original_idx = start_idx_slice + i
|
|
msg_copy['original_index'] = calculated_original_idx
|
|
logger.warning(
|
|
f"ContextManager.get_recent_messages: Re-calculated missing/invalid 'original_index' "
|
|
f"({calculated_original_idx}) for recent message: {str(msg_dict.get('content',''))[:30]}..."
|
|
)
|
|
processed_recent.append(msg_copy)
|
|
|
|
logger.debug(f"Retrieved {len(processed_recent)} messages ({num_messages_to_fetch} requested).")
|
|
return processed_recent
|
|
|
|
def save_context(self) -> None:
|
|
temp_path = self.storage_path.with_suffix(f"{self.storage_path.suffix}.tmp")
|
|
try:
|
|
self.storage_path.parent.mkdir(parents=True, exist_ok=True)
|
|
logger.info(f"Saving full history ({len(self.messages)} messages) to {self.storage_path}")
|
|
# Save role, content, and original_index to allow reconstruction
|
|
messages_to_save = [
|
|
{"role": msg["role"], "content": msg["content"], "original_index": msg.get("original_index", idx)}
|
|
for idx, msg in enumerate(self.messages)
|
|
]
|
|
with temp_path.open("w", encoding="utf-8") as f:
|
|
json.dump(messages_to_save, f, ensure_ascii=False, indent=2)
|
|
shutil.move(str(temp_path), str(self.storage_path))
|
|
logger.info("Full history saved successfully.")
|
|
except Exception as e:
|
|
logger.error(f"Failed to save history '{self.storage_path}': {e}", exc_info=True)
|
|
if temp_path.exists(): temp_path.unlink(missing_ok=True)
|
|
|
|
def clear_context(self) -> None:
|
|
logger.warning("Clearing conversation context (Memory, Vector Store, File)...")
|
|
self.messages = []
|
|
if self.collection:
|
|
try:
|
|
count = self.collection.count()
|
|
if count > 0:
|
|
logger.info(f"Deleting {count} items from Chroma collection '{self.collection_name}'...")
|
|
# Efficient way to clear a Chroma collection (if API supports `delete_collection`)
|
|
# Or, if not, delete all items by IDs.
|
|
# For current chromadb versions, deleting by IDs is standard.
|
|
# If the collection can be deleted and recreated:
|
|
# self.chroma_client.delete_collection(name=self.collection_name)
|
|
# self.collection = self.chroma_client.get_or_create_collection(name=self.collection_name)
|
|
# logger.info(f"Chroma collection '{self.collection_name}' deleted and recreated.")
|
|
# --- OR ---
|
|
all_ids_result = self.collection.get(include=[]) # Only need IDs
|
|
all_ids = all_ids_result.get('ids', [])
|
|
if all_ids:
|
|
self.collection.delete(ids=all_ids)
|
|
logger.info(f"Deleted {len(all_ids)} items from Chroma collection.")
|
|
else:
|
|
logger.info("Chroma collection was already empty (no IDs to delete).")
|
|
else:
|
|
logger.info("Chroma collection already empty.")
|
|
except Exception as e:
|
|
logger.error(f"Failed to clear ChromaDB collection: {e}", exc_info=True)
|
|
self.save_context() # Save the empty state
|
|
logger.warning("Conversation context cleared.")
|
|
|
|
@property
|
|
def history(self) -> List[Dict[str, Any]]:
|
|
return list(self.messages) # Return a shallow copy
|
|
|
|
def shutdown(self) -> None:
|
|
logger.info("ContextManager shutting down...")
|
|
self.embedding_model = None
|
|
self.chroma_client = None # Chroma client usually handles its own persistence
|
|
self.collection = None
|
|
logger.info("ContextManager shutdown complete.") |