lollms-client 0.24.2__tar.gz → 0.25.0__tar.gz

{lollms_client-0.24.2/lollms_client.egg-info → lollms_client-0.25.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lollms_client
-Version: 0.24.2
+Version: 0.25.0
 Summary: A client library for LoLLMs generate endpoint
 Author-email: ParisNeo <parisneoai@gmail.com>
 License: Apache Software License

{lollms_client-0.24.2 → lollms_client-0.25.0}/lollms_client/__init__.py

@@ -8,7 +8,7 @@ from lollms_client.lollms_utilities import PromptReshaper # Keep general utiliti
 from lollms_client.lollms_mcp_binding import LollmsMCPBinding, LollmsMCPBindingManager
 
 
-__version__ = "0.24.2" # Updated version
+__version__ = "0.25.0" # Updated version
 
 # Optionally, you could define __all__ if you want to be explicit about exports
 __all__ = [

{lollms_client-0.24.2 → lollms_client-0.25.0}/lollms_client/llm_bindings/lollms/__init__.py

@@ -4,6 +4,7 @@ from lollms_client.lollms_llm_binding import LollmsLLMBinding
 from lollms_client.lollms_types import MSG_TYPE
 from lollms_client.lollms_utilities import encode_image
 from lollms_client.lollms_types import ELF_COMPLETION_FORMAT
+from lollms_client.lollms_discussion import LollmsDiscussion
 from ascii_colors import ASCIIColors, trace_exception
 from typing import Optional, Callable, List, Union
 import json

{lollms_client-0.24.2 → lollms_client-0.25.0}/lollms_client/llm_bindings/openai/__init__.py

@@ -30,7 +30,8 @@ class OpenAIBinding(LollmsLLMBinding):
                  model_name: str = "",
                  service_key: str = None,
                  verify_ssl_certificate: bool = True,
-                 default_completion_format: ELF_COMPLETION_FORMAT = ELF_COMPLETION_FORMAT.Chat):
+                 default_completion_format: ELF_COMPLETION_FORMAT = ELF_COMPLETION_FORMAT.Chat,
+                 **kwargs):
         """
         Initialize the OpenAI binding.
 
@@ -52,7 +53,7 @@ class OpenAIBinding(LollmsLLMBinding):
 
         if not self.service_key:
             self.service_key = os.getenv("OPENAI_API_KEY", self.service_key)
-        self.client = openai.OpenAI(api_key=self.service_key, base_url=host_address)
+        self.client = openai.OpenAI(api_key=self.service_key, base_url=None if host_address is None else host_address if len(host_address)>0 else None)
         self.completion_format = ELF_COMPLETION_FORMAT.Chat
 
     

{lollms_client-0.24.2 → lollms_client-0.25.0}/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"
@@ -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,
@@ -1444,13 +1445,14 @@ Provide your response as a single JSON object inside a JSON markdown tag. Use th
         system_prompt: str = None,
         reasoning_system_prompt: str = "You are a logical and adaptive AI assistant.",
         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,48 @@ 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:])
+            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,19 +1576,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 a block of code (e.g., Python, SQL) to be used by another tool. It returns a unique 'code_id'. You must then use this 'code_id' as the value for the code parameter in the subsequent tool call. This **does not** execute the code. It only buffers it for future use. Only use it if another tool requires code.",
+            "input_schema": {"type": "object", "properties": {"prompt": {"type": "string", "description": "A detailed natural language description of the code's purpose and requirements."}}, "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"""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.
 
 --- AVAILABLE TOOLS ---
 {formatted_tools_list}
@@ -1577,120 +1617,150 @@ 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 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: {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}")
+                    
+                # --- 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_STEP)
+                _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.
 --- Original User Request ---
@@ -1702,11 +1772,12 @@ 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)
         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 +1787,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 +1865,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,
@@ -2065,7 +2135,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 +2229,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 +2302,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 +2926,3 @@ def chunk_text(text, tokenizer, detokenizer, chunk_size, overlap, use_separators
                 break
 
     return chunks
-
-