xgae 0.1.2__py3-none-any.whl → 0.1.3__py3-none-any.whl

xgae/engine/xga_base.py

@@ -2,7 +2,9 @@ from typing import Union, Optional, Dict, List, Any, Literal
 from dataclasses import dataclass
 from abc import ABC, abstractmethod
 
-
+class XGAError(Exception):
+    """Custom exception for errors in the XGA system."""
+    pass
 
 @dataclass
 class XGAMessage:

xgae/engine/xga_engine.py

@@ -2,68 +2,74 @@
 from typing import List, Any, Dict, Optional, AsyncGenerator
 from uuid import uuid4
 
-from xgae.engine.xga_base import XGAMessage, XGAToolSchema, XGAToolBox
+from xgae.engine.xga_base import XGAMessage, XGAToolBox
 from xgae.utils.llm_client import LLMClient
 from xgae.utils.setup_env import langfuse
 from xga_prompt_builder import XGAPromptBuilder
 from xga_mcp_tool_box import XGAMcpToolBox
 
 
-class XGAEngine():
+class XGATaskEngine():
     def __init__(self,
                  session_id: Optional[str] = None,
-                 trace_id: Optional[str] = None,
+                 task_id: Optional[str] = None,
                  agent_id: Optional[str] = None,
+                 trace_id: Optional[str] = None,
+                 system_prompt: Optional[str] = None,
                  llm_config: Optional[Dict[str, Any]] = None,
                  prompt_builder: Optional[XGAPromptBuilder] = None,
                  tool_box: Optional[XGAToolBox] = None):
         self.session_id = session_id if session_id else f"xga_sid_{uuid4()}"
+        self.task_id = task_id if task_id else f"xga_task_{uuid4()}"
         self.agent_id = agent_id
+        self.trace_id = trace_id if trace_id else langfuse.create_trace_id()
 
         self.messages: List[XGAMessage] = []
         self.llm_client = LLMClient(llm_config)
         self.model_name = self.llm_client.model_name
         self.is_stream = self.llm_client.is_stream
 
-        self.prompt_builder = prompt_builder or XGAPromptBuilder()
+        self.prompt_builder = prompt_builder or XGAPromptBuilder(system_prompt)
         self.tool_box = tool_box or XGAMcpToolBox()
 
-        self.task_id = None
-        self.trace_id = trace_id if trace_id else langfuse.create_trace_id()
 
-    async def __async_init__(self) -> None:
+    async def __async_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)
 
     @classmethod
     async def create(cls,
                      session_id: Optional[str] = None,
-                     trace_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[Dict[str, Any]] = None,
                      prompt_builder: Optional[XGAPromptBuilder] = None,
-                     tool_box: Optional[XGAToolBox] = None) -> 'XGAEngine' :
-        engine: XGAEngine = cls(session_id=session_id,
-                                trace_id=trace_id,
-                                agent_id=agent_id,
-                                llm_config=llm_config,
-                                prompt_builder=prompt_builder,
-                                tool_box=tool_box)
-
-        await engine.__async_init__()
+                     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)
+        general_tools = general_tools or ["*"]
+        custom_tools = custom_tools or []
+        await engine.__async_init__(general_tools, custom_tools)
         return engine
 
 
-    async def run_task(self,
-                       task_messages: List[Dict[str, Any]],
-                       task_id: Optional[str] = None,
-                       system_prompt: Optional[str] = None,
-                       general_tools: Optional[List[str]] = ["*"],
-                       custom_tools: Optional[List[str]] = []) -> AsyncGenerator:
+    async def run_task(self, task_messages: List[Dict[str, Any]]) -> AsyncGenerator:
         try:
-            self.task_id = task_id if task_id else f"xga_task_{uuid4()}"
-            await self.tool_box.creat_task_tool_box(self.task_id, general_tools, custom_tools)
-            task_prompt = self.build_task_prompt(system_prompt)
-            yield task_prompt
+            yield self.task_prompt
 
         finally:
             await self.tool_box.destroy_task_tool_box(self.task_id)
@@ -72,18 +78,6 @@ class XGAEngine():
     def _run_task_once(self):
         pass
 
