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

xgae/engine/engine_base.py

@@ -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"
@@ -46,6 +45,12 @@ class XGAMcpToolBox(XGAToolBox):
                     task_tool_schemas[tool_schema.tool_name] = tool_schema
         task_tool_schemas.pop("end_task", None)
 
+        if len(custom_tools) == 1 and custom_tools[0] == "*":
+            custom_tools = []
+            for server_name in self.mcp_server_names:
+                if server_name != XGAMcpToolBox.GENERAL_MCP_SERVER_NAME:
+                    custom_tools.append(f"{server_name}.*")
+
         for server_tool_name in custom_tools:
             parts = server_tool_name.split(".")
             if len(parts) != 2:
@@ -188,8 +193,8 @@ if __name__ == "__main__":
 
     async def main():
         task_id = "task1"
-        #mcp_tool_box = XGAMcpToolBox(custom_mcp_server_file="mcpservers/custom_servers.json")
-        mcp_tool_box = XGAMcpToolBox()
+        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")
@@ -204,9 +209,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 engine_base import XGAToolSchema, XGAError
+from xgae.engine.engine_base import XGAToolSchema, XGAError
 from xgae.utils.misc import read_file, format_file_with_args
 
 

xgae/engine/responser/non_stream_responser.py

@@ -3,7 +3,6 @@ import logging
 from typing import List, Dict, Any, AsyncGenerator, override,Optional
 
 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):
@@ -16,6 +15,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,8 +24,8 @@ 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_streaming_finish_reason", level="DEFAULT",
-                                     status_message=(f"Non-streaming finish_reason: {finish_reason}"))
+                    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
                 if response_message:
@@ -54,8 +54,6 @@ class NonStreamTaskResponser(TaskResponseProcessor):
             tool_calls_to_execute = [item['tool_call'] for item in parsed_xml_data]
             if  len(tool_calls_to_execute) > 0:
                 logging.info(f"NonStreamTask：Executing {len(tool_calls_to_execute)} tools with strategy: {self.tool_execution_strategy}")
-                langfuse.create_event(trace_context=self.trace_context, name="executing_tools_with_strategy", level="DEFAULT", status_message=(
-                    f"NonStreamTask Executing {len(tool_calls_to_execute)} tools with strategy: {self.tool_execution_strategy}"))
 
                 tool_results = await self._execute_tools(tool_calls_to_execute, self.tool_execution_strategy)
 
@@ -92,7 +90,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)}
@@ -102,7 +100,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

@@ -7,8 +7,9 @@ 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.setup_env 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
@@ -21,7 +22,9 @@ class TaskResponserContext(TypedDict, total=False):
     is_stream: bool
     task_id: str
     task_run_id: str
+    task_no: int
     trace_id: str
+    root_span_id: str
     model_name: str
     max_xml_tool_calls: int             # LLM generate max_xml_tool limit, 0 is no limit
     add_response_msg_func: Callable
@@ -34,7 +37,7 @@ class TaskRunContinuousState(TypedDict, total=False):
     accumulated_content: str
     auto_continue_count: int
     auto_continue: bool
-    max_auto_run: int
+
 
 @dataclass
 class ToolExecutionContext:
@@ -55,15 +58,12 @@ 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": None
-        }
-
+        self.root_span = langfuse.span(trace_id=response_context.get("trace_id"), id=response_context.get("root_span_id"))
         self.add_response_message = response_context.get("add_response_msg_func")
 
         self.tool_box = response_context.get("tool_box")
@@ -173,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
@@ -223,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
 
@@ -249,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
@@ -257,15 +257,12 @@ class TaskResponseProcessor(ABC):
 
     async def _execute_tool(self, tool_call: Dict[str, Any]) -> XGAToolResult:
         """Execute a single tool call and return the result."""
-        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
+        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}")
-            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):
                 try:
@@ -281,13 +278,13 @@ class TaskResponseProcessor(ABC):
                 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}")
-            langfuse.update_current_span(status_message="tool_executed", output=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)
 
-            langfuse.update_current_span(status_message="tool_execution_error", output=f"Error executing tool: {str(e)}", level="ERROR")
+            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(
@@ -296,8 +293,6 @@ 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}")
-        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":
             return await self._execute_tools_sequentially(tool_calls)
@@ -324,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 = []
@@ -340,22 +335,20 @@ 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",
-                                          level="DEFAULT", status_message=(
-                            f"Terminating tool '{tool_name}' executed. Stopping further tool execution."))
+                    # 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",
-                              status_message=(
-                                  f"Sequential execution completed for {len(results)} tools (out of {len(tool_calls)} total)"))
+        # 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
 
 
@@ -366,8 +359,8 @@ 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",
-                             status_message=(f"Executing {len(tool_calls)} tools in parallel: {tool_names}"))
+            # 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
             tasks = [self._execute_tool(tool_call) for tool_call in tool_calls]
@@ -380,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)}")
@@ -389,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",
-                             status_message=(f"Parallel execution completed for {len(tool_calls)} tools"))
+            # 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)}"))
@@ -417,16 +410,11 @@ 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}")
-                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")
-                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})
 
             # For XML and other non-native tools, use the new structured format
             # Determine message role based on strategy
@@ -471,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,
@@ -491,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,
@@ -599,8 +587,8 @@ 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=(
-                f"Marking tool status for '{context.function_name}' with termination signal."))
+            # 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 >>>
 
         return  self.add_response_message(

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)}"))

