Coverage for mcpgateway / plugins / framework / hooks / agents.py: 97%

1# -*- coding: utf-8 -*- 

2"""Location: ./mcpgateway/plugins/models/agents.py 

3Copyright 2025 

4SPDX-License-Identifier: Apache-2.0 

5Authors: Teryl Taylor 

6 

7Pydantic models for agent plugins. 

8This module implements the pydantic models associated with 

9the base plugin layer including configurations, and contexts. 

10""" 

11 

12# Standard 

13from enum import Enum 

14from typing import Any, Dict, List, Optional 

15 

16# Third-Party 

17from pydantic import Field 

18 

19# First-Party 

20from mcpgateway.common.models import Message 

21from mcpgateway.plugins.framework.hooks.http import HttpHeaderPayload 

22from mcpgateway.plugins.framework.models import PluginPayload, PluginResult 

23 

24 

25class AgentHookType(str, Enum): 

26 """Agent hook points. 

27 

28 Attributes: 

29 AGENT_PRE_INVOKE: Before agent invocation. 

30 AGENT_POST_INVOKE: After agent responds. 

31 

32 Examples: 

33 >>> AgentHookType.AGENT_PRE_INVOKE 

34 <AgentHookType.AGENT_PRE_INVOKE: 'agent_pre_invoke'> 

35 >>> AgentHookType.AGENT_PRE_INVOKE.value 

36 'agent_pre_invoke' 

37 >>> AgentHookType('agent_post_invoke') 

38 <AgentHookType.AGENT_POST_INVOKE: 'agent_post_invoke'> 

39 >>> list(AgentHookType) 

40 [<AgentHookType.AGENT_PRE_INVOKE: 'agent_pre_invoke'>, <AgentHookType.AGENT_POST_INVOKE: 'agent_post_invoke'>] 

41 """ 

42 

43 AGENT_PRE_INVOKE = "agent_pre_invoke" 

44 AGENT_POST_INVOKE = "agent_post_invoke" 

45 

46 

47class AgentPreInvokePayload(PluginPayload): 

48 """Agent payload for pre-invoke hook. 

49 

50 Attributes: 

51 agent_id: The agent identifier (can be modified for routing). 

52 messages: Conversation messages (can be filtered/transformed). 

53 tools: Optional list of tools available to agent. 

54 headers: Optional HTTP headers. 

55 model: Optional model override. 

56 system_prompt: Optional system instructions. 

57 parameters: Optional LLM parameters (temperature, max_tokens, etc.). 

58 

59 Examples: 

60 >>> payload = AgentPreInvokePayload(agent_id="agent-123", messages=[]) 

61 >>> payload.agent_id 

62 'agent-123' 

63 >>> payload.messages 

64 [] 

65 >>> payload.tools is None 

66 True 

67 >>> from mcpgateway.common.models import Message, Role, TextContent 

68 >>> msg = Message(role=Role.USER, content=TextContent(type="text", text="Hello")) 

69 >>> payload = AgentPreInvokePayload( 

70 ... agent_id="agent-456", 

71 ... messages=[msg], 

72 ... tools=["search", "calculator"], 

73 ... model="claude-3-5-sonnet-20241022" 

74 ... ) 

75 >>> payload.tools 

76 ['search', 'calculator'] 

77 >>> payload.model 

78 'claude-3-5-sonnet-20241022' 

79 """ 

80 

81 agent_id: str 

82 messages: List[Message] 

83 tools: Optional[List[str]] = None 

84 headers: Optional[HttpHeaderPayload] = None 

85 model: Optional[str] = None 

86 system_prompt: Optional[str] = None 

87 parameters: Optional[Dict[str, Any]] = Field(default_factory=dict) 

88 

89 

90class AgentPostInvokePayload(PluginPayload): 

91 """Agent payload for post-invoke hook. 

92 

93 Attributes: 

94 agent_id: The agent identifier. 

95 messages: Response messages from agent (can be filtered/transformed). 

96 tool_calls: Optional tool invocations made by agent. 

97 

98 Examples: 

99 >>> payload = AgentPostInvokePayload(agent_id="agent-123", messages=[]) 

100 >>> payload.agent_id 

101 'agent-123' 

102 >>> payload.messages 

103 [] 

104 >>> payload.tool_calls is None 

105 True 

106 >>> from mcpgateway.common.models import Message, Role, TextContent 

107 >>> msg = Message(role=Role.ASSISTANT, content=TextContent(type="text", text="Response")) 

108 >>> payload = AgentPostInvokePayload( 

109 ... agent_id="agent-456", 

110 ... messages=[msg], 

111 ... tool_calls=[{"name": "search", "arguments": {"query": "test"}}] 

112 ... ) 

113 >>> payload.tool_calls 

114 [{'name': 'search', 'arguments': {'query': 'test'}}] 

115 """ 

116 

117 agent_id: str 

118 messages: List[Message] 

119 tool_calls: Optional[List[Dict[str, Any]]] = None 

120 

121 

122AgentPreInvokeResult = PluginResult[AgentPreInvokePayload] 

123AgentPostInvokeResult = PluginResult[AgentPostInvokePayload] 

124 

125 

126def _register_agent_hooks() -> None: 

127 """Register agent hooks in the global registry. 

128 

129 This is called lazily to avoid circular import issues. 

130 """ 

131 # Import here to avoid circular dependency at module load time 

132 # First-Party 

133 from mcpgateway.plugins.framework.hooks.registry import get_hook_registry # pylint: disable=import-outside-toplevel 

134 

135 registry = get_hook_registry() 

136 

137 # Only register if not already registered (idempotent) 

138 if not registry.is_registered(AgentHookType.AGENT_PRE_INVOKE): 138 ↛ exitline 138 didn't return from function '_register_agent_hooks' because the condition on line 138 was always true

139 registry.register_hook(AgentHookType.AGENT_PRE_INVOKE, AgentPreInvokePayload, AgentPreInvokeResult) 

140 registry.register_hook(AgentHookType.AGENT_POST_INVOKE, AgentPostInvokePayload, AgentPostInvokeResult) 

141 

142 

143_register_agent_hooks()