xgae 0.2.3__py3-none-any.whl → 0.3.1__py3-none-any.whl

xgae/engine/engine_base.py

@@ -6,7 +6,7 @@ class XGAError(Exception):
     """Custom exception for errors in the XGA system."""
     pass
 
-XGAMsgStatusType = Literal["error", "finish", "tool_started", "tool_completed", "tool_error", "tool_failed"]
+XGAMsgStatusType = Literal["error", "stop", "finish", "tool_started", "tool_completed", "tool_error", "tool_failed"]
 XGAResponseMsgType = Literal["user", "status",  "tool", "assistant", "assistant_chunk"]
 
 class XGAResponseMessage(TypedDict, total=False):
@@ -40,6 +40,10 @@ class XGAToolResult:
 
 
 class XGAToolBox(ABC):
+    @abstractmethod
+    async def init_tool_schemas(self):
+        pass
+
     @abstractmethod
     async def creat_task_tool_box(self, task_id: str, general_tools: List[str], custom_tools: List[str]):
         pass

xgae/engine/mcp_tool_box.py

@@ -2,7 +2,8 @@ import json
 import logging
 import os
 
-from typing import List, Any, Dict, Optional, Literal, override
+from typing import List, Any, Dict, Optional, Literal
+from typing_extensions import override
 
 from langchain_mcp_adapters.client import MultiServerMCPClient
 from langchain_mcp_adapters.tools import load_mcp_tools
@@ -33,7 +34,11 @@ class XGAMcpToolBox(XGAToolBox):
         self.mcp_tool_schemas: Dict[str, List[XGAToolSchema]] = {}
         self.task_tool_schemas: Dict[str, Dict[str,XGAToolSchema]] = {}
 
-        self.is_loaded_tool_schemas = False
+        self._is_loaded_mcp_tool_schemas = False
+
+    @override
+    async def init_tool_schemas(self):
+        await self._load_mcp_tools_schema()
 
     @override
     async def creat_task_tool_box(self, task_id: str, general_tools: List[str], custom_tools: List[str]):
@@ -139,8 +144,8 @@ class XGAMcpToolBox(XGAToolBox):
             return result
 
 
-    async def load_mcp_tools_schema(self)-> None:
-        if not self.is_loaded_tool_schemas:
+    async def _load_mcp_tools_schema(self)-> None:
+        if not self._is_loaded_mcp_tool_schemas:
             for server_name in self.mcp_server_names:
                 self.mcp_tool_schemas[server_name] = []
                 try:
@@ -169,11 +174,12 @@ class XGAMcpToolBox(XGAToolBox):
                                                 input_schema=input_schema,
                                                 metadata=metadata)
                     self.mcp_tool_schemas[server_name].append(tool_schema)
-            self.is_loaded_tool_schemas = True
+
+            self._is_loaded_mcp_tool_schemas = True
 
     async def reload_mcp_tools_schema(self) -> None:
-        self.is_loaded_tool_schemas = False
-        await self.load_mcp_tools_schema()
+        self._is_loaded_mcp_tool_schemas = False
+        await self.init_tool_schemas()
 
 
     def _load_mcp_servers_config(self, mcp_config_path: str) -> Dict[str, Any]:
@@ -219,7 +225,7 @@ if __name__ == "__main__":
         #mcp_tool_box = XGAMcpToolBox()
 
         task_id = "task1"
-        await mcp_tool_box.load_mcp_tools_schema()
+        await mcp_tool_box.init_tool_schemas()
         await mcp_tool_box.creat_task_tool_box(task_id=task_id, general_tools=["*"], custom_tools=["*"])
         tool_schemas = mcp_tool_box.get_task_tool_schemas(task_id, "general")
         print("general_tools_schemas" + "*"*50)

xgae/engine/responser/non_stream_responser.py

@@ -1,6 +1,7 @@
 import logging
 
-from typing import List, Dict, Any, AsyncGenerator, override,Optional
+from typing import List, Dict, Any, AsyncGenerator,Optional
+from typing_extensions import override
 
 from xgae.utils import log_trace
 

