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 tool.name == shared_inputs.selected_tool_name()
183 || tool.name == selected_sanitized
184 || tool.name.starts_with(&selected_prefix)
185 });
186 }
187
188 let mut invocations = Vec::new();
189 let mut should_use_legacy_fallback = false;
190 let mut fallback_draft_text: Option<String> = None;
191 let mut rounds_executed = 0usize;
192 if !tools.is_empty() {
193 for _ in 0..max_tool_rounds {
194 rounds_executed += 1;
195 let chat_request = ChatRequest::new(messages.clone()).with_tools(tools.clone());
196 let completion = match chunk_tx {
197 Some(tx) => {
198 self.prompt_pipeline
199 .complete_chat_stream(chat_request, shared_inputs.context_clone(), Some(tx))
200 .await?
201 }
202 None => {
203 self.prompt_pipeline
204 .complete_chat(chat_request, shared_inputs.context_clone())
205 .await?
206 }
207 };
208 let response = completion.response;
209 let maybe_text = response
210 .first_text()
211 .map(|value| value.trim().to_string())
212 .filter(|value| !value.is_empty());
213 let tool_calls = response.clone().into_tool_calls();
214
215 if tool_calls.is_empty() {
216 if invocations.is_empty() && has_selected_tool {
217 if shared_inputs.tool_call_mode == ToolCallMode::Strict {
218 return Err(StasisError::PortFailure(
219 "policy violation: strict tool-call mode expected model tool call but none was returned"
220 .to_string(),
221 ));
222 }
223
224 should_use_legacy_fallback = true;
225 fallback_draft_text = maybe_text;
226 break;
227 }
228
229 if let Some(text) = maybe_text {
230 let last = invocations.last().cloned().unwrap_or(ToolInvocation {
231 tool_name: shared_inputs.selected_tool_name().to_string(),
232 tool_input: (*shared_inputs.tool_input).clone(),
233 tool_output: Value::Null,
234 });
235
236 return Ok(ToolLoopExecutionResponse {
237 text,
238 metadata: shared_inputs.context_clone(),
239 tool_name: last.tool_name,
240 tool_output: last.tool_output,
241 tool_invocations: invocations,
242 rounds_executed,
243 termination_reason: "model_completed_no_tool_calls".to_string(),
244 });
245 }
246
247 return Err(StasisError::PortFailure(
248 "chat response was empty after tool loop".to_string(),
249 ));
250 }
251
252 messages.push(ChatMessage::from(tool_calls.clone()));
253 for call in tool_calls {
254 let tool_output = self
255 .tool_registry
256 .invoke_tool(&call.fn_name, call.fn_arguments.clone())
257 .await?;
258
259 let tool_output_text = tool_output.to_string();
260 messages.push(ChatMessage::from(ToolResponse::new(
261 call.call_id,
262 tool_output_text,
263 )));
264 invocations.push(ToolInvocation {
265 tool_name: call.fn_name,
266 tool_input: call.fn_arguments,
267 tool_output,
268 });
269 }
270 }
271
272 if !should_use_legacy_fallback {
273 return Err(StasisError::PortFailure(format!(
274 "tool loop exceeded max rounds ({max_tool_rounds}) without final response"
275 )));
276 }
277 }
278
279 if !should_use_legacy_fallback {
280 return Err(StasisError::PortFailure(
281 "no matching tools available for tool loop execution".to_string(),
282 ));
283 }
284
285 let draft_text = if let Some(text) = fallback_draft_text {
286 text
287 } else {
288 let mut first_request =
289 PromptExecutionRequest::from_user_prompt(shared_inputs.user_prompt.to_string())
290 .with_context(shared_inputs.context_clone());
291 if let Some(system_prompt) = shared_inputs.system_prompt.as_ref() {
292 first_request = first_request.with_system_prompt(system_prompt.to_string());
293 }
294 self.prompt_pipeline.execute(first_request).await?.text
295 };
296 let tool_output = self
297 .tool_registry
298 .invoke_tool(shared_inputs.selected_tool_name(), (*shared_inputs.tool_input).clone())
299 .await?;
300
301 let synthesis_prompt = build_fallback_synthesis_prompt(
302 &shared_inputs.user_prompt,
303 &draft_text,
304 shared_inputs.selected_tool_name(),
305 &tool_output,
306 );
307
308 let mut final_request = PromptExecutionRequest::from_user_prompt(synthesis_prompt)
309 .with_context(shared_inputs.context_clone());
310 if let Some(system_prompt) = shared_inputs.system_prompt.as_ref() {
311 final_request = final_request.with_system_prompt(system_prompt.to_string());
312 }
313
314 let final_response = self.prompt_pipeline.execute(final_request).await?;
315
316 let fallback_invocation = ToolInvocation {
317 tool_name: shared_inputs.selected_tool_name().to_string(),
318 tool_input: (*shared_inputs.tool_input).clone(),
319 tool_output: tool_output.clone(),
320 };
321
322 Ok(ToolLoopExecutionResponse {
323 text: final_response.text,
324 metadata: final_response.metadata,
325 tool_name: shared_inputs.selected_tool_name().to_string(),
326 tool_output,
327 tool_invocations: vec![fallback_invocation],
328 rounds_executed,
329 termination_reason: "legacy_fallback_no_model_tool_call".to_string(),
330 })
331 }
332}
333
334fn build_fallback_synthesis_prompt(
335 user_prompt: &str,
336 draft_text: &str,
337 tool_name: &str,
338 tool_output: &Value,
339) -> String {
340 let tool_output_text = tool_output.to_string();
341 let mut prompt = String::with_capacity(
342 user_prompt.len() + draft_text.len() + tool_name.len() + tool_output_text.len() + 128,
343 );
344 prompt.push_str("User request:\n");
345 prompt.push_str(user_prompt);
346 prompt.push_str("\n\nDraft analysis:\n");
347 prompt.push_str(draft_text);
348 prompt.push_str("\n\nTool '");
349 prompt.push_str(tool_name);
350 prompt.push_str("' output JSON:\n");
351 prompt.push_str(&tool_output_text);
352 prompt.push_str("\n\nProduce final answer grounded in the tool output.");
353 prompt
354}
355
356fn sanitize_tool_name_for_model(name: &str) -> String {
357 let mut out = String::with_capacity(name.len());
358 for ch in name.chars() {
359 if ch.is_ascii_alphanumeric() || ch == '_' || ch == '-' {
360 out.push(ch);
361 } else {
362 out.push('_');
363 }
364 }
365
366 let trimmed = out.trim_matches('_');
367 if trimmed.is_empty() {
368 "tool".to_string()
369 } else {
370 trimmed.to_string()
371 }
372}