lollms-client 0.20.10__py3-none-any.whl → 0.22.0__py3-none-any.whl

lollms_client/lollms_discussion.py

@@ -1,412 +1,664 @@
-# lollms_discussion.py
-
 import yaml
-from dataclasses import dataclass, field
-from typing import List, Dict, Optional, Union, Any
+import json
+import base64
+import os
 import uuid
+import shutil
+import re
 from collections import defaultdict
+from datetime import datetime
+from typing import List, Dict, Optional, Union, Any, Type, Callable
+from pathlib import Path
+from types import SimpleNamespace
+
+from sqlalchemy import (create_engine, Column, String, Text, Integer, DateTime,
+                        ForeignKey, JSON, Boolean, LargeBinary, Index, Float)
+from sqlalchemy.orm import sessionmaker, relationship, Session, declarative_base, declared_attr
+from sqlalchemy.types import TypeDecorator
+from sqlalchemy.orm.exc import NoResultFound
+
+try:
+    from cryptography.fernet import Fernet, InvalidToken
+    from cryptography.hazmat.primitives import hashes
+    from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
+    from cryptography.hazmat.backends import default_backend
+    ENCRYPTION_AVAILABLE = True
+except ImportError:
+    ENCRYPTION_AVAILABLE = False
+
+# Type hint placeholders for classes defined externally
+if False: 
+    from lollms_client import LollmsClient
+    from lollms_client.lollms_types import MSG_TYPE
+    from lollms_personality import LollmsPersonality
+
+class EncryptedString(TypeDecorator):
+    """A SQLAlchemy TypeDecorator for field-level database encryption."""
+    impl = LargeBinary
+    cache_ok = True
+
+    def __init__(self, key: str, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        if not ENCRYPTION_AVAILABLE:
+            raise ImportError("'cryptography' is required for DB encryption.")
+        self.salt = b'lollms-fixed-salt-for-db-encryption'
+        kdf = PBKDF2HMAC(
+            algorithm=hashes.SHA256(), length=32, salt=self.salt,
+            iterations=480000, backend=default_backend()
+        )
+        derived_key = base64.urlsafe_b64encode(kdf.derive(key.encode()))
+        self.fernet = Fernet(derived_key)
 
-# It's good practice to forward-declare the type for the client to avoid circular imports.
-if False:
-    from lollms.client import LollmsClient
+    def process_bind_param(self, value: Optional[str], dialect) -> Optional[bytes]:
+        if value is None:
+            return None
+        return self.fernet.encrypt(value.encode('utf-8'))
 
+    def process_result_value(self, value: Optional[bytes], dialect) -> Optional[str]:
+        if value is None:
+            return None
+        try:
+            return self.fernet.decrypt(value).decode('utf-8')
+        except InvalidToken:
+            return "<DECRYPTION_FAILED: Invalid Key or Corrupt Data>"
+
+def create_dynamic_models(discussion_mixin: Optional[Type] = None, message_mixin: Optional[Type] = None, encryption_key: Optional[str] = None):
+    """Factory to dynamically create SQLAlchemy ORM models with custom mixins."""
+    Base = declarative_base()
+    EncryptedText = EncryptedString(encryption_key) if encryption_key else Text
+
+    class DiscussionBase:
+        __abstract__ = True
+        id = Column(String, primary_key=True, default=lambda: str(uuid.uuid4()))
+        system_prompt = Column(EncryptedText, nullable=True)
+        participants = Column(JSON, nullable=True, default=dict)
+        active_branch_id = Column(String, nullable=True)
+        discussion_metadata = Column(JSON, nullable=True, default=dict)
+        created_at = Column(DateTime, default=datetime.utcnow)
+        updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
+        
+        @declared_attr
+        def messages(cls):
+            return relationship("Message", back_populates="discussion", cascade="all, delete-orphan", lazy="joined")
+
+    class MessageBase:
+        __abstract__ = True
+        id = Column(String, primary_key=True, default=lambda: str(uuid.uuid4()))
+        discussion_id = Column(String, ForeignKey('discussions.id'), nullable=False, index=True)
+        parent_id = Column(String, ForeignKey('messages.id'), nullable=True, index=True)
+        sender = Column(String, nullable=False)
+        sender_type = Column(String, nullable=False)
+        
+        raw_content = Column(EncryptedText, nullable=True)
+        thoughts = Column(EncryptedText, nullable=True)
+        content = Column(EncryptedText, nullable=False)
+        scratchpad = Column(EncryptedText, nullable=True)
+        
+        tokens = Column(Integer, nullable=True)
+        binding_name = Column(String, nullable=True)
+        model_name = Column(String, nullable=True)
+        generation_speed = Column(Float, nullable=True)
+        
+        message_metadata = Column(JSON, nullable=True, default=dict)
+        images = Column(JSON, nullable=True, default=list)
+        created_at = Column(DateTime, default=datetime.utcnow)
+        
+        @declared_attr
+        def discussion(cls):
+            return relationship("Discussion", back_populates="messages")
+        
+    discussion_bases = (discussion_mixin, DiscussionBase, Base) if discussion_mixin else (DiscussionBase, Base)
+    DynamicDiscussion = type('Discussion', discussion_bases, {'__tablename__': 'discussions'})
+
+    message_bases = (message_mixin, MessageBase, Base) if message_mixin else (MessageBase, Base)
+    DynamicMessage = type('Message', message_bases, {'__tablename__': 'messages'})
+    
+    return Base, DynamicDiscussion, DynamicMessage
+
+class LollmsDataManager:
+    """Manages database connection, session, and table creation."""
+    def __init__(self, db_path: str, discussion_mixin: Optional[Type] = None, message_mixin: Optional[Type] = None, encryption_key: Optional[str] = None):
+        if not db_path:
+            raise ValueError("Database path cannot be empty.")
+        self.Base, self.DiscussionModel, self.MessageModel = create_dynamic_models(
+            discussion_mixin, message_mixin, encryption_key
+        )
+        self.engine = create_engine(db_path)
+        self.SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=self.engine)
+        self.create_tables()
+
+    def create_tables(self):
+        self.Base.metadata.create_all(bind=self.engine)
+
+    def get_session(self) -> Session:
+        return self.SessionLocal()
+
+    def list_discussions(self) -> List[Dict]:
+        with self.get_session() as session:
+            discussions = session.query(self.DiscussionModel).all()
+            return [{c.name: getattr(disc, c.name) for c in disc.__table__.columns} for disc in discussions]
+
+    def get_discussion(self, lollms_client: 'LollmsClient', discussion_id: str, **kwargs) -> Optional['LollmsDiscussion']:
+        with self.get_session() as session:
+            try:
+                db_disc = session.query(self.DiscussionModel).filter_by(id=discussion_id).one()
+                session.expunge(db_disc)
+                return LollmsDiscussion(lollmsClient=lollms_client, db_manager=self, db_discussion_obj=db_disc, **kwargs)
+            except NoResultFound:
+                return None
+
+    def search_discussions(self, **criteria) -> List[Dict]:
+        with self.get_session() as session:
+            query = session.query(self.DiscussionModel)
+            for key, value in criteria.items():
+                if hasattr(self.DiscussionModel, key):
+                    query = query.filter(getattr(self.DiscussionModel, key).ilike(f"%{value}%"))
+            discussions = query.all()
+            return [{c.name: getattr(disc, c.name) for c in disc.__table__.columns} for disc in discussions]
+
+    def delete_discussion(self, discussion_id: str):
+        with self.get_session() as session:
+            db_disc = session.query(self.DiscussionModel).filter_by(id=discussion_id).first()
+            if db_disc:
+                session.delete(db_disc)
+                session.commit()
 
