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

xgae/engine/xga_base.py

@@ -1,4 +1,4 @@
-from typing import Union, Optional, Dict, List, Any, Literal
+from typing import Union, Optional, Dict, List, Any, Literal, TypedDict
 from dataclasses import dataclass
 from abc import ABC, abstractmethod
 
@@ -6,16 +6,24 @@ class XGAError(Exception):
     """Custom exception for errors in the XGA system."""
     pass
 
-@dataclass
-class XGAMessage:
-    message_id: str
-    type: Literal["status",  "tool", "assistant", "assistant_response_end"]
+
+class XGAContextMsg(TypedDict, total=False):
+    type: Literal["user", "status",  "tool", "assistant", "assistant_response_end"]
+    content: Union[Dict[str, Any], List[Any], str]
     is_llm_message: bool
+    metadata: Dict[str, Any]
+    message_id: str
+    session_id: str
+    agent_id: str
+    task_id: str
+    task_run_id: str
+    trace_id: str
+
+class XGAResponseMsg(TypedDict, total=False):
+    type: Literal["content", "status"]
     content: Union[Dict[str, Any], List[Any], str]
-    metadata: Optional[Dict[str, Any]]
-    session_id: Optional[str]
-    agent_id: Optional[str]
-    task_id: Optional[str]
+    status: Literal["error", "status"]
+    message: str
 
 @dataclass
 class XGAToolSchema:
@@ -44,6 +52,10 @@ class XGAToolBox(ABC):
     def get_task_tool_schemas(self, task_id: str, type: Literal["general_tool",  "custom_tool"]) -> List[XGAToolSchema]:
         pass
 
+    @abstractmethod
+    def get_task_tool_names(self, task_id: str) -> List[str]:
+        pass
+
     @abstractmethod
     async def call_tool(self, task_id: str, tool_name: str, args: Optional[Dict[str, Any]] = None) -> XGAToolResult:
         pass

xgae/engine/xga_engine.py

@@ -1,30 +1,32 @@
 
-from typing import List, Any, Dict, Optional, AsyncGenerator
+import logging
+import json
+
+from typing import List, Any, Dict, Optional, AsyncGenerator, cast, Union, Literal
 from uuid import uuid4
 
-from xgae.engine.xga_base import XGAMessage, XGAToolBox
-from xgae.utils.llm_client import LLMClient
+from xgae.engine.responser.xga_responser_base import TaskResponseContext, TaskResponseProcessor, TaskRunContinuousState
+from xgae.engine.xga_base import XGAContextMsg, XGAToolBox, XGAResponseMsg
+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
 
-
-class XGATaskEngine():
+class XGATaskEngine:
     def __init__(self,
                  session_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,
+                 llm_config: Optional[LLMConfig] = 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
@@ -32,6 +34,10 @@ class XGATaskEngine():
         self.prompt_builder = prompt_builder or XGAPromptBuilder(system_prompt)
         self.tool_box = tool_box or XGAMcpToolBox()
 
+        self.task_context_msgs: List[XGAContextMsg] = []
+        self.task_no = -1
+        self.task_run_id = f"{self.task_id}[{self.task_no}]"
+        self.trace_id = None
 
     async def __async_init__(self, general_tools:List[str], custom_tools: List[str]) -> None:
         await  self.tool_box.load_mcp_tools_schema()
@@ -46,54 +52,227 @@ class XGATaskEngine():
                      session_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,
+                     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,
-                                    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 ["*"]
+
+        general_tools = general_tools or ["complete"]
         custom_tools = custom_tools or []
         await engine.__async_init__(general_tools, custom_tools)
+
+        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}")
+
         return engine
 
 
-    async def run_task(self, task_messages: List[Dict[str, Any]]) -> AsyncGenerator:
+    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:
-            yield self.task_prompt
+            self.trace_id = trace_id or self.trace_id or langfuse.create_trace_id()
 
+            self.task_no += 1
+            self.task_run_id = f"{self.task_id}[{self.task_no}]"
+
+            self.add_context_msg(type="user", content=task_message, is_llm_message=True)
+
+            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)
-        
 
-    def _run_task_once(self):
-        pass
+    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_context_llm_contents()
+        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)
+        response_processor = self._create_response_processer()
+
+        async for chunk in response_processor.process_response(llm_response, llm_messages, continuous_state):
+            yield chunk
+
 
+    async def _run_task_auto(self, max_auto_run: int) -> AsyncGenerator:
+        continuous_state: TaskRunContinuousState = {
+            "accumulated_content": "",
+            "auto_continue_count": 0,
+            "auto_continue": True
+        }
 