xgae/engine/responser/stream_responser.py

@@ -1,6 +1,7 @@
 import logging
-import asyncio
-from typing import List, Dict, Any, Optional, AsyncGenerator, override
+
+from typing import List, Dict, Any, Optional, AsyncGenerator
+from typing_extensions import override
 
 from xgae.utils import log_trace
 

xgae/engine/task_engine.py

@@ -59,6 +59,8 @@ class XGATaskEngine:
 
         self.task_response_msgs: List[XGAResponseMessage] = []
 
+        self.terminate_task = False
+
     async def run_task_with_final_answer(self,
                                          task_input: Dict[str, Any],
                                          trace_id: Optional[str] = None) -> XGATaskResult:
@@ -114,8 +116,7 @@ class XGATaskEngine:
                     general_tools.append("ask")
 
             custom_tools = self.custom_tools or []
-            if isinstance(self.tool_box, XGAMcpToolBox):
-                await  self.tool_box.load_mcp_tools_schema()
+            await  self.tool_box.init_tool_schemas()
 
             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")
@@ -148,6 +149,7 @@ class XGATaskEngine:
         auto_continue = True
         while auto_continue and auto_continue_count < self.max_auto_run:
             auto_continue = False
+            iterations = auto_continue_count
 
             try:
                 async for chunk in self._run_task_once(continuous_state):
@@ -193,11 +195,15 @@ class XGATaskEngine:
                 self.task_langfuse.root_span.event(name="engine_task_run_once_error", level="ERROR",
                                                    status_message=f"Call task_run_once error: {run_error}",
                                                    metadata={"trace": trace})
-
                 status_content = {'status_type': "error", 'role': "system", 'message': "Call run_task_once error"}
                 error_msg = self.add_response_message(type="status", content=status_content, is_llm_message=False)
                 yield error_msg
-
+            finally:
+                if not self.running_task_checkpoint("termination_check", iterations):
+                    status_content = {'status_type': "stop", 'role': "system", 'message': "Task is termiated by Stop Command"}
+                    error_msg = self.add_response_message(type="status", content=status_content, is_llm_message=False)
+                    yield error_msg
+                    break
 
     async def _run_task_once(self, continuous_state: TaskRunContinuousState) -> AsyncGenerator[Dict[str, Any], None]:
         llm_messages = [{"role": "system", "content": self.task_prompt}]
@@ -213,6 +219,9 @@ class XGATaskEngine:
             llm_messages.append(temp_assistant_message)
 
         auto_count = continuous_state.get("auto_continue_count")
+
+        self.running_task_checkpoint("before_completion", auto_count, llm_messages)
+
         langfuse_metadata = self.task_langfuse.create_llm_langfuse_meta(auto_count)
 
         llm_response = await self.llm_client.acompletion(llm_messages, langfuse_metadata)
@@ -220,12 +229,21 @@ class XGATaskEngine:
 
         async for chunk in response_processor.process_response(llm_response, llm_messages, continuous_state):
             self._logging_reponse_chunk(chunk, auto_count)
+
+            if chunk['type'] == "assistant":
+                assis_content = chunk['content']
+                self.running_task_checkpoint("after_completion", auto_count, llm_messages, assis_content)
+
             yield chunk
 
     def parse_final_result(self, chunks: List[Dict[str, Any]]) -> XGATaskResult:
         final_result: XGATaskResult = None
         reverse_chunks = reversed(chunks)
         chunk = None
+
+        # if self.terminate_task:
+        #     return XGATaskResult(type="error", content="LLM Task is terminated !")
+
         try:
             finish_reason = ''
             for chunk in reverse_chunks:
@@ -233,7 +251,7 @@ class XGATaskEngine:
                 if chunk_type == "status":
                     status_content = chunk['content']
                     status_type = status_content['status_type']
-                    if status_type == "error":
+                    if status_type == "error" or  status_type == "stop":
                         error = status_content['message']
                         final_result = XGATaskResult(type="error", content=error)
                     elif status_type == "finish":