xgae/engine/task_engine.py

@@ -1,6 +1,7 @@
 
 import logging
 import json
+import os
 
 from typing import List, Any, Dict, Optional, AsyncGenerator, Union, Literal
 from uuid import uuid4
@@ -8,20 +9,26 @@ from uuid import uuid4
 from xgae.engine.responser.responser_base import TaskResponserContext, TaskResponseProcessor, TaskRunContinuousState
 from xgae.engine.engine_base import XGAResponseMsgType, XGAResponseMessage, XGAToolBox, XGATaskResult
 
-from xgae.utils import langfuse, handle_error
-from xgae.utils.llm_client import LLMClient, LLMConfig
+from xgae.utils import handle_error
+from xgae.utils.setup_env import langfuse
+
+from xgae.utils.llm_client import LLMClient, LLMConfig, LangfuseMetadata
 
 from xgae.utils.json_helpers import format_for_yield
-from prompt_builder import XGAPromptBuilder
-from mcp_tool_box import XGAMcpToolBox
+from xgae.engine.prompt_builder import XGAPromptBuilder
+from xgae.engine.mcp_tool_box import XGAMcpToolBox
+
 
 class XGATaskEngine:
     def __init__(self,
                  session_id: Optional[str] = None,
                  task_id: Optional[str] = None,
                  agent_id: Optional[str] = None,
-                 trace_id: Optional[str] = None,
+                 general_tools: Optional[List[str]] = None,
+                 custom_tools: Optional[List[str]] = None,
                  system_prompt: Optional[str] = None,
+                 max_auto_run: Optional[int] = None,
+                 tool_exec_parallel: Optional[bool] = None,
                  llm_config: Optional[LLMConfig] = None,
                  prompt_builder: Optional[XGAPromptBuilder] = None,
                  tool_box: Optional[XGAToolBox] = None):
@@ -34,125 +41,100 @@ class XGATaskEngine:
         self.is_stream = self.llm_client.is_stream
 
         self.prompt_builder = prompt_builder or XGAPromptBuilder(system_prompt)
-        self.tool_box = tool_box or XGAMcpToolBox()
+        self.tool_box: XGAToolBox = tool_box or XGAMcpToolBox()
 
+        self.general_tools:List[str] = general_tools
+        self.custom_tools:List[str] = custom_tools
         self.task_response_msgs: List[XGAResponseMessage] = []
-        self.task_no = -1
-        self.task_run_id = f"{self.task_id}[{self.task_no}]"
-        self.trace_id :str = trace_id or langfuse.create_trace_id()
 
-    async def _post_init_(self, general_tools:List[str], custom_tools: List[str]) -> None:
-        await  self.tool_box.load_mcp_tools_schema()
-        await self.tool_box.creat_task_tool_box(self.task_id, general_tools, custom_tools)
-        general_tool_schemas = self.tool_box.get_task_tool_schemas(self.task_id, "general_tool")
-        custom_tool_schemas = self.tool_box.get_task_tool_schemas(self.task_id, "custom_tool")
+        max_auto_run = max_auto_run if max_auto_run  else int(os.getenv("MAX_AUTO_RUN", 15))
+        self.max_auto_run: int = 1 if max_auto_run <= 1 else max_auto_run
+        self.tool_exec_parallel = True if tool_exec_parallel is None else tool_exec_parallel
 
-        self.task_prompt = self.prompt_builder.build_task_prompt(self.model_name, general_tool_schemas, custom_tool_schemas)
+        self.task_no = -1
+        self.task_run_id :str = None
 
-    @classmethod
-    async def create(cls,
-                     session_id: Optional[str] = None,
-                     task_id: Optional[str] = None,
-                     agent_id: Optional[str] = None,
-                     trace_id: Optional[str] = None,
-                     system_prompt: Optional[str] = None,
-                     general_tools: Optional[List[str]] = None,
-                     custom_tools: Optional[List[str]] = None,
-                     llm_config: Optional[LLMConfig] = None,
-                     prompt_builder: Optional[XGAPromptBuilder] = None,
-                     tool_box: Optional[XGAToolBox] = None) -> 'XGATaskEngine':
-        engine: XGATaskEngine = cls(session_id=session_id,
-                                    task_id=task_id,
-                                    agent_id=agent_id,
-                                    trace_id=trace_id,
-                                    system_prompt=system_prompt,
-                                    llm_config=llm_config,
-                                    prompt_builder=prompt_builder,
-                                    tool_box=tool_box)
+        self.task_prompt :str = None
+        self.trace_id :str = None
+        self.root_span_id :str = None
+        self.root_span_name :str = None
 
-        general_tools = general_tools or ["complete", "ask"]
-        if "*" not in general_tools:
-            if "complete" not in general_tools:
-                general_tools.append("complete")
-            elif "ask" not in general_tools:
-                general_tools.append("ask")
+    async def run_task_with_final_answer(self,
+                                         task_message: Dict[str, Any],
+                                         trace_id: Optional[str] = None) -> XGATaskResult:
+        final_result:XGATaskResult = None
+        try:
+            self._init_langfuse("run_task_with_final_answer", task_message, trace_id)
+            chunks = []
+            async for chunk in self.run_task(task_message=task_message, trace_id=trace_id):
+                chunks.append(chunk)
 