-    def add_message(self, message: XGAMessage):
-        message.message_id = f"xga_msg_{uuid4()}"
-        message.session_id = self.session_id
-        message.agent_id = self.agent_id
-        self.messages.append(message)
+        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:
+            auto_continue = False
+
+            try:
+                async for chunk in self._run_task_once(continuous_state):
+                    try:
+                        if chunk.get("type") == "status":
+                            content = json.loads(chunk.get('content', '{}'))
+                            status_type = content.get('status_type', None)
+                            if status_type == "error":
+                                logging.error(f"run_task_auto: task_response error: {chunk.get('message', 'Unknown error')}")
+                                yield chunk
+                                return
+                            elif status_type == 'finish':
+                                finish_reason = content.get('finish_reason', None)
+                                if finish_reason == 'stop' :
+                                    auto_continue = True
+                                    auto_continue_count += 1
+                                    update_continuous_state(auto_continue_count, auto_continue)
+                                    logging.info(f"run_task_auto: Detected finish_reason='stop', auto-continuing ({auto_continue_count}/{max_auto_run})")
+                                    continue
+                                elif finish_reason == 'xml_tool_limit_reached':
+                                    logging.info(f"run_task_auto: Detected finish_reason='xml_tool_limit_reached', stopping auto-continue")
+                                    auto_continue = False
+                                    update_continuous_state(auto_continue_count, auto_continue)
+                                elif finish_reason == 'length':
+                                    auto_continue = True
+                                    auto_continue_count += 1
+                                    update_continuous_state(auto_continue_count, auto_continue)
+                                    logging.info(f"run_task_auto: Detected finish_reason='length', auto-continuing ({auto_continue_count}/{max_auto_run})")
+                                    continue
+                    except Exception as parse_error:
+                        logging.error(f"run_task_auto: Error in parse chunk: {str(parse_error)}")
+                        yield {
+                            "type": "status",
+                            "status": "error",
+                            "message": f"Error in parse chunk: {str(parse_error)}"
+                        }
+                        return
+
+                    # Otherwise just yield the chunk normally
+                    yield chunk
+
+                # If not auto-continuing, we're done
+                if not auto_continue:
+                    break
+            except Exception as run_error:
+                logging.error(f"run_task_auto: Call task_run_once error: {str(run_error)}")
+                yield {
+                    "type": "status",
+                    "status": "error",
+                    "message": f"Call task_run_once error: {str(run_error)}"
+                }
+                return
+
+    def add_context_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)-> XGAContextMsg:
+        message = XGAContextMsg(
+            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
+        )
+        self.task_context_msgs.append(message)
+
+        return message
+
+    def _get_context_llm_contents (self) -> List[Dict[str, Any]]:
+        llm_messages = []
+        for message in self.task_context_msgs:
+            if message["is_llm_message"]:
+                llm_messages.append(message)
+
+        cxt_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)
+                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)
+
+        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
+            return StreamTaskResponser(response_context)
+        else:
+            from xgae.engine.responser.xga_non_stream_responser import NonStreamTaskResponser
+            return NonStreamTaskResponser(response_context)
+
+    def _create_response_context(self) -> TaskResponseContext:
+        response_context: TaskResponseContext = {
+            "is_stream": self.is_stream,
+            "task_id": self.task_id,
+            "task_run_id": self.task_run_id,
+            "trace_id": self.trace_id,
+            "model_name": self.model_name,
+            "max_xml_tool_calls": 0,
+            "add_context_msg": self.add_context_msg,
+            "tool_box": self.tool_box,
+            "tool_execution_strategy": "parallel",
+            "xml_adding_strategy": "user_message",
+        }
+        return response_context
 
 if __name__ == "__main__":
     import asyncio
-
+    from xgae.utils.utils import read_file
     async def main():
         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()
+        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()
 
-        async for chunk in engine.run_task(task_messages=[{}]):
+        async for chunk in engine.run_task(task_message={"role": "user", "content": "定位10.0.0.1的故障"},
+                                           max_auto_run=8):
              print(chunk)
 
     asyncio.run(main())

xgae/engine/xga_mcp_tool_box.py

@@ -75,6 +75,12 @@ class XGAMcpToolBox(XGAToolBox):
             await self.call_tool(task_id, "end_task", {"task_id": task_id})
         self.task_tool_schemas.pop(task_id, None)
 
+    @override
+    def get_task_tool_names(self, task_id: str) -> List[str]:
+        task_tool_schema = self.task_tool_schemas.get(task_id, {})
+        task_tool_names =  list(task_tool_schema.keys())
+        return task_tool_names
+
     @override
     def get_task_tool_schemas(self, task_id: str, type: Literal["general_tool",  "custom_tool"]) -> List[XGAToolSchema]:
         task_tool_schemas = []

xgae/utils/json_helpers.py