@@ -284,6 +302,25 @@ class XGATaskEngine:
 
         return final_result
 
+
+    def stop_task(self):
+        logging.warning(f"⚠️ Begin Terminate Task: {self.task_id}")
+        self.task_langfuse.root_span.event(name="stop_task", level="DEFAULT",
+                                           status_message="Begin Terminate Task")
+        self.terminate_task = True
+
+
+    def running_task_checkpoint(self,
+                                task_state: Literal["before_completion", "after_completion", "termination_check"],
+                                iterations: int,
+                                llm_messages: List[Dict[str, Any]] = None,
+                                llm_response: Dict[str, Any] = None
+                                )-> bool:
+        if self.terminate_task and task_state == "termination_check":
+            logging.warning(f"⚠️ TASK: {self.task_id} STOP RUNNING for STOP Command !")
+        return not self.terminate_task
+
+
     def create_response_message(self, type: XGAResponseMsgType,
                              content: Union[Dict[str, Any], List[Any], str],
                              is_llm_message: bool,

xgae/gaia2/are_engine.py

@@ -0,0 +1,123 @@
+import logging
+from typing import Any, Optional, Callable, Literal, Dict, List
+from typing_extensions import override
+
+from xgae.engine.engine_base import XGAToolBox
+from xgae.engine.task_engine import XGATaskEngine
+from xgae.utils.llm_client import LLMConfig
+
+
+class ARETaskEngine(XGATaskEngine):
+    def __init__(self,
+                 agent: Any,
+                 agent_id: str,
+                 system_prompt: str,
+                 max_auto_run: int,
+                 llm_config: Optional[LLMConfig] = None,
+                 tool_box: Optional[XGAToolBox] = None,
+                 pre_run_task_fn : Callable[[Any, int, List[Dict[str, Any]]], Any] = None,
+                 post_run_task_fn : Callable[[Any, int, Dict[str, Any]], Any] = None,
+                 terminate_task_fn : Callable[[Any, int], bool] = None,
+                ):
+        super().__init__(agent_id       = agent_id,
+                         general_tools  = [],
+                         custom_tools   = ["*"],
+                         system_prompt  = system_prompt,
+                         max_auto_run   = max_auto_run,
+                         llm_config     = llm_config,
+                         tool_box       = tool_box,
+                         )
+        self.agent = agent
+        self.pre_run_task_fn = pre_run_task_fn
+        self.post_run_task_fn = post_run_task_fn
+        self.terminate_task_fn = terminate_task_fn
+
+
+    @override
+    def running_task_checkpoint(self,
+                                task_state: Literal["before_completion", "after_completion", "termination_check"],
+                                iterations: int,
+                                llm_messages: List[Dict[str, Any]] = None,
+                                llm_response: Dict[str, Any] = None
+                                )-> bool:
+        is_continue_task = True
+
+        if task_state == "before_completion" and self.pre_run_task_fn:
+            self.pre_run_task_fn(self.agent, iterations, llm_messages)
+        elif task_state == "after_completion" and self.post_run_task_fn:
+            self.post_run_task_fn(self.agent, iterations, llm_response)
+        elif task_state == "termination_check":
+            if self.terminate_task:
+                logging.warning(f"running_task_checkpoint: ⚠️ TASK: {self.task_id} STOP RUNNING for STOP Command !")
+
+            if self.terminate_task_fn:
+                is_terminate = self.terminate_task_fn(self.agent, iterations) if self.terminate_task_fn  else False
+                if is_terminate:
+                    logging.warning(f"running_task_checkpoint: ⚠️ TASK: {self.task_id} STOP RUNNING for Termination Function !")
+                    self.stop_task()
+
+            is_continue_task = not self.terminate_task
+
+        return is_continue_task
+
+
+
+if __name__ == "__main__":
+    import asyncio
+    import os
+    from xgae.utils.misc import read_file
+    from xgae.utils.setup_env import setup_logging
+    from xgae.engine.mcp_tool_box import XGAMcpToolBox
+
+    setup_logging()
+
+    def pre_run_task(agent, iterations:int, llm_messages: List[Dict[str, Any]]):
+        prompt = "\n\n".join([f"{key}: {value}" for d in llm_messages for key, value in d.items()]) if llm_messages else ""
+        logging.info(f"pre_run_task: iterations={iterations}, prompt: \n{prompt}\n")
+
+
+    def post_run_task(agent, iterations: int, llm_response: Dict[str, Any]):
+        logging.info(f"post_run_task: iterations={iterations}, prompt: \n{llm_response}\n")
+
+
+    def terminate_task(agent, iterations: int) -> bool:
+        logging.info(f"terminate_task: iterations={iterations}")
+        return iterations > 3 # can test terminate by > 3
+
+
+    async def main():
+        # Before Run Exec: uv run example-fault-tools
+        # LLAMA_API_KEY , LLAMA_API_BASE
+        tool_box = XGAMcpToolBox(custom_mcp_server_file="mcpservers/custom_servers.json")
+        system_prompt = read_file("templates/example/fault_user_prompt.txt")
+        llm_config = LLMConfig(
+            model            = "openai/qwen-plus",
+            api_key          = os.getenv('LLAMA_API_KEY') ,
+            api_base         = os.getenv('LLAMA_API_BASE'),
+            stream           =True,
+            enable_thinking  = False,
+        )
+
+        engine =  ARETaskEngine(
+            agent               = "AREAgent", # Just for test,ARE use real Agent Object
+            agent_id            = "agent_1",
+            system_prompt       = system_prompt,
+            max_auto_run        = 15,
+            llm_config          = llm_config,
+            tool_box            = tool_box,
+            pre_run_task_fn     = pre_run_task,
+            post_run_task_fn    = post_run_task,
+            terminate_task_fn   = terminate_task
+        )
+
+        user_input = "locate 10.0.0.1 fault and solution"
+        chunks = []
+        async for chunk in engine.run_task(task_input={"role": "user", "content": user_input}):
+            chunks.append(chunk)
+            print(chunk)
+
+        final_result = engine.parse_final_result(chunks)
+        print(f"\n\nFINAL_RESULT: {final_result}")
+
+
+    asyncio.run(main())

{xgae-0.2.3.dist-info → xgae-0.3.1.dist-info}/METADATA

@@ -1,13 +1,13 @@
 Metadata-Version: 2.4
 Name: xgae
-Version: 0.2.3
+Version: 0.3.1
 Summary: Extreme General Agent Engine
-Requires-Python: >=3.13
+Requires-Python: >=3.11
 Requires-Dist: colorlog==6.9.0
 Requires-Dist: langchain-mcp-adapters==0.1.9
 Requires-Dist: langfuse==2.60.9
-Requires-Dist: litellm==1.74.15
-Requires-Dist: mcp==1.13.0
+Requires-Dist: litellm==1.71.1
+Requires-Dist: mcp==1.11.0
 Provides-Extra: examples
 Requires-Dist: chromadb==1.1.0; extra == 'examples'
 Requires-Dist: langchain-community==0.3.29; extra == 'examples'

{xgae-0.2.3.dist-info → xgae-0.3.1.dist-info}/RECORD

@@ -1,13 +1,14 @@
 xgae/__init__.py,sha256=oBX_YzTliM-343BAlR-sD7BUZmsCJ7PY2oYrGBhsdLM,79
 xgae/engine_cli_app.py,sha256=FdmIpq8KDsgyZNfwCDgNX7FEZFeRFyGOt_H1oZF8aKs,2890
-xgae/engine/engine_base.py,sha256=RR1em2wHiM2jP-peHt77SKdHWjnYOjdIIzN93zT61cA,1715
-xgae/engine/mcp_tool_box.py,sha256=G4hKIMguwg1cO4Us2NMfdloYim8kuikVyVTIPucJr7o,10903
+xgae/engine/engine_base.py,sha256=bXwLxli7CK7JndAcOu-Rx4uTuy0pJDh_ZOGyhym76WY,1796
+xgae/engine/mcp_tool_box.py,sha256=7dOiP9cfBU8TfAxzbdw0V31ISF-o2TgdTVUqEemz004,11044
 xgae/engine/prompt_builder.py,sha256=6I5rjgvNJ27QJ8DDuBTplutoPZdGs9LYFv3TSgT7zmc,5045
-xgae/engine/task_engine.py,sha256=bz3s2w6-QHYRJjwO2KLtCv46gFb1p5JZeEW6cX2A7MI,21720
+xgae/engine/task_engine.py,sha256=QhjTPQ_aREsT19QfT1nQwNLKUMwWdSgqlfjfCNq_dho,23470
 xgae/engine/task_langfuse.py,sha256=ifkGrPBv2daLTKE-fCfEtOoI0n4Pd-lCwhyRRL0h308,2850
-xgae/engine/responser/non_stream_responser.py,sha256=zEJjqCgZVe2B8gkHYRFU7tmBV834f7w2a4Ws25P1N-c,5289
+xgae/engine/responser/non_stream_responser.py,sha256=WFzrT0tHGVLi_AR1IrIPPvpQ94ne7slwbFyXaoTpCQc,5318
 xgae/engine/responser/responser_base.py,sha256=jhl1Bdz1Fs3KofGEymThNXlQuCORFTTkTAR_U47krds,24403
-xgae/engine/responser/stream_responser.py,sha256=cv4UGcxj8OksEogW7DUGTCvSJabu-DF6GceFyUwaXI4,7627
+xgae/engine/responser/stream_responser.py,sha256=5n7i2QnE_u7RP1bcJFky8sFPjYMjfYY9NcR5TeFxl9Y,7642
+xgae/gaia2/are_engine.py,sha256=QPvyWEydeiOECQ7WfxjMIurQHD_ET7uWQCnXVCh_ZQk,5094
 xgae/tools/without_general_tools_app.py,sha256=KqsdhxD3hvTpiygaGUVHysRFjvv_1A8zOwMKN1J0J0U,3821
 xgae/utils/__init__.py,sha256=ElaGS-zdeZeu6is41u3Ny7lkvhg7BDSK-jMNg9j6K5A,499
 xgae/utils/json_helpers.py,sha256=WD4G5U9Dh8N6J9O0L5wGyqj-NHi09kcXHGdLD_26nlc,3607
@@ -15,7 +16,7 @@ xgae/utils/llm_client.py,sha256=rqnu_NYXBC0hl4aozP5UOSyf0q-ONB5ywtnrXzA88OE,1505
 xgae/utils/misc.py,sha256=aMWOvJ9VW52q-L9Lkjl1hvXqLwpJAmyxA-Z8jzqFG0U,907
 xgae/utils/setup_env.py,sha256=MqNG0c2QQBDFU1kI8frxr9kB5d08Mmi3QZ1OoorgIa0,2662
 xgae/utils/xml_tool_parser.py,sha256=Mb0d8kBrfyAEvUwW1Nqir-3BgxZRr0ZX3WymQouuFSo,4859
-xgae-0.2.3.dist-info/METADATA,sha256=rgZI_57gbtxNfmYY4LN8AcNn5WPuYd-vExkDXurWlCQ,471
-xgae-0.2.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-xgae-0.2.3.dist-info/entry_points.txt,sha256=wmvgtMQbtzTbDPETS-tbQJD7jVlcs4hp0w6wOB0ooCc,229
-xgae-0.2.3.dist-info/RECORD,,
+xgae-0.3.1.dist-info/METADATA,sha256=vOaoxeqkHDKwSuRUscc9--4KpuOTqM5BwEXqsCPUQ1I,470
+xgae-0.3.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+xgae-0.3.1.dist-info/entry_points.txt,sha256=wmvgtMQbtzTbDPETS-tbQJD7jVlcs4hp0w6wOB0ooCc,229
+xgae-0.3.1.dist-info/RECORD,,