-    def build_task_prompt(self, system_prompt: Optional[str] = None) -> str:
-        task_prompt = self.prompt_builder.build_system_prompt(self.model_name, system_prompt)
-
-        tool_schemas =self.tool_box.get_task_tool_schemas(self.task_id, "general_tool")
-        tool_prompt = self.prompt_builder.build_general_tool_prompt(tool_schemas)
-        task_prompt = task_prompt + "\n" + tool_prompt
-
-        tool_schemas = self.tool_box.get_task_tool_schemas(self.task_id, "custom_tool")
-        tool_prompt = self.prompt_builder.build_custom_tool_prompt(tool_schemas)
-        task_prompt = task_prompt + "\n" + tool_prompt
-
-        return task_prompt
 
     def add_message(self, message: XGAMessage):
         message.message_id = f"xga_msg_{uuid4()}"
@@ -95,11 +89,11 @@ if __name__ == "__main__":
     import asyncio
 
     async def main():
-        #tool_box = XGAMcpToolBox(custom_mcp_server_file="mcpservers/custom_servers.json")
-        tool_box = None
-        engine = await XGAEngine.create(tool_box=tool_box)
-        # async for chunk in engine.run_task(task_messages=[{}], custom_tools=["bomc_fault.*"]):
-        async for chunk in engine.run_task(task_messages=[{}], custom_tools=[]):
+        tool_box = XGAMcpToolBox(custom_mcp_server_file="mcpservers/custom_servers.json")
+        engine = await XGATaskEngine.create(tool_box=tool_box, custom_tools=["bomc_fault.*"])
+        # engine = await XGATaskEngine.create()
+
+        async for chunk in engine.run_task(task_messages=[{}]):
              print(chunk)
 
     asyncio.run(main())

xgae/engine/xga_mcp_tool_box.py

@@ -7,9 +7,8 @@ from typing import List, Any, Dict, Optional, Literal, override
 from langchain_mcp_adapters.client import MultiServerMCPClient
 from langchain_mcp_adapters.tools import load_mcp_tools
 
-from xgae.engine.xga_base import XGAToolSchema, XGAToolBox, XGAToolResult
-from xgae.utils.setup_env import XGAError
-
+from xgae.engine.xga_base import XGAError, XGAToolSchema, XGAToolBox, XGAToolResult
+from xgae.utils.setup_env import langfuse
 
 class XGAMcpToolBox(XGAToolBox):
     GENERAL_MCP_SERVER_NAME = "xga_general"

xgae/engine/xga_prompt_builder.py

@@ -1,25 +1,28 @@
-import datetime
-import sys
 import json
-
+import datetime
 from typing import Optional, List
-from io import StringIO
 
-from xga_base import XGAToolSchema
-from xgae.utils.setup_env import read_file, XGAError
+from xga_base import XGAToolSchema, XGAError
+from xgae.utils.utils import read_file, format_file_with_args
 
 
 class XGAPromptBuilder():
-    def __init__(self, system_prompt_template: Optional[str] = None):
-        self.system_prompt = None
-        if system_prompt_template:
-            self.system_prompt = system_prompt_template
+    def __init__(self, system_prompt: Optional[str] = None):
+        self.system_prompt = system_prompt
+
+    def build_task_prompt(self, model_name: str, general_tool_schemas: List[XGAToolSchema], custom_tool_schemas: List[XGAToolSchema])-> str:
+        if self.system_prompt is None:
+            self.system_prompt = self._load_default_system_prompt(model_name)
+
+        task_prompt = self.system_prompt
+
+        tool_prompt = self.build_general_tool_prompt(general_tool_schemas)
+        task_prompt = task_prompt + "\n" + tool_prompt
+
+        tool_prompt = self.build_custom_tool_prompt(custom_tool_schemas)
+        task_prompt = task_prompt + "\n" + tool_prompt
 
-    def build_system_prompt(self, model_name: str, system_prompt: Optional[str]=None)-> str:
-        task_system_prompt = system_prompt if system_prompt else self.system_prompt
-        if task_system_prompt is None:
-            task_system_prompt = self._load_default_system_prompt(model_name)
-        return task_system_prompt
+        return task_prompt
 
     def build_general_tool_prompt(self, tool_schemas:List[XGAToolSchema])-> str:
         tool_prompt = ""
@@ -41,7 +44,7 @@ class XGAPromptBuilder():
                 input_schema = tool_schema.input_schema
                 openai_function["parameters"] = openai_parameters
                 openai_parameters["type"] = input_schema["type"]