@@ -0,0 +1,174 @@
+"""
+JSON helper utilities for handling both legacy (string) and new (dict/list) formats.
+
+These utilities help with the transition from storing JSON as strings to storing
+them as proper JSONB objects in the database.
+"""
+
+import json
+from typing import Any, Union, Dict, List
+
+
+def ensure_dict(value: Union[str, Dict[str, Any], None], default: Dict[str, Any] = None) -> Dict[str, Any]:
+    """
+    Ensure a value is a dictionary.
+    
+    Handles:
+    - None -> returns default or {}
+    - Dict -> returns as-is
+    - JSON string -> parses and returns dict
+    - Other -> returns default or {}
+    
+    Args:
+        value: The value to ensure is a dict
+        default: Default value if conversion fails
+        
+    Returns:
+        A dictionary
+    """
+    if default is None:
+        default = {}
+        
+    if value is None:
+        return default
+        
+    if isinstance(value, dict):
+        return value
+        
+    if isinstance(value, str):
+        try:
+            parsed = json.loads(value)
+            if isinstance(parsed, dict):
+                return parsed
+            return default
+        except (json.JSONDecodeError, TypeError):
+            return default
+            
+    return default
+
+
+def ensure_list(value: Union[str, List[Any], None], default: List[Any] = None) -> List[Any]:
+    """
+    Ensure a value is a list.
+    
+    Handles:
+    - None -> returns default or []
+    - List -> returns as-is
+    - JSON string -> parses and returns list
+    - Other -> returns default or []
+    
+    Args:
+        value: The value to ensure is a list
+        default: Default value if conversion fails
+        
+    Returns:
+        A list
+    """
+    if default is None:
+        default = []
+        
+    if value is None:
+        return default
+        
+    if isinstance(value, list):
+        return value
+        
+    if isinstance(value, str):
+        try:
+            parsed = json.loads(value)
+            if isinstance(parsed, list):
+                return parsed
+            return default
+        except (json.JSONDecodeError, TypeError):
+            return default
+            
+    return default
+
+
+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.
+    
+    This handles the transition period where some data might be stored as
+    JSON strings (old format) and some as proper objects (new format).
+    
+    Args:
+        value: The value to parse
+        default: Default value if parsing fails
+        
+    Returns:
+        Parsed value or default
+    """
+    if value is None:
+        return default
+        
+    # 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
+
+
+def to_json_string(value: Any) -> str:
+    """
+    Convert a value to a JSON string if needed.
+    
+    This is used for backwards compatibility when yielding data that
+    expects JSON strings.
+    
+    Args:
+        value: The value to convert
+        
+    Returns:
+        JSON string representation
+    """
+    if isinstance(value, str):
+        # If it's already a string, check if it's valid JSON
+        try:
+            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)
+    
+    # For all other types, convert to JSON
+    return json.dumps(value)
+
+
+def format_for_yield(message_object: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Format a message object for yielding, ensuring content and metadata are JSON strings.
+    
+    This maintains backward compatibility with clients expecting JSON strings
+    while the database now stores proper objects.
+    
+    Args:
+        message_object: The message object from the database
+        
+    Returns:
+        Message object with content and metadata as JSON strings
+    """
+    if not message_object:
+        return message_object
+        
+    # Create a copy to avoid modifying the original
+    formatted = message_object.copy()
+    
+    # Ensure content is a JSON string
+    if 'content' in formatted and not isinstance(formatted['content'], str):
+        formatted['content'] = json.dumps(formatted['content'])
+        
+    # Ensure metadata is a JSON string
+    if 'metadata' in formatted and not isinstance(formatted['metadata'], str):
+        formatted['metadata'] = json.dumps(formatted['metadata'])
+        
+    return formatted 

xgae/utils/llm_client.py

@@ -4,11 +4,24 @@ import logging
 import os
 import litellm
 
-from typing import Union, Dict, Any, Optional, List
+from typing import Union, Dict, Any, Optional, List, TypedDict
 
 from litellm.utils import ModelResponse, CustomStreamWrapper
 from openai import OpenAIError
 
+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
 
 class LLMError(Exception):
     """Base exception for LLM-related errors."""
@@ -18,7 +31,7 @@ class LLMClient:
     RATE_LIMIT_DELAY = 30
     RETRY_DELAY = 0.1
 
-    def __init__(self, llm_config: Optional[Dict[str, Any]]=None) -> None:
+    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
@@ -34,7 +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 {}
+        llm_config = llm_config or LLMConfig()
         litellm.modify_params = True
         litellm.drop_params = True
 
@@ -214,9 +227,7 @@ class LLMClient:
 
 if __name__ == "__main__":
     async def llm_completion():
-        llm_client = LLMClient({
-            "stream": False #default is True
-        })
+        llm_client = LLMClient(LLMConfig(stream=False))
         messages = [{"role": "user", "content": "今天是2025年8月15日，北京本周每天温度"}]
         response = await llm_client.create_completion(messages)
         if llm_client.is_stream:

xgae/utils/setup_env.py

@@ -89,8 +89,5 @@ langfuse: Langfuse = Langfuse if _langfuse_initialized else setup_langfuse()
 
 
 if __name__ == "__main__":
-    try:
         trace_id = langfuse.create_trace_id()
-        print(f"trace_id={trace_id}")
-    except Exception as e:
-        handle_error(e)
+        logging.warning(f"trace_id={trace_id}")