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

xgae/engine/task_engine.py

@@ -1,4 +1,3 @@
-
 import logging
 import json
 import os
@@ -6,15 +5,15 @@ import os
 from typing import List, Any, Dict, Optional, AsyncGenerator, Union, Literal
 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 import handle_error
 from xgae.utils.llm_client import LLMClient, LLMConfig
-
 from xgae.utils.json_helpers import format_for_yield
+
+from xgae.engine.engine_base import XGAResponseMsgType, XGAResponseMessage, XGAToolBox, XGATaskResult
+from xgae.engine.task_langfuse import XGATaskLangFuse
 from xgae.engine.prompt_builder import XGAPromptBuilder
 from xgae.engine.mcp_tool_box import XGAMcpToolBox
+from xgae.engine.responser.responser_base import TaskResponserContext, TaskResponseProcessor, TaskRunContinuousState
 
 class XGATaskEngine:
     def __init__(self,
@@ -39,6 +38,7 @@ class XGATaskEngine:
 
         self.prompt_builder = prompt_builder or XGAPromptBuilder(system_prompt)
         self.tool_box: XGAToolBox = tool_box or XGAMcpToolBox()
+        self.task_langfuse: XGATaskLangFuse = None
 
         self.general_tools:List[str] = general_tools
         self.custom_tools:List[str] = custom_tools
@@ -50,21 +50,17 @@ class XGATaskEngine:
 
         self.task_no = -1
         self.task_run_id :str = None
-
         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],
                                          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
+        final_result:XGATaskResult = None
+        try:
+            await self._init_task()
+
+            self.task_langfuse.start_root_span("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):
@@ -75,8 +71,9 @@ class XGATaskEngine:
             else:
                 final_result = XGATaskResult(type="error", content="LLM Answer is Empty")
 
-            root_span.update(output=final_result)
             return final_result
+        finally:
+            self.task_langfuse.end_root_span("run_task_with_final_answer", final_result)
 
 
     async def run_task(self,
@@ -84,60 +81,58 @@ class XGATaskEngine:
                        trace_id: Optional[str] = None) -> AsyncGenerator[Dict[str, Any], None]:
         try:
             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.task_langfuse.start_root_span("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 self.max_auto_run <= 1 else True
-            }
-            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.root_span_id = None
-
+            self.task_langfuse.end_root_span("run_task")
+            self.task_run_id = None
 
     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 self.task_run_id is None:
+            self.task_no = self.task_no + 1
+            self.task_run_id = f"{self.task_id}[{self.task_no}]"
+
+            self.task_langfuse =self._create_task_langfuse()
 
-        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")
+            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")
 
-        custom_tools = self.custom_tools or []
-        if isinstance(self.tool_box, XGAMcpToolBox):
-            await  self.tool_box.load_mcp_tools_schema()
+            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")
+            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.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}")
+            logging.info("*" * 30 + f"   XGATaskEngine Task'{self.task_id}' Initialized   " + "*" * 30)
+            logging.info(f"model_name={self.model_name}, is_stream={self.is_stream}")
+            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]:
+    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 < self.max_auto_run:
@@ -196,14 +191,16 @@ class XGATaskEngine:
             }
             llm_messages.append(temp_assistant_message)
 
-        llm_response = await self.llm_client.create_completion(llm_messages, self.trace_id)
+        llm_count = continuous_state.get("auto_continue_count")
+        langfuse_metadata = self.task_langfuse.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:
@@ -241,7 +238,7 @@ class XGATaskEngine:
                             result_type = "answer" if success else "error"
                             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':
+                elif chunk_type == "assistant_complete" and finish_reason == 'stop':
                     assis_content = chunk.get('content', {})
                     result_content = assis_content.get("content", "LLM output is empty")
                     final_result = XGATaskResult(type="answer", content=result_content)
@@ -263,7 +260,7 @@ class XGATaskEngine:
         metadata = metadata or {}
         metadata["task_id"] = self.task_id
         metadata["task_run_id"] = self.task_run_id
-        metadata["trace_id"] = self.trace_id
+        metadata["trace_id"] = self.task_langfuse.trace_id
         metadata["session_id"] = self.session_id
         metadata["agent_id"] = self.agent_id
 
@@ -300,7 +297,6 @@ class XGATaskEngine:
 
         return response_llm_contents
 
-
     def _create_response_processer(self) -> TaskResponseProcessor:
         response_context = self._create_response_context()
         is_stream = response_context.get("is_stream", False)
