xgae 0.1.4__py3-none-any.whl → 0.1.6__py3-none-any.whl

xgae/engine/responser/{xga_responser_base.py → responser_base.py}

@@ -1,16 +1,14 @@
 import asyncio
 import json
 import logging
-from dataclasses import dataclass
-from typing import List, Dict, Any, Optional, Tuple, Union, Literal, Callable,TypedDict,AsyncGenerator
+
 from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import List, Dict, Any, Optional, Tuple, Union, Literal, Callable, TypedDict, AsyncGenerator
 
-from xgae.engine.xga_base import XGAToolResult, XGAContextMsg, XGAToolBox
-# from xgae.utils.setup_env import langfuse
-from xgae.utils.json_helpers import (
-    safe_json_parse,
-    format_for_yield
-)
+from xgae.engine.engine_base import XGAToolResult, XGAToolBox
+from xgae.utils import langfuse
+from xgae.utils.json_helpers import safe_json_parse,format_for_yield
 from xgae.utils.xml_tool_parser import XMLToolParser
 
 # Type alias for XML result adding strategy
@@ -19,36 +17,24 @@ XmlAddingStrategy = Literal["user_message", "assistant_message", "inline_edit"]
 # Type alias for tool execution strategy
 ToolExecutionStrategy = Literal["sequential", "parallel"]
 
-class TaskResponseContext(TypedDict, total=False):
+class TaskResponserContext(TypedDict, total=False):
     is_stream: bool
     task_id: str
     task_run_id: str
     trace_id: str
     model_name: str
-    max_xml_tool_calls: int
-    add_context_msg: Callable
+    max_xml_tool_calls: int             # LLM generate max_xml_tool limit, 0 is no limit
+    add_response_msg_func: Callable
     tool_box: XGAToolBox
-    tool_execution_strategy: Literal["sequential", "parallel"]
-    xml_adding_strategy: Literal["user_message", "assistant_message", "inline_edit"]
+    tool_execution_strategy: ToolExecutionStrategy
+    xml_adding_strategy: XmlAddingStrategy
 
 
 class TaskRunContinuousState(TypedDict, total=False):
     accumulated_content: str
     auto_continue_count: int
     auto_continue: bool
-
-class Span:
-    def end(self, **kwargs):
-        pass
-
-class Trace:
-    def event(self, **kwargs):
-        pass
-
-    def span(self, **kwargs):
-        return Span()
-
-
+    max_auto_run: int
 
 @dataclass
 class ToolExecutionContext:
@@ -63,62 +49,28 @@ class ToolExecutionContext:
     parsing_details: Optional[Dict[str, Any]] = None
 
 
-@dataclass
-class ProcessorConfig:
-    """
-    Configuration for response processing and tool execution.
-
-    This class controls how the LLM's responses are processed, including how tool calls
-    are detected, executed, and their results handled.
-
-    Attributes:
-        xml_tool_calling: Enable XML-based tool call detection (<tool>...</tool>)
-        native_tool_calling: Enable OpenAI-style function calling format
-        execute_tools: Whether to automatically execute detected tool calls
-        execute_on_stream: For streaming, execute tools as they appear vs. at the end
-        tool_execution_strategy: How to execute multiple tools ("sequential" or "parallel")
-        xml_adding_strategy: How to add XML tool results to the conversation
-        max_xml_tool_calls: Maximum number of XML tool calls to process (0 = no limit)
-    """
-
-    xml_tool_calling: bool = True
-    native_tool_calling: bool = False
-
-    execute_tools: bool = True
-    execute_on_stream: bool = False
-    tool_execution_strategy: ToolExecutionStrategy = "sequential"
-    xml_adding_strategy: XmlAddingStrategy = "assistant_message"
-    max_xml_tool_calls: int = 0  # 0 means no limit
-
-    def __post_init__(self):
-        """Validate configuration after initialization."""
-        if self.xml_tool_calling is False and self.native_tool_calling is False and self.execute_tools:
-            raise ValueError(
-                "At least one tool calling format (XML or native) must be enabled if execute_tools is True")
-
-        if self.xml_adding_strategy not in ["user_message", "assistant_message", "inline_edit"]:
-            raise ValueError("xml_adding_strategy must be 'user_message', 'assistant_message', or 'inline_edit'")
-
-        if self.max_xml_tool_calls < 0:
-            raise ValueError("max_xml_tool_calls must be a non-negative integer (0 = no limit)")
-
-
 class TaskResponseProcessor(ABC):
