Coverage for src / agent_contracts / runtime / executor.py: 98%

1"""Agent Runtime - Unified execution engine.""" 

2from __future__ import annotations 

3 

4from typing import Any, Callable, Awaitable 

5import logging 

6 

7from agent_contracts.runtime.context import RequestContext, ExecutionResult 

8from agent_contracts.runtime.hooks import RuntimeHooks, DefaultHooks 

9from agent_contracts.runtime.session import SessionStore 

10from agent_contracts.state_accessors import ( 

11 Internal, 

12 Request, 

13 Response, 

14 reset_response, 

15) 

16 

17logger = logging.getLogger(__name__) 

18 

19 

20class AgentRuntime: 

21 """Unified execution engine for agent graphs. 

22  

23 Provides a standardized execution lifecycle: 

24 1. Create initial state from request 

25 2. Restore session if resuming 

26 3. Apply prepare_state hook 

27 4. Execute the graph 

28 5. Build result 

29 6. Apply after_execution hook 

30  

31 Example: 

32 >>> from agent_contracts.runtime import AgentRuntime, InMemorySessionStore 

33 >>>  

34 >>> runtime = AgentRuntime( 

35 ... graph=compiled_graph, 

36 ... session_store=InMemorySessionStore(), 

37 ... ) 

38 >>>  

39 >>> result = await runtime.execute(RequestContext( 

40 ... session_id="abc123", 

41 ... action="answer", 

42 ... message="I like casual style", 

43 ... )) 

44 """ 

45 

46 def __init__( 

47 self, 

48 graph: Any, # CompiledGraph from langgraph 

49 hooks: RuntimeHooks | None = None, 

50 session_store: SessionStore | None = None, 

51 slices_to_restore: list[str] | None = None, 

52 ) -> None: 

53 """Initialize the runtime. 

54  

55 Args: 

56 graph: Compiled LangGraph graph 

57 hooks: Custom runtime hooks (optional) 

58 session_store: Session persistence store (optional) 

59 slices_to_restore: Slice names to restore from session (default: common ones) 

60 """ 

61 self.graph = graph 

62 self.hooks = hooks or DefaultHooks() 

63 self.session_store = session_store 

64 self.slices_to_restore = slices_to_restore or ["_internal", "interview", "shopping"] 

65 

66 async def execute(self, request: RequestContext) -> ExecutionResult: 

67 """Execute the agent graph. 

68  

69 Args: 

70 request: Execution request context 

71  

72 Returns: 

73 Execution result with final state and response 

74 """ 

75 try: 

76 # 1. Create initial state 

77 state = self._create_initial_state(request) 

78 logger.debug(f"Created initial state for session {request.session_id}") 

79 

80 # 2. Restore session if resuming 

81 if request.resume_session and self.session_store: 

82 session_data = await self.session_store.load(request.session_id) 

83 if session_data: 

84 state = self._merge_session(state, session_data) 

85 logger.debug(f"Restored session {request.session_id}") 

86 

87 # 3. Apply prepare_state hook 

88 state = await self.hooks.prepare_state(state, request) 

89 

90 # 4. Execute graph 

91 logger.debug(f"Executing graph for action: {request.action}") 

92 result_state = await self.graph.ainvoke(state) 

93 

94 # 5. Build result 

95 result = ExecutionResult.from_state(result_state) 

96 

97 # 6. Apply after_execution hook 

98 await self.hooks.after_execution(result_state, result) 

99 

100 logger.debug(f"Execution complete: {result.response_type}") 

101 return result 

102 

103 except Exception as e: 

104 logger.error(f"Execution failed: {e}", exc_info=True) 

105 return ExecutionResult.error_result(str(e)) 

106 

107 def _create_initial_state(self, request: RequestContext) -> dict[str, Any]: 

108 """Create initial state from request context. 

109  

110 Args: 

111 request: Request context 

112  

113 Returns: 

114 Initial state dictionary 

115 """ 

116 state: dict[str, Any] = {} 

117 

118 # Build request slice 

119 state = Request.session_id.set(state, request.session_id) 

120 state = Request.action.set(state, request.action) 

121 state = Request.params.set(state, request.params) 

122 state = Request.message.set(state, request.message) 

123 state = Request.image.set(state, request.image) 

124 

125 # Initialize response slice (empty) 

126 state = reset_response(state) 

127 

128 # Initialize internal slice 

129 state = Internal.turn_count.set(state, 0) 

130 state = Internal.is_first_turn.set(state, True) 

131 

132 return state 

133 

134 def _merge_session( 

135 self, 

136 state: dict[str, Any], 

137 session_data: dict[str, Any], 

138 ) -> dict[str, Any]: 

139 """Merge session data into state. 

140  

141 Args: 

142 state: Current state 

143 session_data: Session data to merge 

144  

145 Returns: 

146 Merged state 

147 """ 

148 result = dict(state) 

149 

150 for slice_name in self.slices_to_restore: 

151 if slice_name in session_data: 

152 current_slice = result.get(slice_name, {}) 

153 if not isinstance(current_slice, dict): 

154 current_slice = {} 

155 result[slice_name] = {**current_slice, **session_data[slice_name]} 

156 

157 return result