1use std::sync::Arc;
4
5use crate::{
6 BridgeCustomization, BridgeMode, BridgePrimitiveContext, BridgePrimitiveKind,
7 BridgePrimitiveRemapper, BridgeReport, BridgeTarget, BridgeWarning, BridgeWarningKind,
8 RequestBridgeContext, ResponseBridgeContext, StreamBridgeContext,
9};
10use siumai_core::streaming::ChatStreamEvent;
11use siumai_core::types::{
12 ChatMessage, ChatRequest, ChatResponse, ChatStreamPart, ChatStreamReplay, ContentPart,
13 MessageContent, ProviderDefinedTool, Tool, ToolChoice,
14};
15
16type ProviderToolArgsMapper = dyn Fn(&RequestBridgeContext, &ProviderDefinedTool, &mut BridgeReport) -> serde_json::Value
17 + Send
18 + Sync;
19
20#[derive(Clone)]
22pub struct ProviderToolRewriteRule {
23 from_tool_id: String,
24 to_tool_id: String,
25 target_name: Option<String>,
26 args_mapper: Option<Arc<ProviderToolArgsMapper>>,
27}
28
29impl ProviderToolRewriteRule {
30 pub fn new(from_tool_id: impl Into<String>, to_tool_id: impl Into<String>) -> Self {
31 Self {
32 from_tool_id: from_tool_id.into(),
33 to_tool_id: to_tool_id.into(),
34 target_name: None,
35 args_mapper: None,
36 }
37 }
38
39 pub fn with_target_name(mut self, target_name: impl Into<String>) -> Self {
40 self.target_name = Some(target_name.into());
41 self
42 }
43
44 pub fn with_args_mapper(mut self, args_mapper: Arc<ProviderToolArgsMapper>) -> Self {
45 self.args_mapper = Some(args_mapper);
46 self
47 }
48
49 fn matches(&self, tool: &ProviderDefinedTool) -> bool {
50 tool.id == self.from_tool_id
51 }
52
53 fn rewrite(
54 &self,
55 ctx: &RequestBridgeContext,
56 tool: &ProviderDefinedTool,
57 report: &mut BridgeReport,
58 ) -> ProviderDefinedTool {
59 let args = self
60 .args_mapper
61 .as_ref()
62 .map(|mapper| mapper(ctx, tool, report))
63 .unwrap_or_else(|| tool.args.clone());
64
65 ProviderDefinedTool::new(
66 self.to_tool_id.clone(),
67 self.target_name
68 .clone()
69 .unwrap_or_else(|| tool.name.clone()),
70 )
71 .with_args(args)
72 }
73}
74
75#[derive(Clone, Default)]
81pub struct ProviderToolRewriteCustomization {
82 rules: Vec<ProviderToolRewriteRule>,
83}
84
85impl ProviderToolRewriteCustomization {
86 pub fn new() -> Self {
87 Self::default()
88 }
89
90 pub fn with_rule(mut self, rule: ProviderToolRewriteRule) -> Self {
91 self.rules.push(rule);
92 self
93 }
94
95 pub fn map_provider_tool_id(
96 mut self,
97 from_tool_id: impl Into<String>,
98 to_tool_id: impl Into<String>,
99 ) -> Self {
100 self.rules
101 .push(ProviderToolRewriteRule::new(from_tool_id, to_tool_id));
102 self
103 }
104}
105
106impl BridgeCustomization for ProviderToolRewriteCustomization {
107 fn transform_request(
108 &self,
109 ctx: &RequestBridgeContext,
110 request: &mut ChatRequest,
111 report: &mut BridgeReport,
112 ) -> Result<(), siumai_core::LlmError> {
113 let Some(tools) = request.tools.as_mut() else {
114 return Ok(());
115 };
116
117 for tool in tools {
118 let Tool::ProviderDefined(provider_tool) = tool else {
119 continue;
120 };
121
122 let Some(rule) = self.rules.iter().find(|rule| rule.matches(provider_tool)) else {
123 continue;
124 };
125
126 let original_id = provider_tool.id.clone();
127 let rewritten = rule.rewrite(ctx, provider_tool, report);
128 let rewritten_id = rewritten.id.clone();
129 *provider_tool = rewritten;
130
131 report.add_warning(BridgeWarning::new(
132 BridgeWarningKind::Custom,
133 format!(
134 "bridge customization rewrote provider-defined tool `{original_id}` to `{rewritten_id}`"
135 ),
136 ));
137 }
138
139 Ok(())
140 }
141}
142
143fn primitive_context(
144 source: Option<BridgeTarget>,
145 target: BridgeTarget,
146 mode: BridgeMode,
147 route_label: Option<&String>,
148 path_label: Option<&String>,
149 primitive: BridgePrimitiveKind,
150) -> BridgePrimitiveContext {
151 BridgePrimitiveContext::new(
152 source,
153 target,
154 mode,
155 route_label.cloned(),
156 path_label.cloned(),
157 primitive,
158 )
159}
160
161fn remap_tool_name(
162 remapper: &dyn BridgePrimitiveRemapper,
163 ctx: BridgePrimitiveContext,
164 value: &mut String,
165) -> bool {
166 let original = value.clone();
167 if let Some(mapped) = remapper.remap_tool_name(&ctx, value) {
168 *value = mapped;
169 }
170 *value != original
171}
172
173fn remap_tool_call_id(
174 remapper: &dyn BridgePrimitiveRemapper,
175 ctx: BridgePrimitiveContext,
176 value: &mut String,
177) -> bool {
178 let original = value.clone();
179 if let Some(mapped) = remapper.remap_tool_call_id(&ctx, value) {
180 *value = mapped;
181 }
182 *value != original
183}
184
185fn remap_message_content(
186 content: &mut MessageContent,
187 source: Option<BridgeTarget>,
188 target: BridgeTarget,
189 mode: BridgeMode,
190 route_label: Option<&String>,
191 path_label: Option<&String>,
192 remapper: &dyn BridgePrimitiveRemapper,
193) {
194 let MessageContent::MultiModal(parts) = content else {
195 return;
196 };
197
198 for part in parts {
199 match part {
200 ContentPart::ToolCall {
201 tool_call_id,
202 tool_name,
203 ..
204 } => {
205 remap_tool_call_id(
206 remapper,
207 primitive_context(
208 source,
209 target,
210 mode,
211 route_label,
212 path_label,
213 BridgePrimitiveKind::ToolCall,
214 ),
215 tool_call_id,
216 );
217 remap_tool_name(
218 remapper,
219 primitive_context(
220 source,
221 target,
222 mode,
223 route_label,
224 path_label,
225 BridgePrimitiveKind::ToolCall,
226 ),
227 tool_name,
228 );
229 }
230 ContentPart::ToolResult {
231 tool_call_id,
232 tool_name,
233 ..
234 } => {
235 remap_tool_call_id(
236 remapper,
237 primitive_context(
238 source,
239 target,
240 mode,
241 route_label,
242 path_label,
243 BridgePrimitiveKind::ToolResult,
244 ),
245 tool_call_id,
246 );
247 remap_tool_name(
248 remapper,
249 primitive_context(
250 source,
251 target,
252 mode,
253 route_label,
254 path_label,
255 BridgePrimitiveKind::ToolResult,
256 ),
257 tool_name,
258 );
259 }
260 ContentPart::ToolApprovalRequest { tool_call_id, .. } => {
261 remap_tool_call_id(
262 remapper,
263 primitive_context(
264 source,
265 target,
266 mode,
267 route_label,
268 path_label,
269 BridgePrimitiveKind::ToolCall,
270 ),
271 tool_call_id,
272 );
273 }
274 _ => {}
275 }
276 }
277}
278
279fn remap_message(
280 message: &mut ChatMessage,
281 source: Option<BridgeTarget>,
282 target: BridgeTarget,
283 mode: BridgeMode,
284 route_label: Option<&String>,
285 path_label: Option<&String>,
286 remapper: &dyn BridgePrimitiveRemapper,
287) {
288 remap_message_content(
289 &mut message.content,
290 source,
291 target,
292 mode,
293 route_label,
294 path_label,
295 remapper,
296 );
297}
298
299fn remap_stream_part(
300 part: &mut ChatStreamPart,
301 source: Option<BridgeTarget>,
302 target: BridgeTarget,
303 mode: BridgeMode,
304 route_label: Option<&String>,
305 path_label: Option<&String>,
306 remapper: &dyn BridgePrimitiveRemapper,
307) -> bool {
308 match part {
309 ChatStreamPart::ToolInputStart { id, tool_name, .. } => {
310 let mut changed = remap_tool_call_id(
311 remapper,
312 primitive_context(
313 source,
314 target,
315 mode,
316 route_label,
317 path_label,
318 BridgePrimitiveKind::ToolCall,
319 ),
320 id,
321 );
322 changed |= remap_tool_name(
323 remapper,
324 primitive_context(
325 source,
326 target,
327 mode,
328 route_label,
329 path_label,
330 BridgePrimitiveKind::ToolCall,
331 ),
332 tool_name,
333 );
334 changed
335 }
336 ChatStreamPart::ToolInputDelta { id, .. } | ChatStreamPart::ToolInputEnd { id, .. } => {
337 remap_tool_call_id(
338 remapper,
339 primitive_context(
340 source,
341 target,
342 mode,
343 route_label,
344 path_label,
345 BridgePrimitiveKind::ToolCall,
346 ),
347 id,
348 )
349 }
350 ChatStreamPart::ToolApprovalRequest(request) => remap_tool_call_id(
351 remapper,
352 primitive_context(
353 source,
354 target,
355 mode,
356 route_label,
357 path_label,
358 BridgePrimitiveKind::ToolCall,
359 ),
360 &mut request.tool_call_id,
361 ),
362 ChatStreamPart::ToolCall(call) => {
363 let mut changed = remap_tool_call_id(
364 remapper,
365 primitive_context(
366 source,
367 target,
368 mode,
369 route_label,
370 path_label,
371 BridgePrimitiveKind::ToolCall,
372 ),
373 &mut call.tool_call_id,
374 );
375 changed |= remap_tool_name(
376 remapper,
377 primitive_context(
378 source,
379 target,
380 mode,
381 route_label,
382 path_label,
383 BridgePrimitiveKind::ToolCall,
384 ),
385 &mut call.tool_name,
386 );
387 changed
388 }
389 ChatStreamPart::ToolResult(result) => {
390 let mut changed = remap_tool_call_id(
391 remapper,
392 primitive_context(
393 source,
394 target,
395 mode,
396 route_label,
397 path_label,
398 BridgePrimitiveKind::ToolResult,
399 ),
400 &mut result.tool_call_id,
401 );
402 changed |= remap_tool_name(
403 remapper,
404 primitive_context(
405 source,
406 target,
407 mode,
408 route_label,
409 path_label,
410 BridgePrimitiveKind::ToolResult,
411 ),
412 &mut result.tool_name,
413 );
414 changed
415 }
416 _ => false,
417 }
418}
419
420fn drop_openai_responses_raw_item(replay: &mut ChatStreamReplay) {
421 if let Some(openai_responses) = replay.openai_responses.as_mut() {
422 openai_responses.raw_item = None;
423 }
424}
425
426pub(crate) fn apply_request_remapper(
427 request: &mut ChatRequest,
428 ctx: &RequestBridgeContext,
429 remapper: &dyn BridgePrimitiveRemapper,
430) {
431 if let Some(tools) = request.tools.as_mut() {
432 for tool in tools {
433 match tool {
434 Tool::Function { function } => {
435 remap_tool_name(
436 remapper,
437 primitive_context(
438 ctx.source,
439 ctx.target,
440 ctx.mode,
441 ctx.route_label.as_ref(),
442 ctx.path_label.as_ref(),
443 BridgePrimitiveKind::ToolDefinition,
444 ),
445 &mut function.name,
446 );
447 }
448 Tool::ProviderDefined(tool) => {
449 remap_tool_name(
450 remapper,
451 primitive_context(
452 ctx.source,
453 ctx.target,
454 ctx.mode,
455 ctx.route_label.as_ref(),
456 ctx.path_label.as_ref(),
457 BridgePrimitiveKind::ToolDefinition,
458 ),
459 &mut tool.name,
460 );
461 }
462 }
463 }
464 }
465
466 if let Some(ToolChoice::Tool { name }) = request.tool_choice.as_mut() {
467 remap_tool_name(
468 remapper,
469 primitive_context(
470 ctx.source,
471 ctx.target,
472 ctx.mode,
473 ctx.route_label.as_ref(),
474 ctx.path_label.as_ref(),
475 BridgePrimitiveKind::ToolChoice,
476 ),
477 name,
478 );
479 }
480
481 for message in &mut request.messages {
482 remap_message(
483 message,
484 ctx.source,
485 ctx.target,
486 ctx.mode,
487 ctx.route_label.as_ref(),
488 ctx.path_label.as_ref(),
489 remapper,
490 );
491 }
492}
493
494pub(crate) fn apply_response_remapper(
495 response: &mut ChatResponse,
496 ctx: &ResponseBridgeContext,
497 remapper: &dyn BridgePrimitiveRemapper,
498) {
499 remap_message_content(
500 &mut response.content,
501 ctx.source,
502 ctx.target,
503 ctx.mode,
504 ctx.route_label.as_ref(),
505 ctx.path_label.as_ref(),
506 remapper,
507 );
508}
509
510pub(crate) fn remap_stream_event(
511 event: ChatStreamEvent,
512 ctx: &StreamBridgeContext,
513 remapper: &dyn BridgePrimitiveRemapper,
514) -> ChatStreamEvent {
515 match event {
516 ChatStreamEvent::Part { mut part } => {
517 remap_stream_part(
518 &mut part,
519 ctx.source,
520 ctx.target,
521 ctx.mode,
522 ctx.route_label.as_ref(),
523 ctx.path_label.as_ref(),
524 remapper,
525 );
526 ChatStreamEvent::Part { part }
527 }
528 ChatStreamEvent::PartWithReplay {
529 mut part,
530 mut replay,
531 } => {
532 let changed = remap_stream_part(
533 &mut part,
534 ctx.source,
535 ctx.target,
536 ctx.mode,
537 ctx.route_label.as_ref(),
538 ctx.path_label.as_ref(),
539 remapper,
540 );
541 if changed {
542 drop_openai_responses_raw_item(&mut replay);
543 }
544 ChatStreamEvent::PartWithReplay { part, replay }
545 }
546 ChatStreamEvent::StreamEnd { mut response } => {
547 let response_ctx = ResponseBridgeContext::new(
548 ctx.source,
549 ctx.target,
550 ctx.mode,
551 ctx.route_label.clone(),
552 ctx.path_label.clone(),
553 );
554 apply_response_remapper(&mut response, &response_ctx, remapper);
555 ChatStreamEvent::StreamEnd { response }
556 }
557 other => other,
558 }
559}
560
561#[cfg(test)]
562mod tests {
563 #[test]
564 fn bridge_customization_source_stays_primitive_only() {
565 let source = include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/src/customize.rs"));
566 let production_source = source
567 .split("#[cfg(test)]")
568 .next()
569 .expect("production source section");
570
571 for forbidden in [
572 "provider_options",
573 "providerOptions",
574 "provider_metadata",
575 "providerMetadata",
576 ] {
577 assert!(
578 !production_source.contains(forbidden),
579 "bridge customization/remapper helpers should stay primitive-only and must not read or write `{forbidden}`"
580 );
581 }
582 }
583}