feat: 为 agent 添加 instruction args

wangtao2001 · wangtao2001 · commit bceb1cb5b25b · 2025-04-17T19:35:53.000+08:00
diff --git a/docs/tutorials/agent/agent.md b/docs/tutorials/agent/agent.md
@@ -26,6 +26,19 @@ translator = Agent(name='translator',
 
 当然你也可以将具体的指令需求直接写在 `instruction` 中, 这种方式将在 [工作流编排](#工作流编排) 中具体解释。
 
+### 指令参数
+
+在创建 `Agent` 对象时, `instruction` 不仅可以是一个字符串，也可以是一个函数，但这个函数 **必需** 返回一个字符串。该函数可以通过 `instruction_args` 传递参数, 这些参数将会在智能体初始化时传递给 `instruction` 函数。
+
+```python
+def get_instruction(name: str):
+    return f'你的名字是{name}'
+
+agent = Agent(llm=llm,
+              instruction=get_instruction,
+              instruction_args={'name': '张三'})
+```
+
 ## 创建一个控制器
 
 ```python
@@ -194,7 +207,7 @@ controller = Controller(context_variables={'current_time': '2024/09/01'})
 
 #### instruction中使用
 
-在创建 `Agent` 对象时, `instruction` 不仅可以是一个字符串，也可以是一个函数，但这个函数 **必需** 返回一个字符串, 同时可以额外传递一个`ContextVariables` 类型的形参。
+当 `instruction` 类型为函数时, 可以额外传递一个`ContextVariables` 类型的形参, 但 **必需** 标注这个形参的类型为 `ContextVariables` 类型。
 
 ``` python
 from course_graph.agent import ContextVariables
@@ -214,7 +227,7 @@ assistant = Agent(name="assistant",
 
 在定义外部工具函数时, 也可以传递一个 `ContextVariables` 类型的形参。同样的, 控制器也会在调用这些函数的时候自动注入上下文变量。
 
-虽然不需要在文档中描述这个形参, 但是 **必需** 标注这个形参的类型为 `ContextVariables` 类型:
+相同的，虽然不需要在文档中描述这个形参, 但是 **必需** 标注这个形参的类型为 `ContextVariables` 类型:
 
 ```python
 def get_weather(location: str, context_variables: ContextVariables) -> str:
@@ -354,6 +367,7 @@ controller = Controller(trace_callback=pprint)
 - `TOOL_CALL`: 工具调用
 - `TOOL_RESULT`: 工具调用结果
 - `CONTEXT_UPDATE`: 上下文变量更新
+- `MCP_TOOL_CALL`: MCP 工具调用
 
 ## 多智能体编排
 
diff --git a/src/course_graph/agent/agent.py b/src/course_graph/agent/agent.py
@@ -10,10 +10,10 @@
 from openai.types.chat import *
 import inspect
 import docstring_parser
-from typing import Callable, Awaitable
-from typing import Literal
+from typing import Callable, Awaitable, Literal
 from openai import NOT_GIVEN, NotGiven
 from .mcp import MCPServer
+from shortuuid import uuid
 
 
 class Agent:
@@ -23,9 +23,10 @@ def __init__(
             llm: LLMBase,
             name: str = 'Assistant',
             functions: list[Callable | Awaitable] = None,
-            tool_choice: str | NotGiven | Literal['required', 'auto', 'none'] = NOT_GIVEN,
-            parallel_tool_calls: bool | NotGiven = NOT_GIVEN,
-            instruction: str | Callable[[ContextVariables], str] | Callable[[], str] = 'You are a helpful assistant.',
+            tool_choice: str | Literal['required', 'auto', 'none'] = 'auto',
+            parallel_tool_calls: bool = False,
+            instruction: str | Callable[..., str] = 'You are a helpful assistant.',
+            instruction_args: dict = None,
             mcp_server: list[MCPServer] = None,
             mcp_impl: Literal['function_call'] = 'function_call'
     ) -> None:
@@ -36,15 +37,17 @@ def __init__(
             name (str, optional): 名称. Defaults to 'Assistant'.
             functions: (list[Callable | Awaitable], optional): 工具函数. Defaults to None.
             parallel_tool_calls: (bool, optional): 允许工具并行调用. Defaults to False.
-            tool_choice: (Literal['required', 'auto', 'none'] | NotGiven, optional). 强制使用工具函数, 选择模式或提供函数名称. Defaults to NOT_GIVEN.
-            instruction (str | Callable[[ContextVariables], str] | Callable[[], str], optional): 指令. Defaults to 'You are a helpful assistant.'.
+            tool_choice: (Literal['required', 'auto', 'none'], optional). 强制使用工具函数, 选择模式或提供函数名称. Defaults to 'auto'.
+            instruction (str | Callable[Any, str], optional): 指令. Defaults to 'You are a helpful assistant.'.
+            instruction_args: (dict, optional): 指令参数. Defaults to {}.
             mcp_server: (list[MCPServer], optional): MCP 服务器. Defaults to None.
             mcp_impl: (Literal['function_call'] | NotGiven, optional): MCP 协议实现方式, 目前只支持 'function_call'. Defaults to 'function_call'.
         """
+        self.id = str(uuid())
         self.llm = llm
         self.name = name
         self.instruction = instruction
-
+        self.instruction_args = instruction_args
         self.tools: list[ChatCompletionToolParam] = []  # for LLM
         
         self.tool_functions: dict[str, Callable | Awaitable] = {}  # for local function call
@@ -56,8 +59,9 @@ def __init__(
 
         if functions:
             self.add_tool_functions(*functions)
+        
 
-        if tool_choice != NOT_GIVEN and tool_choice not in ['required', 'auto', 'none']:  # 需要注意工具无限循环
+        if tool_choice not in ['required', 'auto', 'none']:  # 需要注意工具无限循环
             self.tool_choice = {
                 "type": "function",
                 "function": {
@@ -281,3 +285,6 @@ def add_tool_functions(self, *functions: Callable | Awaitable) -> 'Agent':
                 }
             })
         return self
+    
+    def __repr__(self) -> str:
+        return f'Agent(id={self.id}, name={self.name})'
diff --git a/src/course_graph/agent/controller.py b/src/course_graph/agent/controller.py
@@ -55,8 +55,13 @@ def set_agent_instruction(self, agent: Agent) -> None:
         match agent.instruction:
             case _ if callable(agent.instruction):
                 parameters = inspect.signature(agent.instruction).parameters
-                args = (self.context_variables,) if len(parameters) == 1 else ()
-                agent.llm.instruction = agent.instruction(*args)
+                args = {}
+                for arg_name, p in parameters.items():
+                    if p.annotation == ContextVariables:
+                        args[arg_name] = self.context_variables
+                    else:
+                        args[arg_name] = agent.instruction_args.get(arg_name, p.default)
+                agent.llm.instruction = agent.instruction(**args)
             case _:
                 agent.llm.instruction = agent.instruction
 
@@ -88,32 +93,26 @@ async def run(self, agent: Agent, message: str = None) -> tuple[Agent, str]:
         self._add_trace_event(TraceEvent(
             timestamp=datetime.now(),
             event_type=TraceEventType.USER_MESSAGE,
-            agent_name=agent.name,
+            agent=agent,
             data={'message': message}
         ))
 
-        assistant_output = agent.chat_completion(message)
-
-        self._add_trace_event(TraceEvent(
-            timestamp=datetime.now(),
-            event_type=TraceEventType.AGENT_MESSAGE,
-            agent_name=agent.name,
-            data={'message': assistant_output.content}
-        ))
+        assistant_output = agent.chat_completion(message)    
 
         while assistant_output.tool_calls:  # None 或者空数组
             functions = assistant_output.tool_calls
             for item in functions:
                 function = item.function
                 args = json.loads(function.arguments)
 
-                self._add_trace_event(TraceEvent(
-                    timestamp=datetime.now(),
-                    event_type=TraceEventType.TOOL_CALL,
-                    agent_name=agent.name,
-                    data={'function': function.name, 'arguments': args}
-                ))
                 if (tool_function := agent.tool_functions.get(function.name)) is not None:
+                    
+                    self._add_trace_event(TraceEvent(
+                        timestamp=datetime.now(),
+                        event_type=TraceEventType.TOOL_CALL,
+                        agent=agent,
+                        data={'function': function.name, 'arguments': args}
+                    ))
 
                     # 自动注入上下文变量
                     if (var_name := agent.use_context_variables.get(function.name)) is not None:
@@ -141,6 +140,13 @@ async def run(self, agent: Agent, message: str = None) -> tuple[Agent, str]:
                             result = Result()
 
                 elif (mcp_sever := agent.mcp_functions.get(function.name)) is not None:
+                    self._add_trace_event(TraceEvent(
+                        timestamp=datetime.now(),
+                        agent=agent,
+                        event_type=TraceEventType.MCP_TOOL_CALL,
+                        data={'function': function.name, 'arguments': args}
+                    ))
+
                     resp = (await mcp_sever.session.call_tool(function.name, args)).content
                     text_contents = []
                     for content in resp:
@@ -158,11 +164,17 @@ async def run(self, agent: Agent, message: str = None) -> tuple[Agent, str]:
                 else:
                     result = Result(content=f'Failed to call tool: {function.name}')
 
+                trace_result = {'content': result.content}
+                if result.context_variables._vars:
+                    trace_result['context_variables'] = result.context_variables._vars
+                if not result.message:
+                    trace_result['message'] = False
+                    
                 self._add_trace_event(TraceEvent(
                     timestamp=datetime.now(),
-                    agent_name=agent.name,
+                    agent=agent,
                     event_type=TraceEventType.TOOL_RESULT,
-                    data={'function': function.name, 'result': result}
+                    data={'function': function.name, 'result': trace_result}
                 ))
 
                 agent.add_tool_call_message(result.content, item.id)
@@ -172,19 +184,19 @@ async def run(self, agent: Agent, message: str = None) -> tuple[Agent, str]:
 
                     self._add_trace_event(TraceEvent(
                         timestamp=datetime.now(),
-                        agent_name=agent.name,
+                        agent=agent,
                         event_type=TraceEventType.AGENT_SWITCH,
                         data={'to_agent': result.agent.name}
                     ))
                     agent = result.agent
-
-                self._add_trace_event(TraceEvent(
-                    timestamp=datetime.now(),
-                    agent_name=agent.name,
-                    event_type=TraceEventType.CONTEXT_UPDATE,
-                    data={'old_context': self.context_variables, 'new_context': result.context_variables}
-                ))
-                self.context_variables.update(result.context_variables)
+                if result.context_variables._vars:
+                    self._add_trace_event(TraceEvent(
+                        timestamp=datetime.now(),
+                        agent=agent,
+                        event_type=TraceEventType.CONTEXT_UPDATE,
+                        data={'old_context': self.context_variables, 'new_context': result.context_variables}
+                    ))
+                    self.context_variables.update(result.context_variables)
 
                 self.set_agent_instruction(agent)
 
@@ -193,13 +205,13 @@ async def run(self, agent: Agent, message: str = None) -> tuple[Agent, str]:
             if turn > self.max_turns:
                 raise MaxTurnsException
 
-        message = assistant_output.content
         self._add_trace_event(TraceEvent(
             timestamp=datetime.now(),
             event_type=TraceEventType.AGENT_MESSAGE,
-            agent_name=agent.name,
-            data={'message': message}
+            agent=agent,
+            data={'message': assistant_output.content}
         ))
+
         self.trace['end_time'] = datetime.now()
 
         return agent, assistant_output.content
diff --git a/src/course_graph/agent/trace.py b/src/course_graph/agent/trace.py
@@ -9,6 +9,7 @@
 from typing import TypedDict, List
 from datetime import datetime
 from course_graph import set_logger, logger
+from course_graph.agent import Agent
 
 set_logger(console=True, file=False)
 
@@ -20,12 +21,13 @@ class TraceEventType(Enum):
     TOOL_CALL = 'tool_call'
     TOOL_RESULT = 'tool_result'
     CONTEXT_UPDATE = 'context_update'
+    MCP_TOOL_CALL = 'mcp_tool_call'
 
 
 @dataclass
 class TraceEvent:
     timestamp: datetime
-    agent_name: str
+    agent: Agent
     event_type: TraceEventType
     data: dict
 
diff --git a/src/course_graph/agent/types.py b/src/course_graph/agent/types.py
@@ -44,6 +44,9 @@ def update(self, other: Union[dict, 'ContextVariables']):
             self._vars.update(other)
         else:
             self._vars.update(other._vars)
+            
+    def get(self, key: Any, default: Any = None) -> Any:
+        return self._vars.get(key, default)
 
 
 @dataclass
