xgae 0.2.2__py3-none-any.whl → 0.3.0__py3-none-any.whl

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:
@@ -148,6 +150,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 +196,12 @@ 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):
+                    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 +217,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 +227,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:
@@ -284,6 +300,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,122 @@
+import logging
+from typing import Any, Optional, override, Callable, Literal, Dict, List
+
+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 > 6 # can test terminate by > 3
+
+
+    async def main():
+        # Before Run Exec: uv run example-fault-tools
+        # LLAMA_API_KEY ,
+        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.2.dist-info → xgae-0.3.0.dist-info}/METADATA

@@ -1,11 +1,10 @@
 Metadata-Version: 2.4
 Name: xgae
-Version: 0.2.2
+Version: 0.3.0
 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: langchain==0.3.27
 Requires-Dist: langfuse==2.60.9
 Requires-Dist: litellm==1.74.15
 Requires-Dist: mcp==1.13.0

{xgae-0.2.2.dist-info → xgae-0.3.0.dist-info}/RECORD

@@ -3,11 +3,12 @@ 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/prompt_builder.py,sha256=6I5rjgvNJ27QJ8DDuBTplutoPZdGs9LYFv3TSgT7zmc,5045
-xgae/engine/task_engine.py,sha256=bz3s2w6-QHYRJjwO2KLtCv46gFb1p5JZeEW6cX2A7MI,21720
+xgae/engine/task_engine.py,sha256=1YgkcsX8LAAIIp1DsVbphwOVO9xA0FioABxzaRvers8,23223
 xgae/engine/task_langfuse.py,sha256=ifkGrPBv2daLTKE-fCfEtOoI0n4Pd-lCwhyRRL0h308,2850
 xgae/engine/responser/non_stream_responser.py,sha256=zEJjqCgZVe2B8gkHYRFU7tmBV834f7w2a4Ws25P1N-c,5289
 xgae/engine/responser/responser_base.py,sha256=jhl1Bdz1Fs3KofGEymThNXlQuCORFTTkTAR_U47krds,24403
 xgae/engine/responser/stream_responser.py,sha256=cv4UGcxj8OksEogW7DUGTCvSJabu-DF6GceFyUwaXI4,7627
+xgae/gaia2/are_engine.py,sha256=GPs6c94CN8jJjmNNsE7XHxdskPJ31p4Ds9QCST7BiI4,5050
 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.2.dist-info/METADATA,sha256=9VaGgRhzSWNmTuYbKAB9knyXKCJdbNP-40vRenAKi-0,504
-xgae-0.2.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-xgae-0.2.2.dist-info/entry_points.txt,sha256=wmvgtMQbtzTbDPETS-tbQJD7jVlcs4hp0w6wOB0ooCc,229
-xgae-0.2.2.dist-info/RECORD,,
+xgae-0.3.0.dist-info/METADATA,sha256=OAdIkEtp1xWqtU2MrOBRJfPyu7TLJk6m5EPYh9WYNdQ,471
+xgae-0.3.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+xgae-0.3.0.dist-info/entry_points.txt,sha256=wmvgtMQbtzTbDPETS-tbQJD7jVlcs4hp0w6wOB0ooCc,229
+xgae-0.3.0.dist-info/RECORD,,