-        custom_tools = custom_tools or []
-        await engine._post_init_(general_tools, custom_tools)
+            if len(chunks) > 0:
+                final_result = self._parse_final_result(chunks)
+            else:
+                final_result = XGATaskResult(type="error", content="LLM Answer is Empty")
 
-        logging.info("*"*30 + f"   XGATaskEngine Task'{engine.task_id}' Initialized   " + "*"*30)
-        logging.info(f"model_name={engine.model_name}, is_stream={engine.is_stream}, trace_id={engine.trace_id}")
-        logging.info(f"general_tools={general_tools}, custom_tools={custom_tools}")
+            return final_result
+        finally:
+            self._end_langfuse("run_task_with_final_answer", final_result)
 
-        return engine
-
-    async def run_task_with_final_answer(self,
-                       task_message: Dict[str, Any],
-                       max_auto_run: int = 25,
-                       trace_id: Optional[str] = None) -> XGATaskResult:
-        chunks = []
-        async for chunk in self.run_task(task_message=task_message, max_auto_run=max_auto_run, trace_id=trace_id):
-            chunks.append(chunk)
-        if len(chunks) > 0:
-            final_result = self._parse_final_result(chunks)
-        else:
-            final_result = XGATaskResult(type="error", content="LLM Answer is Empty")
-        return final_result
 
     async def run_task(self,
                        task_message: Dict[str, Any],
-                       max_auto_run: int = 25,
                        trace_id: Optional[str] = None) -> AsyncGenerator[Dict[str, Any], None]:
         try:
-            self.trace_id = trace_id or self.trace_id or langfuse.create_trace_id()
-
-            self.task_no += 1
-            self.task_run_id = f"{self.task_id}[{self.task_no}]"
+            await self._init_task()
+            self._init_langfuse("run_task", task_message, trace_id)
 
             self.add_response_message(type="user", content=task_message, is_llm_message=True)
 
-            continuous_state: TaskRunContinuousState = {
-                "accumulated_content": "",
-                "auto_continue_count": 0,
-                "auto_continue": False if max_auto_run <= 1 else True,
-                "max_auto_run": max_auto_run
-            }
-            async for chunk in self._run_task_auto(continuous_state):
+            async for chunk in self._run_task_auto():
                 yield chunk
         finally:
             await self.tool_box.destroy_task_tool_box(self.task_id)
+            self._end_langfuse("run_task")
 
-    async def _run_task_once(self, continuous_state: TaskRunContinuousState) -> AsyncGenerator[Dict[str, Any], None]:
-        llm_messages = [{"role": "system", "content": self.task_prompt}]
-        cxt_llm_contents = self.get_history_llm_messages()
-        llm_messages.extend(cxt_llm_contents)
 
-        partial_content = continuous_state.get('accumulated_content', '')
-        if partial_content:
-            temp_assistant_message = {
-                "role": "assistant",
-                "content": partial_content
-            }
-            llm_messages.append(temp_assistant_message)
+    async def _init_task(self) -> None:
+        self.task_no = self.task_no + 1
+        self.task_run_id = f"{self.task_id}[{self.task_no}]"
 
-        llm_response = await self.llm_client.create_completion(llm_messages)
-        response_processor = self._create_response_processer()
+        general_tools = self.general_tools or ["complete", "ask"]
+        if "*" not in general_tools:
+            if "complete" not in general_tools:
+                general_tools.append("complete")
+            elif "ask" not in general_tools:
+                general_tools.append("ask")
 
-        async for chunk in response_processor.process_response(llm_response, llm_messages, continuous_state):
-            self._logging_reponse_chunk(chunk)
-            yield chunk
+        custom_tools = self.custom_tools or []
+        if isinstance(self.tool_box, XGAMcpToolBox):
+            await  self.tool_box.load_mcp_tools_schema()
+
+        await self.tool_box.creat_task_tool_box(self.task_id, general_tools, custom_tools)
+        general_tool_schemas = self.tool_box.get_task_tool_schemas(self.task_id, "general_tool")
+        custom_tool_schemas = self.tool_box.get_task_tool_schemas(self.task_id, "custom_tool")
+
+        self.task_prompt = self.prompt_builder.build_task_prompt(self.model_name, general_tool_schemas, custom_tool_schemas)
+
+        logging.info("*" * 30 + f"   XGATaskEngine Task'{self.task_id}' Initialized   " + "*" * 30)
+        logging.info(f"model_name={self.model_name}, is_stream={self.is_stream}, trace_id={self.trace_id}")
+        logging.info(f"general_tools={general_tools}, custom_tools={custom_tools}")
 
-    async def _run_task_auto(self, continuous_state: TaskRunContinuousState) -> AsyncGenerator[Dict[str, Any], None]:
-        max_auto_run = continuous_state['max_auto_run']
-        max_auto_run = max_auto_run if max_auto_run > 0 else 1
 
+    async def _run_task_auto(self) -> AsyncGenerator[Dict[str, Any], None]:
         def update_continuous_state(_auto_continue_count,  _auto_continue):
             continuous_state["auto_continue_count"] = _auto_continue_count
             continuous_state["auto_continue"] = _auto_continue
 
