tenant.openagents/omega

1use std::sync::Arc;
2
3use super::{ChatLocation, copilot_request_headers};
4use anyhow::{Result, anyhow};
5use futures::{AsyncBufReadExt, AsyncReadExt, StreamExt, io::BufReader, stream::BoxStream};
6use http_client::{AsyncBody, HttpClient, HttpRequestExt, Method, Request as HttpRequest};
7use serde::{Deserialize, Serialize};
8use serde_json::Value;
9pub use settings::OpenAiReasoningEffort as ReasoningEffort;
10
11#[derive(Serialize, Debug)]
12pub struct Request {
13    pub model: String,
14    pub input: Vec<ResponseInputItem>,
15    #[serde(default)]
16    pub stream: bool,
17    #[serde(skip_serializing_if = "Option::is_none")]
18    pub temperature: Option<f32>,
19    #[serde(skip_serializing_if = "Vec::is_empty")]
20    pub tools: Vec<ToolDefinition>,
21    #[serde(skip_serializing_if = "Option::is_none")]
22    pub tool_choice: Option<ToolChoice>,
23    #[serde(skip_serializing_if = "Option::is_none")]
24    pub reasoning: Option<ReasoningConfig>,
25    #[serde(skip_serializing_if = "Option::is_none")]
26    pub include: Option<Vec<ResponseIncludable>>,
27    pub store: bool,
28}
29
30#[derive(Serialize, Deserialize, Debug, Clone)]
31#[serde(rename_all = "snake_case")]
32pub enum ResponseIncludable {
33    #[serde(rename = "reasoning.encrypted_content")]
34    ReasoningEncryptedContent,
35}
36
37#[derive(Serialize, Deserialize, Debug)]
38#[serde(tag = "type", rename_all = "snake_case")]
39pub enum ToolDefinition {
40    Function {
41        name: String,
42        #[serde(skip_serializing_if = "Option::is_none")]
43        description: Option<String>,
44        #[serde(skip_serializing_if = "Option::is_none")]
45        parameters: Option<Value>,
46        #[serde(skip_serializing_if = "Option::is_none")]
47        strict: Option<bool>,
48    },
49}
50
51#[derive(Serialize, Deserialize, Debug)]
52#[serde(rename_all = "lowercase")]
53pub enum ToolChoice {
54    Auto,
55    Required,
56    None,
57    #[serde(untagged)]
58    Other(ToolDefinition),
59}
60
61#[derive(Serialize, Deserialize, Debug)]
62#[serde(rename_all = "lowercase")]
63pub enum ReasoningSummary {
64    Auto,
65    Concise,
66    Detailed,
67}
68
69#[derive(Serialize, Debug)]
70pub struct ReasoningConfig {
71    pub effort: ReasoningEffort,
72    #[serde(skip_serializing_if = "Option::is_none")]
73    pub summary: Option<ReasoningSummary>,
74}
75
76#[derive(Serialize, Deserialize, Debug, Clone, Default)]
77#[serde(rename_all = "snake_case")]
78pub enum ResponseImageDetail {
79    Low,
80    High,
81    #[default]
82    Auto,
83}
84
85#[derive(Serialize, Deserialize, Debug, Clone)]
86#[serde(tag = "type", rename_all = "snake_case")]
87pub enum ResponseInputContent {
88    InputText {
89        text: String,
90    },
91    OutputText {
92        text: String,
93    },
94    InputImage {
95        #[serde(skip_serializing_if = "Option::is_none")]
96        image_url: Option<String>,
97        #[serde(default)]
98        detail: ResponseImageDetail,
99    },
100}
101
102#[derive(Serialize, Deserialize, Debug, Clone)]
103#[serde(rename_all = "snake_case")]
104pub enum ItemStatus {
105    InProgress,
106    Completed,
107    Incomplete,
108}
109
110#[derive(Serialize, Deserialize, Debug, Clone)]
111#[serde(untagged)]
112pub enum ResponseFunctionOutput {
113    Text(String),
114    Content(Vec<ResponseInputContent>),
115}
116
117#[derive(Serialize, Deserialize, Debug, Clone)]
118#[serde(tag = "type", rename_all = "snake_case")]
119pub enum ResponseInputItem {
120    Message {
121        role: String,
122        #[serde(skip_serializing_if = "Option::is_none")]
123        content: Option<Vec<ResponseInputContent>>,
124        #[serde(skip_serializing_if = "Option::is_none")]
125        status: Option<String>,
126    },
127    FunctionCall {
128        call_id: String,
129        name: String,
130        arguments: String,
131        #[serde(skip_serializing_if = "Option::is_none")]
132        status: Option<ItemStatus>,
133        #[serde(default, skip_serializing_if = "Option::is_none")]
134        thought_signature: Option<String>,
135    },
136    FunctionCallOutput {
137        call_id: String,
138        output: ResponseFunctionOutput,
139        #[serde(skip_serializing_if = "Option::is_none")]
140        status: Option<ItemStatus>,
141    },
142    Reasoning(ResponseReasoningInputItem),
143}
144
145#[derive(Deserialize, Debug, Clone)]
146#[serde(rename_all = "snake_case")]
147pub enum IncompleteReason {
148    #[serde(rename = "max_output_tokens")]
149    MaxOutputTokens,
150    #[serde(rename = "content_filter")]
151    ContentFilter,
152}
153
154#[derive(Deserialize, Debug, Clone)]
155pub struct IncompleteDetails {
156    #[serde(skip_serializing_if = "Option::is_none")]
157    pub reason: Option<IncompleteReason>,
158}
159
160#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)]
161pub struct ResponseReasoningInputItem {
162    #[serde(default, skip_serializing_if = "Option::is_none")]
163    pub id: Option<String>,
164    #[serde(default)]
165    pub summary: Vec<ResponseReasoningItem>,
166    #[serde(default, skip_serializing_if = "Option::is_none")]
167    pub encrypted_content: Option<String>,
168}
169
170#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)]
171pub struct ResponseReasoningItem {
172    #[serde(rename = "type")]
173    pub kind: String,
174    pub text: String,
175}
176
177#[derive(Deserialize, Debug)]
178#[serde(tag = "type")]
179pub enum StreamEvent {
180    #[serde(rename = "error")]
181    GenericError { error: ResponseError },
182
183    #[serde(rename = "response.created")]
184    Created { response: Response },
185
186    #[serde(rename = "response.output_item.added")]
187    OutputItemAdded {
188        output_index: usize,
189        #[serde(default)]
190        sequence_number: Option<u64>,
191        item: ResponseOutputItem,
192    },
193
194    #[serde(rename = "response.output_text.delta")]
195    OutputTextDelta {
196        item_id: String,
197        output_index: usize,
198        delta: String,
199    },
200
201    #[serde(rename = "response.output_item.done")]
202    OutputItemDone {
203        output_index: usize,
204        #[serde(default)]
205        sequence_number: Option<u64>,
206        item: ResponseOutputItem,
207    },
208
209    #[serde(rename = "response.incomplete")]
210    Incomplete { response: Response },
211
212    #[serde(rename = "response.completed")]
213    Completed { response: Response },
214
215    #[serde(rename = "response.failed")]
216    Failed { response: Response },
217
218    #[serde(other)]
219    Unknown,
220}
221
222#[derive(Deserialize, Debug, Clone)]
223pub struct ResponseError {
224    pub code: String,
225    pub message: String,
226}
227
228#[derive(Deserialize, Debug, Default, Clone)]
229pub struct Response {
230    pub id: Option<String>,
231    pub status: Option<String>,
232    pub usage: Option<ResponseUsage>,
233    pub output: Vec<ResponseOutputItem>,
234    #[serde(skip_serializing_if = "Option::is_none")]
235    pub incomplete_details: Option<IncompleteDetails>,
236    #[serde(skip_serializing_if = "Option::is_none")]
237    pub error: Option<ResponseError>,
238}
239
240#[derive(Deserialize, Debug, Default, Clone)]
241pub struct ResponseUsage {
242    pub input_tokens: Option<u64>,
243    pub output_tokens: Option<u64>,
244    pub total_tokens: Option<u64>,
245}
246
247#[derive(Deserialize, Debug, Clone)]
248#[serde(tag = "type", rename_all = "snake_case")]
249pub enum ResponseOutputItem {
250    Message {
251        id: String,
252        role: String,
253        #[serde(skip_serializing_if = "Option::is_none")]
254        content: Option<Vec<ResponseOutputContent>>,
255    },
256    FunctionCall {
257        #[serde(skip_serializing_if = "Option::is_none")]
258        id: Option<String>,
259        call_id: String,
260        name: String,
261        arguments: String,
262        #[serde(skip_serializing_if = "Option::is_none")]
263        status: Option<ItemStatus>,
264        #[serde(default, skip_serializing_if = "Option::is_none")]
265        thought_signature: Option<String>,
266    },
267    Reasoning {
268        id: String,
269        #[serde(skip_serializing_if = "Option::is_none")]
270        summary: Option<Vec<ResponseReasoningItem>>,
271        #[serde(skip_serializing_if = "Option::is_none")]
272        encrypted_content: Option<String>,
273    },
274}
275
276#[derive(Deserialize, Debug, Clone)]
277#[serde(tag = "type", rename_all = "snake_case")]
278pub enum ResponseOutputContent {
279    OutputText { text: String },
280    Refusal { refusal: String },
281}
282
283pub async fn stream_response(
284    client: Arc<dyn HttpClient>,
285    oauth_token: String,
286    api_url: String,
287    request: Request,
288    is_user_initiated: bool,
289    location: ChatLocation,
290) -> Result<BoxStream<'static, Result<StreamEvent>>> {
291    let is_vision_request = request.input.iter().any(|item| match item {
292        ResponseInputItem::Message {
293            content: Some(parts),
294            ..
295        } => parts
296            .iter()
297            .any(|p| matches!(p, ResponseInputContent::InputImage { .. })),
298        _ => false,
299    });
300
301    let request_builder = copilot_request_headers(
302        HttpRequest::builder().method(Method::POST).uri(&api_url),
303        &oauth_token,
304        Some(is_user_initiated),
305        Some(location),
306    )
307    .when(is_vision_request, |builder| {
308        builder.header("Copilot-Vision-Request", "true")
309    });
310
311    let is_streaming = request.stream;
312    let json = serde_json::to_string(&request)?;
313    let request = request_builder.body(AsyncBody::from(json))?;
314    let mut response = client.send(request).await?;
315
316    if !response.status().is_success() {
317        let mut body = String::new();
318        response.body_mut().read_to_string(&mut body).await?;
319        anyhow::bail!("Failed to connect to API: {} {}", response.status(), body);
320    }
321
322    if is_streaming {
323        let reader = BufReader::new(response.into_body());
324        Ok(reader
325            .lines()
326            .filter_map(|line| async move {
327                match line {
328                    Ok(line) => {
329                        let line = line.strip_prefix("data: ")?;
330                        if line.starts_with("[DONE]") || line.is_empty() {
331                            return None;
332                        }
333
334                        match serde_json::from_str::<StreamEvent>(line) {
335                            Ok(event) => Some(Ok(event)),
336                            Err(error) => {
337                                log::error!(
338                                    "Failed to parse Copilot responses stream event: `{}`\nResponse: `{}`",
339                                    error,
340                                    line,
341                                );
342                                Some(Err(anyhow!(error)))
343                            }
344                        }
345                    }
346                    Err(error) => Some(Err(anyhow!(error))),
347                }
348            })
349            .boxed())
350    } else {
351        // Simulate streaming this makes the mapping of this function return more straight-forward to handle if all callers assume it streams.
352        // Removes the need of having a method to map StreamEvent and another to map Response to a LanguageCompletionEvent
353        let mut body = String::new();
354        response.body_mut().read_to_string(&mut body).await?;
355
356        match serde_json::from_str::<Response>(&body) {
357            Ok(response) => {
358                let events = vec![StreamEvent::Created {
359                    response: response.clone(),
360                }];
361
362                let mut all_events = events;
363                for (output_index, item) in response.output.iter().enumerate() {
364                    all_events.push(StreamEvent::OutputItemAdded {
365                        output_index,
366                        sequence_number: None,
367                        item: item.clone(),
368                    });
369
370                    if let ResponseOutputItem::Message {
371                        id,
372                        content: Some(content),
373                        ..
374                    } = item
375                    {
376                        for part in content {
377                            if let ResponseOutputContent::OutputText { text } = part {
378                                all_events.push(StreamEvent::OutputTextDelta {
379                                    item_id: id.clone(),
380                                    output_index,
381                                    delta: text.clone(),
382                                });
383                            }
384                        }
385                    }
386
387                    all_events.push(StreamEvent::OutputItemDone {
388                        output_index,
389                        sequence_number: None,
390                        item: item.clone(),
391                    });
392                }
393
394                let final_event = if response.error.is_some() {
395                    StreamEvent::Failed { response }
396                } else if response.incomplete_details.is_some() {
397                    StreamEvent::Incomplete { response }
398                } else {
399                    StreamEvent::Completed { response }
400                };
401                all_events.push(final_event);
402
403                Ok(futures::stream::iter(all_events.into_iter().map(Ok)).boxed())
404            }
405            Err(error) => {
406                log::error!(
407                    "Failed to parse Copilot non-streaming response: `{}`\nResponse: `{}`",
408                    error,
409                    body,
410                );
411                Err(anyhow!(error))
412            }
413        }
414    }
415}
416
417#[cfg(test)]
418mod tests {
419    use super::*;
420
421    #[test]
422    fn test_tool_choice_required_serializes_as_required() {
423        // Regression test: ToolChoice::Required must serialize as "required" (not "any")
424        // for OpenAI Responses API. Reverting the rename would break this.
425        assert_eq!(
426            serde_json::to_string(&ToolChoice::Required).unwrap(),
427            "\"required\""
428        );
429        assert_eq!(
430            serde_json::to_string(&ToolChoice::Auto).unwrap(),
431            "\"auto\""
432        );
433        assert_eq!(
434            serde_json::to_string(&ToolChoice::None).unwrap(),
435            "\"none\""
436        );
437    }
438}
439