stasis/application/orchestration/
tool_loop_pipeline.rs1use std::sync::Arc;
2
3use genai::chat::{ChatMessage, ChatRequest, ToolResponse};
4use serde::Serialize;
5use serde_json::Value;
6use tokio::sync::mpsc;
7
8use crate::application::orchestration::prompt_pipeline::{
9 PromptExecutionContext, PromptExecutionPipeline, PromptExecutionRequest,
10};
11use crate::application::orchestration::tool_registry::ToolRegistry;
12use crate::domain::errors::{Result, StasisError};
13use crate::ports::outbound::ai_chat_client::StreamDelta;
14
15const DEFAULT_MAX_TOOL_ROUNDS: usize = 10;
16
17#[derive(Clone, Debug, Eq, PartialEq, Default)]
18pub enum ToolCallMode {
19 #[default]
20 Auto,
21 Strict,
22}
23
24#[derive(Clone, Debug, Serialize)]
25pub struct ToolInvocation {
26 pub tool_name: String,
27 pub tool_input: Value,
28 pub tool_output: Value,
29}
30
31#[derive(Clone, Debug)]
32pub struct ToolLoopExecutionRequest {
33 pub user_prompt: String,
34 pub system_prompt: Option<String>,
35 pub context: PromptExecutionContext,
36 pub tool_name: String,
37 pub tool_input: Value,
38 pub tool_call_mode: ToolCallMode,
39}
40
41#[derive(Clone, Debug)]
42pub struct ToolLoopExecutionResponse {
43 pub text: String,
44 pub metadata: PromptExecutionContext,
45 pub tool_name: String,
46 pub tool_output: Value,
47 pub tool_invocations: Vec<ToolInvocation>,
48 pub rounds_executed: usize,
49 pub termination_reason: String,
50}
51
52#[derive(Clone)]
53pub struct ToolLoopPipeline {
54 prompt_pipeline: PromptExecutionPipeline,
55 tool_registry: Arc<dyn ToolRegistry>,
56}
57
58#[derive(Clone)]
59struct ToolLoopSharedInputs {
60 user_prompt: Arc<str>,
61 system_prompt: Option<Arc<str>>,
62 context: Arc<PromptExecutionContext>,
63 selected_tool_name: Arc<str>,
64 tool_input: Arc<Value>,
65 tool_call_mode: ToolCallMode,
66}
67
68impl ToolLoopSharedInputs {
69 fn context_clone(&self) -> PromptExecutionContext {
70 (*self.context).clone()
71 }
72
73 fn selected_tool_name(&self) -> &str {
74 &self.selected_tool_name
75 }
76}
77
78impl ToolLoopPipeline {
79 pub fn new(
80 prompt_pipeline: PromptExecutionPipeline,
81 tool_registry: Arc<dyn ToolRegistry>,
82 ) -> Self {
83 Self {
84 prompt_pipeline,
85 tool_registry,
86 }
87 }
88
89 pub async fn execute(
90 &self,
91 request: ToolLoopExecutionRequest,
92 ) -> Result<ToolLoopExecutionResponse> {
93 self.execute_with_defaults(request, Vec::new(), None).await
94 }
95
96 pub async fn execute_with_prior_messages(
97 &self,
98 request: ToolLoopExecutionRequest,
99 prior_messages: Vec<ChatMessage>,
100 ) -> Result<ToolLoopExecutionResponse> {
101 self.execute_with_defaults(request, prior_messages, None).await
102 }
103
104 pub async fn execute_with_stream(
105 &self,
106 request: ToolLoopExecutionRequest,
107 chunk_tx: Option<&mpsc::UnboundedSender<StreamDelta>>,
108 ) -> Result<ToolLoopExecutionResponse> {
109 self.execute_with_defaults(request, Vec::new(), chunk_tx).await
110 }
111
112 pub async fn execute_with_stream_prior_messages(
113 &self,
114 request: ToolLoopExecutionRequest,
115 prior_messages: Vec<ChatMessage>,
116 chunk_tx: Option<&mpsc::UnboundedSender<StreamDelta>>,
117 ) -> Result<ToolLoopExecutionResponse> {
118 self.execute_with_defaults(request, prior_messages, chunk_tx)
119 .await
120 }
121
122 pub async fn execute_with_stream_prior_messages_max_rounds(
123 &self,
124 request: ToolLoopExecutionRequest,
125 prior_messages: Vec<ChatMessage>,
126 chunk_tx: Option<&mpsc::UnboundedSender<StreamDelta>>,
127 max_tool_rounds: usize,
128 ) -> Result<ToolLoopExecutionResponse> {
129 self.execute_internal(request, prior_messages, chunk_tx, max_tool_rounds)
130 .await
131 }
132
133 async fn execute_with_defaults(
134 &self,
135 request: ToolLoopExecutionRequest,
136 prior_messages: Vec<ChatMessage>,
137 chunk_tx: Option<&mpsc::UnboundedSender<StreamDelta>>,
138 ) -> Result<ToolLoopExecutionResponse> {
139 self.execute_internal(request, prior_messages, chunk_tx, DEFAULT_MAX_TOOL_ROUNDS)
140 .await
141 }
142
143 async fn execute_internal(
144 &self,
145 request: ToolLoopExecutionRequest,
146 prior_messages: Vec<ChatMessage>,
147 chunk_tx: Option<&mpsc::UnboundedSender<StreamDelta>>,
148 max_tool_rounds: usize,
149 ) -> Result<ToolLoopExecutionResponse> {
150 let ToolLoopExecutionRequest {
151 user_prompt,
152 system_prompt,
153 context,
154 tool_name,
155 tool_input,
156 tool_call_mode,
157 } = request;
158
159 let max_tool_rounds = max_tool_rounds.max(1);
160 let shared_inputs = ToolLoopSharedInputs {
161 user_prompt: Arc::<str>::from(user_prompt),
162 system_prompt: system_prompt.map(Arc::<str>::from),
163 context: Arc::new(context),
164 selected_tool_name: Arc::<str>::from(tool_name),
165 tool_input: Arc::new(tool_input),
166 tool_call_mode,
167 };
168 let has_selected_tool = !shared_inputs.selected_tool_name().trim().is_empty();
169
170 let mut messages = Vec::with_capacity(2 + prior_messages.len());
171 if let Some(system_prompt) = shared_inputs.system_prompt.as_ref() {
172 messages.push(ChatMessage::system(system_prompt.to_string()));
173 }
174 messages.extend(prior_messages);
175 messages.push(ChatMessage::user(shared_inputs.user_prompt.to_string()));
176
177 let mut tools = self.tool_registry.list_tools().await?;
178 if has_selected_tool {
179 let selected_sanitized = sanitize_tool_name_for_model(shared_inputs.selected_tool_name());
180 let selected_prefix = format!("{selected_sanitized}_");
181 tools.retain(|tool| {
182 let tool_name = tool.name.as_ref();
183 tool_name == shared_inputs.selected_tool_name()
184 || tool_name == selected_sanitized
185 || tool_name.starts_with(&selected_prefix)
186 });
187 }
188
189 let mut invocations = Vec::new();
190 let mut should_use_legacy_fallback = false;
191 let mut fallback_draft_text: Option<String> = None;
192 let mut rounds_executed = 0usize;
193 if !tools.is_empty() {
194 for _ in 0..max_tool_rounds {
195 rounds_executed += 1;
196 let chat_request = ChatRequest::new(messages.clone()).with_tools(tools.clone());
197 let completion = match chunk_tx {
198 Some(tx) => {
199 self.prompt_pipeline
200 .complete_chat_stream(chat_request, shared_inputs.context_clone(), Some(tx))
201 .await?
202 }
203 None => {
204 self.prompt_pipeline
205 .complete_chat(chat_request, shared_inputs.context_clone())
206 .await?
207 }
208 };
209 let response = completion.response;
210 let maybe_text = response
211 .first_text()
212 .map(|value| value.trim().to_string())
213 .filter(|value| !value.is_empty());
214 let tool_calls = response.clone().into_tool_calls();
215
216 if tool_calls.is_empty() {
217 if invocations.is_empty() && has_selected_tool {
218 if shared_inputs.tool_call_mode == ToolCallMode::Strict {
219 return Err(StasisError::PortFailure(
220 "policy violation: strict tool-call mode expected model tool call but none was returned"
221 .to_string(),
222 ));
223 }
224
225 should_use_legacy_fallback = true;
226 fallback_draft_text = maybe_text;
227 break;
228 }
229
230 if let Some(text) = maybe_text {
231 let last = invocations.last().cloned().unwrap_or(ToolInvocation {
232 tool_name: shared_inputs.selected_tool_name().to_string(),
233 tool_input: (*shared_inputs.tool_input).clone(),
234 tool_output: Value::Null,
235 });
236
237 return Ok(ToolLoopExecutionResponse {
238 text,
239 metadata: shared_inputs.context_clone(),
240 tool_name: last.tool_name,
241 tool_output: last.tool_output,
242 tool_invocations: invocations,
243 rounds_executed,
244 termination_reason: "model_completed_no_tool_calls".to_string(),
245 });
246 }
247
248 return Err(StasisError::PortFailure(
249 "chat response was empty after tool loop".to_string(),
250 ));
251 }
252
253 messages.push(ChatMessage::from(tool_calls.clone()));
254 for call in tool_calls {
255 let tool_output = self
256 .tool_registry
257 .invoke_tool(&call.fn_name, call.fn_arguments.clone())
258 .await?;
259
260 let tool_output_text = tool_output.to_string();
261 messages.push(ChatMessage::from(ToolResponse::new(
262 call.call_id,
263 tool_output_text,
264 )));
265 invocations.push(ToolInvocation {
266 tool_name: call.fn_name,
267 tool_input: call.fn_arguments,
268 tool_output,
269 });
270 }
271 }
272
273 if !should_use_legacy_fallback {
274 return Err(StasisError::PortFailure(format!(
275 "tool loop exceeded max rounds ({max_tool_rounds}) without final response"
276 )));
277 }
278 }
279
280 if !should_use_legacy_fallback {
281 return Err(StasisError::PortFailure(
282 "no matching tools available for tool loop execution".to_string(),
283 ));
284 }
285
286 let draft_text = if let Some(text) = fallback_draft_text {
287 text
288 } else {
289 let mut first_request =
290 PromptExecutionRequest::from_user_prompt(shared_inputs.user_prompt.to_string())
291 .with_context(shared_inputs.context_clone());
292 if let Some(system_prompt) = shared_inputs.system_prompt.as_ref() {
293 first_request = first_request.with_system_prompt(system_prompt.to_string());
294 }
295 self.prompt_pipeline.execute(first_request).await?.text
296 };
297 let tool_output = self
298 .tool_registry
299 .invoke_tool(shared_inputs.selected_tool_name(), (*shared_inputs.tool_input).clone())
300 .await?;
301
302 let synthesis_prompt = build_fallback_synthesis_prompt(
303 &shared_inputs.user_prompt,
304 &draft_text,
305 shared_inputs.selected_tool_name(),
306 &tool_output,
307 );
308
309 let mut final_request = PromptExecutionRequest::from_user_prompt(synthesis_prompt)
310 .with_context(shared_inputs.context_clone());
311 if let Some(system_prompt) = shared_inputs.system_prompt.as_ref() {
312 final_request = final_request.with_system_prompt(system_prompt.to_string());
313 }
314
315 let final_response = self.prompt_pipeline.execute(final_request).await?;
316
317 let fallback_invocation = ToolInvocation {
318 tool_name: shared_inputs.selected_tool_name().to_string(),
319 tool_input: (*shared_inputs.tool_input).clone(),
320 tool_output: tool_output.clone(),
321 };
322
323 Ok(ToolLoopExecutionResponse {
324 text: final_response.text,
325 metadata: final_response.metadata,
326 tool_name: shared_inputs.selected_tool_name().to_string(),
327 tool_output,
328 tool_invocations: vec![fallback_invocation],
329 rounds_executed,
330 termination_reason: "legacy_fallback_no_model_tool_call".to_string(),
331 })
332 }
333}
334
335fn build_fallback_synthesis_prompt(
336 user_prompt: &str,
337 draft_text: &str,
338 tool_name: &str,
339 tool_output: &Value,
340) -> String {
341 let tool_output_text = tool_output.to_string();
342 let mut prompt = String::with_capacity(
343 user_prompt.len() + draft_text.len() + tool_name.len() + tool_output_text.len() + 128,
344 );
345 prompt.push_str("User request:\n");
346 prompt.push_str(user_prompt);
347 prompt.push_str("\n\nDraft analysis:\n");
348 prompt.push_str(draft_text);
349 prompt.push_str("\n\nTool '");
350 prompt.push_str(tool_name);
351 prompt.push_str("' output JSON:\n");
352 prompt.push_str(&tool_output_text);
353 prompt.push_str("\n\nProduce final answer grounded in the tool output.");
354 prompt
355}
356
357fn sanitize_tool_name_for_model(name: &str) -> String {
358 let mut out = String::with_capacity(name.len());
359 for ch in name.chars() {
360 if ch.is_ascii_alphanumeric() || ch == '_' || ch == '-' {
361 out.push(ch);
362 } else {
363 out.push('_');
364 }
365 }
366
367 let trimmed = out.trim_matches('_');
368 if trimmed.is_empty() {
369 "tool".to_string()
370 } else {
371 trimmed.to_string()
372 }
373}