+        continuous_state: TaskRunContinuousState = {
+            "accumulated_content": "",
+            "auto_continue_count": 0,
+            "auto_continue": False if self.max_auto_run <= 1 else True
+        }
+
         auto_continue_count = 0
         auto_continue = True
-        while auto_continue and auto_continue_count < max_auto_run:
+        while auto_continue and auto_continue_count < self.max_auto_run:
             auto_continue = False
 
             try:
@@ -180,7 +162,7 @@ class XGATaskEngine:
                                     auto_continue = True
                                     auto_continue_count += 1
                                     update_continuous_state(auto_continue_count, auto_continue)
-                                    logging.info(f"run_task_auto: Detected finish_reason='{finish_reason}', auto-continuing ({auto_continue_count}/{max_auto_run})")
+                                    logging.info(f"run_task_auto: Detected finish_reason='{finish_reason}', auto-continuing ({auto_continue_count}/{self.max_auto_run})")
                     except Exception as parse_error:
                         logging.error(f"run_task_auto: Error in parse chunk: {str(parse_error)}")
                         content = {"role": "system", "status_type": "error", "message": "Parse response chunk Error"}
@@ -194,6 +176,29 @@ class XGATaskEngine:
                 error_msg = self.add_response_message(type="status", content=content, is_llm_message=False)
                 yield format_for_yield(error_msg)
 
+
+    async def _run_task_once(self, continuous_state: TaskRunContinuousState) -> AsyncGenerator[Dict[str, Any], None]:
+        llm_messages = [{"role": "system", "content": self.task_prompt}]
+        cxt_llm_contents = self.get_history_llm_messages()
+        llm_messages.extend(cxt_llm_contents)
+
+        partial_content = continuous_state.get('accumulated_content', '')
+        if partial_content:
+            temp_assistant_message = {
+                "role": "assistant",
+                "content": partial_content
+            }
+            llm_messages.append(temp_assistant_message)
+
+        llm_count = continuous_state.get("auto_continue_count")
+        langfuse_metadata = self._create_llm_langfuse_meta(llm_count)
+        llm_response = await self.llm_client.create_completion(llm_messages, langfuse_metadata)
+        response_processor = self._create_response_processer()
+
+        async for chunk in response_processor.process_response(llm_response, llm_messages, continuous_state):
+            self._logging_reponse_chunk(chunk)
+            yield chunk
+
     def _parse_final_result(self, chunks: List[Dict[str, Any]]) -> XGATaskResult:
         final_result: XGATaskResult = None
         try:
@@ -245,6 +250,7 @@ class XGATaskEngine:
 
         return final_result
 
+
     def add_response_message(self, type: XGAResponseMsgType,
                              content: Union[Dict[str, Any], List[Any], str],
                              is_llm_message: bool,
@@ -289,6 +295,41 @@ class XGATaskEngine:
 
         return response_llm_contents
 
+
+    def _create_llm_langfuse_meta(self, llm_count:int)-> LangfuseMetadata:
+        generation_name = f"xga_task_engine_llm_completion[{self.task_no}]({llm_count})"
+        generation_id = f"{self.task_run_id}({llm_count})"
+        return LangfuseMetadata(
+            generation_name=generation_name,
+            generation_id=generation_id,
+            existing_trace_id=self.trace_id,
+            session_id=self.session_id,
+        )
+
+    def _init_langfuse(self,
+                       root_span_name: str,
+                       task_message: Dict[str, Any],
+                       trace_id: Optional[str] = None):
+
+        if self.root_span_id is None:
+            trace = None
+            if trace_id:
+                self.trace_id = trace_id
+                trace = langfuse.trace(id=trace_id)
+            else:
+                trace = langfuse.trace(name="xga_task_engine")
+                self.trace_id = trace.id
+
+            span = trace.span(name=root_span_name, input=task_message,metadata={"task_id": self.task_id})
+            self.root_span_id = span.id
+            self.root_span_name = root_span_name
+
+    def _end_langfuse(self, root_span_name:str, output: Optional[XGATaskResult]=None):
+        if self.root_span_id and self.root_span_name == root_span_name:
+            langfuse.span(trace_id=self.trace_id, id=self.root_span_id).end(output=output)
+            self.root_span_id = None
+            self.root_span_name = None
+
     def _create_response_processer(self) -> TaskResponseProcessor:
         response_context = self._create_response_context()
         is_stream = response_context.get("is_stream", False)
@@ -304,12 +345,14 @@ class XGATaskEngine:
             "is_stream": self.is_stream,
             "task_id": self.task_id,
             "task_run_id": self.task_run_id,
+            "task_no": self.task_no,
             "trace_id": self.trace_id,
+            "root_span_id": self.root_span_id,
             "model_name": self.model_name,
             "max_xml_tool_calls": 0,
             "add_response_msg_func": self.add_response_message,
             "tool_box": self.tool_box,
-            "tool_execution_strategy": "sequential" ,#"parallel",
+            "tool_execution_strategy": "parallel" if self.tool_exec_parallel else "sequential" ,#,
             "xml_adding_strategy": "user_message",
         }
         return response_context
@@ -338,23 +381,17 @@ if __name__ == "__main__":
 
     async def main():
         tool_box = XGAMcpToolBox(custom_mcp_server_file="mcpservers/custom_servers.json")