-    def __init__(self, response_context: TaskResponseContext):
+    def __init__(self, response_context: TaskResponserContext):
         self.response_context = response_context
 
-        self.task_id = response_context.get("task_id", "")
-        self.task_run_id = response_context.get("task_run_id", "")
+        self.task_id = response_context.get("task_id")
+        self.task_run_id = response_context.get("task_run_id")
+        self.tool_execution_strategy = self.response_context.get("tool_execution_strategy", "parallel")
+        self.xml_adding_strategy = self.response_context.get("xml_adding_strategy", "user_message")
+        self.max_xml_tool_calls = self.response_context.get("max_xml_tool_calls", 0)
 
-        # Initialize the XML parser
-        self.trace = Trace()
+        self.trace_context = {
+            "trace_id": self.response_context.get("trace_id"),
+            "parent_span_id": None
+        }
 
-        self.add_message = response_context.get("add_context_msg")
-        self._add_message_with_agent_info =  self.add_message
-        self.tool_box = response_context.get("tool_box")
+        self.add_response_message = response_context.get("add_response_msg_func")
 
+        self.tool_box = response_context.get("tool_box")
         self.xml_parser = XMLToolParser()
 
+
+
     @abstractmethod
     async def process_response(self,
                                llm_response: AsyncGenerator,
@@ -127,16 +79,6 @@ class TaskResponseProcessor(ABC):
                                ) -> AsyncGenerator[Dict[str, Any], None]:
         pass
 
