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

xgae/engine/{xga_engine.py → task_engine.py}

@@ -1,25 +1,31 @@
 
 import logging
 import json
+import os
 
-from typing import List, Any, Dict, Optional, AsyncGenerator, cast, Union, Literal
+from typing import List, Any, Dict, Optional, AsyncGenerator, Union, Literal
 from uuid import uuid4
 
-from xgae.engine.responser.xga_responser_base import TaskResponseContext, TaskResponseProcessor, TaskRunContinuousState
-from xgae.engine.xga_base import XGAResponseMsg, XGAToolBox, XGATaskResult
+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.setup_env import langfuse
-from xgae.utils.utils import handle_error
 
-from xga_prompt_builder import XGAPromptBuilder
-from xga_mcp_tool_box import XGAMcpToolBox
+from xgae.utils.json_helpers import format_for_yield
+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,
+                 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):
@@ -32,121 +38,109 @@ 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.task_response_msgs: List[XGAResponseMsg] = []
-        self.task_no = -1
-        self.task_run_id = f"{self.task_id}[{self.task_no}]"
-        self.trace_id = None
-
-    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")
-
-        self.task_prompt = self.prompt_builder.build_task_prompt(self.model_name, general_tool_schemas, custom_tool_schemas)
+        self.tool_box: XGAToolBox = tool_box or XGAMcpToolBox()
 
-    @classmethod
-    async def create(cls,
-                     session_id: Optional[str] = None,
-                     task_id: Optional[str] = None,
-                     agent_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,
-                                    system_prompt=system_prompt,
-                                    llm_config=llm_config,
-                                    prompt_builder=prompt_builder,
-                                    tool_box=tool_box)
+        self.general_tools:List[str] = general_tools
+        self.custom_tools:List[str] = custom_tools
+        self.task_response_msgs: List[XGAResponseMessage] = []
 
-        general_tools = general_tools or ["complete", "ask"]
-        custom_tools = custom_tools or []
-        await engine._post_init_(general_tools, custom_tools)
+        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
 
-        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}")
+        self.task_no = -1
+        self.task_run_id :str = None
 
-        return engine
+        self.task_prompt :str = None
+        self.trace_id :str = None
+        self.root_span_id :str = None
 
     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)
+                                         task_message: Dict[str, Any],
+                                         trace_id: Optional[str] = None) -> XGATaskResult:
+        self.trace_id = trace_id or langfuse.create_trace_id()
+        with langfuse.start_as_current_span(trace_context={"trace_id": self.trace_id},
+                                            name="run_task_with_final_answer",
+                                            input=task_message,
+                                            metadata={"task_id": self.task_id},
+                                            ) as root_span:
+            self.root_span_id = root_span.id
+
+            chunks = []
+            async for chunk in self.run_task(task_message=task_message, 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")
+
+            root_span.update(output=final_result)
+            return final_result
 
-        final_result = self._parse_final_result(chunks)
-        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()
+            await self._init_task()
+            if self.root_span_id is None:
+                self.trace_id = trace_id or langfuse.create_trace_id()
+                with langfuse.start_as_current_span(trace_context={"trace_id": self.trace_id},
+                                                    name="run_task",
+                                                    input=task_message
+                                                    ) as root_span:
+                    self.root_span_id = root_span.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 self.max_auto_run <= 1 else True
+            }
+            async for chunk in self._run_task_auto(continuous_state):
+                yield chunk
+        finally:
+            await self.tool_box.destroy_task_tool_box(self.task_id)
+            self.root_span_id = None
 
-            self.task_no += 1
-            self.task_run_id = f"{self.task_id}[{self.task_no}]"
 
-            self.add_response_msg(type="user", content=task_message, is_llm_message=True)
+    async def _init_task(self) -> None:
+        self.task_no = self.task_no + 1
+        self.task_run_id = f"{self.task_id}[{self.task_no}]"
 
-            if max_auto_run <= 1:
-                continuous_state:TaskRunContinuousState  = {
-                    "accumulated_content": "",
-                    "auto_continue_count": 0,
-                    "auto_continue": False
-                }
-                async for chunk in self._run_task_once(continuous_state):
-                    yield chunk
-            else:
-                async for chunk in self._run_task_auto(max_auto_run):
-                    yield chunk
-        finally:
-            await self.tool_box.destroy_task_tool_box(self.task_id)
+        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 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_response_llm_contents()
-        llm_messages.extend(cxt_llm_contents)
+        custom_tools = self.custom_tools or []
+        if isinstance(self.tool_box, XGAMcpToolBox):
+            await  self.tool_box.load_mcp_tools_schema()
 
-        partial_content = continuous_state.get('accumulated_content', '')
-        if partial_content:
-            temp_assistant_message = {
-                "role": "assistant",
-                "content": partial_content
-            }
-            llm_messages.append(temp_assistant_message)
+        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")
 
-        llm_response = await self.llm_client.create_completion(llm_messages)
-        response_processor = self._create_response_processer()
+        self.task_prompt = self.prompt_builder.build_task_prompt(self.model_name, general_tool_schemas, custom_tool_schemas)
 
-        async for chunk in response_processor.process_response(llm_response, llm_messages, continuous_state):
-            self._reponse_chunk_log(chunk)
-            yield chunk
+        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, max_auto_run: int) -> AsyncGenerator[Dict[str, Any], None]:
-        continuous_state: TaskRunContinuousState = {
-            "accumulated_content": "",
-            "auto_continue_count": 0,
-            "auto_continue": True
-        }
 
+    async def _run_task_auto(self, continuous_state: TaskRunContinuousState) -> 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
 
         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:
@@ -163,30 +157,52 @@ class XGATaskEngine:
                             elif status_type == 'finish':
                                 finish_reason = content.get('finish_reason', None)
                                 if finish_reason == 'completed':
-                                    logging.warning(f"run_task_auto: Detected finish_reason='completed', Task Completed Success !")
+                                    logging.info(f"run_task_auto: Detected finish_reason='completed', TASK_COMPLETE Success !")
                                     auto_continue = False
                                     break
                                 elif finish_reason == 'xml_tool_limit_reached':
-                                    logging.warning(f"run_task_auto: Detected finish_reason='xml_tool_limit_reached', stopping auto-continue")
+                                    logging.warning(f"run_task_auto: Detected finish_reason='xml_tool_limit_reached', stop auto-continue")
                                     auto_continue = False
                                     break
                                 elif finish_reason == 'stop' or finish_reason == 'length': # 'length' never occur
                                     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})")
-                    except StopAsyncIteration:
-                        pass
+                                    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"}
-                        error_msg = self.add_response_msg(type="status", content=content, is_llm_message=False)
-                        yield error_msg
+                        handle_error(parse_error)
+                        error_msg = self.add_response_message(type="status", content=content, is_llm_message=False)
+                        yield format_for_yield(error_msg)
             except Exception as run_error:
                 logging.error(f"run_task_auto: Call task_run_once error: {str(run_error)}")
                 content = {"role": "system", "status_type": "error", "message": "Call task_run_once error"}
-                error_msg = self.add_response_msg(type="status", content=content, is_llm_message=False)
-                yield error_msg
+                handle_error(run_error)
+                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_response = await self.llm_client.create_completion(llm_messages, self.trace_id)
+        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
@@ -200,14 +216,11 @@ class XGATaskEngine:
                     if status_type == "error":
                         error = status_content.get('message', 'Unknown error')
                         final_result = XGATaskResult(type="error", content=error)
-                        break
                     elif status_type == "finish":
                         finish_reason = status_content.get('finish_reason', None)
                         if finish_reason == 'xml_tool_limit_reached':
                             error = "Completed due to over task max_auto_run limit !"
                             final_result = XGATaskResult(type="error", content=error)
-                            break
-                        continue
                 elif chunk_type == "tool" and finish_reason in ['completed', 'stop']:
                     tool_content = json.loads(chunk.get('content', '{}'))
                     tool_execution = tool_content.get('tool_execution')
@@ -229,9 +242,10 @@ class XGATaskEngine:
                             result_content = f"Task execute '{tool_name}' {result_type}: {output}"
                             final_result = XGATaskResult(type=result_type, content=result_content)
                 elif chunk_type == "assistant" and finish_reason == 'stop':
-                    assis_content = chunk.get('content', '{}')
+                    assis_content = chunk.get('content', {})
                     result_content = assis_content.get("content", "LLM output is empty")
                     final_result = XGATaskResult(type="answer", content=result_content)
+
                 if final_result is not None:
                     break
         except Exception as e:
@@ -241,74 +255,80 @@ class XGATaskEngine:
 
         return final_result
 
