Coverage for src / agent_contracts / state.py: 90%

1"""State Slices - Domain-separated State. 

2 

3Following LangGraph best practices, State is separated by domain. 

4Each node declares and accesses only the slices it needs via Contract. 

5""" 

6from __future__ import annotations 

7 

8from typing import Any, TypedDict 

9 

10 

11# ============================================================================= 

12# Base Slices (Generic - Extend in your project) 

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

14 

15class BaseRequestSlice(TypedDict, total=False): 

16 """Base request information (read-only). 

17  

18 Set from API request, nodes typically only read. 

19 """ 

20 session_id: str 

21 action: str 

22 params: dict | None 

23 message: str | None 

24 image: str | None 

25 

26 

27class BaseResponseSlice(TypedDict, total=False): 

28 """Base response information. 

29  

30 Holds information to return as API response. 

31 """ 

32 response_type: str | None 

33 response_data: dict | None 

34 response_message: str | None 

35 

36 

37class BaseInternalSlice(TypedDict, total=False): 

38 """Internal flags (for Supervisor/Router). 

39  

40 Nodes should not access these directly. 

41 """ 

42 active_mode: str | None 

43 turn_count: int 

44 is_first_turn: bool 

45 next_node: str | None 

46 decision: str | None 

47 error: str | None 

48 

49 

50# ============================================================================= 

51# Base Agent State 

52# ============================================================================= 

53 

54class BaseAgentState(TypedDict, total=False): 

55 """Base Agent State - Extend in your project. 

56  

57 This provides the minimal state structure. 

58 Add your domain slices by subclassing. 

59  

60 Example: 

61 class MyAgentState(BaseAgentState): 

62 user: UserSlice 

63 shopping: ShoppingSlice 

64 """ 

65 request: BaseRequestSlice 

66 response: BaseResponseSlice 

67 _internal: BaseInternalSlice 

68 

69 

70# ============================================================================= 

71# Helper Functions 

72# ============================================================================= 

73 

74def get_slice(state: dict, slice_name: str) -> dict: 

75 """Get specified slice from state. 

76  

77 Args: 

78 state: AgentState or dict 

79 slice_name: Slice name (request, response, _internal, etc.) 

80  

81 Returns: 

82 Slice dict (empty dict if not found) 

83 """ 

84 return state.get(slice_name, {}) 

85 

86 

87def merge_slice_updates(state: dict, updates: dict[str, Any] | None) -> dict[str, Any]: 

88 """Merge slice-level updates for LangGraph. 

89  

90 Args: 

91 state: Current state 

92 updates: {slice_name: {field: value, ...}, ...} 

93  

94 Returns: 

95 Merged update dict (full slices) 

96 """ 

97 if not updates: 

98 return {} 

99 

100 merged: dict[str, Any] = {} 

101 for slice_name, slice_updates in updates.items(): 

102 current_slice = state.get(slice_name, {}) 

103 if isinstance(current_slice, dict) and isinstance(slice_updates, dict): 

104 merged[slice_name] = {**current_slice, **slice_updates} 

105 else: 

106 merged[slice_name] = slice_updates 

107 return merged 

108 

109 

110def apply_slice_updates(state: dict, updates: dict[str, Any] | None) -> dict: 

111 """Apply updates to state and return new state.""" 

112 merged_updates = merge_slice_updates(state, updates) 

113 if not merged_updates: 

114 return dict(state) 

115 new_state = dict(state) 

116 new_state.update(merged_updates) 

117 return new_state