-        system_prompt = read_file("templates/scp_test_prompt.txt")
-        engine = await XGATaskEngine.create(tool_box=tool_box,
-                                            general_tools=[],
-                                            custom_tools=["bomc_fault.*"],
-                                            llm_config=LLMConfig(stream=False),
-                                            system_prompt=system_prompt)
-
-        final_result = await engine.run_task_with_final_answer(task_message={"role": "user", "content": "定位10.0.1.1故障"},max_auto_run=8)
+        system_prompt = read_file("templates/example_user_prompt.txt")
+        engine =  XGATaskEngine(tool_box=tool_box,
+                                    general_tools=[],
+                                    custom_tools=["*"],
+                                    llm_config=LLMConfig(stream=False),
+                                    system_prompt=system_prompt,
+                                    max_auto_run=8)
+
+        final_result = await engine.run_task_with_final_answer(task_message={"role": "user",
+                                                                             "content": "locate 10.0.0.1 fault and solution"})
         print("FINAL RESULT:", final_result)
 
-        # ====  test streaming response   ========
-        #chunks = []
-        # async for chunk in engine.run_task(task_message={"role": "user", "content": "定位10.0.0.1的故障"}, max_auto_run=8):
-        #     print(chunk)
 
-        # ====  test no tool call   ========
-        # engine = await XGATaskEngine.create(llm_config=LLMConfig(stream=False))
-        # final_result = await engine.run_task_with_final_answer(task_message={"role": "user", "content": "1+1"}, max_auto_run=2)
-        # print("FINAL RESULT:", final_result)
     asyncio.run(main())

xgae/tools/without_general_tools_app.py

@@ -0,0 +1,48 @@
+from typing import Annotated, Optional
+from pydantic import Field
+
+from mcp.server.fastmcp import FastMCP
+
+from xgae.engine.engine_base import  XGAToolResult
+
+mcp = FastMCP(name="XGAE Message Tools")
+
+@mcp.tool(
+    description="""A special tool to indicate you have completed all tasks and are about to enter complete state. Use ONLY when: 1) All tasks in todo.md are marked complete [x], 2) The user's original request has been fully addressed, 3) There are no pending actions or follow-ups required, 4) You've delivered all final outputs and results to the user. IMPORTANT: This is the ONLY way to properly terminate execution. Never use this tool unless ALL tasks are complete and verified. Always ensure you've provided all necessary outputs and references before using this tool. Include relevant attachments when the completion relates to specific files or resources."""
+)
+async def complete(task_id: str,
+                   text: Annotated[Optional[str], Field(default=None,
+                       description="Completion summary. Include: 1) Task summary 2) Key deliverables 3) Next steps 4) Impact achieved")],
+                   attachments: Annotated[Optional[str], Field(default=None,
+                       description="Comma-separated list of final outputs. Use when: 1) Completion relates to files 2) User needs to review outputs 3) Deliverables in files")]
+                   ):
+    print(f"<XGAETools-complete>: task_id={task_id}, text={text}, attachments={attachments}")
+    return XGAToolResult(success=True, output=str({"status": "complete"}))
+
+
+@mcp.tool(
+    description="""Ask user a question and wait for response. Use for: 1) Requesting clarification on ambiguous requirements, 2) Seeking confirmation before proceeding with high-impact changes, 3) Gathering additional information needed to complete a task, 4) Offering options and requesting user preference, 5) Validating assumptions when critical to task success, 6) When encountering unclear or ambiguous results during task execution, 7) When tool results don't match expectations, 8) For natural conversation and follow-up questions, 9) When research reveals multiple entities with the same name, 10) When user requirements are unclear or could be interpreted differently. IMPORTANT: Use this tool when user input is essential to proceed. Always provide clear context and options when applicable. Use natural, conversational language that feels like talking with a helpful friend. Include relevant attachments when the question relates to specific files or resources. CRITICAL: When you discover ambiguity (like multiple people with the same name), immediately stop and ask for clarification rather than making assumptions."""
+)
+async def ask(task_id: str,
+              text: Annotated[str, Field(
+                  description="Question text to present to user. Include: 1) Clear question/request 2) Context why input is needed 3) Available options 4) Impact of choices 5) Relevant constraints")],
+              attachments: Annotated[Optional[str], Field(default=None,
+                  description="Comma-separated list of files/URLs to attach. Use when: 1) Question relates to files/configs 2) User needs to review content 3) Options documented in files 4) Supporting evidence needed")]
+              ):
+    print(f"<XGAETools-ask>: task_id={task_id}, text={text}, attachments={attachments}")
+    return XGAToolResult(success=True, output=str({"status": "Awaiting user response..."}))
+
+@mcp.tool(
+    description="end task, destroy sandbox"
+)
+async def end_task(task_id: str) :
+    print(f"<XGAETools-end_task> task_id: {task_id}")
+
+
+
+def main():
+    #print("="*20 + "   XGAE Message Tools Sever Started in Stdio mode   " + "="*20)
+    mcp.run(transport="stdio")
+
+if __name__ == "__main__":
+    main()

xgae/utils/__init__.py

@@ -1,13 +1,9 @@
 import logging
 
-from .setup_env import setup_langfuse, setup_logging
-
-setup_logging()
-langfuse = setup_langfuse()
-
 def handle_error(e: Exception) -> None:
     import traceback
 
     logging.error("An error occurred: %s", str(e))
     logging.error("Traceback details:\n%s", traceback.format_exc())
     raise (e) from e
