Coverage for src / agent_contracts / contracts.py: 88%

1"""NodeContract - Node I/O contracts. 

2 

3Each node defines a CONTRACT class variable to declare its 

4inputs, outputs, dependencies, and trigger conditions. 

5""" 

6from __future__ import annotations 

7 

8from typing import Any 

9 

10from pydantic import BaseModel, Field, ConfigDict 

11 

12 

13# ============================================================================= 

14# Trigger Conditions 

15# ============================================================================= 

16 

17class TriggerCondition(BaseModel): 

18 """Condition for when a node should be triggered. 

19  

20 The Supervisor collects matching conditions as hints for LLM decision. 

21 If no LLM is available, rule-based matching is used as fallback. 

22 """ 

23 model_config = ConfigDict(frozen=True) 

24 

25 # Priority (higher = evaluated first) 

26 priority: int = Field(default=0, description="Evaluation priority. Higher values are evaluated first.") 

27 

28 # Rule-based conditions 

29 when: dict[str, Any] | None = Field( 

30 default=None, 

31 description="Match conditions. {slice.field: expected_value, ...}", 

32 ) 

33 when_not: dict[str, Any] | None = Field( 

34 default=None, 

35 description="Non-match conditions. Matches when these are NOT true.", 

36 ) 

37 

38 # LLM decision hint 

39 llm_hint: str | None = Field( 

40 default=None, 

41 description="Hint for LLM when making routing decisions.", 

42 ) 

43 

44 

45# ============================================================================= 

46# Node Contract 

47# ============================================================================= 

48 

49class NodeContract(BaseModel): 

50 """Node I/O contract. 

51  

52 Each node defines this as a CONTRACT class variable. 

53 Registry uses this for I/O validation, routing, and dependency analysis. 

54  

55 Example: 

56 class LikeHandlerNode(ModularNode): 

57 CONTRACT = NodeContract( 

58 name="like_handler", 

59 description="Records product LIKE to card", 

60 reads=["request", "card", "shopping"], 

61 writes=["card", "shopping", "response"], 

62 ... 

63 ) 

64 """ 

65 model_config = ConfigDict(frozen=True) 

66 

67 # === Identification === 

68 name: str = Field(description="Node name (key for graph registration)") 

69 description: str = Field(description="Node role description") 

70 

71 # === I/O Definition (by slice) === 

72 reads: list[str] = Field( 

73 description="List of slice names to read from" 

74 ) 

75 writes: list[str] = Field( 

76 description="List of slice names to write to" 

77 ) 

78 

79 # === Dependencies === 

80 requires_llm: bool = Field( 

81 default=False, 

82 description="Whether LLM client is required", 

83 ) 

84 services: list[str] = Field( 

85 default_factory=list, 

86 description="Required service names (card_service, shopify_service, etc.)", 

87 ) 

88 

89 # === Supervisor === 

90 supervisor: str = Field( 

91 description="Supervisor this node belongs to", 

92 ) 

93 

94 # === Trigger Conditions === 

95 trigger_conditions: list[TriggerCondition] = Field( 

96 default_factory=list, 

97 description="List of conditions that trigger this node", 

98 ) 

99 

100 # === Terminal Condition === 

101 is_terminal: bool = Field( 

102 default=False, 

103 description="Whether this node should transition to END after execution", 

104 ) 

105 

106 def get_highest_priority_condition(self) -> TriggerCondition | None: 

107 """Get the highest priority trigger condition.""" 

108 if not self.trigger_conditions: 

109 return None 

110 return max(self.trigger_conditions, key=lambda c: c.priority) 

111 

112 def get_llm_hints(self) -> list[str]: 

113 """Get all LLM hints.""" 

114 return [c.llm_hint for c in self.trigger_conditions if c.llm_hint] 

115 

116 

117# ============================================================================= 

118# I/O Types 

119# ============================================================================= 

120 

121class NodeInputs(BaseModel): 

122 """Node inputs. 

123  

124 Holds slices extracted based on Contract.reads. 

125 """ 

126 model_config = ConfigDict(extra="allow") # Allow dynamic slice addition 

127 

128 def get_slice(self, name: str) -> dict: 

129 """Get specified slice.""" 

130 return getattr(self, name, {}) 

131 

132 

133class NodeOutputs(BaseModel): 

134 """Node outputs. 

135  

136 Holds slices to update based on Contract.writes. 

137 """ 

138 model_config = ConfigDict(extra="allow") # Allow dynamic slice addition 

139 

140 def to_state_updates(self) -> dict[str, dict]: 

141 """Convert to State update dict.""" 

142 return {k: v for k, v in self.model_dump().items() if v is not None}