xgae 0.1.3__py3-none-any.whl → 0.1.5__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,23 @@ 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"]
-    is_llm_message: bool
+
+class XGAResponseMsg(TypedDict, total=False):
+    type: Literal["user", "status",  "tool", "assistant", "assistant_response_end"]
     content: Union[Dict[str, Any], List[Any], str]
-    metadata: Optional[Dict[str, Any]]
+    is_llm_message: bool
+    metadata: Dict[str, Any]
+    message_id: str
+    task_id: str
+    task_run_id: str
+    trace_id: str
     session_id: Optional[str]
     agent_id: Optional[str]
-    task_id: Optional[str]
+
+class XGATaskResult(TypedDict, total=False):
+    type: Literal["ask", "answer", "error"]
+    content: str
+    attachments: Optional[List[str]]
 
 @dataclass
 class XGAToolSchema:
@@ -31,6 +38,7 @@ class XGAToolResult:
     success: bool
     output: str
 
+
 class XGAToolBox(ABC):
     @abstractmethod
     async def creat_task_tool_box(self, task_id: str, general_tools: List[str], custom_tools: List[str]):
@@ -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 XGAResponseMsg, XGAToolBox, XGATaskResult
+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.session_id = session_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,8 +34,12 @@ class XGATaskEngine():
         self.prompt_builder = prompt_builder or XGAPromptBuilder(system_prompt)
         self.tool_box = tool_box or XGAMcpToolBox()
 
+        self.task_response_msgs: List[XGAResponseMsg] = []
+        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:
+    async def _post_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")
@@ -46,54 +52,296 @@ 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", "ask"]
         custom_tools = custom_tools or []
-        await engine.__async_init__(general_tools, custom_tools)
+        await engine._post_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_with_final_answer(self,
+                       task_message: Dict[str, Any],
+                       max_auto_run: int = 25,
+                       trace_id: Optional[str] = None) -> XGATaskResult:
+        chunks = []
+        async for chunk in self.run_task(task_message=task_message, max_auto_run=max_auto_run, trace_id=trace_id):
+            chunks.append(chunk)
 