+

xgae/utils/llm_client.py

@@ -9,19 +9,29 @@ from typing import Union, Dict, Any, Optional, List, TypedDict
 from litellm.utils import ModelResponse, CustomStreamWrapper
 from openai import OpenAIError
 
+from xgae.utils.setup_env import langfuse
+
 class LLMConfig(TypedDict, total=False):
-    model: str
-    model_name: str
-    model_id: str
-    api_key: str
-    api_base: str
-    temperature: float
-    max_tokens: int
-    stream: bool
-    enable_thinking: bool
-    reasoning_effort: str
-    response_format: str
-    top_p: int
+    model: str              # Optional Name of the model to use , Override .env LLM_MODEL
+    model_name: str         # Optional Name of the model to use , use model if empty
+    model_id: str           # Optional ARN for Bedrock inference profiles, default is None
+    api_key: str            # Optional  API key, Override .env LLM_API_KEY or OS env variable
+    api_base: str           # Optional API base URL, Override .env LLM_API_BASE
+    temperature: float      # temperature: Optional Sampling temperature (0-1), Override .env LLM_TEMPERATURE
+    max_tokens: int         # max_tokens: Optional Maximum tokens in the response, Override .env LLM_MAX_TOKENS
+    stream: bool            # stream: Optional whether to stream the response, default is True
+    enable_thinking: bool   # Optional whether to enable thinking, default is  False
+    reasoning_effort: str   # Optional level of reasoning effort, default is  ‘low’
+    response_format: str    # response_format: Optional desired format for the response, default is  None
+    top_p: int              # Optional Top-p sampling parameter, default is None
+
+
+class LangfuseMetadata(TypedDict, total=False):
+    generation_name: str
+    generation_id: str
+    existing_trace_id: str
+    session_id: str
+
 
 class LLMError(Exception):
     """Base exception for LLM-related errors."""
@@ -31,26 +41,15 @@ class LLMClient:
     RATE_LIMIT_DELAY = 30
     RETRY_DELAY = 0.1
 
-    def __init__(self, llm_config: LLMConfig=None) -> None:
-        """
-        Arg: llm_config (Optional[Dict[str, Any]], optional)
-            model: Override default model to use, default set by .env LLM_MODEL
-            model_name: Optional Name of the model to use , use model if empty
-            model_id: Optional ARN for Bedrock inference profiles, default is None
-            api_key: Optional  API key, Override .env LLM_API_KEY or OS environment variable
-            api_base: Optional API base URL, Override .env LLM_API_BASE
-            temperature: Optional Sampling temperature (0-1), Override .env LLM_TEMPERATURE
-            max_tokens: Optional Maximum tokens in the response, Override .env LLM_MAX_TOKENS
-            stream: Optional whether to stream the response, default is True
-            response_format: Optional desired format for the response, default is  None
-            enable_thinking: Optional whether to enable thinking, default is  False
-            reasoning_effort: Optional level of reasoning effort, default is  ‘low’
-            top_p: Optional Top-p sampling parameter, default is None
-        """
-        llm_config = llm_config or LLMConfig()
+    langfuse_inited = False
+    langfuse_enabled = False
+
+    def __init__(self, llm_config: LLMConfig=None):
         litellm.modify_params = True
         litellm.drop_params = True
+        self._init_langfuse()
 
+        llm_config = llm_config or LLMConfig()
         self.max_retries = int(os.getenv("LLM_MAX_RETRIES", 1))
 
         env_llm_model = os.getenv("LLM_MODEL", "openai/qwen3-235b-a22b")
@@ -82,6 +81,20 @@ class LLMClient:
         self.lite_llm_params = self._prepare_llm_params(llm_config_params)
         logging.info(f"📡 LLMClient initialed : model={self.model_name}, is_stream={self.is_stream}, enable thinking={self.lite_llm_params['enable_thinking']}")
 
+    @staticmethod
+    def _init_langfuse():
+        if not LLMClient.langfuse_inited:
+            LLMClient.langfuse_inited =True
+            if langfuse and langfuse.enabled:
+                litellm.success_callback = ["langfuse"]
+                litellm.failure_callback = ["langfuse"]
+                LLMClient.langfuse_enabled = True
+                logging.info("=== LiteLLM Langfuse is enable !")
+            else:
+                LLMClient.langfuse_enabled = False
+                logging.warning("*** LiteLLM Langfuse is disable !")
+
+
 
     def _prepare_llm_params(self, llm_config_params: Dict[str, Any]) -> Dict[str, Any]:
         prepared_llm_params = llm_config_params.copy()
@@ -205,9 +218,10 @@ class LLMClient:
             logging.debug(f"LLMClient: Waiting {delay} seconds before retry llm completion...")
             await asyncio.sleep(delay)
 
-
-    async def create_completion(self, messages: List[Dict[str, Any]]) -> Union[ModelResponse, CustomStreamWrapper]:
+    async def create_completion(self, messages: List[Dict[str, Any]], langfuse_metadata: Optional[LangfuseMetadata]=None) -> Union[ModelResponse, CustomStreamWrapper]:
         complete_params = self._prepare_complete_params(messages)
+        if LLMClient.langfuse_enabled and langfuse_metadata:
+            complete_params["metadata"] = langfuse_metadata
 
         last_error = None
         for attempt in range(self.max_retries):