-@dataclass
 class LollmsMessage:
-    """
-    Represents a single message in a LollmsDiscussion, including its content,
-    sender, and relationship within the discussion tree.
-    """
-    sender: str
-    sender_type: str
-    content: str
-    id: str = field(default_factory=lambda: str(uuid.uuid4()))
-    parent_id: Optional[str] = None
-    metadata: str = "{}"
-    images: List[Dict[str, str]] = field(default_factory=list)
-
-    def to_dict(self) -> Dict[str, Any]:
-        """Serializes the message object to a dictionary."""
-        return {
-            'sender': self.sender,
-            'sender_type': self.sender_type,
-            'content': self.content,
-            'id': self.id,
-            'parent_id': self.parent_id,
-            'metadata': self.metadata,
-            'images': self.images
-        }
+    """A wrapper for a message ORM object, providing direct attribute access."""
+    def __init__(self, discussion: 'LollmsDiscussion', db_message: Any):
+        object.__setattr__(self, '_discussion', discussion)
+        object.__setattr__(self, '_db_message', db_message)
+
+    def __getattr__(self, name: str) -> Any:
+        if name == 'metadata':
+            return getattr(self._db_message, 'message_metadata', None)
+        return getattr(self._db_message, name)
+
+    def __setattr__(self, name: str, value: Any):
+        if name == 'metadata':
+            setattr(self._db_message, 'message_metadata', value)
+        else:
+            setattr(self._db_message, name, value)
+        self._discussion.touch()
 
+    def __repr__(self) -> str:
+        return f"<LollmsMessage id={self.id} sender='{self.sender}'>"
 
 class LollmsDiscussion:
-    """
-    Manages a branching conversation tree, including system prompts, participants,
-    an internal knowledge scratchpad, and context pruning capabilities.
-    """
-
-    def __init__(self, lollmsClient: 'LollmsClient'):
-        """
-        Initializes a new LollmsDiscussion instance.
-
-        Args:
-            lollmsClient: An instance of LollmsClient, required for tokenization.
-        """
-        self.lollmsClient = lollmsClient
-        self.version: int = 3  # Current version of the format with scratchpad support
-        self._reset_state()
-
-    def _reset_state(self):
-        """Helper to reset all discussion attributes to their defaults."""
-        self.messages: List[LollmsMessage] = []
-        self.active_branch_id: Optional[str] = None
-        self.message_index: Dict[str, LollmsMessage] = {}
-        self.children_index: Dict[Optional[str], List[str]] = defaultdict(list)
-        self.participants: Dict[str, str] = {}
-        self.system_prompt: Optional[str] = None
-        self.scratchpad: Optional[str] = None
-
-    # --- Scratchpad Management Methods ---
-    def set_scratchpad(self, content: str):
-        """Sets or replaces the entire content of the internal scratchpad."""
-        self.scratchpad = content
-
-    def update_scratchpad(self, new_content: str, append: bool = True):
-        """
-        Updates the scratchpad. By default, it appends with a newline separator.
-
-        Args:
-            new_content: The new text to add to the scratchpad.
-            append: If True, appends to existing content. If False, replaces it.
-        """
-        if append and self.scratchpad:
-            self.scratchpad += f"\n{new_content}"
-        else:
-            self.scratchpad = new_content
-
-    def get_scratchpad(self) -> Optional[str]:
-        """Returns the current content of the scratchpad."""
-        return self.scratchpad
-
-    def clear_scratchpad(self):
-        """Clears the scratchpad content."""
-        self.scratchpad = None
-
-    # --- Configuration Methods ---
-    def set_system_prompt(self, prompt: str):
-        """Sets the main system prompt for the discussion."""
-        self.system_prompt = prompt
-
-    def set_participants(self, participants: Dict[str, str]):
-        """
-        Defines the participants and their roles ('user' or 'assistant').
-
-        Args:
-            participants: A dictionary mapping sender names to roles.
-        """
-        for name, role in participants.items():
-            if role not in ["user", "assistant"]:
-                raise ValueError(f"Invalid role '{role}' for participant '{name}'")
-        self.participants = participants
-
-    # --- Core Message Tree Methods ---
-    def add_message(
-        self,
-        sender: str,
-        sender_type: str,
-        content: str,
-        metadata: Optional[Dict] = None,
-        parent_id: Optional[str] = None,
-        images: Optional[List[Dict[str, str]]] = None,
-        override_id: Optional[str] = None
-    ) -> str:
-        """
-        Adds a new message to the discussion tree.
-        """
-        if parent_id is None:
-            parent_id = self.active_branch_id
-        if parent_id is None:
-            parent_id = "main_root"
-
-        message = LollmsMessage(
-            sender=sender, sender_type=sender_type, content=content,
-            parent_id=parent_id, metadata=str(metadata or {}), images=images or []
-        )
-        if override_id:
-            message.id = override_id
-
-        self.messages.append(message)
-        self.message_index[message.id] = message
-        self.children_index[parent_id].append(message.id)
-        self.active_branch_id = message.id
-        return message.id
-
-    def get_branch(self, leaf_id: str) -> List[LollmsMessage]:
-        """Gets the full branch of messages from the root to the specified leaf."""
-        branch = []
-        current_id: Optional[str] = leaf_id
-        while current_id and current_id in self.message_index:
-            msg = self.message_index[current_id]
-            branch.append(msg)
-            current_id = msg.parent_id
-        return list(reversed(branch))
-
-    def set_active_branch(self, message_id: str):
-        """Sets the active message, effectively switching to a different branch."""
-        if message_id not in self.message_index:
-            raise ValueError(f"Message ID {message_id} not found in discussion.")
-        self.active_branch_id = message_id
-
-    # --- Persistence ---
-    def save_to_disk(self, file_path: str):
-        """Saves the entire discussion state to a YAML file."""
-        data = {
-            'version': self.version, 'active_branch_id': self.active_branch_id,
-            'system_prompt': self.system_prompt, 'participants': self.participants,
-            'scratchpad': self.scratchpad, 'messages': [m.to_dict() for m in self.messages]
-        }
-        with open(file_path, 'w', encoding='utf-8') as file:
-            yaml.dump(data, file, allow_unicode=True, sort_keys=False)
-
-    def load_from_disk(self, file_path: str):
-        """Loads a discussion state from a YAML file."""
-        with open(file_path, 'r', encoding='utf-8') as file:
-            data = yaml.safe_load(file)
-
-        self._reset_state()
-        version = data.get("version", 1)
-        if version > self.version:
-            raise ValueError(f"File version {version} is newer than supported version {self.version}.")
-
-        self.active_branch_id = data.get('active_branch_id')
-        self.system_prompt = data.get('system_prompt', None)
-        self.participants = data.get('participants', {})
-        self.scratchpad = data.get('scratchpad', None)
-
-        for msg_data in data.get('messages', []):
-            msg = LollmsMessage(
-                sender=msg_data['sender'], sender_type=msg_data.get('sender_type', 'user'),
-                content=msg_data['content'], parent_id=msg_data.get('parent_id'),
-                id=msg_data.get('id', str(uuid.uuid4())), metadata=msg_data.get('metadata', '{}'),
-                images=msg_data.get('images', [])
-            )
-            self.messages.append(msg)
-            self.message_index[msg.id] = msg
-            self.children_index[msg.parent_id].append(msg.id)
-
-    # --- Context Management and Formatting ---
-    def _get_full_system_prompt(self) -> Optional[str]:
-        """Combines the scratchpad and system prompt into a single string for the LLM."""
-        full_sys_prompt_parts = []
-        if self.scratchpad and self.scratchpad.strip():
-            full_sys_prompt_parts.append("--- KNOWLEDGE SCRATCHPAD ---")
-            full_sys_prompt_parts.append(self.scratchpad.strip())
-            full_sys_prompt_parts.append("--- END SCRATCHPAD ---")
+    """Represents and manages a single discussion, acting as a high-level interface."""
+    def __init__(self, lollmsClient: 'LollmsClient', db_manager: Optional[LollmsDataManager] = None, 
+                 discussion_id: Optional[str] = None, db_discussion_obj: Optional[Any] = None,
+                 autosave: bool = False, max_context_size: Optional[int] = None):
         
