lollms-client 1.5.6__py3-none-any.whl → 1.7.13__py3-none-any.whl

lollms_client/lollms_discussion.py

@@ -34,8 +34,7 @@ if False:
 
 from lollms_client.lollms_utilities import build_image_dicts, robust_json_parser
 from ascii_colors import ASCIIColors, trace_exception
-# Assuming lollms_types is a local module providing MSG_TYPE enum/constants
-# from .lollms_types import MSG_TYPE 
+from lollms_client.lollms_types import MSG_TYPE 
 
 class EncryptedString(TypeDecorator):
     """A SQLAlchemy TypeDecorator for field-level database encryption.
@@ -1054,7 +1053,7 @@ class LollmsDiscussion:
         debug: bool = False,
         remove_thinking_blocks:bool = True,
         **kwargs
-    ) -> Dict[str, 'LollmsMessage']:
+    ) -> Dict[str, Any]:
         """Main interaction method that can invoke the dynamic, multi-modal agent.
 
         This method orchestrates the entire response generation process. It can
@@ -1095,14 +1094,35 @@ class LollmsDiscussion:
             where the 'ai_message' will contain rich metadata if an agentic turn was used.
         """
         callback = kwargs.get("streaming_callback")
+        collected_sources = []
+        
+
+        # Step 1: Add user message, now including any images.
+        if add_user_message:
+            user_msg = self.add_message(
+                sender=kwargs.get("user_name", "user"), 
+                sender_type="user", 
+                content=user_message,
+                images=images,
+                **kwargs
+            )
+        else: # Regeneration logic
+            # _validate_and_set_active_branch ensures active_branch_id is valid and a leaf.
+            # So, if we are regenerating, active_branch_id must be valid.
+            if self.active_branch_id not in self._message_index: # Redundant check, but safe
+                 raise ValueError("Regeneration failed: active branch tip not found or is invalid.")
+            user_msg_orm = self._message_index[self.active_branch_id]
+            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)
+            images = user_msg.images
+                    
         # 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:
-                # Placeholder for MSG_TYPE if not imported
-                MSG_TYPE = SimpleNamespace(MSG_TYPE_STEP="step", MSG_TYPE_STEP_START="step_start", MSG_TYPE_STEP_END="step_end", MSG_TYPE_EXCEPTION="exception")
 
                 if isinstance(personality.data_source, str):
                     # --- Static Data Source ---
@@ -1117,7 +1137,7 @@ class LollmsDiscussion:
                     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')
+                    context_for_query = self.export('markdown', suppress_system_prompt=True)
                     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"
@@ -1154,6 +1174,15 @@ class LollmsDiscussion:
                                 
                                 if retrieved_data:
                                     self.personality_data_zone = retrieved_data.strip()
+                                    source_item = {
+                                        "title": "Personality Data Source",
+                                        "content": retrieved_data,
+                                        "source": personality.name if hasattr(personality, 'name') else "Personality",
+                                        "query": generated_query
+                                    }
+                                    collected_sources.append(source_item)
+                                    if callback:
+                                        callback([source_item], MSG_TYPE.MSG_TYPE_SOURCES_LIST)
 
                             except Exception as e:
                                 trace_exception(e)
@@ -1175,26 +1204,6 @@ class LollmsDiscussion:
         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:
-            user_msg = self.add_message(
-                sender=kwargs.get("user_name", "user"), 
-                sender_type="user", 
-                content=user_message,
-                images=images,
-                **kwargs
-            )
-        else: # Regeneration logic
-            # _validate_and_set_active_branch ensures active_branch_id is valid and a leaf.
-            # So, if we are regenerating, active_branch_id must be valid.
-            if self.active_branch_id not in self._message_index: # Redundant check, but safe
-                 raise ValueError("Regeneration failed: active branch tip not found or is invalid.")
-            user_msg_orm = self._message_index[self.active_branch_id]
-            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)
-            images = user_msg.images
-
         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()
@@ -1205,17 +1214,30 @@ class LollmsDiscussion:
         final_content = ""
 
         if is_agentic_turn:
-            prompt_for_agent = self.export("markdown", branch_tip_id if branch_tip_id else self.active_branch_id)
+            prompt_for_agent = self.export("markdown", branch_tip_id if branch_tip_id else self.active_branch_id, suppress_system_prompt=True)
             if debug:
                 ASCIIColors.cyan("\n" + "="*50 + "\n--- DEBUG: AGENTIC TURN TRIGGERED ---\n" + f"--- PROMPT FOR AGENT (from discussion history) ---\n{prompt_for_agent}\n" + "="*50 + "\n")
-
+            
+            
+            # 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
             agent_result = self.lollmsClient.generate_with_mcp_rag(
                 prompt=prompt_for_agent,
                 use_mcps=effective_use_mcps,
                 use_data_store=use_data_store,
                 max_reasoning_steps=max_reasoning_steps,
                 images=images,
-                system_prompt = self._system_prompt,
+                system_prompt = full_system_prompt,
                 debug=debug,
                 **kwargs
             )
@@ -1249,9 +1271,12 @@ class LollmsDiscussion:
         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 "sources" in agent_result: collected_sources.extend(agent_result["sources"])
             if agent_result.get("clarification_required", False): message_meta["clarification_required"] = True
 
+        if collected_sources:
+             message_meta["sources"] = collected_sources
+
         ai_message_obj = self.add_message(
             sender=personality.name if personality else "assistant",
             sender_type="assistant",
@@ -1269,9 +1294,9 @@ class LollmsDiscussion:
         if self._is_db_backed and self.autosave:
             self.commit()
             
-        return {"user_message": user_msg, "ai_message": ai_message_obj}
+        return {"user_message": user_msg, "ai_message": ai_message_obj, "sources": collected_sources}
 
-    def regenerate_branch(self, branch_tip_id: Optional[str] = None, **kwargs) -> Dict[str, 'LollmsMessage']:
+    def regenerate_branch(self, branch_tip_id: Optional[str] = None, **kwargs) -> Dict[str, Any]:
         """Regenerates the AI response for a given message or the active branch's AI response.
 
         If the target is an AI message, it's deleted and its children are re-parented to its parent
@@ -1426,7 +1451,7 @@ class LollmsDiscussion:
         self.touch() # Mark for update and auto-save if configured
         print(f"Branch starting at {message_id} ({len(messages_to_delete_ids)} messages) removed. New active branch: {self.active_branch_id}")
         
-    def export(self, format_type: str, branch_tip_id: Optional[str] = None, max_allowed_tokens: Optional[int] = None) -> Union[List[Dict], str]:
+    def export(self, format_type: str, branch_tip_id: Optional[str] = None, max_allowed_tokens: Optional[int] = None, suppress_system_prompt=False) -> Union[List[Dict], str]:
         """Exports the discussion history into a specified format.
 
         This method can format the conversation for different backends like OpenAI,
@@ -1464,12 +1489,13 @@ class LollmsDiscussion:
         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
+        if not suppress_system_prompt:
+            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 {}