Coverage for mcpgateway / plugins / framework / hooks / prompts.py: 96%

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

2"""Location: ./mcpgateway/plugins/hooks/prompts.py 

3Copyright 2025 

4SPDX-License-Identifier: Apache-2.0 

5Authors: Teryl Taylor 

6 

7Pydantic models for prompt 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 Optional 

15 

16# Third-Party 

17from pydantic import Field 

18 

19# First-Party 

20from mcpgateway.common.models import PromptResult 

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

22 

23 

24class PromptHookType(str, Enum): 

25 """MCP Forge Gateway hook points. 

26 

27 Attributes: 

28 prompt_pre_fetch: The prompt pre hook. 

29 prompt_post_fetch: The prompt post hook. 

30 tool_pre_invoke: The tool pre invoke hook. 

31 tool_post_invoke: The tool post invoke hook. 

32 resource_pre_fetch: The resource pre fetch hook. 

33 resource_post_fetch: The resource post fetch hook. 

34 

35 Examples: 

36 >>> PromptHookType.PROMPT_PRE_FETCH 

37 <PromptHookType.PROMPT_PRE_FETCH: 'prompt_pre_fetch'> 

38 >>> PromptHookType.PROMPT_PRE_FETCH.value 

39 'prompt_pre_fetch' 

40 >>> PromptHookType('prompt_post_fetch') 

41 <PromptHookType.PROMPT_POST_FETCH: 'prompt_post_fetch'> 

42 >>> list(PromptHookType) 

43 [<PromptHookType.PROMPT_PRE_FETCH: 'prompt_pre_fetch'>, <PromptHookType.PROMPT_POST_FETCH: 'prompt_post_fetch'>] 

44 """ 

45 

46 PROMPT_PRE_FETCH = "prompt_pre_fetch" 

47 PROMPT_POST_FETCH = "prompt_post_fetch" 

48 

49 

50class PromptPrehookPayload(PluginPayload): 

51 """A prompt payload for a prompt prehook. 

52 

53 Attributes: 

54 prompt_id (str): The ID of the prompt template. 

55 args (dic[str,str]): The prompt template arguments. 

56 

57 Examples: 

58 >>> payload = PromptPrehookPayload(prompt_id="123", args={"user": "alice"}) 

59 >>> payload.prompt_id 

60 '123' 

61 >>> payload.args 

62 {'user': 'alice'} 

63 >>> payload2 = PromptPrehookPayload(prompt_id="empty") 

64 >>> payload2.args 

65 {} 

66 >>> p = PromptPrehookPayload(prompt_id="123", args={"name": "Bob", "time": "morning"}) 

67 >>> p.prompt_id 

68 '123' 

69 >>> p.args["name"] 

70 'Bob' 

71 """ 

72 

73 prompt_id: str 

74 args: Optional[dict[str, str]] = Field(default_factory=dict) 

75 

76 

77class PromptPosthookPayload(PluginPayload): 

78 """A prompt payload for a prompt posthook. 

79 

80 Attributes: 

81 prompt_id (str): The prompt ID. 

82 result (PromptResult): The prompt after its template is rendered. 

83 

84 Examples: 

85 >>> from mcpgateway.common.models import PromptResult, Message, TextContent 

86 >>> msg = Message(role="user", content=TextContent(type="text", text="Hello World")) 

87 >>> result = PromptResult(messages=[msg]) 

88 >>> payload = PromptPosthookPayload(prompt_id="123", result=result) 

89 >>> payload.prompt_id 

90 '123' 

91 >>> payload.result.messages[0].content.text 

92 'Hello World' 

93 >>> from mcpgateway.common.models import PromptResult, Message, TextContent 

94 >>> msg = Message(role="assistant", content=TextContent(type="text", text="Test output")) 

95 >>> r = PromptResult(messages=[msg]) 

96 >>> p = PromptPosthookPayload(prompt_id="123", result=r) 

97 >>> p.prompt_id 

98 '123' 

99 """ 

100 

101 prompt_id: str 

102 result: PromptResult 

103 

104 

105PromptPrehookResult = PluginResult[PromptPrehookPayload] 

106PromptPosthookResult = PluginResult[PromptPosthookPayload] 

107 

108 

109def _register_prompt_hooks() -> None: 

110 """Register prompt hooks in the global registry. 

111 

112 This is called lazily to avoid circular import issues. 

113 """ 

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

115 # First-Party 

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

117 

118 registry = get_hook_registry() 

119 

120 # Only register if not already registered (idempotent) 

121 if not registry.is_registered(PromptHookType.PROMPT_PRE_FETCH): 121 ↛ exitline 121 didn't return from function '_register_prompt_hooks' because the condition on line 121 was always true

122 registry.register_hook(PromptHookType.PROMPT_PRE_FETCH, PromptPrehookPayload, PromptPrehookResult) 

123 registry.register_hook(PromptHookType.PROMPT_POST_FETCH, PromptPosthookPayload, PromptPosthookResult) 

124 

125 

126_register_prompt_hooks()