Skip to main content

rig_core/providers/internal/
mod.rs

1//! Shared provider infrastructure: the wire-adapter contract, its
2//! single-policy-site driver, and the decode-then-validate classify layer.
3//!
4//! [`adapter`], [`wire`], [`tool_call_bridge`], and [`chunk_lifecycle`] are
5//! public so out-of-tree providers implement [`adapter::WireAdapter`] and
6//! inherit the shared driver, frame-triage policy, index→identity tool-call
7//! bridging, and the boundary-less reasoning lifecycle derivation instead of
8//! hand-rolling per-provider assemblers; the remaining helpers are
9//! crate-private.
10
11pub mod adapter;
12pub(crate) mod anthropic_compatible;
13#[cfg(feature = "audio")]
14pub(crate) mod audio_generation;
15pub(crate) mod auth;
16pub mod chunk_lifecycle;
17pub(crate) mod completion_send;
18#[cfg(not(target_family = "wasm"))]
19pub(crate) mod device_auth;
20pub(crate) mod envelope;
21#[cfg(feature = "image")]
22pub(crate) mod image_generation;
23pub(crate) mod model_listing;
24pub(crate) mod openai_chat_completions_compatible;
25pub(crate) mod schema;
26#[cfg(any(test, debug_assertions))]
27pub(crate) mod sequence_law;
28pub(crate) mod sse_transport;
29pub mod tool_call_bridge;
30pub(crate) mod transcription;
31pub mod wire;
32
33/// Fill empty [`ToolResult::name`](crate::message::ToolResult::name)s from
34/// the calls they answer, for wires that key the replay on the tool name
35/// (Gemini `functionResponse.name`, Ollama tool messages, Vertex AI,
36/// gemini-grpc, Interactions).
37///
38/// `ToolResult::name` is required data, but rig's own inbound converters
39/// cannot supply it: Anthropic, OpenAI-chat, Cohere, and Bedrock tool
40/// messages carry no name on their wires, so a cross-provider ingested
41/// transcript arrives with `name: ""`. The name lives on the paired
42/// assistant call in the same history — match by rig's correlation handle
43/// first, then by provider identifiers. A result matching no call keeps
44/// its empty name: the transcript genuinely lacks the data, and the wire's
45/// own rejection is the honest failure.
46///
47/// `pub` (not `pub(crate)`) because sibling serializer crates that speak a
48/// name-keyed tool-result wire (rig-vertexai, rig-gemini-grpc) carry the
49/// same contract; it is not part of rig-core's stable public API.
50pub fn resolve_empty_tool_result_names(history: &mut [crate::message::Message]) {
51    use std::collections::HashMap;
52
53    let mut names_by_id: HashMap<String, String> = HashMap::new();
54    for message in history.iter() {
55        let crate::message::Message::Assistant { content, .. } = message else {
56            continue;
57        };
58        for item in content.iter() {
59            let crate::message::AssistantContent::ToolCall(call) = item else {
60                continue;
61            };
62            names_by_id.insert(call.id.as_str().to_owned(), call.function.name.clone());
63            if let Some(provider) = &call.provider {
64                names_by_id.insert(provider.call_id.clone(), call.function.name.clone());
65                if let Some(item_id) = &provider.item_id {
66                    names_by_id.insert(item_id.clone(), call.function.name.clone());
67                }
68            }
69        }
70    }
71    if names_by_id.is_empty() {
72        return;
73    }
74
75    for message in history.iter_mut() {
76        let crate::message::Message::User { content } = message else {
77            continue;
78        };
79        for item in content.iter_mut() {
80            let crate::message::UserContent::ToolResult(result) = item else {
81                continue;
82            };
83            if !result.name.is_empty() {
84                continue;
85            }
86            let resolved = names_by_id.get(result.call.as_str()).or_else(|| {
87                result.provider.as_ref().and_then(|provider| {
88                    names_by_id.get(&provider.call_id).or_else(|| {
89                        provider
90                            .item_id
91                            .as_ref()
92                            .and_then(|item_id| names_by_id.get(item_id))
93                    })
94                })
95            });
96            if let Some(name) = resolved {
97                result.name = name.clone();
98            }
99        }
100    }
101}
102
103/// A rig logging target for [`trace_json`]. An enum (not a `&str`) because
104/// `tracing` targets must be literals, so the dispatch is total by
105/// construction.
106#[derive(Clone, Copy)]
107pub(crate) enum LogTarget {
108    Completions,
109    Streaming,
110}
111
112/// Trace-log `value` as pretty-printed JSON under one of rig's logging
113/// targets. Infallible: does nothing when TRACE is disabled for the target or
114/// the value fails to serialize.
115pub(crate) fn trace_json(target: LogTarget, label: &str, value: &impl serde::Serialize) {
116    macro_rules! emit {
117        ($target:literal) => {
118            if tracing::enabled!(target: $target, tracing::Level::TRACE) {
119                if let Ok(json) = serde_json::to_string_pretty(value) {
120                    tracing::trace!(target: $target, "{label}: {json}");
121                }
122            }
123        };
124    }
125    match target {
126        LogTarget::Streaming => emit!("rig::streaming"),
127        LogTarget::Completions => emit!("rig::completions"),
128    }
129}
130
131pub(crate) fn completion_usage(
132    input_tokens: u64,
133    output_tokens: u64,
134    total_tokens: u64,
135    cached_input_tokens: u64,
136) -> crate::completion::Usage {
137    crate::completion::Usage {
138        input_tokens,
139        output_tokens,
140        total_tokens,
141        cached_input_tokens,
142        cache_creation_input_tokens: 0,
143        tool_use_prompt_tokens: 0,
144        reasoning_tokens: 0,
145    }
146}
147
148#[cfg(test)]
149mod tests {
150    use crate::message::{
151        AssistantContent, Message, ToolCall, ToolFunction, ToolResultContent, UserContent,
152    };
153
154    fn call(wire_id: &str, name: &str) -> Message {
155        Message::Assistant {
156            id: None,
157            content: vec![AssistantContent::ToolCall(ToolCall::from_wire(
158                wire_id,
159                ToolFunction {
160                    name: name.to_owned(),
161                    arguments: serde_json::json!({}),
162                },
163            ))],
164        }
165    }
166
167    fn nameless_result(wire_id: &str) -> Message {
168        Message::User {
169            content: vec![UserContent::tool_result_from_wire(
170                wire_id,
171                "",
172                vec![ToolResultContent::text("out")],
173            )],
174        }
175    }
176
177    fn result_names(history: &[Message]) -> Vec<String> {
178        history
179            .iter()
180            .filter_map(|message| match message {
181                Message::User { content } => content.iter().next().and_then(|item| match item {
182                    UserContent::ToolResult(result) => Some(result.name.clone()),
183                    _ => None,
184                }),
185                _ => None,
186            })
187            .collect()
188    }
189
190    /// An ingested cross-provider transcript (converters stamp `name: ""`)
191    /// resolves each result's name from its paired call — by provider id
192    /// here, since `from_wire` on both sides shares it.
193    #[test]
194    fn empty_names_resolve_from_the_paired_call() {
195        let mut history = vec![
196            call("toolu_1", "get_weather"),
197            nameless_result("toolu_1"),
198            call("toolu_2", "get_time"),
199            nameless_result("toolu_2"),
200        ];
201        super::resolve_empty_tool_result_names(&mut history);
202        assert_eq!(result_names(&history), ["get_weather", "get_time"]);
203    }
204
205    /// A result no call in the history answers keeps its empty name: the
206    /// transcript genuinely lacks the data, and inventing one would ship a
207    /// fabricated name to a name-keyed wire.
208    #[test]
209    fn an_unmatched_result_keeps_its_empty_name() {
210        let mut history = vec![call("toolu_1", "get_weather"), nameless_result("toolu_9")];
211        super::resolve_empty_tool_result_names(&mut history);
212        assert_eq!(result_names(&history), [""]);
213    }
214
215    /// An established name is data, never overwritten — a repair hook may
216    /// have renamed the executed tool relative to the model's call.
217    #[test]
218    fn an_established_name_is_never_overwritten() {
219        let mut history = vec![
220            call("toolu_1", "add"),
221            Message::User {
222                content: vec![UserContent::tool_result_from_wire(
223                    "toolu_1",
224                    "sum",
225                    vec![ToolResultContent::text("3")],
226                )],
227            },
228        ];
229        super::resolve_empty_tool_result_names(&mut history);
230        assert_eq!(result_names(&history), ["sum"]);
231    }
232
233    /// Matching falls through the identifier tiers: rig's correlation
234    /// handle first (a driver-built result answering an id-less call),
235    /// then the provider identifiers.
236    #[test]
237    fn a_handle_only_result_resolves_from_an_id_less_call() {
238        let id_less = ToolCall::new(
239            crate::message::ToolCallId::mint(),
240            ToolFunction {
241                name: "lookup".to_owned(),
242                arguments: serde_json::json!({}),
243            },
244        );
245        let handle = id_less.id.as_str().to_owned();
246        let mut history = vec![
247            Message::Assistant {
248                id: None,
249                content: vec![AssistantContent::ToolCall(id_less)],
250            },
251            Message::User {
252                content: vec![UserContent::tool_result(
253                    handle,
254                    "",
255                    vec![ToolResultContent::text("out")],
256                )],
257            },
258        ];
259        super::resolve_empty_tool_result_names(&mut history);
260        assert_eq!(result_names(&history), ["lookup"]);
261    }
262}