-        if self.system_prompt and self.system_prompt.strip():
-            full_sys_prompt_parts.append(self.system_prompt.strip())
-        return "\n\n".join(full_sys_prompt_parts) if full_sys_prompt_parts else None
-
-    def summarize_and_prune(self, max_tokens: int, preserve_last_n: int = 4, branch_tip_id: Optional[str] = None) -> Dict[str, Any]:
-        """
-        Checks context size and, if exceeded, summarizes the oldest messages
-        into the scratchpad and prunes them to free up token space.
-        """
-        if branch_tip_id is None: branch_tip_id = self.active_branch_id
-        if not branch_tip_id: return {"pruned": False, "reason": "No active branch."}
-
-        full_prompt_text = self.export("lollms_text", branch_tip_id)
-        current_tokens = len(self.lollmsClient.binding.tokenize(full_prompt_text))
-        if current_tokens <= max_tokens: return {"pruned": False, "reason": "Token count within limit."}
-
-        branch = self.get_branch(branch_tip_id)
-        if len(branch) <= preserve_last_n: return {"pruned": False, "reason": "Not enough messages to prune."}
-
-        messages_to_prune = branch[:-preserve_last_n]
-        messages_to_keep = branch[-preserve_last_n:]
-        text_to_summarize = "\n\n".join([f"{self.participants.get(m.sender, 'user').capitalize()}: {m.content}" for m in messages_to_prune])
+        object.__setattr__(self, 'lollmsClient', lollmsClient)
+        object.__setattr__(self, 'db_manager', db_manager)
+        object.__setattr__(self, 'autosave', autosave)
+        object.__setattr__(self, 'max_context_size', max_context_size)
+        object.__setattr__(self, 'scratchpad', "")
+        object.__setattr__(self, 'show_thoughts', False)
+        object.__setattr__(self, 'include_thoughts_in_context', False)
+        object.__setattr__(self, 'thought_placeholder', "<thought process hidden>")
         