-    async def run_task(self, task_messages: List[Dict[str, Any]]) -> AsyncGenerator:
+        final_result = self._parse_final_result(chunks)
+        return final_result
+
+    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_response_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_response_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):
+            self._reponse_chunk_log(chunk)
+            yield chunk
+
+    async def _run_task_auto(self, max_auto_run: int) -> AsyncGenerator[Dict[str, Any], None]:
+        continuous_state: TaskRunContinuousState = {
+            "accumulated_content": "",
+            "auto_continue_count": 0,
+            "auto_continue": True
+        }
+
+        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):
+                    yield chunk
+                    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')}")
+                                auto_continue = False
+                                break
+                            elif status_type == 'finish':
+                                finish_reason = content.get('finish_reason', None)
+                                if finish_reason == 'completed':
+                                    logging.warning(f"run_task_auto: Detected finish_reason='completed', Task Completed Success !")
+                                    auto_continue = False
+                                    break
+                                elif finish_reason == 'xml_tool_limit_reached':
+                                    logging.warning(f"run_task_auto: Detected finish_reason='xml_tool_limit_reached', stopping auto-continue")
+                                    auto_continue = False
+                                    break
+                                elif finish_reason == 'stop' or finish_reason == 'length': # 'length' never occur
+                                    auto_continue = True
+                                    auto_continue_count += 1
+                                    update_continuous_state(auto_continue_count, auto_continue)
+                                    logging.info(f"run_task_auto: Detected finish_reason='{finish_reason}', auto-continuing ({auto_continue_count}/{max_auto_run})")
+                    except StopAsyncIteration:
+                        pass
+                    except Exception as parse_error:
+                        logging.error(f"run_task_auto: Error in parse chunk: {str(parse_error)}")
+                        content = {"role": "system", "status_type": "error", "message": "Parse response chunk Error"}
+                        error_msg = self.add_response_msg(type="status", content=content, is_llm_message=False)
+                        yield error_msg
+            except Exception as run_error:
+                logging.error(f"run_task_auto: Call task_run_once error: {str(run_error)}")
+                content = {"role": "system", "status_type": "error", "message": "Call task_run_once error"}
+                error_msg = self.add_response_msg(type="status", content=content, is_llm_message=False)
+                yield error_msg
+
+    def _parse_final_result(self, chunks: List[Dict[str, Any]]) -> XGATaskResult:
+        final_result: XGATaskResult = None
+        try:
+            finish_reason = ''
+            for chunk in reversed(chunks):
+                chunk_type = chunk.get("type")
+                if chunk_type == "status":
+                    status_content = json.loads(chunk.get('content', '{}'))
+                    status_type = status_content.get('status_type', None)
+                    if status_type == "error":
+                        error = status_content.get('message', 'Unknown error')
+                        final_result = XGATaskResult(type="error", content=error)
+                        break
+                    elif status_type == "finish":
+                        finish_reason = status_content.get('finish_reason', None)
+                        if finish_reason == 'xml_tool_limit_reached':
+                            error = "Completed due to over task max_auto_run limit !"
+                            final_result = XGATaskResult(type="error", content=error)
+                            break
+                        continue
+                elif chunk_type == "tool" and finish_reason in ['completed', 'stop']:
+                    tool_content = json.loads(chunk.get('content', '{}'))
+                    tool_execution = tool_content.get('tool_execution')
+                    tool_name = tool_execution.get('function_name')
+                    if tool_name == "complete":
+                        result_content = tool_execution["arguments"].get("text", "Task completed with no answer")
+                        attachments = tool_execution["arguments"].get("attachments", None)
+                        final_result = XGATaskResult(type="answer", content=result_content, attachments=attachments)
+                    elif tool_name == "ask":
+                        result_content = tool_execution["arguments"].get("text", "Task ask for more info")
+                        attachments = tool_execution["arguments"].get("attachments", None)
+                        final_result = XGATaskResult(type="ask", content=result_content, attachments=attachments)
+                    else:
+                        tool_result = tool_execution.get("result", None)
+                        if tool_result is not None:
+                            success = tool_result.get("success")
+                            output = tool_result.get("output")
+                            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':
+                    assis_content = chunk.get('content', '{}')
+                    result_content = assis_content.get("content", "LLM output is empty")
+                    final_result = XGATaskResult(type="answer", content=result_content)
+                if final_result is not None:
+                    break
+        except Exception as e:
+            logging.error(f"parse_final_result: Final result pass error: {str(e)}")
+            final_result = XGATaskResult(type="error", content="Parse final result failed!")
+            handle_error(e)
+
+        return final_result
+
+    def add_response_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)-> XGAResponseMsg:
+        message = XGAResponseMsg(
+            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_response_msgs.append(message)
+
+        return message
+
+    def _get_response_llm_contents (self) -> List[Dict[str, Any]]:
+        llm_messages = []
+        for message in self.task_response_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_response_msg,
+            "tool_box": self.tool_box,
+            "tool_execution_strategy": "parallel",
+            "xml_adding_strategy": "user_message",
+        }
+        return response_context
+
+    def _reponse_chunk_log(self, chunk):
+        chunk_type = chunk.get('type')
+        prefix = ""
+
+        if chunk_type == 'status':
+            content = json.loads(chunk.get('content', '{}'))
+            status_type = content.get('status_type', "empty")
+            prefix = "-" + status_type
+        elif chunk_type == 'tool':
+            tool_content = json.loads(chunk.get('content', '{}'))
+            tool_execution = tool_content.get('tool_execution')
+            tool_name = tool_execution.get('function_name')
+            prefix = "-" + tool_name
+
+        logging.info(f"TASK_RESP_CHUNK[{chunk_type}{prefix}]: {chunk}")
 
-    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)
 
 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()
-
-        async for chunk in engine.run_task(task_messages=[{}]):
-             print(chunk)
+        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(llm_config=LLMConfig(stream=False))
+        #chunks = []
+        # async for chunk in engine.run_task(task_message={"role": "user", "content": "定位10.0.0.1的故障"},max_auto_run=8):
+        #     print(chunk)
+        #final_result = await engine.run_task_with_final_answer(task_message={"role": "user", "content": "1+1"}, max_auto_run=2)
 
+        final_result = await engine.run_task_with_final_answer(task_message={"role": "user", "content": "定位10.0.1.1故障"},max_auto_run=8)
+        print("FINAL RESULT:", final_result)
     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 = []
@@ -106,10 +112,7 @@ class XGAMcpToolBox(XGAToolBox):
             if mcp_tool:
                 tool_args = args or {}
                 if server_name == self.GENERAL_MCP_SERVER_NAME:
-                    pass
-                    #tool_args["task_id"] = task_id #xga general tool, first param must be task_id
-                else:
-                    tool_args = args
+                    tool_args = dict({"task_id": task_id}, **tool_args)
 
                 try:
                     tool_result = await mcp_tool.arun(tool_args)

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: