lollms-client 0.28.0__py3-none-any.whl → 0.29.1__py3-none-any.whl

lollms_client/lollms_discussion.py

@@ -31,6 +31,7 @@ if False:
 
 from lollms_client.lollms_utilities import build_image_dicts, robust_json_parser
 from ascii_colors import ASCIIColors, trace_exception
+from .lollms_types import MSG_TYPE
 
 class EncryptedString(TypeDecorator):
     """A SQLAlchemy TypeDecorator for field-level database encryption.
@@ -127,6 +128,11 @@ def create_dynamic_models(
         __abstract__ = True
         id = Column(String, primary_key=True, default=lambda: str(uuid.uuid4()))
         system_prompt = Column(EncryptedText, nullable=True)
+        user_data_zone = Column(EncryptedText, nullable=True) # Field for persistent user-specific data
+        discussion_data_zone = Column(EncryptedText, nullable=True) # Field for persistent discussion-specific data
+        personality_data_zone = Column(EncryptedText, nullable=True) # Field for persistent personality-specific data
+        memory = Column(EncryptedText, nullable=True) # New field for long-term memory across discussions
+        
         participants = Column(JSON, nullable=True, default=dict)
         active_branch_id = Column(String, nullable=True)
         discussion_metadata = Column(JSON, nullable=True, default=dict)
@@ -211,8 +217,6 @@ class LollmsDataManager:
             with self.engine.connect() as connection:
                 print("Checking for database schema upgrades...")
                 
-                # --- THIS IS THE FIX ---
-                # We must wrap raw SQL strings in the `text()` function for direct execution.
                 cursor = connection.execute(text("PRAGMA table_info(discussions)"))
                 columns = [row[1] for row in cursor.fetchall()]
                 
@@ -223,9 +227,28 @@ class LollmsDataManager:
                 if 'pruning_point_id' not in columns:
                     print("  -> Upgrading 'discussions' table: Adding 'pruning_point_id' column.")
                     connection.execute(text("ALTER TABLE discussions ADD COLUMN pruning_point_id VARCHAR"))
+
+                if 'data_zone' in columns:
+                    print("  -> Upgrading 'discussions' table: Removing 'data_zone' column.")
+                    connection.execute(text("ALTER TABLE discussions DROP COLUMN data_zone"))
+
+                if 'user_data_zone' not in columns:
+                    print("  -> Upgrading 'discussions' table: Adding 'user_data_zone' column.")
+                    connection.execute(text("ALTER TABLE discussions ADD COLUMN user_data_zone TEXT"))
+
+                if 'discussion_data_zone' not in columns:
+                    print("  -> Upgrading 'discussions' table: Adding 'discussion_data_zone' column.")
+                    connection.execute(text("ALTER TABLE discussions ADD COLUMN discussion_data_zone TEXT"))
+
+                if 'personality_data_zone' not in columns:
+                    print("  -> Upgrading 'discussions' table: Adding 'personality_data_zone' column.")
+                    connection.execute(text("ALTER TABLE discussions ADD COLUMN personality_data_zone TEXT"))
+
+                if 'memory' not in columns:
+                    print("  -> Upgrading 'discussions' table: Adding 'memory' column.")
+                    connection.execute(text("ALTER TABLE discussions ADD COLUMN memory TEXT"))
                 
                 print("Database schema is up to date.")
-                # This is important to apply the ALTER TABLE statements
                 connection.commit() 
 
         except Exception as e:
@@ -382,7 +405,6 @@ class LollmsDiscussion:
         object.__setattr__(self, '_is_db_backed', db_manager is not None)
         
         object.__setattr__(self, '_system_prompt', 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.")
@@ -484,6 +506,10 @@ class LollmsDiscussion:
         proxy = SimpleNamespace()
         proxy.id = id or str(uuid.uuid4())
         proxy.system_prompt = None
+        proxy.user_data_zone = None
+        proxy.discussion_data_zone = None
+        proxy.personality_data_zone = None
+        proxy.memory = None
         proxy.participants = {}
         proxy.active_branch_id = None
         proxy.discussion_metadata = {}
@@ -594,7 +620,23 @@ class LollmsDiscussion:
             
         return [LollmsMessage(self, orm) for orm in reversed(branch_orms)]
 
-
+    def get_full_data_zone(self):
+        """Assembles all data zones into a single, formatted string for the prompt."""
+        parts = []
+        # If memory is not empty, add it to the list of zones.
+        if self.memory and self.memory.strip():
+            parts.append(f"-- Memory --\n{self.memory.strip()}")
+        if self.user_data_zone and self.user_data_zone.strip():
+            parts.append(f"-- User Data Zone --\n{self.user_data_zone.strip()}")
+        if self.discussion_data_zone and self.discussion_data_zone.strip():
+            parts.append(f"-- Discussion Data Zone --\n{self.discussion_data_zone.strip()}")
+        if self.personality_data_zone and self.personality_data_zone.strip():
+            parts.append(f"-- Personality Data Zone --\n{self.personality_data_zone.strip()}")
+        
+        # Join the zones with double newlines for clear separation in the prompt.
+        return "\n\n".join(parts)
+    
+    
     def chat(
         self,
         user_message: str,
@@ -642,21 +684,92 @@ class LollmsDiscussion:
             A dictionary with 'user_message' and 'ai_message' LollmsMessage objects,
             where the 'ai_message' will contain rich metadata if an agentic turn was used.
         """
+        callback = kwargs.get("streaming_callback")
+        # extract personality data
         if personality is not None:
             object.__setattr__(self, '_system_prompt', personality.system_prompt)
+
+            # --- New Data Source Handling Logic ---
+            if hasattr(personality, 'data_source') and personality.data_source is not None:
+                if isinstance(personality.data_source, str):
+                    # --- Static Data Source ---
+                    if callback:
+                        callback("Loading static personality data...", MSG_TYPE.MSG_TYPE_STEP, {"id": "static_data_loading"})
+                    if personality.data_source:
+                        self.personality_data_zone = personality.data_source.strip()
+
+                elif callable(personality.data_source):
+                    # --- Dynamic Data Source ---
+                    qg_id = None
+                    if callback:
+                        qg_id = callback("Generating query for dynamic personality data...", MSG_TYPE.MSG_TYPE_STEP_START, {"id": "dynamic_data_query_gen"})
+
+                    context_for_query = self.export('markdown')
+                    query_prompt = (
+                        "You are an expert query generator. Based on the current conversation, formulate a concise and specific query to retrieve relevant information from a knowledge base. "
+                        "The query will be used to fetch data that will help you answer the user's latest request.\n\n"
+                        f"--- Conversation History ---\n{context_for_query}\n\n"
+                        "--- Instructions ---\n"
+                        "Generate a single query string."
+                    )
+                    
+                    try:
+                        query_json = self.lollmsClient.generate_structured_content(
+                            prompt=query_prompt,
+                            output_format={"query": "Your generated search query here."},
+                            system_prompt="You are an AI assistant that generates search queries in JSON format.",
+                            temperature=0.0
+                        )
+
+                        if not query_json or "query" not in query_json:
+                            if callback:
+                                callback("Failed to generate data query.", MSG_TYPE.MSG_TYPE_EXCEPTION, {"id": qg_id})
+                        else:
+                            generated_query = query_json["query"]
+                            if callback:
+                                callback(f"Generated query: '{generated_query}'", MSG_TYPE.MSG_TYPE_STEP_END, {"id": qg_id, "query": generated_query})
+                            
+                            dr_id = None
+                            if callback:
+                                dr_id = callback("Retrieving dynamic data from personality source...", MSG_TYPE.MSG_TYPE_STEP_START, {"id": "dynamic_data_retrieval"})
+                            
+                            try:
+                                retrieved_data = personality.data_source(generated_query)
+                                if callback:
+                                    callback(f"Retrieved data successfully.", MSG_TYPE.MSG_TYPE_STEP_END, {"id": dr_id, "data_snippet": retrieved_data[:200]})
+                                
+                                
+                                if retrieved_data:
+                                    self.personality_data_zone = retrieved_data.strip()
+
+                            except Exception as e:
+                                trace_exception(e)
+                                if callback:
+                                    callback(f"Error retrieving dynamic data: {e}", MSG_TYPE.MSG_TYPE_EXCEPTION, {"id": dr_id})
+                    except Exception as e:
+                        trace_exception(e)
+                        if callback:
+                            callback(f"An error occurred during query generation: {e}", MSG_TYPE.MSG_TYPE_EXCEPTION, {"id": qg_id})
+
+        # Determine effective MCPs by combining personality defaults and turn-specific overrides
+        effective_use_mcps = use_mcps
+        if personality and hasattr(personality, 'active_mcps') and personality.active_mcps:
+            if effective_use_mcps in [None, False]:
+                effective_use_mcps = personality.active_mcps
+            elif isinstance(effective_use_mcps, list):
+                effective_use_mcps = list(set(effective_use_mcps + personality.active_mcps))
             
         if self.max_context_size is not None:
             self.summarize_and_prune(self.max_context_size)
 
         # Step 1: Add user message, now including any images.
         if add_user_message:
-            # Pass kwargs through to capture images and other potential message attributes
             user_msg = self.add_message(
                 sender="user", 
                 sender_type="user", 
                 content=user_message,
                 images=images,
-                **kwargs # Use kwargs to allow other fields to be set from the caller
+                **kwargs
             )
         else: # Regeneration logic
             if self.active_branch_id not in self._message_index:
@@ -665,11 +778,9 @@ class LollmsDiscussion:
             if user_msg_orm.sender_type != 'user':
                 raise ValueError(f"Regeneration failed: active branch tip is a '{user_msg_orm.sender_type}' message, not 'user'.")
             user_msg = LollmsMessage(self, user_msg_orm)
-            # For regeneration, we use the images from the original user message
             images = user_msg.images
 
-        # Step 2: Determine if this is a simple chat or a complex agentic turn.
-        is_agentic_turn = (use_mcps is not None and use_mcps) or (use_data_store is not None and use_data_store)
+        is_agentic_turn = (effective_use_mcps is not None and effective_use_mcps) or (use_data_store is not None and use_data_store)
         
         start_time = datetime.now()
         
@@ -678,79 +789,56 @@ class LollmsDiscussion:
         final_raw_response = ""
         final_content = ""
 
-        # Step 3: Execute the appropriate generation logic.
         if is_agentic_turn:
-            # --- AGENTIC TURN ---
             prompt_for_agent = self.export("markdown", branch_tip_id if branch_tip_id else self.active_branch_id)
             if debug:
-                ASCIIColors.cyan("\n" + "="*50)
-                ASCIIColors.cyan("--- DEBUG: AGENTIC TURN TRIGGERED ---")
-                ASCIIColors.cyan(f"--- PROMPT FOR AGENT (from discussion history) ---")
-                ASCIIColors.magenta(prompt_for_agent)
-                ASCIIColors.cyan("="*50 + "\n")
+                ASCIIColors.cyan("\n" + "="*50 + "\n--- DEBUG: AGENTIC TURN TRIGGERED ---\n" + f"--- PROMPT FOR AGENT (from discussion history) ---\n{prompt_for_agent}\n" + "="*50 + "\n")
 
             agent_result = self.lollmsClient.generate_with_mcp_rag(
                 prompt=prompt_for_agent,
-                use_mcps=use_mcps,
+                use_mcps=effective_use_mcps,
                 use_data_store=use_data_store,
                 max_reasoning_steps=max_reasoning_steps,
                 images=images,
                 system_prompt = self._system_prompt,
-                debug=debug, # Pass the debug flag down
+                debug=debug,
                 **kwargs
             )
             final_content = agent_result.get("final_answer", "The agent did not produce a final answer.")
             final_scratchpad = agent_result.get("final_scratchpad", "")
             final_raw_response = json.dumps(agent_result, indent=2)
-
         else:
-            # --- SIMPLE CHAT TURN ---
             if debug:
                 prompt_for_chat = self.export("markdown", branch_tip_id if branch_tip_id else self.active_branch_id)