-        summary_prompt = (
-            "You are a summarization expert. Read the following conversation excerpt and create a "
-            "concise, factual summary of all key information, decisions, and outcomes. This summary "
-            "will be placed in a knowledge scratchpad for future reference. Omit conversational filler.\n\n"
-            f"CONVERSATION EXCERPT:\n---\n{text_to_summarize}\n---\n\nCONCISE SUMMARY:"
-        )
+        object.__setattr__(self, '_session', None)
+        object.__setattr__(self, '_db_discussion', None)
+        object.__setattr__(self, '_message_index', None)
+        object.__setattr__(self, '_messages_to_delete_from_db', set())
+        object.__setattr__(self, '_is_db_backed', db_manager is not None)
+
+        if self._is_db_backed:
+            if not db_discussion_obj and not discussion_id:
+                raise ValueError("Either discussion_id or db_discussion_obj must be provided for DB-backed discussions.")
+            
+            self._session = db_manager.get_session()
+            if db_discussion_obj:
+                self._db_discussion = self._session.merge(db_discussion_obj)
+            else:
+                try:
+                    self._db_discussion = self._session.query(db_manager.DiscussionModel).filter_by(id=discussion_id).one()
+                except NoResultFound:
+                    self._session.close()
+                    raise ValueError(f"No discussion found with ID: {discussion_id}")
+        else:
+            self._create_in_memory_proxy(id=discussion_id)
+        self._rebuild_message_index()
+    
+    @property
+    def remaining_tokens(self) -> Optional[int]:
+        """Calculates the remaining tokens available in the context window."""
+        binding = self.lollmsClient.binding
+        if not binding or not hasattr(binding, 'ctx_size') or not binding.ctx_size:
+            return None
+        max_ctx = binding.ctx_size
+        current_prompt = self.format_discussion(max_ctx)
+        current_tokens = self.lollmsClient.count_tokens(current_prompt)
+        return max_ctx - current_tokens
+
+    @classmethod
+    def create_new(cls, lollms_client: 'LollmsClient', db_manager: Optional[LollmsDataManager] = None, **kwargs) -> 'LollmsDiscussion':
+        init_args = {
+            'autosave': kwargs.pop('autosave', False),
+            'max_context_size': kwargs.pop('max_context_size', None)
+        }
+        if db_manager:
+            with db_manager.get_session() as session:
+                valid_keys = db_manager.DiscussionModel.__table__.columns.keys()
+                db_creation_args = {k: v for k, v in kwargs.items() if k in valid_keys}
+                db_discussion_orm = db_manager.DiscussionModel(**db_creation_args)
+                session.add(db_discussion_orm)
+                session.commit()
+                session.expunge(db_discussion_orm)
+            return cls(lollmsClient=lollms_client, db_manager=db_manager, db_discussion_obj=db_discussion_orm, **init_args)
+        else:
+            return cls(lollmsClient=lollms_client, discussion_id=kwargs.get('id'), **init_args)
+
+    def __getattr__(self, name: str) -> Any:
+        if name == 'metadata':
+            return getattr(self._db_discussion, 'discussion_metadata', None)
+        if name == 'messages':
+            return [LollmsMessage(self, msg) for msg in self._db_discussion.messages]
+        return getattr(self._db_discussion, name)
+
+    def __setattr__(self, name: str, value: Any):
+        internal_attrs = [
+            'lollmsClient','db_manager','autosave','max_context_size','scratchpad',
+            'show_thoughts', 'include_thoughts_in_context', 'thought_placeholder',
+            '_session','_db_discussion','_message_index','_messages_to_delete_from_db', '_is_db_backed'
+        ]
+        if name in internal_attrs:
+            object.__setattr__(self, name, value)
+        else:
+            if name == 'metadata':
+                setattr(self._db_discussion, 'discussion_metadata', value)
+            else:
+                setattr(self._db_discussion, name, value)
+            self.touch()
+    
+    def _create_in_memory_proxy(self, id: Optional[str] = None):
+        proxy = SimpleNamespace()
+        proxy.id, proxy.system_prompt, proxy.participants = id or str(uuid.uuid4()), None, {}
+        proxy.active_branch_id, proxy.discussion_metadata = None, {}
+        proxy.created_at, proxy.updated_at = datetime.utcnow(), datetime.utcnow()
+        proxy.messages = []
+        object.__setattr__(self, '_db_discussion', proxy)
+    
+    def _rebuild_message_index(self):
+        if self._is_db_backed and self._session.is_active and self._db_discussion in self._session:
+            self._session.refresh(self._db_discussion, ['messages'])
+        self._message_index = {msg.id: msg for msg in self._db_discussion.messages}
+
+    def touch(self):
+        setattr(self._db_discussion, 'updated_at', datetime.utcnow())
+        if self._is_db_backed and self.autosave:
+            self.commit()
+
+    def commit(self):
+        if not self._is_db_backed or not self._session:
+            return
+        if self._messages_to_delete_from_db:
+            for msg_id in self._messages_to_delete_from_db:
+                msg_to_del = self._session.get(self.db_manager.MessageModel, msg_id)
+                if msg_to_del:
+                    self._session.delete(msg_to_del)
+            self._messages_to_delete_from_db.clear()
         try:
-            summary = self.lollmsClient.generate_text(summary_prompt, max_new_tokens=300, temperature=0.1)
+            self._session.commit()
+            self._rebuild_message_index()
         except Exception as e:
-            return {"pruned": False, "reason": f"Failed to generate summary: {e}"}
-
-        summary_block = f"--- Summary of earlier conversation (pruned on {uuid.uuid4().hex[:8]}) ---\n{summary.strip()}"
-        self.update_scratchpad(summary_block, append=True)
-
-        ids_to_prune = {msg.id for msg in messages_to_prune}
-        new_root_of_branch = messages_to_keep[0]
-        original_parent_id = messages_to_prune[0].parent_id
-
-        self.message_index[new_root_of_branch.id].parent_id = original_parent_id
-        if original_parent_id in self.children_index:
-            self.children_index[original_parent_id] = [mid for mid in self.children_index[original_parent_id] if mid != messages_to_prune[0].id]
-            self.children_index[original_parent_id].append(new_root_of_branch.id)
-
-        for msg_id in ids_to_prune:
-            self.message_index.pop(msg_id, None)
-            self.children_index.pop(msg_id, None)
-        self.messages = [m for m in self.messages if m.id not in ids_to_prune]
-
-        new_prompt_text = self.export("lollms_text", branch_tip_id)
-        new_tokens = len(self.lollmsClient.binding.tokenize(new_prompt_text))
-        return {"pruned": True, "tokens_saved": current_tokens - new_tokens, "summary_added": True}
-
-    def format_discussion(self, max_allowed_tokens: int, splitter_text: str = "!@>", branch_tip_id: Optional[str] = None) -> str:
-        """
-        Formats the discussion into a single string for instruct models,
-        truncating from the start to respect the token limit.
-
-        Args:
-            max_allowed_tokens: The maximum token limit for the final prompt.
-            splitter_text: The separator token to use (e.g., '!@>').
-            branch_tip_id: The ID of the branch to format. Defaults to active.
-
-        Returns:
-            A single, truncated prompt string.
-        """
-        if branch_tip_id is None:
-            branch_tip_id = self.active_branch_id
+            self._session.rollback()
+            raise e
+
+    def close(self):
+        if self._session:
+            self.commit()
+            self._session.close()
+
+    def add_message(self, **kwargs) -> LollmsMessage:
+        msg_id, parent_id = kwargs.get('id', str(uuid.uuid4())), kwargs.get('parent_id', self.active_branch_id)
+        message_data = {'id': msg_id, 'parent_id': parent_id, 'discussion_id': self.id, 'created_at': datetime.utcnow(), **kwargs}
+        if 'metadata' in message_data:
+            message_data['message_metadata'] = message_data.pop('metadata')
+        if self._is_db_backed:
+            valid_keys = {c.name for c in self.db_manager.MessageModel.__table__.columns}
+            filtered_data = {k: v for k, v in message_data.items() if k in valid_keys}
+            new_msg_orm = self.db_manager.MessageModel(**filtered_data)
+            self._db_discussion.messages.append(new_msg_orm)
+            if new_msg_orm not in self._session:
+                self._session.add(new_msg_orm)
+        else:
+            new_msg_orm = SimpleNamespace(**message_data)
+            self._db_discussion.messages.append(new_msg_orm)
+        self._message_index[msg_id], self.active_branch_id = new_msg_orm, msg_id
+        self.touch()
+        return LollmsMessage(self, new_msg_orm)
         