@@ -316,18 +312,22 @@ class XGATaskEngine:
             "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,
+            "task_no": self.task_no,
             "model_name": self.model_name,
             "max_xml_tool_calls": 0,
+            "tool_execution_strategy": "parallel" if self.tool_exec_parallel else "sequential",  # ,
+            "xml_adding_strategy": "user_message",
             "add_response_msg_func": self.add_response_message,
             "tool_box": self.tool_box,
-            "tool_execution_strategy": "parallel" if self.tool_exec_parallel else "sequential" ,#,
-            "xml_adding_strategy": "user_message",
+            "task_langfuse": self.task_langfuse,
         }
         return response_context
 
 
+    def _create_task_langfuse(self)-> XGATaskLangFuse:
+        return XGATaskLangFuse(self.session_id, self.task_id, self.task_run_id, self.task_no, self.agent_id)
+
+
     def _logging_reponse_chunk(self, chunk):
         chunk_type = chunk.get('type')
         prefix = ""
@@ -350,6 +350,7 @@ if __name__ == "__main__":
     from xgae.utils.misc import read_file
 
     async def main():
+        # Before Run Exec: uv run custom_fault_tools
         tool_box = XGAMcpToolBox(custom_mcp_server_file="mcpservers/custom_servers.json")
         system_prompt = read_file("templates/example_user_prompt.txt")
         engine =  XGATaskEngine(tool_box=tool_box,
@@ -357,7 +358,9 @@ if __name__ == "__main__":
                                     custom_tools=["*"],
                                     llm_config=LLMConfig(stream=False),
                                     system_prompt=system_prompt,
-                                    max_auto_run=8)
+                                    max_auto_run=8,
+                                    session_id="session_1",
+                                    agent_id="agent_1",)
 
         final_result = await engine.run_task_with_final_answer(task_message={"role": "user",
                                                                              "content": "locate 10.0.0.1 fault and solution"})

xgae/engine/task_langfuse.py

@@ -0,0 +1,63 @@
+
+from typing import Any, Dict, Optional
+from langfuse import Langfuse
+
+from xgae.utils.setup_env import setup_langfuse, setup_env_logging
+from xgae.utils.llm_client import LangfuseMetadata
+from xgae.engine.engine_base import XGATaskResult
+
+setup_env_logging()
+langfuse:Langfuse = setup_langfuse()
+
+class XGATaskLangFuse:
+    def __init__(self,
+                 session_id: str,
+                 task_id:str,
+                 task_run_id: str,
+                 task_no: int,
+                 agent_id: str) -> None:
+        self.session_id = session_id
+        self.task_id = task_id
+        self.task_run_id = task_run_id
+        self.task_no = task_no
+        self.agent_id = agent_id
+
+        self.trace_id = None
+        self.root_span = None
+        self.root_span_name = None
+
+
+    def start_root_span(self,
+                        root_span_name: str,
+                        task_message: Dict[str, Any],
+                        trace_id: Optional[str] = None):
+        if self.root_span 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
+
+            metadata = {"task_id": self.task_id, "session_id": self.session_id, "agent_id": self.agent_id}
+            self.root_span = trace.span(id=self.task_run_id, name=root_span_name, input=task_message,metadata=metadata)
+            self.root_span_name = root_span_name
+
+
+    def end_root_span(self, root_span_name:str, output: Optional[XGATaskResult]=None):
+        if self.root_span and self.root_span_name == root_span_name:
+            self.root_span.end(output=output)
+            self.root_span = None
+            self.root_span_name = None
+
+
+    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,
+        )

xgae/tools/without_general_tools_app.py

@@ -41,7 +41,7 @@ async def end_task(task_id: str) :
 
 
 def main():
-    #print("="*20 + "   XGAE Message Tools Sever Started in Stdio mode   " + "="*20)
+    print("="*20 + "   XGAE Message Tools Sever Started in Stdio mode   " + "="*20)
     mcp.run(transport="stdio")
 
 if __name__ == "__main__":

xgae/utils/__init__.py

@@ -1,13 +1,15 @@
 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
+
+
+def to_bool(value: str) -> bool:
+    if value is None:
+        return False
+
+    return True if value.lower() == "true" else False

xgae/utils/json_helpers.py