-                ASCIIColors.cyan("\n" + "="*50)
-                ASCIIColors.cyan("--- DEBUG: SIMPLE CHAT PROMPT ---")
-                ASCIIColors.magenta(prompt_for_chat)
-                ASCIIColors.cyan("="*50 + "\n")
+                ASCIIColors.cyan("\n" + "="*50 + f"\n--- DEBUG: SIMPLE CHAT PROMPT ---\n{prompt_for_chat}\n" + "="*50 + "\n")
 
-            # For simple chat, we also need to consider images if the model is multi-modal
             final_raw_response = self.lollmsClient.chat(self, images=images, **kwargs) or ""
             
             if debug:
-                ASCIIColors.cyan("\n" + "="*50)
-                ASCIIColors.cyan("--- DEBUG: RAW SIMPLE CHAT RESPONSE ---")
-                ASCIIColors.magenta(final_raw_response)
-                ASCIIColors.cyan("="*50 + "\n")
+                ASCIIColors.cyan("\n" + "="*50 + f"\n--- DEBUG: RAW SIMPLE CHAT RESPONSE ---\n{final_raw_response}\n" + "="*50 + "\n")
 
             if isinstance(final_raw_response, dict) and final_raw_response.get("status") == "error":
                 raise Exception(final_raw_response.get("message", "Unknown error from lollmsClient.chat"))
             else:
                 final_content = self.lollmsClient.remove_thinking_blocks(final_raw_response)
+            final_scratchpad = None
 
-            final_scratchpad = None # No agentic scratchpad in a simple turn
-
-        # Step 4: Post-generation processing and statistics.
         end_time = datetime.now()
         duration = (end_time - start_time).total_seconds()
         token_count = self.lollmsClient.count_tokens(final_content)
         tok_per_sec = (token_count / duration) if duration > 0 else 0
 
-        # Step 5: Collect metadata from the agentic turn for storage.
         message_meta = {}
         if is_agentic_turn and isinstance(agent_result, dict):