@@ -228,8 +242,10 @@ class LLMClient:
 if __name__ == "__main__":
     async def llm_completion():
         llm_client = LLMClient(LLMConfig(stream=False))
-        messages = [{"role": "user", "content": "今天是2025年8月15日，北京本周每天温度"}]
-        response = await llm_client.create_completion(messages)
+        messages = [{"role": "user", "content": "1+1="}]
+        trace_id = langfuse.trace(name = "xgae_litellm_test").trace_id
+
+        response = await llm_client.create_completion(messages, trace_id)
         if llm_client.is_stream:
             async for chunk in response:
                 choices = chunk.get("choices", [{}])
@@ -242,6 +258,7 @@ if __name__ == "__main__":
         else:
             print(response.choices[0].message.content)
 
+
     asyncio.run(llm_completion())
 
 

xgae/utils/setup_env.py

@@ -1,22 +1,27 @@
 import logging
 import os
 
+from dotenv import load_dotenv
 from langfuse import Langfuse
 
-def setup_logging() -> None:
+load_dotenv()
+
+def setup_logging() :
+    log_enable = bool(os.getenv("LOG_ENABLE", True))
+    if not log_enable :
+        return
+
     import colorlog
-    from dotenv import load_dotenv
-    load_dotenv()
 
     env_log_level = os.getenv("LOG_LEVEL", "INFO")
-    env_log_file = os.getenv("LOG_FILE", "log/xga.log")
+    log_file = os.getenv("LOG_FILE", "log/xga.log")
     log_level = getattr(logging, env_log_level.upper(), logging.INFO)
 
-    log_dir = os.path.dirname(env_log_file)
+    log_dir = os.path.dirname(log_file)
     if log_dir and not os.path.exists(log_dir):
         os.makedirs(log_dir, exist_ok=True)
     else:
-        os.remove(env_log_file)
+        os.remove(log_file)
 
     logger = logging.getLogger()
     for handler in logger.handlers[:]:
@@ -43,7 +48,7 @@ def setup_logging() -> None:
     console_handler = logging.StreamHandler()
     console_handler.setFormatter(console_formatter)
 
-    file_handler = logging.FileHandler(env_log_file, encoding='utf-8')
+    file_handler = logging.FileHandler(log_file, encoding='utf-8')
     file_handler.setFormatter(file_formatter)
 
     logger.addHandler(console_handler)
@@ -51,28 +56,25 @@ def setup_logging() -> None:
 
     logger.setLevel(log_level)
 
-    logging.info(f"Logger is initialized, log_level={env_log_level}, log_file={env_log_file}")
+    logging.info(f"📡 XGAE_LOGGING is initialized, log_level={env_log_level}, log_file={log_file}")
 
+setup_logging()
 
 def setup_langfuse() -> Langfuse:
-    _langfuse = None
     env_public_key = os.getenv("LANGFUSE_PUBLIC_KEY")
     env_secret_key = os.getenv("LANGFUSE_SECRET_KEY")
     env_host = os.getenv("LANGFUSE_HOST", "https://cloud.langfuse.com")
     if env_public_key and env_secret_key:
-        _langfuse = Langfuse(tracing_enabled=True,
+        _langfuse = Langfuse(enabled=True,
                             public_key=env_public_key,
                             secret_key=env_secret_key,
                             host=env_host)
-        logging.info("Langfuse initialized Successfully by Key !")
+
+        logging.info("📡 XGAE_LANGFUSE initialized Successfully by Key !")
     else:
-        _langfuse = Langfuse(tracing_enabled=False)
-        logging.warning("Not set key, Langfuse is disabled!")
+        _langfuse = Langfuse(enabled=False)
+        logging.warning("📡 XGAE_LANGFUSE Not set key, Langfuse is disabled!")
 
     return _langfuse
 
-
-if __name__ == "__main__":
-        from xgae.utils import langfuse
-        trace_id = langfuse.create_trace_id()
-        logging.warning(f"trace_id={trace_id}")
+langfuse = setup_langfuse()

xgae-0.1.8.dist-info/METADATA

@@ -0,0 +1,11 @@
+Metadata-Version: 2.4
+Name: xgae
+Version: 0.1.8
+Summary: Extreme General Agent Engine
+Requires-Python: >=3.13
+Requires-Dist: colorlog==6.9.0
+Requires-Dist: langchain-mcp-adapters==0.1.9
+Requires-Dist: langfuse==2.60.9
+Requires-Dist: langgraph==0.6.5
+Requires-Dist: litellm==1.74.15
+Requires-Dist: mcp==1.13.0

xgae-0.1.8.dist-info/RECORD

@@ -0,0 +1,19 @@
+xgae/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+xgae/engine/engine_base.py,sha256=MExJxXHkufv_a-n9tm3FZAbJyxVAwEkRC3C_WAGdu4Q,1746
+xgae/engine/mcp_tool_box.py,sha256=6bsORGK2HE-NS8vWNnbbgUQ9uGr_MI7JGO9j_FYKUKA,9903
+xgae/engine/prompt_builder.py,sha256=8_rNRJksn2QLV_K98S0x0qNeHcmxhU0kB_53IZJTGOU,4366
+xgae/engine/task_engine.py,sha256=kOntdzVtTjjakIhcBNK5vBMiBuAXPqB4R5ustY-BNfU,18931
+xgae/engine/responser/non_stream_responser.py,sha256=tXvkGHFGm2oO0IE8Bz9DFDwqemBkO92Nb6MbX1ukHeE,6352
+xgae/engine/responser/responser_base.py,sha256=FdR0yX7VfU5kW-4JmTXBfnn9AwyX1EbOAnslkC7Vcyg,29874
+xgae/engine/responser/stream_responser.py,sha256=oPGtrT1nedGMjiBAwPzUlu6Z_rPWeVSODC1xQ6D8cTY,52055
+xgae/tools/without_general_tools_app.py,sha256=QknIF4OW9xvOad8gx-F_sCBwQYXqMalnNFvYvZXkQ_I,3789
+xgae/utils/__init__.py,sha256=GPubcIs2XFPiWKnuCpevAtYEmVWKJuXlmGkmsH9qoXA,219
+xgae/utils/json_helpers.py,sha256=K1ja6GJCatrAheW9bEWAYSQbDI42__boBCZgtsv1gtk,4865
+xgae/utils/llm_client.py,sha256=Y-o26VW1MOhJYsWJ0zR4u_YXsHSEbvVPY6r90zLQJXU,13213
+xgae/utils/misc.py,sha256=EK94YesZp8AmRUqWfN-CjTxyEHPWdIIWpFNO17dzm9g,915
+xgae/utils/setup_env.py,sha256=P_p74q3nroBdTkAElfGr4QLm4fu7ZRP0R9BMGNUL010,2352
+xgae/utils/xml_tool_parser.py,sha256=EJ6BjpD4CSdmS_LqViUJ6P8H9GY2R1e4Dh8rLCR6nSE,7474
+xgae-0.1.8.dist-info/METADATA,sha256=mk0D1208wUJnbIZxyteCAjE2tlKtQixumHo0mcqD3RI,309
+xgae-0.1.8.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+xgae-0.1.8.dist-info/entry_points.txt,sha256=rhQ9Vksnu8nA78lPTjiJxOCZ5k6sH6s5YNMR68y7C-A,73
+xgae-0.1.8.dist-info/RECORD,,

xgae-0.1.8.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+xgae-tools = xgae.tools.without_general_tools_app:main

xgae-0.1.6.dist-info/METADATA

@@ -1,11 +0,0 @@
-Metadata-Version: 2.4
-Name: xgae
-Version: 0.1.6
-Summary: Extreme General Agent Engine
-Requires-Python: >=3.13
-Requires-Dist: colorlog>=6.9.0
-Requires-Dist: langchain-mcp-adapters>=0.1.4
-Requires-Dist: langfuse>=2.60.5
-Requires-Dist: langgraph>=0.3.21
-Requires-Dist: litellm>=1.74.8
-Requires-Dist: mcp>=1.12.1

xgae-0.1.6.dist-info/RECORD

@@ -1,17 +0,0 @@
-xgae/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-xgae/engine/engine_base.py,sha256=ySERuLy1YWsf-3s0NFKcyTnXQ4g69wR-cQhtnG0OFmU,1747
-xgae/engine/mcp_tool_box.py,sha256=6mdvu9-aquyLJEwebTtpa_bfGmgT1jPszKE90NIpR5c,9852
-xgae/engine/prompt_builder.py,sha256=ygFAIc4p3opIMyl6g1JeBuSiMjNVxwRloKeF2eX8R5I,4354
-xgae/engine/task_engine.py,sha256=xxAWtPfKgSpf6L7wOc243U-7YP8AC2WYoCI-FUdDpOc,18132
-xgae/engine/responser/non_stream_responser.py,sha256=QEFE4JGYVaIbFeMUMJa1Mt1uBblU_hAOywAhyp9V1k4,6634
-xgae/engine/responser/responser_base.py,sha256=aHKJ880B1ezfBWzyHoOSNVDb-CJY4ujH2MGm61aJLy8,31468
-xgae/engine/responser/stream_responser.py,sha256=5KzCHApiPplZ-zN_sbbEbSvj2rtvKWBshJKe_-x7RDI,52927
-xgae/utils/__init__.py,sha256=jChvD-p_p5gsrCZUVYPUGJs4CS9gIdNFcSOpkRpcM4Y,317
-xgae/utils/json_helpers.py,sha256=K1ja6GJCatrAheW9bEWAYSQbDI42__boBCZgtsv1gtk,4865
-xgae/utils/llm_client.py,sha256=mgzn8heUyRm92HTLEYGdfsGEpFtD-xLFr39P98_JP0s,12402
-xgae/utils/misc.py,sha256=EK94YesZp8AmRUqWfN-CjTxyEHPWdIIWpFNO17dzm9g,915
-xgae/utils/setup_env.py,sha256=Nc0HCQOnK-EGNLTWCQ9-iYysNRdIvwGhcHdqpNeV910,2407
-xgae/utils/xml_tool_parser.py,sha256=EJ6BjpD4CSdmS_LqViUJ6P8H9GY2R1e4Dh8rLCR6nSE,7474
-xgae-0.1.6.dist-info/METADATA,sha256=Q5OiPe5W3H7ym2TDPaM1x3k6jSTIol3QDyWI0dsQetw,309
-xgae-0.1.6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-xgae-0.1.6.dist-info/RECORD,,