-    def add_response_msg(self, type: Literal["user", "status", "tool", "assistant", "assistant_response_end"],
-                         content: Union[Dict[str, Any], List[Any], str],
-                         is_llm_message: bool,
-                         metadata: Optional[Dict[str, Any]]=None)-> XGAResponseMsg:
-        message = XGAResponseMsg(
+
+    def add_response_message(self, type: XGAResponseMsgType,
+                             content: Union[Dict[str, Any], List[Any], str],
+                             is_llm_message: bool,
+                             metadata: Optional[Dict[str, Any]]=None)-> XGAResponseMessage:
+        metadata = metadata or {}
+        metadata["task_id"] = self.task_id
+        metadata["task_run_id"] = self.task_run_id
+        metadata["trace_id"] = self.trace_id
+        metadata["session_id"] = self.session_id
+        metadata["agent_id"] = self.agent_id
+
+        message = XGAResponseMessage(
             message_id = f"xga_msg_{uuid4()}",
             type = type,
-            content = content,
             is_llm_message=is_llm_message,
-            metadata = metadata,
-            session_id = self.session_id,
-            agent_id = self.agent_id,
-            task_id = self.task_id,
-            task_run_id = self.task_run_id,
-            trace_id = self.trace_id
+            content = content,
+            metadata = metadata
         )
         self.task_response_msgs.append(message)
 
         return message
 
-    def _get_response_llm_contents (self) -> List[Dict[str, Any]]:
+    def get_history_llm_messages (self) -> List[Dict[str, Any]]:
         llm_messages = []
         for message in self.task_response_msgs:
             if message["is_llm_message"]:
                 llm_messages.append(message)
 
-        cxt_llm_contents = []
+        response_llm_contents = []
         for llm_message in llm_messages:
             content = llm_message["content"]
             # @todo content List type
             if isinstance(content, str):
                 try:
                     _content = json.loads(content)
-                    cxt_llm_contents.append(_content)
+                    response_llm_contents.append(_content)
                 except json.JSONDecodeError as e:
                     logging.error(f"get_context_llm_contents: Failed to decode json, content=:{content}")
                     handle_error(e)
             else:
-                cxt_llm_contents.append(content)
+                response_llm_contents.append(content)
+
+        return response_llm_contents
 
-        return cxt_llm_contents
 
     def _create_response_processer(self) -> TaskResponseProcessor:
         response_context = self._create_response_context()
         is_stream = response_context.get("is_stream", False)
         if is_stream:
-            from xgae.engine.responser.xga_stream_responser import StreamTaskResponser
+            from xgae.engine.responser.stream_responser import StreamTaskResponser
             return StreamTaskResponser(response_context)
         else:
-            from xgae.engine.responser.xga_non_stream_responser import NonStreamTaskResponser
+            from xgae.engine.responser.non_stream_responser import NonStreamTaskResponser
             return NonStreamTaskResponser(response_context)
 
-    def _create_response_context(self) -> TaskResponseContext:
-        response_context: TaskResponseContext = {
+    def _create_response_context(self) -> TaskResponserContext:
+        response_context: TaskResponserContext = {
             "is_stream": self.is_stream,
             "task_id": self.task_id,
             "task_run_id": self.task_run_id,
             "trace_id": self.trace_id,
+            "root_span_id": self.root_span_id,
             "model_name": self.model_name,
             "max_xml_tool_calls": 0,
-            "add_context_msg": self.add_response_msg,
+            "add_response_msg_func": self.add_response_message,
             "tool_box": self.tool_box,
-            "tool_execution_strategy": "parallel",
+            "tool_execution_strategy": "parallel" if self.tool_exec_parallel else "sequential" ,#,
             "xml_adding_strategy": "user_message",
         }
         return response_context
 
-    def _reponse_chunk_log(self, chunk):
+
+    def _logging_reponse_chunk(self, chunk):
         chunk_type = chunk.get('type')
         prefix = ""
 
@@ -327,21 +347,21 @@ class XGATaskEngine:
 
 if __name__ == "__main__":
     import asyncio
-    from xgae.utils.utils import read_file
+    from xgae.utils.misc import read_file
+
     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)
-        # engine = await XGATaskEngine.create(llm_config=LLMConfig(stream=False))
-        #chunks = []
-        # async for chunk in engine.run_task(task_message={"role": "user", "content": "定位10.0.0.1的故障"},max_auto_run=8):
-        #     print(chunk)
-        #final_result = await engine.run_task_with_final_answer(task_message={"role": "user", "content": "1+1"}, max_auto_run=2)
-
-        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)
+
+
     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

@@ -0,0 +1,13 @@
+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

@@ -47,6 +47,7 @@ class LLMClient:
             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()
         litellm.modify_params = True
         litellm.drop_params = True
@@ -205,9 +206,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]], trace_id: Optional[str]=None) -> Union[ModelResponse, CustomStreamWrapper]:
         complete_params = self._prepare_complete_params(messages)
+        if trace_id:
+            complete_params["litellm_trace_id"] = trace_id
 
         last_error = None
         for attempt in range(self.max_retries):
@@ -226,10 +228,13 @@ class LLMClient:
         raise LLMError(f"LLM completion failed after {self.max_retries} attempts !")
 
 if __name__ == "__main__":
+    from xgae.utils import langfuse
+
     async def llm_completion():
         llm_client = LLMClient(LLMConfig(stream=False))
         messages = [{"role": "user", "content": "今天是2025年8月15日，北京本周每天温度"}]
-        response = await llm_client.create_completion(messages)
+        trace_id = langfuse.create_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", [{}])

xgae/utils/{utils.py → misc.py}

@@ -1,17 +1,9 @@
 import logging
 import os
 import sys
-import datetime
 
 from typing import Any, Dict
 
-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
-
 def read_file(file_path: str) -> str:
     if not os.path.exists(file_path):
         logging.error(f"File '{file_path}' not found")