-            if "tool_calls" in agent_result:
-                message_meta["tool_calls"] = agent_result["tool_calls"]
-            if "sources" in agent_result:
-                message_meta["sources"] = agent_result["sources"]
-            if agent_result.get("clarification_required", False):
-                message_meta["clarification_required"] = True
-
-        # Step 6: Add the final AI message to the discussion.
+            if "tool_calls" in agent_result: message_meta["tool_calls"] = agent_result["tool_calls"]
+            if "sources" in agent_result: message_meta["sources"] = agent_result["sources"]
+            if agent_result.get("clarification_required", False): message_meta["clarification_required"] = True
+
         ai_message_obj = self.add_message(
             sender=personality.name if personality else "assistant",
             sender_type="assistant",
             content=final_content,
             raw_content=final_raw_response,
-            # Store the agent's full reasoning log in the message's dedicated scratchpad field
             scratchpad=final_scratchpad,
             tokens=token_count,
             generation_speed=tok_per_sec,
@@ -892,7 +980,21 @@ class LollmsDiscussion:
             return "" if format_type == "lollms_text" else []
 
         branch = self.get_branch(branch_tip_id)
-        full_system_prompt = self._system_prompt # Simplified for clarity
+        
+        # Combine system prompt and data zones
+        system_prompt_part = (self._system_prompt or "").strip()
+        data_zone_part = self.get_full_data_zone() # This now returns a clean, multi-part block or an empty string
+        full_system_prompt = ""
+
+        # Combine them intelligently
+        if system_prompt_part and data_zone_part:
+            full_system_prompt = f"{system_prompt_part}\n\n{data_zone_part}"
+        elif system_prompt_part:
+            full_system_prompt = system_prompt_part
+        else:
+            full_system_prompt = data_zone_part
+
+
         participants = self.participants or {}
 
         def get_full_content(msg: 'LollmsMessage') -> str:
@@ -1049,6 +1151,73 @@ class LollmsDiscussion:
         self.touch()
         print(f"[INFO] Discussion auto-pruned. {len(messages_to_prune)} messages summarized. History preserved.")
 
+    def memorize(self, branch_tip_id: Optional[str] = None):
+        """
+        Analyzes the current discussion, extracts key information suitable for long-term
+        memory, and appends it to the discussion's 'memory' field.
+
+        This is intended to build a persistent knowledge base about user preferences,
+        facts, and context that can be useful across different future discussions.
+
+        Args:
+            branch_tip_id: The ID of the message to use as the end of the context
+                           for memory extraction. Defaults to the active branch.
+        """
+        try:
+            # 1. Get the current conversation context
+            discussion_context = self.export("markdown", branch_tip_id=branch_tip_id)
+            if not discussion_context.strip():
+                print("[INFO] Memorize: Discussion is empty, nothing to memorize.")
+                return
+
+            # 2. Formulate the prompt for the LLM
+            system_prompt = (
+                "You are a Memory Extractor AI. Your task is to analyze a conversation "
+                "and extract only the most critical pieces of information that would be "
+                "valuable for a future, unrelated conversation with the same user. "
+                "Focus on: \n"
+                "- Explicit user preferences, goals, or facts about themselves.\n"
+                "- Key decisions or conclusions reached.\n"
+                "- Important entities, projects, or topics mentioned that are likely to recur.\n"
+                "Format the output as a concise list of bullet points. Be brief and factual. "
+                "If no new, significant long-term information is present, output the single word: 'NOTHING'."
+            )
+            
+            prompt = (
+                "Analyze the following discussion and extract key information for long-term memory:\n\n"
+                f"--- Conversation ---\n{discussion_context}\n\n"
+                "--- Extracted Memory Points (as a bulleted list) ---"
+            )
+
+            # 3. Call the LLM to extract information
+            print("[INFO] Memorize: Extracting key information from discussion...")
+            extracted_info = self.lollmsClient.generate_text(
+                prompt,
+                system_prompt=system_prompt,
+                n_predict=512, # A reasonable length for a summary
+                temperature=0.1, # Low temperature for factual extraction
+                top_k=10,
+            )
+
+            # 4. Process and append the information
+            if extracted_info and "NOTHING" not in extracted_info.upper():
+                new_memory_entry = extracted_info.strip()
+                
+                # Format with a timestamp for context
+                timestamp = datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S UTC")
+                formatted_entry = f"\n\n--- Memory entry from {timestamp} ---\n{new_memory_entry}"
+
+                current_memory = self.memory or ""
+                self.memory = (current_memory + formatted_entry).strip()
+                self.touch() # Mark as updated and save if autosave is on
+                print(f"[INFO] Memorize: New information added to long-term memory.")
+            else:
+                print("[INFO] Memorize: No new significant information found to add to memory.")
+                
+        except Exception as e:
+            trace_exception(e)
+            print(f"[ERROR] Memorize: Failed to extract memory. {e}")
+
     def count_discussion_tokens(self, format_type: str, branch_tip_id: Optional[str] = None) -> int:
         """Counts the number of tokens in the exported discussion content.
 
@@ -1083,27 +1252,134 @@ class LollmsDiscussion:
 
         return self.lollmsClient.count_tokens(text_to_count)
 
-    def get_context_status(self, branch_tip_id: Optional[str] = None) -> Dict[str, Optional[int]]:
-        """Returns the current token count and the maximum context size.
+    def get_context_status(self, branch_tip_id: Optional[str] = None) -> Dict[str, Any]:
+        """
+        Returns a detailed breakdown of the context size and its components.
 
-        This provides a snapshot of the context usage, taking into account
-        any non-destructive pruning that has occurred. The token count is
-        based on the "lollms_text" export format, which is the format used
-        for pruning calculations.
+        This provides a comprehensive snapshot of the context usage, including the
+        content and token count for each part of the prompt (system prompt, data zones,
+        pruning summary, and message history). The token counts are based on the
+        "lollms_text" export format, which is the format used for pruning calculations.
 
         Args:
             branch_tip_id: The ID of the message branch to measure. Defaults
                            to the active branch.
 
         Returns:
-            A dictionary with 'current_tokens' and 'max_tokens'.
+            A dictionary with a detailed breakdown:
+            {
+                "max_tokens": int | None,
+                "current_tokens": int,
+                "zones": {
+                    "system_prompt": {"content": str, "tokens": int},
+                    "memory": {"content": str, "tokens": int},
+                    "user_data_zone": {"content": str, "tokens": int},
+                    "discussion_data_zone": {"content": str, "tokens": int},
+                    "personality_data_zone": {"content": str, "tokens": int},
+                    "pruning_summary": {"content": str, "tokens": int},
+                    "message_history": {"content": str, "tokens": int, "message_count": int}
+                }
+            }
+            Zones are only included if they contain content.
         """
-        current_tokens = self.count_discussion_tokens("lollms_text", branch_tip_id)
-        return {
-            "current_tokens": current_tokens,
-            "max_tokens": self.max_context_size
+        result = {
+            "max_tokens": self.max_context_size,
+            "current_tokens": 0,
+            "zones": {}
+        }
+        total_tokens = 0
+
+        # 1. System Prompt
+        system_prompt_text = (self._system_prompt or "").strip()
+        if system_prompt_text:
+            # We count tokens for the full block as it would appear in the prompt
+            full_block = f"!@>system:\n{system_prompt_text}\n"
+            tokens = self.lollmsClient.count_tokens(full_block)
+            result["zones"]["system_prompt"] = {
+                "content": system_prompt_text,
+                "tokens": tokens
+            }
+            total_tokens += tokens
+            
+        # 2. All Data Zones
+        zones_to_process = {
+            "memory": self.memory,
+            "user_data_zone": self.user_data_zone,
+            "discussion_data_zone": self.discussion_data_zone,
+            "personality_data_zone": self.personality_data_zone,
         }
 
+        for name, content in zones_to_process.items():
+            content_text = (content or "").strip()
+            if content_text:
+                # Mimic the formatting from get_full_data_zone for accurate token counting
+                header = f"-- {name.replace('_', ' ').title()} --\n"
+                full_block = f"{header}{content_text}"
+                # In lollms_text format, zones are part of the system message, so we add separators
+                # This counts the standalone block.
+                tokens = self.lollmsClient.count_tokens(full_block)
+                result["zones"][name] = {
+                    "content": content_text,
+                    "tokens": tokens
+                }
+                # Note: The 'export' method combines these into one system prompt.
+                # For this breakdown, we count them separately. The total will be a close approximation.
+
+        # 3. Pruning Summary
+        pruning_summary_text = (self.pruning_summary or "").strip()
+        if pruning_summary_text and self.pruning_point_id:
+            full_block = f"!@>system:\n--- Conversation Summary ---\n{pruning_summary_text}\n"
+            tokens = self.lollmsClient.count_tokens(full_block)
+            result["zones"]["pruning_summary"] = {
+                "content": pruning_summary_text,
+                "tokens": tokens
+            }
+            total_tokens += tokens
+        
+        # 4. Message History
+        branch_tip_id = branch_tip_id or self.active_branch_id
+        messages_text = ""
+        message_count = 0
+        if branch_tip_id:
+            branch = self.get_branch(branch_tip_id)
+            messages_to_render = branch
+            
+            # Adjust for pruning to get the active set of messages
+            if self.pruning_summary and self.pruning_point_id:
+                pruning_index = -1
+                for i, msg in enumerate(branch):
+                    if msg.id == self.pruning_point_id:
+                        pruning_index = i
+                        break
+                if pruning_index != -1:
+                    messages_to_render = branch[pruning_index:]
+
+            message_parts = []
+            for msg in messages_to_render:
+                sender_str = msg.sender.replace(':', '').replace('!@>', '')
+                content = msg.content.strip()
+                if msg.images:
+                    content += f"\n({len(msg.images)} image(s) attached)"
+                msg_text = f"!@>{sender_str}:\n{content}\n"
+                message_parts.append(msg_text)
+            
+            messages_text = "".join(message_parts)
+            message_count = len(messages_to_render)
+
+        if messages_text:
+            tokens = self.lollmsClient.count_tokens(messages_text)
+            result["zones"]["message_history"] = {
+                "content": messages_text,
+                "tokens": tokens,
+                "message_count": message_count
+            }
+            total_tokens += tokens
+
+        # Finalize the total count. This re-calculates based on the actual export format
+        # for maximum accuracy, as combining zones can slightly change tokenization.
+        result["current_tokens"] = self.count_discussion_tokens("lollms_text", branch_tip_id)
+
+        return result
     def switch_to_branch(self, branch_id):
         self.active_branch_id = branch_id
 
@@ -1112,15 +1388,21 @@ class LollmsDiscussion:
             if self.metadata is None:
                 self.metadata = {}
             discussion = self.export("markdown")[0:1000]
-            prompt = f"""You are a title builder. Your oibjective is to build a title for the following discussion:
+            system_prompt="You are a title builder out of a discussion."
+            prompt = f"""Build a title for the following discussion:
 {discussion}
 ...
 """
-            template = """{
-    "title": "An short but comprehensive discussion title"
-}"""
-            infos = self.lollmsClient.generate_code(prompt = prompt, template = template)
-            discussion_title = robust_json_parser(infos)["title"]
+            title_generation_schema = {
+                "type": "object",
+                "properties": {
+                    "title": {"type": "string", "description": "Short, catchy title for the discussion."},
+                },
+                "required": ["title"],
+                "description": "JSON object as title of the discussion."
+            }
+            infos = self.lollmsClient.generate_structured_content(prompt = prompt, system_prompt=system_prompt, schema = title_generation_schema)
+            discussion_title = infos["title"]
             new_metadata = (self.metadata or {}).copy()
             new_metadata['title'] = discussion_title
             

lollms_client/lollms_personality.py

@@ -22,6 +22,8 @@ class LollmsPersonality:
         # Core behavioral instruction
         system_prompt: str,        
         icon: Optional[str] = None,  # Base64 encoded image string
+        active_mcps: Optional[List[str]] = None, # The list of MCPs to activate with this personality
+        data_source: Optional[Union[str, Callable[[str], str]]] = None, # Static string data or a callable for dynamic data retrieval
 
 
         # RAG - Data Files and Application-provided Callbacks
@@ -46,6 +48,8 @@ class LollmsPersonality:
             description: A brief description of what the personality does.
             icon: An optional base64 encoded string for a display icon.
             system_prompt: The core system prompt that defines the AI's behavior.
+            active_mcps: An optional list of MCP (tool) names to be automatically activated with this personality.
+            data_source: A source of knowledge. Can be a static string or a callable function that takes a query and returns a string.
             data_files: A list of file paths to be used as a knowledge base for RAG.
             vectorize_chunk_callback: A function provided by the host app to vectorize and store a text chunk.
             is_vectorized_callback: A function provided by the host app to check if a chunk is already vectorized.
@@ -59,6 +63,8 @@ class LollmsPersonality:
         self.description = description
         self.icon = icon
         self.system_prompt = system_prompt
+        self.active_mcps = active_mcps or []
+        self.data_source = data_source
         self.data_files = [Path(f) for f in data_files] if data_files else []
         
         # RAG Callbacks provided by the host application
@@ -177,6 +183,8 @@ class LollmsPersonality:
             "category": self.category,
             "description": self.description,
             "system_prompt": self.system_prompt,
+            "active_mcps": self.active_mcps,
+            "has_data_source": self.data_source is not None,
             "data_files": [str(p) for p in self.data_files],
             "has_script": self.script is not None
         }

lollms_client/lollms_utilities.py

@@ -93,10 +93,18 @@ def robust_json_parser(json_string: str) -> dict:
         ValueError: If parsing fails after all correction attempts.
     """
 
-    # STEP 0: Remove code block wrappers if present (e.g., ```json ... ```)
+
+    # STEP 0: Attempt to parse directly
+    try:
+        return json.loads(json_string)
+    except json.JSONDecodeError as e:
+        trace_exception(e)
+        pass
+    
+    # STEP 1: Remove code block wrappers if present (e.g., ```json ... ```)
     json_string = re.sub(r"^```(?:json)?\s*|\s*```$", '', json_string.strip())
 
-    # STEP 1: Attempt to parse directly
+    # STEP 2: Attempt to parse directly
     try:
         return json.loads(json_string)
     except json.JSONDecodeError: