lollms-client 0.24.2__py3-none-any.whl → 0.27.0__py3-none-any.whl

lollms_client/lollms_core.py

@@ -13,6 +13,8 @@ from lollms_client.lollms_ttm_binding import LollmsTTMBinding, LollmsTTMBindingM
 from lollms_client.lollms_mcp_binding import LollmsMCPBinding, LollmsMCPBindingManager
 
 from lollms_client.lollms_discussion import LollmsDiscussion
+
+from lollms_client.lollms_utilities import build_image_dicts, dict_to_markdown
 import json, re
 from enum import Enum
 import base64
@@ -846,7 +848,7 @@ Don't forget encapsulate the code inside a html code tag. This is mandatory.
             "2.  **Check for a Single-Step Solution:** Scrutinize the available tools. Can a single tool call directly achieve the user's current goal? \n"
             "3.  **Formulate a Plan:** Based on your analysis, create a concise, numbered list of steps to achieve the goal. If the goal is simple, this may be only one step. If it is complex or multi-turn, it may be several steps.\n\n"
             "**CRITICAL RULES:**\n"
-            "*   **MANDATORY: NEVER add steps the user did not ask for.** Do not embellish or add 'nice-to-have' features.\n"
+            "*   **MANDATORY: Be helpful, curious and creative.\n"
             "*   **Focus on the Goal:** Your plan should directly address the user's request as it stands now in the conversation.\n\n"
             "---\n"
             "**Available Tools:**\n"
@@ -877,7 +879,7 @@ Don't forget encapsulate the code inside a html code tag. This is mandatory.
             
             formatted_agent_history = "No actions taken yet in this turn."
             if agent_work_history:
-                history_parts = [ f"### Step {i+1}:\n**Thought:** {entry['thought']}\n**Action:** Called tool `{entry['tool_name']}` with parameters `{json.dumps(entry['tool_params'])}`\n**Observation (Tool Output):**\n```json\n{json.dumps(entry['tool_result'], indent=2)}\n```" for i, entry in enumerate(agent_work_history)]
+                history_parts = [ f"### Step {i+1}:\n**Thought:**\n{entry['thought']}\n**Action:** Called tool `{entry['tool_name']}` with parameters `{json.dumps(entry['tool_params'])}`\n**Observation (Tool Output):**\n```json\n{json.dumps(entry['tool_result'], indent=2)}\n```" for i, entry in enumerate(agent_work_history)]
                 formatted_agent_history = "\n\n".join(history_parts)
 
             llm_decision = None
@@ -1078,7 +1080,7 @@ Provide your response as a single JSON object with one key, "query".
 """
             try:
                 raw_initial_query_response = self.generate_code(initial_query_gen_prompt, system_prompt="You are a query generation expert.", temperature=0.0)
-                initial_plan = json.loads(raw_initial_query_response)
+                initial_plan = robust_json_parser(raw_initial_query_response)
                 current_query_for_rag = initial_plan.get("query")
                 if not current_query_for_rag:
                     raise ValueError("LLM returned an empty initial query.")
@@ -1434,7 +1436,6 @@ Provide your response as a single JSON object inside a JSON markdown tag. Use th
         new_scratchpad_text = self.generate_text(prompt=synthesis_prompt, n_predict=1024, temperature=0.0)
         return self.remove_thinking_blocks(new_scratchpad_text).strip()
 
-# In lollms_client/lollms_discussion.py -> LollmsClient class
 
     def generate_with_mcp_rag(
         self,
@@ -1442,15 +1443,16 @@ Provide your response as a single JSON object inside a JSON markdown tag. Use th
         use_mcps: Union[None, bool, List[str]] = None,
         use_data_store: Union[None, Dict[str, Callable]] = None,
         system_prompt: str = None,
-        reasoning_system_prompt: str = "You are a logical and adaptive AI assistant.",
+        reasoning_system_prompt: str = "You are a logical AI assistant. Your task is to achieve the user's goal by thinking step-by-step and using the available tools.",
         images: Optional[List[str]] = None,
-        max_reasoning_steps: int = 10,
-        decision_temperature: float = 0.0,
+        max_reasoning_steps: int = None,
+        decision_temperature: float = None,
         final_answer_temperature: float = None,
         streaming_callback: Optional[Callable[[str, 'MSG_TYPE', Optional[Dict], Optional[List]], bool]] = None,
-        rag_top_k: int = 5,
-        rag_min_similarity_percent: float = 70.0,
-        output_summarization_threshold: int = 500, # In tokens
+        rag_top_k: int = None,
+        rag_min_similarity_percent: float = None,
+        output_summarization_threshold: int = None, # In tokens
+        debug: bool = False,
         **llm_generation_kwargs
     ) -> Dict[str, Any]:
         """Generates a response using a dynamic agent with stateful, ID-based step tracking.
@@ -1483,6 +1485,7 @@ Provide your response as a single JSON object inside a JSON markdown tag. Use th
             rag_min_similarity_percent: Minimum similarity for RAG results.
             output_summarization_threshold: The token count that triggers automatic
                                             summarization of a tool's text output.
+            debug : If true, we'll report the detailed promptin and response information
             **llm_generation_kwargs: Additional keyword arguments for LLM calls.
 
         Returns:
@@ -1490,12 +1493,28 @@ Provide your response as a single JSON object inside a JSON markdown tag. Use th
             answer, the complete internal scratchpad, a log of tool calls,
             any retrieved RAG sources, and other metadata.
         """
+        reasoning_step_id = None
         if not self.binding:
             return {"final_answer": "", "tool_calls": [], "sources": [], "error": "LLM binding not initialized."}
 
+        if not max_reasoning_steps:
+            max_reasoning_steps= 10
+        if not rag_min_similarity_percent:
+            rag_min_similarity_percent= 50
+        if not rag_top_k:
+            rag_top_k = 5
+        if not decision_temperature:
+            decision_temperature = 0.7
+        if not output_summarization_threshold:
+            output_summarization_threshold = 500
+
+        events = []
+            
+            
         # --- Initialize Agent State ---
         sources_this_turn: List[Dict[str, Any]] = []
         tool_calls_this_turn: List[Dict[str, Any]] = []
+        generated_code_store: Dict[str, str] = {} # NEW: Store for UUID -> code
         original_user_prompt = prompt
         
         initial_state_parts = [
@@ -1507,41 +1526,50 @@ Provide your response as a single JSON object inside a JSON markdown tag. Use th
             initial_state_parts.append(f"- The user has provided {len(images)} image(s) for context.")
         current_scratchpad = "\n".join(initial_state_parts)
 
-        # --- Define Inner Helper Function for Stateful Step Logging ---
-        def log_step(
+        def log_prompt(prompt, type="prompt"):
+            ASCIIColors.cyan(f"** DEBUG: {type} **")
+            ASCIIColors.magenta(prompt[-15000:])
+            prompt_size = self.count_tokens(prompt)
+            ASCIIColors.red(f"Prompt size:{prompt_size}/{self.default_ctx_size}")
+            ASCIIColors.cyan(f"** DEBUG: DONE **")
+
+        # --- Define Inner Helper Functions ---
+        def log_event(
             description: str,
-            step_type: str,
+            event_type: MSG_TYPE = MSG_TYPE.MSG_TYPE_CHUNK,
             metadata: Optional[Dict] = None,
-            is_start: bool = True
+            event_id=None
         ) -> Optional[str]:
-            """
-            Logs a step start or end, generating a unique ID for correlation.
-            This is an inner function that has access to the `streaming_callback`.
-            
-            Returns the ID for start events so it can be used for the end event.
-            """
-            if not streaming_callback:
-                return None
-
-            event_id = str(uuid.uuid4()) if is_start else None
-            
-            params = {"type": step_type, "description": description, **(metadata or {})}
-            
-            if is_start:
-                params["id"] = event_id
-                streaming_callback(description, MSG_TYPE.MSG_TYPE_STEP_START, params)
-                return event_id
-            else:
-                if 'id' in params:
-                    streaming_callback(description, MSG_TYPE.MSG_TYPE_STEP_END, params)
-                else: # Fallback for simple, non-duration steps
-                    streaming_callback(description, MSG_TYPE.MSG_TYPE_STEP, params)
-                return None
+            if not streaming_callback: return None
+            event_id = str(uuid.uuid4()) if event_type==MSG_TYPE.MSG_TYPE_STEP_START else event_id
+            params = {"type": event_type, "description": description, **(metadata or {})}
+            params["id"] = event_id
+            streaming_callback(description, event_type, params)
+            return event_id
+
+        def _substitute_code_uuids_recursive(data: Any, code_store: Dict[str, str]):
+            """Recursively finds and replaces code UUIDs in tool parameters."""
+            if isinstance(data, dict):
+                for key, value in data.items():
+                    if isinstance(value, str) and value in code_store:
+                        data[key] = code_store[value]
+                    else:
+                        _substitute_code_uuids_recursive(value, code_store)
+            elif isinstance(data, list):
+                for i, item in enumerate(data):
+                    if isinstance(item, str) and item in code_store:
+                        data[i] = code_store[item]
+                    else:
+                        _substitute_code_uuids_recursive(item, code_store)
 
+        discovery_step_id = log_event("**Discovering tools**",MSG_TYPE.MSG_TYPE_STEP_START)
         # --- 1. Discover Available Tools ---
         available_tools = []
         if use_mcps and self.mcp:
-            available_tools.extend(self.mcp.discover_tools(force_refresh=True))
+            discovered_tools = self.mcp.discover_tools(force_refresh=True)
+            if isinstance(use_mcps, list):
+                available_tools.extend([t for t in discovered_tools if t["name"] in use_mcps])
+            
         if use_data_store:
             for store_name in use_data_store:
                 available_tools.append({
@@ -1550,20 +1578,33 @@ Provide your response as a single JSON object inside a JSON markdown tag. Use th
                     "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"]}
                 })
         
-        formatted_tools_list = "\n".join([f"- {t['name']}: {t['description']}" for t in available_tools])
-        formatted_tools_list += "\n- request_clarification: Use if the user's request is ambiguous."
-        formatted_tools_list += "\n- final_answer: Use when you are ready to respond to the user."
+        # Add the new put_code_in_buffer tool definition
+        available_tools.append({
+            "name": "put_code_in_buffer",
+            "description": """Generates and stores code into a buffer to be used by another tool. You can put the uuid of the generated code into the fields that require long code among the tools. If no tool requires code as input do not use put_code_in_buffer. put_code_in_buffer do not execute the code nor does it audit it.""",
+            "input_schema": {"type": "object", "properties": {"prompt": {"type": "string", "description": "A detailed natural language description of the code's purpose and requirements."}, "language": {"type": "string", "description": "The programming language of the generated code. By default it uses python."}}, "required": ["prompt"]}
+        })
+        # Add the new refactor_scratchpad tool definition
+        available_tools.append({
+            "name": "refactor_scratchpad",
+            "description": "Rewrites the scratchpad content to clean it and reorganize it. Only use if the scratchpad is messy or contains too much information compared to what you need.",
+            "input_schema": {"type": "object", "properties": {}}
+        })
+
+        formatted_tools_list = "\n".join([f"**{t['name']}**:\n{t['description']}\ninput schema:\n{json.dumps(t['input_schema'])}" for t in available_tools])
+        formatted_tools_list += "\n**request_clarification**:\nUse if the user's request is ambiguous and you can not infer a clear idea of his intent. this tool has no parameters."
+        formatted_tools_list += "\n**final_answer**:\nUse when you are ready to respond to the user. this tool has no parameters."
+
+        if discovery_step_id: log_event("**Discovering tools**",MSG_TYPE.MSG_TYPE_STEP_END, event_id=discovery_step_id)
 
         # --- 2. Dynamic Reasoning Loop ---
         for i in range(max_reasoning_steps):
-            reasoning_step_id = log_step(f"Reasoning Step {i+1}/{max_reasoning_steps}", "reasoning_step", is_start=True)
-
-            user_context = f'Original User Request: "{original_user_prompt}"'
-            if images:
-                user_context += f'\n(Note: {len(images)} image(s) were provided with this request.)'
-            
-            reasoning_prompt_template = f"""You are a logical AI assistant. Your task is to achieve the user's goal by thinking step-by-step and using the available tools.
-
+            try:
+                reasoning_step_id = log_event(f"**Reasoning Step {i+1}/{max_reasoning_steps}**", MSG_TYPE.MSG_TYPE_STEP_START)
+                user_context = f'Original User Request: "{original_user_prompt}"'
+                if images: user_context += f'\n(Note: {len(images)} image(s) were provided with this request.)'
+                
+                reasoning_prompt_template = f"""
 --- AVAILABLE TOOLS ---
 {formatted_tools_list}
 --- CONTEXT ---
@@ -1577,122 +1618,155 @@ Provide your response as a single JSON object inside a JSON markdown tag. Use th
 2.  **THINK:**
     - Does the latest observation completely fulfill the user's original request?
     - If YES, your next action MUST be to use the `final_answer` tool.
-    - If NO, what is the single next logical step needed?
+    - If NO, what is the single next logical step needed? This may involve writing code first with `put_code_in_buffer`, then using another tool.
     - If you are stuck or the request is ambiguous, use `request_clarification`.
 3.  **ACT:** Formulate your decision as a JSON object.
 """
-            action_template = {
-                "thought": "My detailed analysis of the last observation and my reasoning for the next action.",
-                "action": {
-                    "tool_name": "The single tool to use (e.g., 'time_machine::get_current_time', 'final_answer').",
-                    "tool_params": {"param1": "value1"},
-                    "clarification_question": "(string, ONLY if tool_name is 'request_clarification')"
+                action_template = {
+                    "thought": "My detailed analysis of the last observation and my reasoning for the next action and how it integrates with my global plan.",
+                    "action": {
+                        "tool_name": "The single tool to use (e.g., 'put_code_in_buffer', 'time_machine::get_current_time', 'final_answer').",
+                        "tool_params": {"param1": "value1"},
+                        "clarification_question": "(string, ONLY if tool_name is 'request_clarification')"
+                    }
                 }
-            }
-            
-            structured_action_response = self.generate_code(
-                prompt=reasoning_prompt_template,
-                template=json.dumps(action_template, indent=2),
-                system_prompt=reasoning_system_prompt,
-                temperature=decision_temperature,
-                images=images if i == 0 else None
-            )
-
-            try:
-                action_data = json.loads(structured_action_response)
-                thought = action_data.get("thought", "No thought was generated.")
-                action = action_data.get("action", {})
-                tool_name = action.get("tool_name")
-                tool_params = action.get("tool_params", {})
-            except (json.JSONDecodeError, TypeError) as e:
-                current_scratchpad += f"\n\n### Step {i+1} Failure\n- **Error:** Failed to generate a valid JSON action: {e}"
-                log_step(f"\n\n### Step {i+1} Failure\n- **Error:** Failed to generate a valid JSON action: {e}", "scratchpad", is_start=False)
-                if reasoning_step_id:
-                    log_step(f"Reasoning Step {i+1}/{max_reasoning_steps}", "reasoning_step", metadata={"id": reasoning_step_id, "error": str(e)}, is_start=False)
-                break
+                if debug: log_prompt(reasoning_prompt_template, f"REASONING PROMPT (Step {i+1})")
+                structured_action_response = self.generate_code(
+                    prompt=reasoning_prompt_template, template=json.dumps(action_template, indent=2),
+                    system_prompt=reasoning_system_prompt, temperature=decision_temperature,
+                    images=images if i == 0 else None
+                )
+                if structured_action_response is None:
+                    log_event("**Error generating thought.** Retrying..", MSG_TYPE.MSG_TYPE_EXCEPTION)
+                    continue
+                if debug: log_prompt(structured_action_response, f"RAW REASONING RESPONSE (Step {i+1})")
 
-            current_scratchpad += f"\n\n### Step {i+1}: Thought\n{thought}"
-            log_step(f"\n\n### Step {i+1}: Thought\n{thought}", "scratchpad", is_start=False)
+                try:
+                    action_data = robust_json_parser(structured_action_response)
+                    thought = action_data.get("thought", "No thought was generated.")
+                    action = action_data.get("action", {})
+                    if isinstance(action,str):
+                        tool_name = action
+                        tool_params = {}
+                    else:   
+                        tool_name = action.get("tool_name")
+                        tool_params = action.get("tool_params", {})
+                except (json.JSONDecodeError, TypeError) as e:
+                    current_scratchpad += f"\n\n### Step {i+1} Failure\n- **Error:** Failed to generate a valid JSON action: {e}"
+                    log_event(f"Step Failure: Invalid JSON action.", MSG_TYPE.MSG_TYPE_EXCEPTION, metadata={"details": str(e)})
+                    if reasoning_step_id: log_event(f"**Reasoning Step {i+1}/{max_reasoning_steps}**", MSG_TYPE.MSG_TYPE_STEP_END, metadata={"error": str(e)}, event_id=reasoning_step_id)
+                    
 
-            if not tool_name:
-                current_scratchpad += f"\n\n### Step {i+1} Failure\n- **Error:** Did not specify a tool name."
-                log_step(f"\n\n### Step {i+1} Failure\n- **Error:** Did not specify a tool name.", "scratchpad", is_start=False)
-                if reasoning_step_id:
-                    log_step(f"Reasoning Step {i+1}/{max_reasoning_steps}", "reasoning_step", metadata={"id": reasoning_step_id}, is_start=False)
-                break
+                current_scratchpad += f"\n\n### Step {i+1}: Thought\n{thought}"
+                log_event(f"**Thought**:\n{thought}", MSG_TYPE.MSG_TYPE_THOUGHT_CONTENT)
 
-            if tool_name == "request_clarification":
-                clarification_question = action.get("clarification_question", "Could you please provide more details?")
-                current_scratchpad += f"\n\n### Step {i+1}: Action\n- **Action:** Decided to request clarification.\n- **Question:** {clarification_question}"
-                log_step(f"\n\n### Step {i+1}: Action\n- **Action:** Decided to request clarification.\n- **Question:** {clarification_question}", "scratchpad", is_start=False)
-                if reasoning_step_id:
-                    log_step(f"Reasoning Step {i+1}/{max_reasoning_steps}", "reasoning_step", metadata={"id": reasoning_step_id}, is_start=False)
-                return {"final_answer": clarification_question, "final_scratchpad": current_scratchpad, "tool_calls": tool_calls_this_turn, "sources": sources_this_turn, "clarification_required": True, "error": None}
-
-            if tool_name == "final_answer":
-                current_scratchpad += f"\n\n### Step {i+1}: Action\n- **Action:** Decided to formulate the final answer."
-                log_step(f"\n\n### Step {i+1}: Action\n- **Action:** Decided to formulate the final answer.", "scratchpad", is_start=False)
-                if reasoning_step_id:
-                    log_step(f"Reasoning Step {i+1}/{max_reasoning_steps}", "reasoning_step", metadata={"id": reasoning_step_id}, is_start=False)
-                break
+                if not tool_name:
+                    # Handle error...
+                    break
+                
+                # --- Handle special, non-executing tools ---
+                if tool_name == "request_clarification":
+                    # Handle clarification...
+                    return {"final_answer": action.get("clarification_question", "Could you please provide more details?"), "final_scratchpad": current_scratchpad, "tool_calls": tool_calls_this_turn, "sources": sources_this_turn, "clarification_required": True, "error": None}
+
+                if tool_name == "final_answer":
+                    current_scratchpad += f"\n\n### Step {i+1}: Action\n- **Action:** Decided to formulate the final answer."
+                    log_event("**Action**: Formulate final answer.", MSG_TYPE.MSG_TYPE_THOUGHT_CHUNK)
+                    if reasoning_step_id: log_event(f"**Reasoning Step {i+1}/{max_reasoning_steps}**",MSG_TYPE.MSG_TYPE_STEP_END, event_id=reasoning_step_id)
+                    break
 
-            tool_call_id = log_step(f"Executing tool: {tool_name}", "tool_call", metadata={"name": tool_name, "parameters": tool_params}, is_start=True)
-            tool_result = None
-            try:
-                if tool_name.startswith("research::") and use_data_store:
-                    store_name = tool_name.split("::")[1]
-                    rag_callable = use_data_store.get(store_name, {}).get("callable")
-                    query = tool_params.get("query", "")
-                    retrieved_chunks = rag_callable(query, rag_top_k=rag_top_k, rag_min_similarity_percent=rag_min_similarity_percent)
-                    if retrieved_chunks:
-                        sources_this_turn.extend(retrieved_chunks)
-                        tool_result = {"status": "success", "summary": f"Found {len(retrieved_chunks)} relevant chunks.", "chunks": retrieved_chunks}
+                # --- Handle the `put_code_in_buffer` tool specifically ---
+                if tool_name == 'put_code_in_buffer':
+                    code_gen_id = log_event(f"Generating code...", MSG_TYPE.MSG_TYPE_STEP_START, metadata={"name": "put_code_in_buffer", "id": "gencode"})
+                    code_prompt = tool_params.get("prompt", "Generate the requested code.")
+                    
+                    # Use a specific system prompt to get raw code
+                    code_generation_system_prompt = "You are a code generation assistant. Generate ONLY the raw code based on the user's request. Do not add any explanations, markdown code fences, or other text outside of the code itself."
+                    generated_code = self.generate_code(prompt=code_prompt, system_prompt=code_generation_system_prompt + "\n----\n" + reasoning_prompt_template, **llm_generation_kwargs)
+                    
+                    code_uuid = str(uuid.uuid4())
+                    generated_code_store[code_uuid] = generated_code
+                    
+                    tool_result = {"status": "success", "code_id": code_uuid, "summary": f"Code generated successfully. Use this ID in the next tool call that requires code."}
+                    tool_calls_this_turn.append({"name": "put_code_in_buffer", "params": tool_params, "result": tool_result})
+                    observation_text = f"```json\n{json.dumps(tool_result, indent=2)}\n```"
+                    current_scratchpad += f"\n\n### Step {i+1}: Observation\n- **Action:** Called `{tool_name}`\n- **Result:**\n{observation_text}"
+                    log_event(f"**Observation**:Code generated with ID: {code_uuid}", MSG_TYPE.MSG_TYPE_OBSERVATION)
+                    if code_gen_id: log_event(f"Generating code...", MSG_TYPE.MSG_TYPE_TOOL_CALL, metadata={"id": code_gen_id, "result": tool_result})
+                    if reasoning_step_id: log_event(f"**Reasoning Step {i+1}/{max_reasoning_steps}**", MSG_TYPE.MSG_TYPE_STEP_END, event_id= reasoning_step_id)
+                    continue # Go to the next reasoning step immediately
+                if tool_name == 'refactor_scratchpad':
+                    scratchpad_cleaning_prompt = f"""Enhance this scratchpad content to be more organized and comprehensive. Keep relevant experience information and remove any useless redundancies. Try to log learned things from the context so that you won't make the same mistakes again. Do not remove the main objective information or any crucial information that may be useful for the next iterations. Answer directly with the new scratchpad content without any comments.
+--- YOUR INTERNAL SCRATCHPAD (Work History & Analysis) ---
+{current_scratchpad}
+--- END OF SCRATCHPAD ---"""
+                    current_scratchpad = self.generate_text(scratchpad_cleaning_prompt)
+                    log_event(f"**New scratchpad**:\n{current_scratchpad}", MSG_TYPE.MSG_TYPE_SCRATCHPAD)
+                    
+                # --- Substitute UUIDs and Execute Standard Tools ---
+                log_event(f"**Calling tool**: `{tool_name}` with params:\n{dict_to_markdown(tool_params)}", MSG_TYPE.MSG_TYPE_TOOL_CALL)
+                _substitute_code_uuids_recursive(tool_params, generated_code_store)
+                
+                tool_call_id = log_event(f"**Executing tool**: {tool_name}",MSG_TYPE.MSG_TYPE_STEP_START, metadata={"name": tool_name, "parameters": tool_params, "id":"executing tool"})
+                tool_result = None
+                try:
+                    if tool_name.startswith("research::") and use_data_store:
+                        store_name = tool_name.split("::")[1]
+                        rag_callable = use_data_store.get(store_name, {}).get("callable")
+                        query = tool_params.get("query", "")
+                        retrieved_chunks = rag_callable(query, rag_top_k=rag_top_k, rag_min_similarity_percent=rag_min_similarity_percent)
+                        if retrieved_chunks:
+                            sources_this_turn.extend(retrieved_chunks)
+                            tool_result = {"status": "success", "summary": f"Found {len(retrieved_chunks)} relevant chunks.", "chunks": retrieved_chunks}
+                        else:
+                            tool_result = {"status": "success", "summary": "No relevant documents found."}
+                    elif use_mcps and self.mcp:
+                        mcp_result = self.mcp.execute_tool(tool_name, tool_params, lollms_client_instance=self)
+                        tool_result = {"status": "success", "output": mcp_result} if not (isinstance(mcp_result, dict) and "error" in mcp_result) else {"status": "failure", **mcp_result}
                     else:
-                        tool_result = {"status": "success", "summary": "No relevant documents found."}
-                elif use_mcps and self.mcp:
-                    mcp_result = self.mcp.execute_tool(tool_name, tool_params, lollms_client_instance=self)
-                    tool_result = {"status": "success", "output": mcp_result} if not (isinstance(mcp_result, dict) and "error" in mcp_result) else {"status": "failure", **mcp_result}
+                        tool_result = {"status": "failure", "error": f"Tool '{tool_name}' not found."}
+                except Exception as e:
+                    trace_exception(e)
+                    tool_result = {"status": "failure", "error": f"Exception executing tool: {str(e)}"}
+                
+                if tool_call_id: log_event(f"**Executing tool**: {tool_name}", MSG_TYPE.MSG_TYPE_STEP_END, metadata={"result": tool_result}, event_id= tool_call_id)
+
+                observation_text = ""
+                sanitized_result = {}
+                if isinstance(tool_result, dict):
+                    sanitized_result = tool_result.copy()
+                    summarized_fields = {}
+                    for key, value in tool_result.items():
+                        if isinstance(value, str) and key.endswith("_base64") and len(value) > 256:
+                            sanitized_result[key] = f"[Image was generated. Size: {len(value)} bytes]"
+                            continue
+                        if isinstance(value, str) and len(self.tokenize(value)) > output_summarization_threshold:
+                            if streaming_callback: streaming_callback(f"Summarizing long output from field '{key}'...", MSG_TYPE.MSG_TYPE_STEP, {"type": "summarization"})
+                            summary = self.sequential_summarize(text=value, chunk_processing_prompt=f"Summarize key info from this chunk of '{key}'.", callback=streaming_callback)
+                            summarized_fields[key] = summary
+                            sanitized_result[key] = f"[Content summarized, see summary below. Original length: {len(value)} chars]"
+                    observation_text = f"```json\n{json.dumps(sanitized_result, indent=2)}\n```"
+                    if summarized_fields:
+                        observation_text += "\n\n**Summaries of Long Outputs:**"
+                        for key, summary in summarized_fields.items():
+                            observation_text += f"\n- **Summary of '{key}':**\n{summary}"
                 else:
-                    tool_result = {"status": "failure", "error": f"Tool '{tool_name}' not found."}
-            except Exception as e:
-                trace_exception(e)
-                tool_result = {"status": "failure", "error": f"Exception executing tool: {str(e)}"}
-            
-            if tool_call_id:
-                log_step(f"Executing tool: {tool_name}", "tool_call", metadata={"id": tool_call_id, "result": tool_result}, is_start=False)
-
-            observation_text = ""
-            if isinstance(tool_result, dict):
-                sanitized_result = tool_result.copy()
-                summarized_fields = {}
-                for key, value in tool_result.items():
-                    if isinstance(value, str) and key.endswith("_base64") and len(value) > 256:
-                        sanitized_result[key] = f"[Image was generated. Size: {len(value)} bytes]"
-                        continue
-                    if isinstance(value, str) and len(self.tokenize(value)) > output_summarization_threshold:
-                        if streaming_callback: streaming_callback(f"Summarizing long output from field '{key}'...", MSG_TYPE.MSG_TYPE_STEP, {"type": "summarization"})
-                        summary = self.sequential_summarize(text=value, chunk_processing_prompt=f"Summarize key info from this chunk of '{key}'.", callback=streaming_callback)
-                        summarized_fields[key] = summary
-                        sanitized_result[key] = f"[Content summarized, see summary below. Original length: {len(value)} chars]"
-                observation_text = f"```json\n{json.dumps(sanitized_result, indent=2)}\n```"
-                if summarized_fields:
-                    observation_text += "\n\n**Summaries of Long Outputs:**"
-                    for key, summary in summarized_fields.items():
-                        observation_text += f"\n- **Summary of '{key}':**\n{summary}"
-            else:
-                observation_text = f"Tool returned non-dictionary output: {str(tool_result)}"
-
-            tool_calls_this_turn.append({"name": tool_name, "params": tool_params, "result": tool_result})
-            current_scratchpad += f"\n\n### Step {i+1}: Observation\n- **Action:** Called `{tool_name}`\n- **Result:**\n{observation_text}"
-            log_step(f"### Step {i+1}: Observation\n- **Action:** Called `{tool_name}`\n", "scratchpad", is_start=False)
-            
-            if reasoning_step_id:
-                log_step(f"Reasoning Step {i+1}/{max_reasoning_steps}", "reasoning_step", metadata={"id": reasoning_step_id}, is_start=False)
-
+                    observation_text = f"Tool returned non-dictionary output: {str(tool_result)}"
+                
+                tool_calls_this_turn.append({"name": tool_name, "params": tool_params, "result": tool_result})
+                current_scratchpad += f"\n\n### Step {i+1}: Observation\n- **Action:** Called `{tool_name}`\n- **Result:**\n{observation_text}"
+                log_event(f"**Observation**: Result from `{tool_name}`:\n{dict_to_markdown(sanitized_result)}", MSG_TYPE.MSG_TYPE_OBSERVATION)
+                
+                if reasoning_step_id: log_event(f"**Reasoning Step {i+1}/{max_reasoning_steps}**", MSG_TYPE.MSG_TYPE_STEP_END, event_id = reasoning_step_id)
+            except Exception as ex:
+                trace_exception(ex)
+                current_scratchpad += f"\n\n### Error : {ex}"
+                if reasoning_step_id: log_event(f"**Reasoning Step {i+1}/{max_reasoning_steps}**", MSG_TYPE.MSG_TYPE_STEP_END, event_id = reasoning_step_id)
+                
         # --- Final Answer Synthesis ---
-        synthesis_id = log_step("Synthesizing final answer...", "final_answer_synthesis", is_start=True)
+        synthesis_id = log_event("Synthesizing final answer...", MSG_TYPE.MSG_TYPE_STEP_START)
         
-        final_answer_prompt = f"""You are an AI assistant. Provide a final, comprehensive answer based on your work.
+        final_answer_prompt = f"""
 --- Original User Request ---
 "{original_user_prompt}"
 --- Your Internal Scratchpad (Actions Taken & Findings) ---
@@ -1702,11 +1776,25 @@ Provide your response as a single JSON object inside a JSON markdown tag. Use th
 - If images were provided by the user, incorporate your analysis of them into the answer.
 - Do not talk about your internal process unless it's necessary to explain why you couldn't find an answer.
 """
+        if debug: log_prompt(final_answer_prompt, "FINAL ANSWER SYNTHESIS PROMPT")
+
+            
         final_answer_text = self.generate_text(prompt=final_answer_prompt, system_prompt=system_prompt, images=images, stream=streaming_callback is not None, streaming_callback=streaming_callback, temperature=final_answer_temperature, **llm_generation_kwargs)
+        if type(final_answer_text) is dict:
+            if streaming_callback:
+                streaming_callback(final_answer_text["error"], MSG_TYPE.MSG_TYPE_EXCEPTION)
+            return {
+            "final_answer": "",
+            "final_scratchpad": current_scratchpad,
+            "tool_calls": tool_calls_this_turn,
+            "sources": sources_this_turn,
+            "clarification_required": False,
+            "error": final_answer_text["error"]
+        }
         final_answer = self.remove_thinking_blocks(final_answer_text)
+        if debug: log_prompt(final_answer_text, "FINAL ANSWER RESPONSE")
 
-        if synthesis_id:
-            log_step("Synthesizing final answer...", "final_answer_synthesis", metadata={"id": synthesis_id}, is_start=False)
+        if synthesis_id: log_event("Synthesizing final answer...", MSG_TYPE.MSG_TYPE_STEP_END, event_id= synthesis_id)
 
         return {
             "final_answer": final_answer,
@@ -1716,7 +1804,6 @@ Provide your response as a single JSON object inside a JSON markdown tag. Use th
             "clarification_required": False,
             "error": None
         }
-    
     def generate_code(
                         self,
                         prompt,
@@ -1795,7 +1882,7 @@ Do not split the code in multiple tags.
             while not last_code["is_complete"] and retries < max_retries:
                 retries += 1
                 ASCIIColors.info(f"Code block seems incomplete. Attempting continuation ({retries}/{max_retries})...")
-                continuation_prompt = f"{full_prompt}{code_content}\n\n{self.user_full_header}The previous code block was incomplete. Continue the code exactly from where it left off. Do not repeat the previous part. Only provide the continuation inside a single {code_tag_format} code tag.\n{self.ai_full_header}"
+                continuation_prompt = f"{prompt}\n\nAssistant:\n{code_content}\n\n{self.user_full_header}The previous code block was incomplete. Continue the code exactly from where it left off. Do not repeat the previous part. Only provide the continuation inside a single {code_tag_format} code tag.\n{self.ai_full_header}"
 
                 continuation_response = self.generate_text(
                     continuation_prompt,
@@ -1831,6 +1918,97 @@ Do not split the code in multiple tags.
 
         return code_content # Return the (potentially completed) code content or None
 
+    def generate_structured_content(
+        self,
+        prompt,
+        output_format,
+        extra_system_prompt=None,
+        **kwargs
+    ):
+        """
+        Generates structured data (a dict) from a prompt using a JSON template.
+
+        This method is a high-level wrapper around `generate_code`, specializing it 
+        for JSON output. It ensures the LLM sticks to a predefined structure,
+        and then parses the output into a Python dictionary.
+
+        Args:
+            prompt (str): 
+                The user's request (e.g., "Extract the name, age, and city of the person described").
+            output_format (dict or str): 
+                A Python dictionary or a JSON string representing the desired output 
+                structure. This will be used as a template for the LLM.
+                Example: {"name": "string", "age": "integer", "city": "string"}
+            extra_system_prompt (str, optional): 
+                Additional instructions for the system prompt, to be appended to the 
+                main instructions. Defaults to None.
+            **kwargs: 
+                Additional keyword arguments to be passed directly to the 
+                `generate_code` method (e.g., temperature, max_size, top_k, debug).
+
+        Returns:
+            dict: The parsed JSON data as a Python dictionary, or None if
+                generation or parsing fails.
+        """
+        # 1. Validate and prepare the template string from the output_format
+        if isinstance(output_format, dict):
+            # Convert the dictionary to a nicely formatted JSON string for the template
+            template_str = json.dumps(output_format, indent=2)
+        elif isinstance(output_format, str):
+            # Assume it's already a valid JSON string template
+            template_str = output_format
+        else:
+            # It's good practice to fail early for invalid input types
+            raise TypeError("output_format must be a dict or a JSON string.")
+
+        # 2. Construct a specialized system prompt for structured data generation
+        system_prompt = (
+            "You are a highly skilled AI assistant that processes user requests "
+            "and returns structured data in JSON format. You must strictly adhere "
+            "to the provided JSON template, filling in the values accurately based "
+            "on the user's prompt. Do not add any commentary, explanations, or text "
+            "outside of the final JSON code block. Your entire response must be a single "
+            "valid JSON object within a markdown code block."
+        )
+        if extra_system_prompt:
+            system_prompt += f"\n\nAdditional instructions:\n{extra_system_prompt}"
+
+        # 3. Call the underlying generate_code method with JSON-specific settings
+        if kwargs.get('debug'):
+            ASCIIColors.info("Generating structured content...")
+
+        json_string = self.generate_code(
+            prompt=prompt,
+            system_prompt=system_prompt,
+            template=template_str,
+            language="json",
+            code_tag_format="markdown", # Sticking to markdown is generally more reliable
+            **kwargs # Pass other params like temperature, top_k, etc.
+        )
+
+        # 4. Parse the result and return
+        if not json_string:
+            # generate_code already logs the error, so no need for another message
+            return None
+
+        if kwargs.get('debug'):
+            ASCIIColors.info("Parsing generated JSON string...")
+            print(f"--- Raw JSON String ---\n{json_string}\n-----------------------")
+
+        try:
+            # Use the provided robust parser
+            parsed_json = self.robust_json_parser(json_string)
+            
+            if parsed_json is None:
+                ASCIIColors.warning("Failed to robustly parse the generated JSON.")
+                return None
+                
+            return parsed_json
+
+        except Exception as e:
+            ASCIIColors.error(f"An unexpected error occurred during JSON parsing: {e}")
+            return None
+
 
     def extract_code_blocks(self, text: str, format: str = "markdown") -> List[dict]:
         """
@@ -2065,7 +2243,7 @@ Do not split the code in multiple tags.
             response_json_str = re.sub(r",\s*}", "}", response_json_str)
             response_json_str = re.sub(r",\s*]", "]", response_json_str)
 
-            parsed_response = json.loads(response_json_str)
+            parsed_response = robust_json_parser(response_json_str)
             answer = parsed_response.get("answer")
             explanation = parsed_response.get("explanation", "")
 
@@ -2159,7 +2337,7 @@ Do not split the code in multiple tags.
             response_json_str = re.sub(r",\s*}", "}", response_json_str)
             response_json_str = re.sub(r",\s*]", "]", response_json_str)
 
-            result = json.loads(response_json_str)
+            result = robust_json_parser(response_json_str)
             index = result.get("index")
             explanation = result.get("explanation", "")
 
@@ -2232,7 +2410,7 @@ Do not split the code in multiple tags.
             response_json_str = re.sub(r",\s*}", "}", response_json_str)
             response_json_str = re.sub(r",\s*]", "]", response_json_str)
 
-            result = json.loads(response_json_str)
+            result = robust_json_parser(response_json_str)
             ranking = result.get("ranking")
             explanations = result.get("explanations", []) if return_explanation else None
 
@@ -2856,5 +3034,3 @@ def chunk_text(text, tokenizer, detokenizer, chunk_size, overlap, use_separators
                 break
 
     return chunks
-
-