@@ -26,8 +26,7 @@ def ensure_dict(value: Union[str, Dict[str, Any], None], default: Dict[str, Any]
     Returns:
         A dictionary
     """
-    if default is None:
-        default = {}
+    default = default or {}
         
     if value is None:
         return default
@@ -64,8 +63,7 @@ def ensure_list(value: Union[str, List[Any], None], default: List[Any] = None) -
     Returns:
         A list
     """
-    if default is None:
-        default = []
+    default = default or []
         
     if value is None:
         return default
@@ -84,7 +82,7 @@ def ensure_list(value: Union[str, List[Any], None], default: List[Any] = None) -
             
     return default
 
-
+# @todo if all call value is str, delete useless code
 def safe_json_parse(value: Union[str, Dict, List, Any], default: Any = None) -> Any:
     """
     Safely parse a value that might be JSON string or already parsed.
@@ -105,16 +103,13 @@ def safe_json_parse(value: Union[str, Dict, List, Any], default: Any = None) ->
     # If it's already a dict or list, return as-is
     if isinstance(value, (dict, list)):
         return value
-        
-    # If it's a string, try to parse it
+
     if isinstance(value, str):
         try:
             return json.loads(value)
         except (json.JSONDecodeError, TypeError):
-            # If it's not valid JSON, return the string itself
             return value
-            
-    # For any other type, return as-is
+
     return value
 
 
@@ -137,9 +132,8 @@ def to_json_string(value: Any) -> str:
             json.loads(value)
             return value  # It's already a JSON string
         except (json.JSONDecodeError, TypeError):
-            # It's a plain string, encode it as JSON
-            return json.dumps(value)
-    
+            pass
+
     # For all other types, convert to JSON
     return json.dumps(value)
 

xgae/utils/llm_client.py

@@ -1,27 +1,36 @@
 import asyncio
 import json
-import logging
 import os
+import logging
 import litellm
 
 from typing import Union, Dict, Any, Optional, List, TypedDict
-
-from litellm.utils import ModelResponse, CustomStreamWrapper
 from openai import OpenAIError
+from litellm.utils import ModelResponse, CustomStreamWrapper
+
+from xgae.utils.setup_env import setup_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,27 +40,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
-        """
+    langfuse_inited = False
+    langfuse_enabled = False
 
-        llm_config = llm_config or LLMConfig()
+    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")
@@ -83,6 +80,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
+            env_langfuse = setup_langfuse()
+            if env_langfuse and env_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()
@@ -206,10 +217,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]], trace_id: Optional[str]=None) -> 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 trace_id:
-            complete_params["litellm_trace_id"] = trace_id
+        if LLMClient.langfuse_enabled and langfuse_metadata:
+            complete_params["metadata"] = langfuse_metadata
 
         last_error = None
         for attempt in range(self.max_retries):
@@ -228,13 +239,24 @@ class LLMClient:
         raise LLMError(f"LLM completion failed after {self.max_retries} attempts !")
 
 if __name__ == "__main__":
-    from xgae.utils import langfuse
+    from xgae.utils.setup_env import setup_logging
+
+    setup_logging()
+    langfuse = setup_langfuse()
 
     async def llm_completion():
         llm_client = LLMClient(LLMConfig(stream=False))
-        messages = [{"role": "user", "content": "今天是2025年8月15日，北京本周每天温度"}]
-        trace_id = langfuse.create_trace_id()
-        response = await llm_client.create_completion(messages, trace_id)
+
+        messages = [{"role": "user", "content": "1+1="}]
+        trace_id = langfuse.trace(name = "xgae_litellm_test").trace_id
+        meta = LangfuseMetadata(
+                generation_name="llm_completion_test",
+                generation_id="generation_id",
+                existing_trace_id=trace_id,
+                session_id="session_0",
+        )
+
+        response = await llm_client.create_completion(messages, meta)
         if llm_client.is_stream:
             async for chunk in response:
                 choices = chunk.get("choices", [{}])
@@ -247,6 +269,7 @@ if __name__ == "__main__":
         else:
             print(response.choices[0].message.content)
 
+
     asyncio.run(llm_completion())
 
 

xgae/utils/misc.py

@@ -4,6 +4,8 @@ import sys
 
 from typing import Any, Dict
 
+from xgae.utils import handle_error
+
 def read_file(file_path: str) -> str:
     if not os.path.exists(file_path):
         logging.error(f"File '{file_path}' not found")
@@ -31,4 +33,4 @@ def format_file_with_args(file_content:str, args: Dict[str, Any])-> str:
     finally:
         sys.stdout = original_stdout
 
-    return formated
+    return formated