xgae 0.1.7__py3-none-any.whl → 0.1.9__py3-none-any.whl

xgae/engine/engine_base.py

@@ -7,7 +7,7 @@ class XGAError(Exception):
     pass
 
 XGAMsgStatusType = Literal["error", "finish", "tool_error", "tool_started", "tool_completed", "tool_failed", "thread_run_start", "thread_run_end", "assistant_response_start", "assistant_response_end"]
-XGAResponseMsgType = Literal["user", "status",  "tool", "assistant"]
+XGAResponseMsgType = Literal["user", "status",  "tool", "assistant", "assistant_complete"]
 
 class XGAResponseMessage(TypedDict, total=False):
     message_id: str
@@ -55,4 +55,4 @@ class XGAToolBox(ABC):
 
     @abstractmethod
     async def call_tool(self, task_id: str, tool_name: str, args: Optional[Dict[str, Any]] = None) -> XGAToolResult:
-        pass
+        pass

xgae/engine/mcp_tool_box.py

@@ -8,7 +8,6 @@ from langchain_mcp_adapters.client import MultiServerMCPClient
 from langchain_mcp_adapters.tools import load_mcp_tools
 
 from xgae.engine.engine_base import XGAError, XGAToolSchema, XGAToolBox, XGAToolResult
-from xgae.utils import langfuse
 
 class XGAMcpToolBox(XGAToolBox):
     GENERAL_MCP_SERVER_NAME = "xga_general"
@@ -191,11 +190,16 @@ class XGAMcpToolBox(XGAToolBox):
 if __name__ == "__main__":
     import asyncio
     from dataclasses import asdict
+    from xgae.utils.setup_env import setup_logging
+
+    setup_logging()
 
     async def main():
+        ## Before Run Exec: uv run custom_fault_tools
+        mcp_tool_box = XGAMcpToolBox(custom_mcp_server_file="mcpservers/custom_servers.json")
+        #mcp_tool_box = XGAMcpToolBox()
+
         task_id = "task1"
-        #mcp_tool_box = XGAMcpToolBox(custom_mcp_server_file="mcpservers/custom_servers.json")
-        mcp_tool_box = XGAMcpToolBox()
         await mcp_tool_box.load_mcp_tools_schema()
         await mcp_tool_box.creat_task_tool_box(task_id=task_id, general_tools=["*"], custom_tools=["bomc_fault.*"])
         tool_schemas = mcp_tool_box.get_task_tool_schemas(task_id, "general_tool")
@@ -210,9 +214,6 @@ if __name__ == "__main__":
             print(asdict(tool_schema))
         print()
 
-        result = await mcp_tool_box.call_tool(task_id=task_id, tool_name="web_search", args={"task_id": task_id, "query": "查询天津天气"})
-        print(f"call web_search result: {result}")
-
         result = await mcp_tool_box.call_tool(task_id=task_id, tool_name="complete", args={"task_id": task_id})
         print(f"call complete result: {result}")
 

xgae/engine/prompt_builder.py

@@ -3,7 +3,7 @@ import datetime
 
 from typing import Optional, List
 
-from xgae.engine.engine_base import XGAToolSchema, XGAError
+from xgae.engine.engine_base import XGAToolSchema
 from xgae.utils.misc import read_file, format_file_with_args
 
 

xgae/engine/responser/non_stream_responser.py

@@ -1,10 +1,10 @@
 import logging
 
 from typing import List, Dict, Any, AsyncGenerator, override,Optional
+from xgae.utils.json_helpers import format_for_yield
 
 from xgae.engine.responser.responser_base import TaskResponseProcessor, TaskResponserContext, TaskRunContinuousState
-from xgae.utils import langfuse
-from xgae.utils.json_helpers import format_for_yield
+
 
 class NonStreamTaskResponser(TaskResponseProcessor):
     def __init__(self, response_context: TaskResponserContext):
@@ -16,6 +16,7 @@ class NonStreamTaskResponser(TaskResponseProcessor):
         llm_content = ""
         parsed_xml_data = []
         finish_reason = None