-        branch_msgs = self.get_branch(branch_tip_id) if branch_tip_id else []
-        full_system_prompt = self._get_full_system_prompt()
+    def get_branch(self, leaf_id: Optional[str]) -> List[LollmsMessage]:
+        if not leaf_id:
+            return []
+        branch_orms, current_id = [], leaf_id
+        while current_id and current_id in self._message_index:
+            msg_orm = self._message_index[current_id]
+            branch_orms.append(msg_orm)
+            current_id = msg_orm.parent_id
+        return [LollmsMessage(self, orm) for orm in reversed(branch_orms)]
+
+    def chat(self, user_message: str, personality: Optional['LollmsPersonality'] = None, **kwargs) -> LollmsMessage:
+        if self.max_context_size is not None:
+            self.summarize_and_prune(self.max_context_size)
+
+        if user_message:
+            self.add_message(sender="user", sender_type="user", content=user_message)
+
+        rag_context = None
+        original_system_prompt = self.system_prompt
+        if personality:
+            self.system_prompt = personality.system_prompt
+            if user_message:
+                rag_context = personality.get_rag_context(user_message)
         
-        prompt_parts = []
-        current_tokens = 0
+        if rag_context:
+            self.system_prompt = f"{original_system_prompt or ''}\n\n--- Relevant Information ---\n{rag_context}\n---"
+
+        from lollms_client.lollms_types import MSG_TYPE
+        is_streaming = "streaming_callback" in kwargs and kwargs.get("streaming_callback") is not None
         
-        # Start with the system prompt if defined
-        if full_system_prompt:
-            sys_msg_text = f"{splitter_text}system:\n{full_system_prompt}\n"
-            sys_tokens = len(self.lollmsClient.binding.tokenize(sys_msg_text))
-            if sys_tokens <= max_allowed_tokens:
-                prompt_parts.append(sys_msg_text)
-                current_tokens += sys_tokens
+        final_raw_response = ""
+        start_time = datetime.now()
+
+        if personality and personality.script_module and hasattr(personality.script_module, 'run'):
+            try:
+                print(f"[{personality.name}] Running custom script...")
+                final_raw_response = personality.script_module.run(self, kwargs.get("streaming_callback"))
+            except Exception as e:
+                print(f"[{personality.name}] Error in custom script: {e}")
+                final_raw_response = f"Error executing personality script: {e}"
+        else:
+            raw_response_accumulator = []
+            if is_streaming:
+                full_response_parts, token_buffer, in_thought_block = [], "", False
+                original_callback = kwargs.get("streaming_callback")
+                def accumulating_callback(token: str, msg_type: MSG_TYPE = MSG_TYPE.MSG_TYPE_CHUNK):
+                    nonlocal token_buffer, in_thought_block
+                    raw_response_accumulator.append(token)
+                    continue_streaming = True
+                    if token: token_buffer += token
+                    while True:
+                        if in_thought_block:
+                            end_tag_pos = token_buffer.find("</think>")
+                            if end_tag_pos != -1:
+                                thought_chunk = token_buffer[:end_tag_pos]
+                                if self.show_thoughts and original_callback and thought_chunk:
+                                    if not original_callback(thought_chunk, MSG_TYPE.MSG_TYPE_THOUGHT_CHUNK): continue_streaming = False
+                                in_thought_block, token_buffer = False, token_buffer[end_tag_pos + len("</think>"):]
+                            else:
+                                if self.show_thoughts and original_callback and token_buffer:
+                                    if not original_callback(token_buffer, MSG_TYPE.MSG_TYPE_THOUGHT_CHUNK): continue_streaming = False
+                                token_buffer = ""; break
+                        else:
+                            start_tag_pos = token_buffer.find("<think>")
+                            if start_tag_pos != -1:
+                                response_chunk = token_buffer[:start_tag_pos]
+                                if response_chunk:
+                                    full_response_parts.append(response_chunk)
+                                    if original_callback:
+                                        if not original_callback(response_chunk, MSG_TYPE.MSG_TYPE_CHUNK): continue_streaming = False
+                                in_thought_block, token_buffer = True, token_buffer[start_tag_pos + len("<think>"):]
+                            else:
+                                if token_buffer:
+                                    full_response_parts.append(token_buffer)
+                                    if original_callback:
+                                        if not original_callback(token_buffer, MSG_TYPE.MSG_TYPE_CHUNK): continue_streaming = False
+                                token_buffer = ""; break
+                    return continue_streaming
+                kwargs["streaming_callback"], kwargs["stream"] = accumulating_callback, True
+                self.lollmsClient.chat(self, **kwargs)
+                final_raw_response = "".join(raw_response_accumulator)
+            else:
+                kwargs["stream"] = False
+                final_raw_response = self.lollmsClient.chat(self, **kwargs) or ""
+
+        end_time = datetime.now()
+        if rag_context:
+            self.system_prompt = original_system_prompt
+        
+        duration = (end_time - start_time).total_seconds()
+        thoughts_match = re.search(r"<think>(.*?)</think>", final_raw_response, re.DOTALL)
+        thoughts_text = thoughts_match.group(1).strip() if thoughts_match else None
+        final_content = self.lollmsClient.remove_thinking_blocks(final_raw_response)
+        token_count = self.lollmsClient.count_tokens(final_content)
+        tok_per_sec = (token_count / duration) if duration > 0 else 0
         
-        # Iterate from newest to oldest to fill the remaining context
-        for msg in reversed(branch_msgs):
-            sender_str = msg.sender.replace(':', '').replace(splitter_text, '')
-            content = msg.content.strip()
-            if msg.images:
-                content += f"\n({len(msg.images)} image(s) attached)"
+        ai_message_obj = self.add_message(
+            sender="assistant", sender_type="assistant", content=final_content,
+            raw_content=final_raw_response, thoughts=thoughts_text, tokens=token_count,
+            binding_name=self.lollmsClient.binding.binding_name, model_name=self.lollmsClient.binding.model_name,
+            generation_speed=tok_per_sec
+        )
 
-            msg_text = f"{splitter_text}{sender_str}:\n{content}\n"
-            msg_tokens = len(self.lollmsClient.binding.tokenize(msg_text))
+        if self._is_db_backed and not self.autosave:
+            self.commit()
+        return ai_message_obj
+
+    def process_and_summarize(self, large_text: str, user_prompt: str, chunk_size: int = 4096, **kwargs) -> LollmsMessage:
+        user_msg = self.add_message(sender="user", sender_type="user", content=user_prompt)
+        chunks = [large_text[i:i + chunk_size] for i in range(0, len(large_text), chunk_size)]
+        current_summary, total_chunks = "", len(chunks)
+        for i, chunk in enumerate(chunks):
+            print(f"\nProcessing chunk {i+1}/{total_chunks}...")
+            if i == 0:
+                prompt = f"""The user wants to know: "{user_prompt}"\nHere is the first part of the document (chunk 1 of {total_chunks}). \nRead it and create a detailed summary of all information relevant to the user's prompt.\n\nDOCUMENT CHUNK:\n---\n{chunk}\n---\nSUMMARY:"""
+            else:
+                prompt = f"""The user wants to know: "{user_prompt}"\nYou are processing a large document sequentially. Here is the summary of the previous chunks and the content of the next chunk ({i+1} of {total_chunks}).\nUpdate your summary by integrating new relevant information from the new chunk. Do not repeat information you already have. Output ONLY the new, updated, complete summary.\n\nPREVIOUS SUMMARY:\n---\n{current_summary}\n---\n\nNEW DOCUMENT CHUNK:\n---\n{chunk}\n---\nUPDATED SUMMARY:"""
+            current_summary = self.lollmsClient.generate_text(prompt, **kwargs).strip()
+        final_prompt = f"""Based on the following comprehensive summary of a document, provide a final answer to the user's original prompt.\nUser's prompt: "{user_prompt}"\n\nCOMPREHENSIVE SUMMARY:\n---\n{current_summary}\n---\nFINAL ANSWER:"""
+        final_answer = self.lollmsClient.generate_text(final_prompt, **kwargs).strip()
+        ai_message_obj = self.add_message(
+            sender="assistant", sender_type="assistant", content=final_answer,
+            scratchpad=current_summary, parent_id=user_msg.id
+        )
+        if self._is_db_backed and not self.autosave:
+            self.commit()
+        return ai_message_obj
+
+    def regenerate_branch(self, **kwargs) -> LollmsMessage:
+        if not self.active_branch_id or self.active_branch_id not in self._message_index:
+            raise ValueError("No active message to regenerate from.")
+        last_message_orm = self._message_index[self.active_branch_id]
+        if last_message_orm.sender_type != 'assistant':
+            raise ValueError("Can only regenerate from an assistant's message.")
+        parent_id, last_message_id = last_message_orm.parent_id, last_message_orm.id
+        self._db_discussion.messages.remove(last_message_orm)
+        del self._message_index[last_message_id]
+        if self._is_db_backed:
+            self._messages_to_delete_from_db.add(last_message_id)
+        self.active_branch_id = parent_id
+        self.touch()
+        return self.chat("", **kwargs)
+
+    def delete_branch(self, message_id: str):
+        if not self._is_db_backed:
+            raise NotImplementedError("Branch deletion is only supported for database-backed discussions.")
+        if message_id not in self._message_index:
+            raise ValueError("Message not found.")
+        msg_to_delete = self._session.query(self.db_manager.MessageModel).filter_by(id=message_id).first()
+        if msg_to_delete:
+            self.active_branch_id = msg_to_delete.parent_id
+            self._session.delete(msg_to_delete)
+            self.commit()
+
+    def switch_to_branch(self, message_id: str):
+        if message_id not in self._message_index:
+            raise ValueError(f"Message ID '{message_id}' not found in the current discussion.")
+        self.active_branch_id = message_id
+        self.touch()
 
-            if current_tokens + msg_tokens > max_allowed_tokens:
-                break # Stop if adding the next message exceeds the limit
+    def format_discussion(self, max_allowed_tokens: int, branch_tip_id: Optional[str] = None) -> str:
+        return self.export("lollms_text", branch_tip_id, max_allowed_tokens)
 
-            prompt_parts.insert(1 if full_system_prompt else 0, msg_text) # Prepend after system prompt
-            current_tokens += msg_tokens
+    def _get_full_system_prompt(self) -> Optional[str]:
+        parts = []
+        if self.scratchpad:
+            parts.extend(["--- KNOWLEDGE SCRATCHPAD ---", self.scratchpad.strip(), "--- END SCRATCHPAD ---"])
+        if self.system_prompt and self.system_prompt.strip():
+            parts.append(self.system_prompt.strip())
+        return "\n\n".join(parts) if parts else None
+
+    def export(self, format_type: str, branch_tip_id: Optional[str] = None, max_allowed_tokens: Optional[int] = None) -> Union[List[Dict], str]:
+        branch_tip_id = branch_tip_id or self.active_branch_id
+        if not branch_tip_id and format_type in ["lollms_text", "openai_chat", "ollama_chat"]:
+            return "" if format_type == "lollms_text" else []
+        branch, full_system_prompt, participants = self.get_branch(branch_tip_id), self._get_full_system_prompt(), self.participants or {}
+
+        def get_full_content(msg: LollmsMessage) -> str:
+            content_to_use = msg.content
+            if self.include_thoughts_in_context and msg.sender_type == 'assistant' and msg.raw_content:
+                if self.thought_placeholder:
+                    content_to_use = re.sub(r"<think>.*?</think>", f"<think>{self.thought_placeholder}</think>", msg.raw_content, flags=re.DOTALL)
+                else:
+                    content_to_use = msg.raw_content
             
-        return "".join(prompt_parts).strip()
-
-
-    def export(self, format_type: str, branch_tip_id: Optional[str] = None) -> Union[List[Dict], str]:
-        """
-        Exports the full, untruncated discussion history in a specific format.
-        """
-        if branch_tip_id is None: branch_tip_id = self.active_branch_id
-        if branch_tip_id is None and not self._get_full_system_prompt(): return "" if format_type in ["lollms_text", "openai_completion"] else []
-
-        branch = self.get_branch(branch_tip_id) if branch_tip_id else []
-        full_system_prompt = self._get_full_system_prompt()
-
-        if format_type == "openai_chat":
-            messages = []
-            if full_system_prompt: messages.append({"role": "system", "content": full_system_prompt})
-            def openai_image_block(image: Dict[str, str]) -> Dict:
-                image_url = image['data'] if image['type'] == 'url' else f"data:image/jpeg;base64,{image['data']}"
-                return {"type": "image_url", "image_url": {"url": image_url, "detail": "auto"}}
-            for msg in branch:
-                role = self.participants.get(msg.sender, "user")
+            parts = [f"--- Internal Scratchpad ---\n{msg.scratchpad.strip()}\n---"] if msg.scratchpad and msg.scratchpad.strip() else []
+            parts.append(content_to_use.strip())
+            return "\n".join(parts)
+
+        if format_type == "lollms_text":
+            prompt_parts, current_tokens = [], 0
+            if full_system_prompt:
+                sys_msg_text = f"!@>system:\n{full_system_prompt}\n"
+                sys_tokens = self.lollmsClient.count_tokens(sys_msg_text)
+                if max_allowed_tokens is None or sys_tokens <= max_allowed_tokens:
+                    prompt_parts.append(sys_msg_text)
+                    current_tokens += sys_tokens
+            for msg in reversed(branch):
+                sender_str = msg.sender.replace(':', '').replace('!@>', '')
+                content = get_full_content(msg)
                 if msg.images:
-                    content_parts = [{"type": "text", "text": msg.content.strip()}] if msg.content.strip() else []
-                    content_parts.extend(openai_image_block(img) for img in msg.images)
+                    content += f"\n({len(msg.images)} image(s) attached)"
+                msg_text = f"!@>{sender_str}:\n{content}\n"
+                msg_tokens = self.lollmsClient.count_tokens(msg_text)
+                if max_allowed_tokens is not None and current_tokens + msg_tokens > max_allowed_tokens:
+                    break
+                prompt_parts.insert(1 if full_system_prompt else 0, msg_text)
+                current_tokens += msg_tokens
+            return "".join(prompt_parts).strip()
+        
+        messages = []
+        if full_system_prompt:
+            messages.append({"role": "system", "content": full_system_prompt})
+        for msg in branch:
+            role, content, images = participants.get(msg.sender, "user"), get_full_content(msg), msg.images or []
+            if format_type == "openai_chat":
+                if images:
+                    content_parts = [{"type": "text", "text": content}] if content else []
+                    for img in images:
+                        content_parts.append({"type": "image_url", "image_url": {"url": img['data'] if img['type'] == 'url' else f"data:image/jpeg;base64,{img['data']}", "detail": "auto"}})
                     messages.append({"role": role, "content": content_parts})
-                else: messages.append({"role": role, "content": msg.content.strip()})
-            return messages
-
-        elif format_type == "ollama_chat":
-            messages = []
-            if full_system_prompt: messages.append({"role": "system", "content": full_system_prompt})
-            for msg in branch:
-                role = self.participants.get(msg.sender, "user")
-                message_dict = {"role": role, "content": msg.content.strip()}
-                ollama_images = [img['data'] for img in msg.images if img['type'] == 'base64']
-                if ollama_images: message_dict["images"] = ollama_images
+                else:
+                    messages.append({"role": role, "content": content})
+            elif format_type == "ollama_chat":
+                message_dict = {"role": role, "content": content}
+                base64_images = [img['data'] for img in images if img['type'] == 'base64']
+                if base64_images:
+                    message_dict["images"] = base64_images
                 messages.append(message_dict)
-            return messages
+            else:
+                raise ValueError(f"Unsupported export format_type: {format_type}")
+        return messages
+
+    def summarize_and_prune(self, max_tokens: int, preserve_last_n: int = 4):
+        branch_tip_id = self.active_branch_id
+        if not branch_tip_id:
+            return
+        current_tokens = self.lollmsClient.count_tokens(self.format_discussion(999999, branch_tip_id))
+        if current_tokens <= max_tokens:
+            return
+        branch = self.get_branch(branch_tip_id)
+        if len(branch) <= preserve_last_n:
+            return
+        messages_to_prune = branch[:-preserve_last_n]
+        text_to_summarize = "\n\n".join([f"{m.sender}: {m.content}" for m in messages_to_prune])
+        summary_prompt = f"Concisely summarize this conversation excerpt:\n---\n{text_to_summarize}\n---\nSUMMARY:"
+        try:
+            summary = self.lollmsClient.generate_text(summary_prompt, n_predict=300, temperature=0.1)
+        except Exception as e:
+            print(f"\n[WARNING] Pruning failed, couldn't generate summary: {e}")
+            return
+        self.scratchpad = f"{self.scratchpad}\n\n--- Summary of earlier conversation ---\n{summary.strip()}".strip()
+        pruned_ids = {msg.id for msg in messages_to_prune}
+        if self._is_db_backed:
+            self._messages_to_delete_from_db.update(pruned_ids)
+            self._db_discussion.messages = [m for m in self._db_discussion.messages if m.id not in pruned_ids]
+        else:
+            self._db_discussion.messages = [m for m in self._db_discussion.messages if m.id not in pruned_ids]
+        self._rebuild_message_index()
+        self.touch()
+        print(f"\n[INFO] Discussion auto-pruned. {len(messages_to_prune)} messages summarized.")
 
-        elif format_type == "lollms_text":
-            full_prompt_parts = []
-            if full_system_prompt: full_prompt_parts.append(f"!@>system:\n{full_system_prompt}")
-            for msg in branch:
-                sender_str = msg.sender.replace(':', '').replace('!@>', '')
-                content = msg.content.strip()
-                if msg.images: content += f"\n({len(msg.images)} image(s) attached)"
-                full_prompt_parts.append(f"!@>{sender_str}:\n{content}")
-            return "\n".join(full_prompt_parts)
-
-        elif format_type == "openai_completion":
-            full_prompt_parts = []
-            if full_system_prompt: full_prompt_parts.append(f"System:\n{full_system_prompt}")
-            for msg in branch:
-                role_label = self.participants.get(msg.sender, "user").capitalize()
-                content = msg.content.strip()
-                if msg.images: content += f"\n({len(msg.images)} image(s) attached)"
-                full_prompt_parts.append(f"{role_label}:\n{content}")
-            return "\n\n".join(full_prompt_parts)
-
-        else: raise ValueError(f"Unsupported export format_type: {format_type}")
-
-
-if __name__ == "__main__":
-    class MockBinding:
-        def tokenize(self, text: str) -> List[int]: return text.split()
-    class MockLollmsClient:
-        def __init__(self): self.binding = MockBinding()
-        def generate(self, prompt: str, max_new_tokens: int, temperature: float) -> str: return "This is a generated summary."
-
-    print("--- Initializing Mock Client and Discussion ---")
-    mock_client = MockLollmsClient()
-    discussion = LollmsDiscussion(mock_client)
-    discussion.set_participants({"User": "user", "Project Lead": "assistant"})
-    discussion.set_system_prompt("This is a formal discussion about Project Phoenix.")
-    discussion.set_scratchpad("Initial State: Project Phoenix is in the planning phase.")
-
-    print("\n--- Creating a long discussion history ---")
-    parent_id = None
-    long_text = "extra text to increase token count"
-    for i in range(10):
-        user_msg = f"Message #{i*2+1}: Update on task {i+1}? {long_text}"
-        user_id = discussion.add_message("User", "user", user_msg, parent_id=parent_id)
-        assistant_msg = f"Message #{i*2+2}: Task {i+1} status is blocked. {long_text}"
-        assistant_id = discussion.add_message("Project Lead", "assistant", assistant_msg, parent_id=user_id)
-        parent_id = assistant_id
-    
-    initial_tokens = len(mock_client.binding.tokenize(discussion.export("lollms_text")))
-    print(f"Initial message count: {len(discussion.messages)}, Initial tokens: {initial_tokens}")
-
-    print("\n--- Testing Pruning ---")
-    prune_result = discussion.summarize_and_prune(max_tokens=200, preserve_last_n=4)
-    if prune_result.get("pruned"):
-        print("✅ Pruning was successful!")
-        assert "Summary" in discussion.get_scratchpad()
-    else: print(f"❌ Pruning failed: {prune_result.get('reason')}")
-
-    print("\n--- Testing format_discussion (Instruct Model Format) ---")
-    truncated_prompt = discussion.format_discussion(max_allowed_tokens=80)
-    truncated_tokens = len(mock_client.binding.tokenize(truncated_prompt))
-    print(f"Truncated prompt tokens: {truncated_tokens}")
-    print("Truncated Prompt:\n" + "="*20 + f"\n{truncated_prompt}\n" + "="*20)
-
-    # Verification
-    assert truncated_tokens <= 80
-    # Check that it contains the newest message that fits
-    assert "Message #19" in truncated_prompt or "Message #20" in truncated_prompt 
-    print("✅ format_discussion correctly truncated the prompt.")
+    def to_dict(self):
+        return {
+            "id": self.id, "system_prompt": self.system_prompt, "participants": self.participants,
+            "active_branch_id": self.active_branch_id, "metadata": self.metadata, "scratchpad": self.scratchpad,
+            "messages": [{ 'id': m.id, 'parent_id': m.parent_id, 'discussion_id': m.discussion_id, 'sender': m.sender,
+                           'sender_type': m.sender_type, 'content': m.content, 'scratchpad': m.scratchpad, 'images': m.images,
+                           'created_at': m.created_at.isoformat(), 'metadata': m.metadata } for m in self.messages],
+            "created_at": self.created_at.isoformat() if self.created_at else None,
+            "updated_at": self.updated_at.isoformat() if self.updated_at else None
+        }
+
+    def load_from_dict(self, data: Dict):
+        self._create_in_memory_proxy(id=data.get("id"))
+        self.system_prompt, self.participants = data.get("system_prompt"), data.get("participants", {})
+        self.active_branch_id, self.metadata = data.get("active_branch_id"), data.get("metadata", {})
+        self.scratchpad = data.get("scratchpad", "")
+        for msg_data in data.get("messages", []):
+            if 'created_at' in msg_data and isinstance(msg_data['created_at'], str):
+                try:
+                    msg_data['created_at'] = datetime.fromisoformat(msg_data['created_at'])
+                except ValueError:
+                    msg_data['created_at'] = datetime.utcnow()
+            self.add_message(**msg_data)
+        self.created_at = datetime.fromisoformat(data['created_at']) if data.get('created_at') else datetime.utcnow()
+        self.updated_at = datetime.fromisoformat(data['updated_at']) if data.get('updated_at') else self.created_at
+
+    @staticmethod
+    def migrate(lollms_client: 'LollmsClient', db_manager: LollmsDataManager, folder_path: Union[str, Path]):
+        folder = Path(folder_path)
+        if not folder.is_dir():
+            print(f"Error: Path '{folder}' is not a valid directory.")
+            return
+        print(f"\n--- Starting Migration from '{folder}' ---")
+        files = list(folder.glob("*.json")) + list(folder.glob("*.yaml"))
+        with db_manager.get_session() as session:
+            valid_disc_keys = {c.name for c in db_manager.DiscussionModel.__table__.columns}
+            valid_msg_keys = {c.name for c in db_manager.MessageModel.__table__.columns}
+            for i, file_path in enumerate(files):
+                print(f"Migrating file {i+1}/{len(files)}: {file_path.name} ... ", end="")
+                try:
+                    data = yaml.safe_load(file_path.read_text(encoding='utf-8'))
+                    discussion_id = data.get("id", str(uuid.uuid4()))
+                    if session.query(db_manager.DiscussionModel).filter_by(id=discussion_id).first():
+                        print("SKIPPED (already exists)")
+                        continue
+                    discussion_data = data.copy()
+                    if 'metadata' in discussion_data:
+                        discussion_data['discussion_metadata'] = discussion_data.pop('metadata')
+                    for key in ['created_at', 'updated_at']:
+                        if key in discussion_data and isinstance(discussion_data[key], str):
+                            discussion_data[key] = datetime.fromisoformat(discussion_data[key])
+                    db_discussion = db_manager.DiscussionModel(**{k: v for k, v in discussion_data.items() if k in valid_disc_keys})
+                    session.add(db_discussion)
+                    for msg_data in data.get("messages", []):
+                        msg_data['discussion_id'] = db_discussion.id
+                        if 'metadata' in msg_data:
+                            msg_data['message_metadata'] = msg_data.pop('metadata')
+                        if 'created_at' in msg_data and isinstance(msg_data['created_at'], str):
+                            msg_data['created_at'] = datetime.fromisoformat(msg_data['created_at'])
+                        msg_orm = db_manager.MessageModel(**{k: v for k, v in msg_data.items() if k in valid_msg_keys})
+                        session.add(msg_orm)
+                    session.flush()
+                    print("OK")
+                except Exception as e:
+                    print(f"FAILED. Error: {e}")
+                    session.rollback()
+                    continue
+            session.commit()
+        print("--- Migration Finished ---")