rig_core/providers/openai/wire/
route.rs1use serde::{Deserialize, Serialize};
11
12use crate::catalog::ModelFacts;
13use crate::completion::CompletionRequest;
14use crate::error::{EncodeError, ProviderError};
15use crate::operation::Completion;
16use crate::providers::openai::responses_api::streaming::{ResponsesDecoder, ResponsesEvent};
17use crate::providers::openai::responses_api::wire::Responses;
18use crate::providers::openai::responses_api::{
19 ResponsesToolDefinition, SystemInstructionsPlacement,
20};
21use crate::wire::document::Reassemble;
22use crate::wire::{Decoder, Descriptor, Encoded, Flow, Mode, Out, Wire, WireEvent, WireFrame};
23
24use super::OpenAIConfig;
25use super::chat::{Chat, ChatDecoder, ChatEvent};
26
27macro_rules! on_route {
29 ($chosen:expr, $wire:ident => $ask:expr) => {
30 match $chosen {
31 Self::Chat($wire) => $ask,
32 Self::Responses($wire) => $ask,
33 }
34 };
35}
36
37#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
39pub enum Route {
40 Chat,
42 Responses,
44}
45
46#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
49pub enum OpenAiWire {
50 Chat(Chat),
52 Responses(Responses),
54}
55
56impl OpenAiWire {
57 pub fn new(provider: OpenAIConfig, model: impl Into<String>) -> Self {
60 let model = model.into();
61 let route = provider.route.unwrap_or_else(|| {
62 provider
63 .dialect
64 .quirks
65 .hooks
66 .and_then(|hooks| hooks.model_route)
67 .map_or_else(|| provider.completion_route(), |route| route(&model))
68 });
69 match route {
70 Route::Chat => Self::Chat(Chat::new(provider, model)),
71 Route::Responses => Self::Responses(Responses::new(provider, model)),
72 }
73 }
74
75 pub fn provider(&self) -> &OpenAIConfig {
77 on_route!(self, wire => &wire.provider)
78 }
79
80 pub fn with_strict_tools(self) -> Self {
83 self.on_chat(Chat::with_strict_tools)
84 .on_responses(Responses::with_strict_tools)
85 }
86
87 pub fn with_tool_result_array_content(self) -> Self {
91 self.on_chat(Chat::with_tool_result_array_content)
92 }
93
94 pub fn with_tool(self, tool: impl Into<ResponsesToolDefinition>) -> Self {
97 self.on_responses(|wire| wire.with_tool(tool))
98 }
99
100 pub fn with_tools<I, Tool>(self, tools: I) -> Self
103 where
104 I: IntoIterator<Item = Tool>,
105 Tool: Into<ResponsesToolDefinition>,
106 {
107 self.on_responses(|wire| wire.with_tools(tools))
108 }
109
110 pub fn with_system_instructions_placement(
114 self,
115 placement: SystemInstructionsPlacement,
116 ) -> Self {
117 self.on_responses(|wire| wire.with_system_instructions_placement(placement))
118 }
119
120 pub fn with_system_instructions_as_messages(self) -> Self {
123 self.on_responses(Responses::with_system_instructions_as_messages)
124 }
125
126 pub fn with_facts(self, facts: ModelFacts) -> Self {
129 match self {
130 Self::Chat(wire) => Self::Chat(wire.with_facts(facts)),
131 Self::Responses(wire) => Self::Responses(wire.with_facts(facts)),
132 }
133 }
134
135 fn on_chat(self, option: impl FnOnce(Chat) -> Chat) -> Self {
137 match self {
138 Self::Chat(wire) => Self::Chat(option(wire)),
139 responses => responses,
140 }
141 }
142
143 fn on_responses(self, option: impl FnOnce(Responses) -> Responses) -> Self {
145 match self {
146 Self::Responses(wire) => Self::Responses(option(wire)),
147 chat => chat,
148 }
149 }
150}
151
152impl From<Chat> for OpenAiWire {
153 fn from(wire: Chat) -> Self {
154 Self::Chat(wire)
155 }
156}
157
158impl From<Responses> for OpenAiWire {
159 fn from(wire: Responses) -> Self {
160 Self::Responses(wire)
161 }
162}
163
164impl Wire for OpenAiWire {
165 type Op = Completion;
166 type Payload = crate::wire::Encoded;
167 type Frame = crate::wire::WireFrame;
168 type Event = OpenAiEvent;
169 type Decoder<'id> = OpenAiDecoder;
170 type Reassembler = OpenAiReassembler;
171
172 fn describe(&self) -> Descriptor<'_> {
173 on_route!(self, wire => wire.describe())
174 }
175
176 fn encode(&self, request: CompletionRequest, mode: Mode) -> Result<Encoded, EncodeError> {
177 on_route!(self, wire => wire.encode(request, mode))
178 }
179
180 fn decoder<'id>(&self) -> Self::Decoder<'id> {
181 match self {
182 Self::Chat(wire) => OpenAiDecoder::Chat(wire.decoder()),
183 Self::Responses(wire) => OpenAiDecoder::Responses(wire.decoder()),
184 }
185 }
186
187 fn reassembler(&self) -> Self::Reassembler {
188 match self {
189 Self::Chat(wire) => OpenAiReassembler::Chat(wire.reassembler()),
190 Self::Responses(wire) => OpenAiReassembler::Responses(wire.reassembler()),
191 }
192 }
193}
194
195pub enum OpenAiReassembler {
197 Chat(<Chat as Wire>::Reassembler),
199 Responses(<Responses as Wire>::Reassembler),
201}
202
203impl Default for OpenAiReassembler {
205 fn default() -> Self {
206 Self::Chat(Default::default())
207 }
208}
209
210impl crate::wire::document::Serves<crate::operation::Completion> for OpenAiReassembler {}
211
212impl Reassemble<OpenAiEvent> for OpenAiReassembler {
213 fn absorb(&mut self, event: &OpenAiEvent) {
214 match (self, event) {
215 (Self::Chat(document), OpenAiEvent::Chat(event)) => document.absorb(event),
216 (Self::Responses(document), OpenAiEvent::Responses(event)) => document.absorb(event),
217 _ => {}
219 }
220 }
221
222 fn finish(self) -> serde_json::Value {
223 on_route!(self, document => document.finish())
224 }
225}
226
227pub enum OpenAiEvent {
229 Chat(ChatEvent),
231 Responses(ResponsesEvent),
233}
234
235pub enum OpenAiDecoder {
237 Chat(ChatDecoder),
239 Responses(ResponsesDecoder),
241}
242
243impl<'id> Decoder<'id, Completion> for OpenAiDecoder {
244 type Event = OpenAiEvent;
245
246 fn classify(&self, frame: WireFrame) -> WireEvent<OpenAiEvent> {
247 match self {
248 Self::Chat(decoder) => decoder.classify(frame).map(OpenAiEvent::Chat),
249 Self::Responses(decoder) => decoder.classify(frame).map(OpenAiEvent::Responses),
250 }
251 }
252
253 fn decode(
254 &mut self,
255 event: OpenAiEvent,
256 out: Out<'id, Completion>,
257 ) -> Result<Flow, ProviderError> {
258 match (self, event) {
259 (Self::Chat(decoder), OpenAiEvent::Chat(event)) => decoder.decode(event, out),
260 (Self::Responses(decoder), OpenAiEvent::Responses(event)) => decoder.decode(event, out),
261 (Self::Chat(_), OpenAiEvent::Responses(_))
264 | (Self::Responses(_), OpenAiEvent::Chat(_)) => Ok(Flow::More),
265 }
266 }
267
268 fn eof(&mut self, out: Out<'id, Completion>) -> Result<Flow, ProviderError> {
269 on_route!(self, decoder => decoder.eof(out))
270 }
271}
272
273#[cfg(test)]
274mod tests;