-                openai_parameters["properties"] = input_schema.get("properties", [])
+                openai_parameters["properties"] = input_schema.get("properties", {})
                 openai_parameters["required"] = input_schema["required"]
 
                 openai_schemas.append(openai_schema)
@@ -65,32 +68,21 @@ class XGAPromptBuilder():
             for tool_schema in tool_schemas:
                 description = tool_schema.description if tool_schema.description else 'No description available'
                 tool_info += f"- **{tool_schema.tool_name}**: {description}\n"
-                tool_info += f"  Parameters: {tool_schema.input_schema}\n"
+                parameters = tool_schema.input_schema.get("properties", {})
+                tool_info += f"  Parameters: {parameters}\n"
             tool_prompt = tool_prompt.replace("{tool_schemas}", tool_info)
 
         return tool_prompt
 
     def _load_default_system_prompt(self, model_name) -> Optional[str]:
         if "gemini-2.5-flash" in model_name.lower() and "gemini-2.5-pro" not in model_name.lower():
-            org_prompt_template = read_file("templates/gemini_system_prompt_template.txt")
+            system_prompt_template = read_file("templates/gemini_system_prompt_template.txt")
         else:
-            org_prompt_template = read_file("templates/system_prompt_template.txt")
-
-        original_stdout = sys.stdout
-        buffer = StringIO()
-        sys.stdout = buffer
-        try:
-            namespace = {
-                "datetime": datetime,
-                "__builtins__": __builtins__
-            }
-            code = f"print(f\"\"\"{org_prompt_template}\"\"\")"
-            exec(code, namespace)
-            system_prompt_template = buffer.getvalue()
-        finally:
-            sys.stdout = original_stdout
-
-        system_prompt = system_prompt_template.format(
+            system_prompt_template = read_file("templates/system_prompt_template.txt")
+
+        system_prompt = format_file_with_args(system_prompt_template, {"datetime": datetime})
+
+        system_prompt = system_prompt.format(
             current_date=datetime.datetime.now(datetime.timezone.utc).strftime('%Y-%m-%d'),
             current_time=datetime.datetime.now(datetime.timezone.utc).strftime('%H:%M:%S'),
             current_year=datetime.datetime.now(datetime.timezone.utc).strftime('%Y')
@@ -102,14 +94,4 @@ class XGAPromptBuilder():
 
         return system_prompt
 
-if __name__ == "__main__":
-
-    prompt_builder = XGAPromptBuilder()
-    prompt = prompt_builder.build_system_prompt("openai/qwen3-235b-a22b")
-
-    # system_prompt = read_file("templates/scp_test_prompt.txt")
-    # prompt = prompt_builder.build_system_prompt("openai/qwen3-235b-a22b", system_prompt=system_prompt)
-
-    print(prompt)
-
 

xgae/utils/setup_env.py

@@ -1,10 +1,8 @@
 import logging
 import os
-import traceback
 
 from langfuse import Langfuse
 
-
 _log_initialized = False
 
 def setup_logging() -> None:
@@ -65,17 +63,6 @@ def setup_logging() -> None:
 
 setup_logging()
 
-
-class XGAError(Exception):
-    """Custom exception for errors in the XGA system."""
-    pass
-
-def handle_error(e: Exception) -> None:
-    logging.error("An error occurred: %s", str(e))
-    logging.error("Traceback details:\n%s", traceback.format_exc())
-    raise (e) from e
-
-
 _langfuse_initialized = False
 
 def setup_langfuse() -> Langfuse:
@@ -101,20 +88,6 @@ def setup_langfuse() -> Langfuse:
 langfuse: Langfuse = Langfuse if _langfuse_initialized else setup_langfuse()
 
 
-def read_file(file_path: str) -> str:
-    if not os.path.exists(file_path):
-        logging.error(f"File '{file_path}' not found")
-        raise XGAError(f"File '{file_path}' not found")
-
-    try:
-        with open(file_path, "r", encoding="utf-8") as template_file:
-            content = template_file.read()
-        return content
-    except Exception as e:
-        logging.error(f"Read file '{file_path}' failed")
-        handle_error(e)
-
-
 if __name__ == "__main__":
     try:
         trace_id = langfuse.create_trace_id()

xgae/utils/utils.py

@@ -0,0 +1,42 @@
+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")
+        raise Exception(f"File '{file_path}' not found")
+
+    try:
+        with open(file_path, "r", encoding="utf-8") as template_file:
+            content = template_file.read()
+        return content
+    except Exception as e:
+        logging.error(f"Read file '{file_path}' failed")
+        handle_error(e)
+
+def format_file_with_args(file_content:str, args: Dict[str, Any])-> str:
+    from io import StringIO
+
+    formated = file_content
+    original_stdout = sys.stdout
+    buffer = StringIO()
+    sys.stdout = buffer
+    try:
+        code = f"print(f\"\"\"{file_content}\"\"\")"
+        exec(code, args)
+        formated = buffer.getvalue()
+    finally:
+        sys.stdout = original_stdout
+
+    return formated

{xgae-0.1.2.dist-info → xgae-0.1.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xgae
-Version: 0.1.2
+Version: 0.1.3
 Summary: Extreme General Agent Engine
 Requires-Python: >=3.13
 Requires-Dist: colorlog>=6.9.0

xgae-0.1.3.dist-info/RECORD

@@ -0,0 +1,14 @@
+xgae/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+xgae/engine/xga_base.py,sha256=b9FnBbCuO4lCAIBBhnedtRk49Dl3M6cBn5h_Wkzw4o0,1345
+xgae/engine/xga_engine.py,sha256=kci7pvF8NwAYJ3Y6DS0JKLkPEQERZwcwRzjcAFQx_ec,4169
+xgae/engine/xga_mcp_tool_box.py,sha256=coyQy7tj2YTBpOc0JKmVBJO8ItO4Z6s1pdbucHd85qg,9762
+xgae/engine/xga_prompt_builder.py,sha256=RuTvQCNufqxDwVvSOPXR0qxAc42cG7NuIaUy9amu66A,4351
+xgae/engine/responser/xga_non_stream_responser.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+xgae/engine/responser/xga_responser_utils.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+xgae/engine/responser/xga_stream_reponser.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+xgae/utils/llm_client.py,sha256=UMMK84psLsx36-Nn6Q8X1hl9wd-OdaS9ZhxbRjwNCr0,12149
+xgae/utils/setup_env.py,sha256=vY2wnq5KMXQ2dx2wAPeTpudp7b-fvgzenkwHpgd8vuI,2951
+xgae/utils/utils.py,sha256=cCYmWjKFksZ8BRD1YYnaM_jTLVHAg1ibEdjsczEUO6k,1134
+xgae-0.1.3.dist-info/METADATA,sha256=t4sr4k9pbVAcicIRjoqt327eWtSC0uAeo2pDB5LPjs0,309
+xgae-0.1.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+xgae-0.1.3.dist-info/RECORD,,

xgae-0.1.2.dist-info/RECORD

@@ -1,13 +0,0 @@
-xgae/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-xgae/engine/xga_base.py,sha256=vOiYWoBze7VrmndYnaNhMzoB1zCT80WtIlOWdp2IrUc,1253
-xgae/engine/xga_engine.py,sha256=AgklNAimyErvn74xUu6aardZvgjqbIR99gCojEmZfO0,4267
-xgae/engine/xga_mcp_tool_box.py,sha256=jRYvZ8eY72KRU0leD-o15fe_E8wjWJtxQZbsIqGVjPU,9753
-xgae/engine/xga_prompt_builder.py,sha256=Q3nrmYAofRkFQFdfSWkP00FZKjWMUQs2cvMs7bvyQgQ,4819
-xgae/engine/responser/xga_non_stream_responser.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-xgae/engine/responser/xga_responser_utils.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-xgae/engine/responser/xga_stream_reponser.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-xgae/utils/llm_client.py,sha256=UMMK84psLsx36-Nn6Q8X1hl9wd-OdaS9ZhxbRjwNCr0,12149
-xgae/utils/setup_env.py,sha256=qZmVfqaYHtjq5t-zwRt_G3i-QqGPxSb4D2FnFHbhdcY,3689
-xgae-0.1.2.dist-info/METADATA,sha256=iBwWKlpbp7z0ZQgydnieP_2NBf0Z7H0D6AtavbIN3Zc,309
-xgae-0.1.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-xgae-0.1.2.dist-info/RECORD,,