+        llm_count = continuous_state.get("auto_continue_count")
 
         try:
             # Extract finish_reason, content, tool calls
@@ -24,7 +25,7 @@ class NonStreamTaskResponser(TaskResponseProcessor):
                     finish_reason = llm_response.choices[0].finish_reason
                     logging.info(f"NonStreamTask：LLM response finish_reason={finish_reason}")
 
-                    langfuse.create_event(trace_context=self.trace_context, name="non_stream_processor_start", level="DEFAULT",
+                    self.root_span.event(name=f"non_stream_processor_start[{self.task_no}]({llm_count})", level="DEFAULT",
                                      status_message=(f"finish_reason={finish_reason}, tool_exec_strategy={self.tool_execution_strategy}"))
 
                 response_message = llm_response.choices[0].message if hasattr(llm_response.choices[0], 'message') else None
@@ -47,8 +48,8 @@ class NonStreamTaskResponser(TaskResponseProcessor):
                 else:
                     logging.warning(f"NonStreamTask：LLM response_message is empty")
 
-            message_data = {"role": "assistant", "content": llm_content, "index": -1} # index=-1, full llm_content
-            assistant_msg = self.add_response_message(type="assistant", content=message_data, is_llm_message=True)
+            message_data = {"role": "assistant", "content": llm_content} # index=-1, full llm_content
+            assistant_msg = self.add_response_message(type="assistant_complete", content=message_data, is_llm_message=True)
             yield assistant_msg
 
             tool_calls_to_execute = [item['tool_call'] for item in parsed_xml_data]
@@ -90,7 +91,7 @@ class NonStreamTaskResponser(TaskResponseProcessor):
 
         except Exception as e:
             logging.error(f"NonStreamTask: Error processing non-streaming response: {llm_content}")
-            langfuse.create_event(trace_context=self.trace_context, name="error_processing_non_streaming_response", level="ERROR",
+            self.root_span.event(name="error_processing_non_streaming_response", level="ERROR",
                              status_message=(f"Error processing non-streaming response: {str(e)}"))
 
             content = {"role": "system", "status_type": "error", "message": str(e)}
@@ -100,7 +101,7 @@ class NonStreamTaskResponser(TaskResponseProcessor):
 
             # Re-raise the same exception (not a new one) to ensure proper error propagation
             logging.critical(f"NonStreamTask: Re-raising error to stop further processing: {str(e)}")
-            langfuse.create_event(trace_context=self.trace_context, name="re_raising_error_to_stop_further_processing", level="CRITICAL",
+            self.root_span.event(name="re_raising_error_to_stop_further_processing", level="CRITICAL",
                              status_message=(f"Re-raising error to stop further processing: {str(e)}"))
             raise  # Use bare 'raise' to preserve the original exception with its traceback
 

xgae/engine/responser/responser_base.py

@@ -6,11 +6,13 @@ 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.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.json_helpers import safe_json_parse
 from xgae.utils.xml_tool_parser import XMLToolParser
 
+from xgae.engine.engine_base import XGAToolResult, XGAToolBox
+from xgae.engine.task_langfuse import XGATaskLangFuse
+
+
 # Type alias for XML result adding strategy
 XmlAddingStrategy = Literal["user_message", "assistant_message", "inline_edit"]
 
@@ -21,14 +23,14 @@ class TaskResponserContext(TypedDict, total=False):
     is_stream: bool
     task_id: str
     task_run_id: str
-    trace_id: str
-    root_span_id: str
+    task_no: int
     model_name: str
     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: ToolExecutionStrategy
     xml_adding_strategy: XmlAddingStrategy
+    add_response_msg_func: Callable
+    tool_box: XGAToolBox
+    task_langfuse: XGATaskLangFuse
 
 
 class TaskRunContinuousState(TypedDict, total=False):
@@ -56,22 +58,19 @@ class TaskResponseProcessor(ABC):
 
         self.task_id = response_context.get("task_id")
         self.task_run_id = response_context.get("task_run_id")
+        self.task_no = response_context.get("task_no")
         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)
 
-        self.trace_context = {
-            "trace_id": self.response_context.get("trace_id"),
-            "parent_span_id": self.response_context.get("root_span_id"),
-        }
-
+        task_langfuse = response_context.get("task_langfuse")
+        self.root_span = task_langfuse.root_span
         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,
@@ -174,7 +173,7 @@ class TaskResponseProcessor(ABC):
         except Exception as e:
             logging.error(f"Error extracting XML chunks: {e}")
             logging.error(f"Content was: {content}")
-            langfuse.create_event(trace_context=self.trace_context, name="error_extracting_xml_chunks", level="ERROR",
+            self.root_span.event(name="error_extracting_xml_chunks", level="ERROR",
                               status_message=(f"Error extracting XML chunks: {e}"), metadata={"content": content})
 
         return chunks
@@ -224,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}")
-            langfuse.create_event(trace_context=self.trace_context, name="error_parsing_xml_chunk", level="ERROR",
+            self.root_span.event(name="error_parsing_xml_chunk", level="ERROR",
                               status_message=(f"Error parsing XML chunk: {e}"), metadata={"xml_chunk": xml_chunk})
             return None
 
@@ -250,7 +249,7 @@ class TaskResponseProcessor(ABC):
 
         except Exception as e:
             logging.error(f"Error parsing XML tool calls: {e}", exc_info=True)
-            langfuse.create_event(trace_context=self.trace_context, name="error_parsing_xml_tool_calls", level="ERROR",
+            self.root_span.event(name="error_parsing_xml_tool_calls", level="ERROR",
                              status_message=(f"Error parsing XML tool calls: {e}"), metadata={"content": content})
 
         return parsed_data
@@ -258,36 +257,35 @@ class TaskResponseProcessor(ABC):
 
     async def _execute_tool(self, tool_call: Dict[str, Any]) -> XGAToolResult:
         """Execute a single tool call and return the result."""
-        with langfuse.start_as_current_span(trace_context=self.trace_context, name=f"execute_tool.{tool_call['function_name']}", input=tool_call["arguments"]
-                                            ) as exec_tool_span:
-            try:
-                function_name = tool_call["function_name"]
-                arguments = tool_call["arguments"]
-
-                logging.info(f"Executing tool: {function_name} with arguments: {arguments}")
-
-                if isinstance(arguments, str):
-                    try:
-                        arguments = safe_json_parse(arguments)
-                    except json.JSONDecodeError:
-                        arguments = {"text": arguments} # @todo modify
-
-                result = None
-                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}")
-                exec_tool_span.update(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)
+        exec_tool_span = self.root_span.span(name=f"execute_tool.{tool_call['function_name']}", input=tool_call["arguments"])
+        try:
+            function_name = tool_call["function_name"]
+            arguments = tool_call["arguments"]
+
+            logging.info(f"Executing tool: {function_name} with arguments: {arguments}")
+
+            if isinstance(arguments, str):
+                try:
+                    arguments = safe_json_parse(arguments)
+                except json.JSONDecodeError:
+                    arguments = {"text": arguments} # @todo modify
+
+            result = None
+            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}")
+            exec_tool_span.update(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)
 
-                exec_tool_span.update(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)}")
+            exec_tool_span.update(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(
             self,
@@ -321,7 +319,7 @@ class TaskResponseProcessor(ABC):
             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",
+        self.root_span.event(name="executing_tools_sequentially", level="DEFAULT",
                               status_message=(f"Executing {len(tool_calls)} tools sequentially: {tool_names}"))
 
         results = []
@@ -337,19 +335,19 @@ class TaskResponseProcessor(ABC):
                 # 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",
+                    # self.root_span.event(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)}")
-                langfuse.create_event(trace_context=self.trace_context, name="error_executing_tool", level="ERROR",
+                self.root_span.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)")
-        # langfuse.create_event(trace_context=self.trace_context, name="sequential_execution_completed", level="DEFAULT",
+        # self.root_span.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
 
@@ -361,7 +359,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}")
-            # langfuse.create_event(trace_context=self.trace_context, name="executing_tools_in_parallel", level="DEFAULT",
+            # self.root_span.event(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
@@ -375,7 +373,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)}")
-                    langfuse.create_event(trace_context=self.trace_context, name="error_executing_tool", level="ERROR", status_message=(
+                    self.root_span.event(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)}")
@@ -384,13 +382,13 @@ class TaskResponseProcessor(ABC):
                     processed_results.append((tool_call, result))
 
             logging.info(f"Parallel execution completed for {len(tool_calls)} tools")
-            # langfuse.create_event(trace_context=self.trace_context, name="parallel_execution_completed", level="DEFAULT",
+            # self.root_span.event(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)
-            langfuse.create_event(trace_context=self.trace_context, name="error_in_parallel_tool_execution", level="ERROR",
+            self.root_span.event(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)}"))
@@ -461,7 +459,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)
-            langfuse.create_event(trace_context=self.trace_context, name="error_adding_tool_result", level="ERROR",
+            self.root_span.event(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,
@@ -481,7 +479,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)
-                langfuse.create_event(trace_context=self.trace_context, name="failed_even_with_fallback_message", level="ERROR",
+                self.root_span.event(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,
@@ -589,7 +587,7 @@ 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.")
-            # langfuse.create_event(trace_context=self.trace_context, name="marking_tool_status_for_termination", level="DEFAULT", status_message=(
+            # self.root_span.event(name="marking_tool_status_for_termination", level="DEFAULT", status_message=(
             #     f"Marking tool status for '{context.function_name}' with termination signal."))
         # <<< END ADDED >>>
 

xgae/engine/responser/stream_responser.py

@@ -7,7 +7,6 @@ from dataclasses import dataclass
 from datetime import datetime, timezone
 from typing import List, Dict, Any, Optional, AsyncGenerator, override, Literal
 
-from xgae.utils import langfuse
 from xgae.engine.responser.responser_base import TaskResponseProcessor, TaskResponserContext,TaskRunContinuousState,XmlAddingStrategy,ToolExecutionStrategy
 from xgae.utils.json_helpers import (
     ensure_dict, safe_json_parse,
@@ -205,7 +204,7 @@ class StreamTaskResponser(TaskResponseProcessor):
                             __sequence += 1
                         else:
                             logging.info("XML tool call limit reached - not yielding more content chunks")
-                            langfuse.create_event(trace_context=self.trace_context, name="xml_tool_call_limit_reached", level="DEFAULT", status_message=(
+                            self.root_span.event(name="xml_tool_call_limit_reached", level="DEFAULT", status_message=(
                                 f"XML tool call limit reached - not yielding more content chunks"))
 
                         # --- Process XML Tool Calls (if enabled and limit not reached) ---
@@ -317,7 +316,7 @@ class StreamTaskResponser(TaskResponseProcessor):
 
                 if finish_reason == "xml_tool_limit_reached":
                     logging.info("Stopping stream processing after loop due to XML tool call limit")
-                    langfuse.create_event(trace_context=self.trace_context, name="stopping_stream_processing_after_loop_due_to_xml_tool_call_limit",
+                    self.root_span.event(name="stopping_stream_processing_after_loop_due_to_xml_tool_call_limit",
                                      level="DEFAULT", status_message=(
                             f"Stopping stream processing after loop due to XML tool call limit"))
                     break
@@ -352,18 +351,18 @@ class StreamTaskResponser(TaskResponseProcessor):
                     #     f"🔥 Estimated tokens – prompt: {prompt_tokens}, "
                     #     f"completion: {completion_tokens}, total: {prompt_tokens + completion_tokens}"
                     # )
-                    langfuse.create_event(trace_context=self.trace_context, name="usage_calculated_with_litellm_token_counter", level="DEFAULT",
+                    self.root_span.event(name="usage_calculated_with_litellm_token_counter", level="DEFAULT",
                                      status_message=(f"Usage calculated with litellm.token_counter"))
                 except Exception as e:
                     logging.warning(f"Failed to calculate usage: {str(e)}")
-                    langfuse.create_event(trace_context=self.trace_context, name="failed_to_calculate_usage", level="WARNING",
+                    self.root_span.event(name="failed_to_calculate_usage", level="WARNING",
                                      status_message=(f"Failed to calculate usage: {str(e)}"))
 
             # Wait for pending tool executions from streaming phase
             tool_results_buffer = []  # Stores (tool_call, result, tool_index, context)
             if pending_tool_executions:
                 logging.info(f"Waiting for {len(pending_tool_executions)} pending streamed tool executions")
-                langfuse.create_event(trace_context=self.trace_context, name="waiting_for_pending_streamed_tool_executions", level="DEFAULT", status_message=(
+                self.root_span.event(name="waiting_for_pending_streamed_tool_executions", level="DEFAULT", status_message=(
                     f"Waiting for {len(pending_tool_executions)} pending streamed tool executions"))
                 # ... (asyncio.wait logic) ...
                 pending_tasks = [execution["task"] for execution in pending_tool_executions]
@@ -387,19 +386,19 @@ class StreamTaskResponser(TaskResponseProcessor):
                                 if tool_name in ['ask', 'complete']:
                                     logging.info(
                                         f"Terminating tool '{tool_name}' completed during streaming. Setting termination flag.")
-                                    langfuse.create_event(trace_context=self.trace_context, name="terminating_tool_completed_during_streaming",
+                                    self.root_span.event(name="terminating_tool_completed_during_streaming",
                                                      level="DEFAULT", status_message=(
                                             f"Terminating tool '{tool_name}' completed during streaming. Setting termination flag."))
                                     agent_should_terminate = True
 
                             else:  # Should not happen with asyncio.wait
                                 logging.warning(f"Task for tool index {tool_idx} not done after wait.")
-                                langfuse.create_event(trace_context=self.trace_context, name="task_for_tool_index_not_done_after_wait", level="WARNING",
+                                self.root_span.event(name="task_for_tool_index_not_done_after_wait", level="WARNING",
                                                  status_message=(
                                                      f"Task for tool index {tool_idx} not done after wait."))
                         except Exception as e:
                             logging.error(f"Error getting result for pending tool execution {tool_idx}: {str(e)}")
-                            langfuse.create_event(trace_context=self.trace_context, name="error_getting_result_for_pending_tool_execution", level="ERROR",
+                            self.root_span.event(name="error_getting_result_for_pending_tool_execution", level="ERROR",
                                              status_message=(
                                                  f"Error getting result for pending tool execution {tool_idx}: {str(e)}"))
                             context.error = e
@@ -419,7 +418,7 @@ class StreamTaskResponser(TaskResponseProcessor):
                             if tool_name in ['ask', 'complete']:
                                 logging.info(
                                     f"Terminating tool '{tool_name}' completed during streaming. Setting termination flag.")
-                                langfuse.create_event(trace_context=self.trace_context, name="terminating_tool_completed_during_streaming", level="DEFAULT",
+                                self.root_span.event(name="terminating_tool_completed_during_streaming", level="DEFAULT",
                                                  status_message=(
                                                      f"Terminating tool '{tool_name}' completed during streaming. Setting termination flag."))
                                 agent_should_terminate = True
@@ -432,7 +431,7 @@ class StreamTaskResponser(TaskResponseProcessor):
                     except Exception as e:
                         logging.error(
                             f"Error getting result/yielding status for pending tool execution {tool_idx}: {str(e)}")
-                        langfuse.create_event(trace_context=self.trace_context, name="error_getting_result_yielding_status_for_pending_tool_execution",
+                        self.root_span.event(name="error_getting_result_yielding_status_for_pending_tool_execution",
                                          level="ERROR", status_message=(
                                 f"Error getting result/yielding status for pending tool execution {tool_idx}: {str(e)}"))
                         context.error = e
@@ -451,7 +450,7 @@ class StreamTaskResponser(TaskResponseProcessor):
                 if finish_msg_obj: yield format_for_yield(finish_msg_obj)
                 logging.info(
                     f"Stream finished with reason: xml_tool_limit_reached after {xml_tool_call_count} XML tool calls")
-                langfuse.create_event(trace_context=self.trace_context, name="stream_finished_with_reason_xml_tool_limit_reached_after_xml_tool_calls",
+                self.root_span.event(name="stream_finished_with_reason_xml_tool_limit_reached_after_xml_tool_calls",
                                  level="DEFAULT", status_message=(
                         f"Stream finished with reason: xml_tool_limit_reached after {xml_tool_call_count} XML tool calls"))
 
@@ -501,7 +500,7 @@ class StreamTaskResponser(TaskResponseProcessor):
                     yield format_for_yield(yield_message)
                 else:
                     logging.error(f"Failed to save final assistant message for thread {thread_id}")
-                    langfuse.create_event(trace_context=self.trace_context, name="failed_to_save_final_assistant_message_for_thread", level="ERROR",
+                    self.root_span.event(name="failed_to_save_final_assistant_message_for_thread", level="ERROR",
                                      status_message=(f"Failed to save final assistant message for thread {thread_id}"))
                     # Save and yield an error status
                     err_content = {"role": "system", "status_type": "error",
@@ -569,7 +568,7 @@ class StreamTaskResponser(TaskResponseProcessor):
                 # Populate from buffer if executed on stream
                 if config.execute_on_stream and tool_results_buffer:
                     logging.info(f"Processing {len(tool_results_buffer)} buffered tool results")
-                    langfuse.create_event(trace_context=self.trace_context, name="processing_buffered_tool_results", level="DEFAULT",
+                    self.root_span.event(name="processing_buffered_tool_results", level="DEFAULT",
                                      status_message=(f"Processing {len(tool_results_buffer)} buffered tool results"))
                     for tool_call, result, tool_idx, context in tool_results_buffer:
                         if last_assistant_message_object: context.assistant_message_id = last_assistant_message_object[
@@ -580,7 +579,7 @@ class StreamTaskResponser(TaskResponseProcessor):
                 elif final_tool_calls_to_process and not config.execute_on_stream:
                     logging.info(
                         f"Executing {len(final_tool_calls_to_process)} tools ({config.tool_execution_strategy}) after stream")
-                    langfuse.create_event(trace_context=self.trace_context, name="executing_tools_after_stream", level="DEFAULT", status_message=(
+                    self.root_span.event(name="executing_tools_after_stream", level="DEFAULT", status_message=(
                         f"Executing {len(final_tool_calls_to_process)} tools ({config.tool_execution_strategy}) after stream"))
                     results_list = await self._execute_tools(final_tool_calls_to_process,
                                                              config.tool_execution_strategy)
@@ -598,14 +597,14 @@ class StreamTaskResponser(TaskResponseProcessor):
                             tool_results_map[current_tool_idx] = (tc, res, context)
                         else:
                             logging.warning(f"Could not map result for tool index {current_tool_idx}")
-                            langfuse.create_event(trace_context=self.trace_context, name="could_not_map_result_for_tool_index", level="WARNING",
+                            self.root_span.event(name="could_not_map_result_for_tool_index", level="WARNING",
                                              status_message=(f"Could not map result for tool index {current_tool_idx}"))
                         current_tool_idx += 1
 
                 # Save and Yield each result message
                 if tool_results_map:
                     logging.info(f"Saving and yielding {len(tool_results_map)} final tool result messages")
-                    langfuse.create_event(trace_context=self.trace_context, name="saving_and_yielding_final_tool_result_messages", level="DEFAULT",
+                    self.root_span.event(name="saving_and_yielding_final_tool_result_messages", level="DEFAULT",
                                      status_message=(
                                          f"Saving and yielding {len(tool_results_map)} final tool result messages"))
                     for tool_idx in sorted(tool_results_map.keys()):
@@ -640,7 +639,7 @@ class StreamTaskResponser(TaskResponseProcessor):
                         else:
                             logging.error(
                                 f"Failed to save tool result for index {tool_idx}, not yielding result message.")
-                            langfuse.create_event(trace_context=self.trace_context, name="failed_to_save_tool_result_for_index", level="ERROR",
+                            self.root_span.event(name="failed_to_save_tool_result_for_index", level="ERROR",
                                              status_message=(
                                                  f"Failed to save tool result for index {tool_idx}, not yielding result message."))
                             # Optionally yield error status for saving failure?
@@ -658,7 +657,7 @@ class StreamTaskResponser(TaskResponseProcessor):
             if agent_should_terminate:
                 logging.info(
                     "Agent termination requested after executing ask/complete tool. Stopping further processing.")
-                langfuse.create_event(trace_context=self.trace_context, name="agent_termination_requested", level="DEFAULT",
+                self.root_span.event(name="agent_termination_requested", level="DEFAULT",
                                  status_message="Agent termination requested after executing ask/complete tool. Stopping further processing.")
 
                 # Set finish reason to indicate termination
@@ -719,7 +718,7 @@ class StreamTaskResponser(TaskResponseProcessor):
                         logging.info("Assistant response end saved for stream (before termination)")
                     except Exception as e:
                         logging.error(f"Error saving assistant response end for stream (before termination): {str(e)}")
-                        langfuse.create_event(trace_context=self.trace_context, name="error_saving_assistant_response_end_for_stream_before_termination",
+                        self.root_span.event(name="error_saving_assistant_response_end_for_stream_before_termination",
                                          level="ERROR", status_message=(
                                 f"Error saving assistant response end for stream (before termination): {str(e)}"))
 
@@ -775,12 +774,12 @@ class StreamTaskResponser(TaskResponseProcessor):
                         logging.info("Assistant response end saved for stream")
                     except Exception as e:
                         logging.error(f"Error saving assistant response end for stream: {str(e)}")
-                        langfuse.create_event(trace_context=self.trace_context, name="error_saving_assistant_response_end_for_stream", level="ERROR",
+                        self.root_span.event(name="error_saving_assistant_response_end_for_stream", level="ERROR",
                                          status_message=(f"Error saving assistant response end for stream: {str(e)}"))
 
         except Exception as e:
             logging.error(f"Error processing stream: {str(e)}", exc_info=True)
-            langfuse.create_event(trace_context=self.trace_context, name="error_processing_stream", level="ERROR",
+            self.root_span.event(name="error_processing_stream", level="ERROR",
                              status_message=(f"Error processing stream: {str(e)}"))
             # Save and yield error status message
 
@@ -794,12 +793,12 @@ class StreamTaskResponser(TaskResponseProcessor):
                 if err_msg_obj: yield format_for_yield(err_msg_obj)  # Yield the saved error message
                 # Re-raise the same exception (not a new one) to ensure proper error propagation
                 logging.critical(f"Re-raising error to stop further processing: {str(e)}")
-                langfuse.create_event(trace_context=self.trace_context, name="re_raising_error_to_stop_further_processing", level="ERROR",
+                self.root_span.event(name="re_raising_error_to_stop_further_processing", level="ERROR",
                                  status_message=(f"Re-raising error to stop further processing: {str(e)}"))
             else:
                 logging.error(f"AnthropicException - Overloaded detected - Falling back to OpenRouter: {str(e)}",
                              exc_info=True)
-                langfuse.create_event(trace_context=self.trace_context, name="anthropic_exception_overloaded_detected", level="ERROR", status_message=(
+                self.root_span.event(name="anthropic_exception_overloaded_detected", level="ERROR", status_message=(
                     f"AnthropicException - Overloaded detected - Falling back to OpenRouter: {str(e)}"))
             raise  # Use bare 'raise' to preserve the original exception with its traceback
 
@@ -822,5 +821,5 @@ class StreamTaskResponser(TaskResponseProcessor):
                     if end_msg_obj: yield format_for_yield(end_msg_obj)
                 except Exception as final_e:
                     logging.error(f"Error in finally block: {str(final_e)}", exc_info=True)
-                    langfuse.create_event(trace_context=self.trace_context, name="error_in_finally_block", level="ERROR",
+                    self.root_span.event(name="error_in_finally_block", level="ERROR",
                                      status_message=(f"Error in finally block: {str(final_e)}"))