-    async def _yield_message(self, message_obj: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
-        """Helper to yield a message with proper formatting.
-
-        Ensures that content and metadata are JSON strings for client compatibility.
-        """
-        if message_obj:
-            return format_for_yield(message_obj)
-        return None
-
-
 
     def _extract_xml_chunks(self, content: str) -> List[str]:
         """Extract complete XML chunks using start and end pattern matching."""
@@ -231,12 +173,12 @@ class TaskResponseProcessor(ABC):
         except Exception as e:
             logging.error(f"Error extracting XML chunks: {e}")
             logging.error(f"Content was: {content}")
-            self.trace.event(name="error_extracting_xml_chunks", level="ERROR",
+            langfuse.create_event(trace_context=self.trace_context, name="error_extracting_xml_chunks", level="ERROR",
                               status_message=(f"Error extracting XML chunks: {e}"), metadata={"content": content})
 
         return chunks
 
-    def _parse_xml_tool_call(self, xml_chunk: str) -> Optional[Tuple[Dict[str, Any], Dict[str, Any]]]:
+    def _parse_xml_tool_call(self, xml_chunk: str) -> Tuple[Dict[str, Any], Dict[str, Any]]:
         """Parse XML chunk into tool call format and return parsing details.
 
         Returns:
@@ -256,6 +198,9 @@ class TaskResponseProcessor(ABC):
 
                 # Take the first tool call (should only be one per chunk)
                 xml_tool_call = parsed_calls[0]
+                if not xml_tool_call.function_name:
+                    logging.error(f"xml_tool_call function name is empty: {xml_tool_call}")
+                    return None
 
                 # Convert to the expected format
                 tool_call = {
@@ -278,7 +223,7 @@ class TaskResponseProcessor(ABC):
         except Exception as e:
             logging.error(f"Error parsing XML chunk: {e}")
             logging.error(f"XML chunk was: {xml_chunk}")
-            self.trace.event(name="error_parsing_xml_chunk", level="ERROR",
+            langfuse.create_event(trace_context=self.trace_context, name="error_parsing_xml_chunk", level="ERROR",
                               status_message=(f"Error parsing XML chunk: {e}"), metadata={"xml_chunk": xml_chunk})
             return None
 
@@ -304,21 +249,22 @@ class TaskResponseProcessor(ABC):
 
         except Exception as e:
             logging.error(f"Error parsing XML tool calls: {e}", exc_info=True)
-            self.trace.event(name="error_parsing_xml_tool_calls", level="ERROR",
+            langfuse.create_event(trace_context=self.trace_context, name="error_parsing_xml_tool_calls", level="ERROR",
                              status_message=(f"Error parsing XML tool calls: {e}"), metadata={"content": content})
 
         return parsed_data
 
-    # Tool execution methods
+
     async def _execute_tool(self, tool_call: Dict[str, Any]) -> XGAToolResult:
         """Execute a single tool call and return the result."""
-        span = self.trace.span(name=f"execute_tool.{tool_call['function_name']}", input=tool_call["arguments"])
+        span = langfuse.start_span(trace_context=self.trace_context, name=f"execute_tool.{tool_call['function_name']}", input=tool_call["arguments"])
+        self.trace_context["parent_span_id"] = span.id
         try:
             function_name = tool_call["function_name"]
             arguments = tool_call["arguments"]
 
             logging.info(f"Executing tool: {function_name} with arguments: {arguments}")
-            self.trace.event(name="executing_tool", level="DEFAULT",
+            langfuse.create_event(trace_context=self.trace_context, name="executing_tool", level="DEFAULT",
                              status_message=(f"Executing tool: {function_name} with arguments: {arguments}"))
 
             if isinstance(arguments, str):
@@ -327,21 +273,21 @@ class TaskResponseProcessor(ABC):
                 except json.JSONDecodeError:
                     arguments = {"text": arguments} # @todo modify
 
-            # Get available functions from tool registry
-            #available_functions = self.tool_registry.get_available_functions()
             result = None
-            available_function_names = self.tool_box.get_task_tool_names(self.task_id)
-            if function_name in available_function_names:
+            available_tool_names = self.tool_box.get_task_tool_names(self.task_id)
+            if function_name in available_tool_names:
                 result = await self.tool_box.call_tool(self.task_id, function_name, arguments)
             else:
                 logging.error(f"Tool function '{function_name}' not found in registry")
                 result = XGAToolResult(success=False, output=f"Tool function '{function_name}' not found")
             logging.info(f"Tool execution complete: {function_name} -> {result}")
-            span.end(status_message="tool_executed", output=result)
+            langfuse.update_current_span(status_message="tool_executed", output=result)
+
             return result
         except Exception as e:
             logging.error(f"Error executing tool {tool_call['function_name']}: {str(e)}", exc_info=True)
-            span.end(status_message="tool_execution_error", output=f"Error executing tool: {str(e)}", level="ERROR")
+
+            langfuse.update_current_span(status_message="tool_execution_error", output=f"Error executing tool: {str(e)}", level="ERROR")
             return XGAToolResult(success=False, output=f"Error executing tool: {str(e)}")
 
     async def _execute_tools(
@@ -350,7 +296,7 @@ class TaskResponseProcessor(ABC):
             execution_strategy: ToolExecutionStrategy = "sequential"
     ) -> List[Tuple[Dict[str, Any], XGAToolResult]]:
         logging.info(f"Executing {len(tool_calls)} tools with strategy: {execution_strategy}")
-        self.trace.event(name="executing_tools_with_strategy", level="DEFAULT",
+        langfuse.create_event(trace_context=self.trace_context, name="executing_tools_with_strategy", level="DEFAULT",
                          status_message=(f"Executing {len(tool_calls)} tools with strategy: {execution_strategy}"))
 
         if execution_strategy == "sequential":
@@ -376,54 +322,42 @@ class TaskResponseProcessor(ABC):
         """
         if not tool_calls:
             return []
+        tool_names = [t.get('function_name', 'unknown') for t in tool_calls]
+        logging.info(f"Executing {len(tool_calls)} tools sequentially: {tool_names}")
+        langfuse.create_event(trace_context=self.trace_context, name="executing_tools_sequentially", level="DEFAULT",
+                              status_message=(f"Executing {len(tool_calls)} tools sequentially: {tool_names}"))
 
-        try:
-            tool_names = [t.get('function_name', 'unknown') for t in tool_calls]
-            logging.info(f"Executing {len(tool_calls)} tools sequentially: {tool_names}")
-            self.trace.event(name="executing_tools_sequentially", level="DEFAULT",
-                             status_message=(f"Executing {len(tool_calls)} tools sequentially: {tool_names}"))
-
-            results = []
-            for index, tool_call in enumerate(tool_calls):
-                tool_name = tool_call.get('function_name', 'unknown')
-                logging.debug(f"Executing tool {index + 1}/{len(tool_calls)}: {tool_name}")
+        results = []
+        for index, tool_call in enumerate(tool_calls):
+            tool_name = tool_call.get('function_name', 'unknown')
+            logging.debug(f"Executing tool {index + 1}/{len(tool_calls)}: {tool_name}")
 
-                try:
-                    result = await self._execute_tool(tool_call)
-                    results.append((tool_call, result))
-                    logging.debug(f"Completed tool {tool_name} with success={result.success}")
-
-                    # Check if this is a terminating tool (ask or complete)
-                    if tool_name in ['ask', 'complete']:
-                        logging.info(f"Terminating tool '{tool_name}' executed. Stopping further tool execution.")
-                        self.trace.event(name="terminating_tool_executed", level="DEFAULT", status_message=(
+            try:
+                result = await self._execute_tool(tool_call)
+                results.append((tool_call, result))
+                logging.debug(f"Completed tool {tool_name} with success={result.success}")
+
+                # Check if this is a terminating tool (ask or complete)
+                if tool_name in ['ask', 'complete']:
+                    logging.info(f"Terminating tool '{tool_name}' executed. Stopping further tool execution.")
+                    langfuse.create_event(trace_context=self.trace_context, name="terminating_tool_executed",
+                                          level="DEFAULT", status_message=(
                             f"Terminating tool '{tool_name}' executed. Stopping further tool execution."))
-                        break  # Stop executing remaining tools
-
-                except Exception as e:
-                    logging.error(f"Error executing tool {tool_name}: {str(e)}")
-                    self.trace.event(name="error_executing_tool", level="ERROR",
-                                     status_message=(f"Error executing tool {tool_name}: {str(e)}"))
-                    error_result = XGAToolResult(success=False, output=f"Error executing tool: {str(e)}")
-                    results.append((tool_call, error_result))
-
-            logging.info(f"Sequential execution completed for {len(results)} tools (out of {len(tool_calls)} total)")
-            self.trace.event(name="sequential_execution_completed", level="DEFAULT", status_message=(
-                f"Sequential execution completed for {len(results)} tools (out of {len(tool_calls)} total)"))
-            return results
+                    break  # Stop executing remaining tools
 
-        except Exception as e:
-            logging.error(f"Error in sequential tool execution: {str(e)}", exc_info=True)
-            # Return partial results plus error results for remaining tools
-            completed_results = results if 'results' in locals() else []
-            completed_tool_names = [r[0].get('function_name', 'unknown') for r in completed_results]
-            remaining_tools = [t for t in tool_calls if t.get('function_name', 'unknown') not in completed_tool_names]
+            except Exception as e:
+                logging.error(f"Error executing tool {tool_name}: {str(e)}")
+                langfuse.create_event(trace_context=self.trace_context, name="error_executing_tool", level="ERROR",
+                                      status_message=(f"Error executing tool {tool_name}: {str(e)}"))
+                error_result = XGAToolResult(success=False, output=f"Error executing tool: {str(e)}")
+                results.append((tool_call, error_result))
 
-            # Add error results for remaining tools
-            error_results = [(tool, XGAToolResult(success=False, output=f"Execution error: {str(e)}"))
-                             for tool in remaining_tools]
+        logging.info(f"Sequential execution completed for {len(results)} tools (out of {len(tool_calls)} total)")
+        langfuse.create_event(trace_context=self.trace_context, name="sequential_execution_completed", level="DEFAULT",
+                              status_message=(
+                                  f"Sequential execution completed for {len(results)} tools (out of {len(tool_calls)} total)"))
+        return results
 
-            return completed_results + error_results
 
     async def _execute_tools_in_parallel(self, tool_calls: List[Dict[str, Any]]) -> List[Tuple[Dict[str, Any], XGAToolResult]]:
         if not tool_calls:
@@ -432,7 +366,7 @@ class TaskResponseProcessor(ABC):
         try:
             tool_names = [t.get('function_name', 'unknown') for t in tool_calls]
             logging.info(f"Executing {len(tool_calls)} tools in parallel: {tool_names}")
-            self.trace.event(name="executing_tools_in_parallel", level="DEFAULT",
+            langfuse.create_event(trace_context=self.trace_context, name="executing_tools_in_parallel", level="DEFAULT",
                              status_message=(f"Executing {len(tool_calls)} tools in parallel: {tool_names}"))
 
             # Create tasks for all tool calls
@@ -446,7 +380,7 @@ class TaskResponseProcessor(ABC):
             for i, (tool_call, result) in enumerate(zip(tool_calls, results)):
                 if isinstance(result, Exception):
                     logging.error(f"Error executing tool {tool_call.get('function_name', 'unknown')}: {str(result)}")
-                    self.trace.event(name="error_executing_tool", level="ERROR", status_message=(
+                    langfuse.create_event(trace_context=self.trace_context, name="error_executing_tool", level="ERROR", status_message=(
                         f"Error executing tool {tool_call.get('function_name', 'unknown')}: {str(result)}"))
                     # Create error result
                     error_result = XGAToolResult(success=False, output=f"Error executing tool: {str(result)}")
@@ -455,24 +389,23 @@ class TaskResponseProcessor(ABC):
                     processed_results.append((tool_call, result))
 
             logging.info(f"Parallel execution completed for {len(tool_calls)} tools")
-            self.trace.event(name="parallel_execution_completed", level="DEFAULT",
+            langfuse.create_event(trace_context=self.trace_context, name="parallel_execution_completed", level="DEFAULT",
                              status_message=(f"Parallel execution completed for {len(tool_calls)} tools"))
             return processed_results
 
         except Exception as e:
             logging.error(f"Error in parallel tool execution: {str(e)}", exc_info=True)
-            self.trace.event(name="error_in_parallel_tool_execution", level="ERROR",
+            langfuse.create_event(trace_context=self.trace_context, name="error_in_parallel_tool_execution", level="ERROR",
                              status_message=(f"Error in parallel tool execution: {str(e)}"))
             # Return error results for all tools if the gather itself fails
             return [(tool_call, XGAToolResult(success=False, output=f"Execution error: {str(e)}"))
                     for tool_call in tool_calls]
 
-    def _add_tool_result(
+    def _add_tool_messsage(
             self,
-            thread_id: str,
             tool_call: Dict[str, Any],
             result: XGAToolResult,
-            strategy: Union[XmlAddingStrategy, str] = "assistant_message",
+            strategy: XmlAddingStrategy = "assistant_message",
             assistant_message_id: Optional[str] = None,
             parsing_details: Optional[Dict[str, Any]] = None
     ) -> Optional[Dict[str, Any]]:  # Return the full message object
@@ -484,63 +417,16 @@ class TaskResponseProcessor(ABC):
             if assistant_message_id:
                 metadata["assistant_message_id"] = assistant_message_id
                 logging.info(f"Linking tool result to assistant message: {assistant_message_id}")
-                self.trace.event(name="linking_tool_result_to_assistant_message", level="DEFAULT",
+                langfuse.create_event(trace_context=self.trace_context, name="linking_tool_result_to_assistant_message", level="DEFAULT",
                                  status_message=(f"Linking tool result to assistant message: {assistant_message_id}"))
 
             # --- Add parsing details to metadata if available ---
             if parsing_details:
                 metadata["parsing_details"] = parsing_details
                 logging.info("Adding parsing_details to tool result metadata")
-                self.trace.event(name="adding_parsing_details_to_tool_result_metadata", level="DEFAULT",
+                langfuse.create_event(trace_context=self.trace_context, name="adding_parsing_details_to_tool_result_metadata", level="DEFAULT",
                                  status_message=(f"Adding parsing_details to tool result metadata"),
                                  metadata={"parsing_details": parsing_details})
-            # ---
-
-            # Check if this is a native function call (has id field)
-            if "id" in tool_call:
-                # Format as a proper tool message according to OpenAI spec
-                function_name = tool_call.get("function_name", "")
-
-                # Format the tool result content - tool role needs string content
-                if isinstance(result, str):
-                    content = result
-                elif hasattr(result, 'output'):
-                    # If it's a XGAToolResult object
-                    if isinstance(result.output, dict) or isinstance(result.output, list):
-                        # If output is already a dict or list, convert to JSON string
-                        content = json.dumps(result.output)
-                    else:
-                        # Otherwise just use the string representation
-                        content = str(result.output)
-                else:
-                    # Fallback to string representation of the whole result
-                    content = str(result)
-
-                logging.info(f"Formatted tool result content: {content[:100]}...")
-                self.trace.event(name="formatted_tool_result_content", level="DEFAULT",
-                                 status_message=(f"Formatted tool result content: {content[:100]}..."))
-
-                # Create the tool response message with proper format
-                tool_message = {
-                    "role": "tool",
-                    "tool_call_id": tool_call["id"],
-                    "name": function_name,
-                    "content": content
-                }
-
-                logging.info(f"Adding native tool result for tool_call_id={tool_call['id']} with role=tool")
-                self.trace.event(name="adding_native_tool_result_for_tool_call_id", level="DEFAULT", status_message=(
-                    f"Adding native tool result for tool_call_id={tool_call['id']} with role=tool"))
-
-                # Add as a tool message to the conversation history
-                # This makes the result visible to the LLM in the next turn
-                message_obj =  self.add_message(
-                    type="tool",  # Special type for tool responses
-                    content=tool_message,
-                    is_llm_message=True,
-                    metadata=metadata
-                )
-                return message_obj  # Return the full message object
 
             # For XML and other non-native tools, use the new structured format
             # Determine message role based on strategy
@@ -566,7 +452,7 @@ class TaskResponseProcessor(ABC):
                 metadata = {}
             metadata['frontend_content'] = structured_result_for_frontend
 
-            message_obj =  self._add_message_with_agent_info(
+            message_obj =  self.add_response_message(
                 type="tool",
                 content=result_message_for_llm,  # Save the LLM-friendly version
                 is_llm_message=True,
@@ -585,7 +471,7 @@ class TaskResponseProcessor(ABC):
             return message_obj  # Return the modified message object
         except Exception as e:
             logging.error(f"Error adding tool result: {str(e)}", exc_info=True)
-            self.trace.event(name="error_adding_tool_result", level="ERROR",
+            langfuse.create_event(trace_context=self.trace_context, name="error_adding_tool_result", level="ERROR",
                              status_message=(f"Error adding tool result: {str(e)}"),
                              metadata={"tool_call": tool_call, "result": result, "strategy": strategy,
                                        "assistant_message_id": assistant_message_id,
@@ -596,7 +482,7 @@ class TaskResponseProcessor(ABC):
                     "role": "user",
                     "content": str(result)
                 }
-                message_obj = self.add_message(
+                message_obj = self.add_response_message(
                     type="tool",
                     content=fallback_message,
                     is_llm_message=True,
@@ -605,7 +491,7 @@ class TaskResponseProcessor(ABC):
                 return message_obj  # Return the full message object
             except Exception as e2:
                 logging.error(f"Failed even with fallback message: {str(e2)}", exc_info=True)
-                self.trace.event(name="failed_even_with_fallback_message", level="ERROR",
+                langfuse.create_event(trace_context=self.trace_context, name="failed_even_with_fallback_message", level="ERROR",
                                  status_message=(f"Failed even with fallback message: {str(e2)}"),
                                  metadata={"tool_call": tool_call, "result": result, "strategy": strategy,
                                            "assistant_message_id": assistant_message_id,
@@ -670,38 +556,29 @@ class TaskResponseProcessor(ABC):
         )
 
         # Set function_name and xml_tag_name fields
-        if "xml_tag_name" in tool_call:
-            context.xml_tag_name = tool_call["xml_tag_name"]
-            context.function_name = tool_call.get("function_name", tool_call["xml_tag_name"])
-        else:
-            # For non-XML tools, use function name directly
-            context.function_name = tool_call.get("function_name", "unknown")
-            context.xml_tag_name = None
+        context.xml_tag_name = tool_call["xml_tag_name"]
+        context.function_name = tool_call["function_name"]
 
         return context
 
-    def _yield_and_save_tool_started(self, context: ToolExecutionContext, thread_id: str, thread_run_id: str) -> \
-    Optional[Dict[str, Any]]:
+    def _add_tool_start_message(self, context: ToolExecutionContext) -> Optional[Dict[str, Any]]:
         """Formats, saves, and returns a tool started status message."""
         tool_name = context.xml_tag_name or context.function_name
         content = {
             "role": "assistant", "status_type": "tool_started",
             "function_name": context.function_name, "xml_tag_name": context.xml_tag_name,
-            "message": f"Starting execution of {tool_name}", "tool_index": context.tool_index,
-            "tool_call_id": context.tool_call.get("id")  # Include tool_call ID if native
+            "message": f"Starting execution of {tool_name}", "tool_index": context.tool_index # Include tool_call ID if native
         }
-        metadata = {"thread_run_id": thread_run_id}
-        saved_message_obj =  self.add_message(
-             type="status", content=content, is_llm_message=False, metadata=metadata
+
+        return  self.add_response_message(
+             type="status", content=content, is_llm_message=False
         )
-        return saved_message_obj  # Return the full object (or None if saving failed)
 
-    def _yield_and_save_tool_completed(self, context: ToolExecutionContext, tool_message_id: Optional[str],
-                                             thread_id: str, thread_run_id: str) -> Optional[Dict[str, Any]]:
+    def _add_tool_completed_message(self, context: ToolExecutionContext, tool_message_id: Optional[str]) -> Optional[Dict[str, Any]]:
         """Formats, saves, and returns a tool completed/failed status message."""
         if not context.result:
             # Delegate to error saving if result is missing (e.g., execution failed)
-            return  self._yield_and_save_tool_error(context, thread_id, thread_run_id)
+            return  self._add_tool_error_message(context)
 
         tool_name = context.xml_tag_name or context.function_name
         status_type = "tool_completed" if context.result.success else "tool_failed"
@@ -713,7 +590,7 @@ class TaskResponseProcessor(ABC):
             "message": message_text, "tool_index": context.tool_index,
             "tool_call_id": context.tool_call.get("id")
         }
-        metadata = {"thread_run_id": thread_run_id}
+        metadata = {}
         # Add the *actual* tool result message ID to the metadata if available and successful
         if context.result.success and tool_message_id:
             metadata["linked_tool_result_message_id"] = tool_message_id
@@ -722,17 +599,15 @@ class TaskResponseProcessor(ABC):
         if context.function_name in ['ask', 'complete']:
             metadata["agent_should_terminate"] = "true"
             logging.info(f"Marking tool status for '{context.function_name}' with termination signal.")
-            self.trace.event(name="marking_tool_status_for_termination", level="DEFAULT", status_message=(
+            langfuse.create_event(trace_context=self.trace_context, name="marking_tool_status_for_termination", level="DEFAULT", status_message=(
                 f"Marking tool status for '{context.function_name}' with termination signal."))
         # <<< END ADDED >>>
 
-        saved_message_obj =  self.add_message(
+        return  self.add_response_message(
              type="status", content=content, is_llm_message=False, metadata=metadata
         )
-        return saved_message_obj
 
-    def _yield_and_save_tool_error(self, context: ToolExecutionContext, thread_id: str, thread_run_id: str) -> \
-    Optional[Dict[str, Any]]:
+    def _add_tool_error_message(self, context: ToolExecutionContext) -> Optional[Dict[str, Any]]:
         """Formats, saves, and returns a tool error status message."""
         error_msg = str(context.error) if context.error else "Unknown error during tool execution"
         tool_name = context.xml_tag_name or context.function_name
@@ -743,9 +618,9 @@ class TaskResponseProcessor(ABC):
             "tool_index": context.tool_index,
             "tool_call_id": context.tool_call.get("id")
         }
-        metadata = {"thread_run_id": thread_run_id}
+
         # Save the status message with is_llm_message=False
-        saved_message_obj =  self.add_message(
-            type="status", content=content, is_llm_message=False, metadata=metadata
+        return  self.add_response_message(
+            type="status", content=content, is_llm_message=False
         )
-        return saved_message_obj
+