agent_framework_openai/
lib.rs

1//! # agent-framework-openai
2//!
3//! An OpenAI (and OpenAI-compatible) [`ChatClient`] for `agent-framework-rs`.
4//!
5//! Works against the OpenAI Chat Completions API and any compatible endpoint
6//! (Azure OpenAI, Ollama, together.ai, local servers, …) by overriding the
7//! base URL.
8//!
9//! ```no_run
10//! use agent_framework_openai::OpenAIChatCompletionClient;
11//! use agent_framework_core::prelude::*;
12//!
13//! # async fn demo() -> Result<()> {
14//! let client = OpenAIChatCompletionClient::new("sk-...", "gpt-4o-mini");
15//! let agent = Agent::builder(client)
16//!     .instructions("You are concise.")
17//!     .build();
18//! let reply = agent.run_once("Say hi").await?;
19//! println!("{}", reply.text());
20//! # Ok(())
21//! # }
22//! ```
23
24/// Conversion between framework types and the OpenAI chat-completions wire
25/// format.
26///
27/// This module is public (but hidden from docs) so that other
28/// OpenAI-wire-compatible clients in this workspace — currently
29/// `agent-framework-azure` — can reuse request/response conversion instead of
30/// duplicating it. It is not intended as a stable external API.
31#[doc(hidden)]
32pub mod convert;
33pub mod embeddings;
34
35pub mod responses;
36pub use embeddings::OpenAIEmbeddingClient;
37pub use responses::OpenAIChatClient;
38
39use std::collections::{HashMap, VecDeque};
40use std::sync::Arc;
41
42use agent_framework_core::client::{ChatClient, ChatStream};
43use agent_framework_core::error::{Error, Result};
44use agent_framework_core::streaming::Utf8StreamDecoder;
45use agent_framework_core::types::{
46    ChatOptions, ChatResponse, ChatResponseUpdate, Content, FinishReason, FunctionArguments,
47    FunctionCallContent, Message, Role, TextContent, UsageContent,
48};
49use futures::StreamExt;
50use serde_json::{json, Map, Value};
51
52pub(crate) const DEFAULT_BASE_URL: &str = "https://api.openai.com/v1";
53
54/// Parse a `Retry-After` header into a delay in seconds.
55///
56/// HTTP allows either a delay in seconds or an HTTP-date; OpenAI (and other
57/// rate limiters) use the integer/decimal-seconds form for `429`/`503`, which
58/// is what we honor. A date-form or unparseable value is treated as absent.
59/// Shared with [`responses`](crate::responses) so both OpenAI clients attach
60/// the same retry hint to [`Error::ServiceStatus`].
61pub(crate) fn parse_retry_after(headers: &reqwest::header::HeaderMap) -> Option<f64> {
62    headers
63        .get(reqwest::header::RETRY_AFTER)
64        .and_then(|v| v.to_str().ok())
65        .and_then(|s| s.trim().parse::<f64>().ok())
66        .filter(|s| s.is_finite() && *s >= 0.0)
67}
68
69/// Classify a non-success OpenAI-wire (Chat Completions / Responses) HTTP
70/// response into a granular [`Error`].
71///
72/// The single point of truth for status/body interpretation, used by every
73/// endpoint in this crate (`OpenAIChatCompletionClient::post` and [`responses`]) and reused
74/// by `agent-framework-azure` (Azure OpenAI is
75/// wire-compatible for Chat Completions and Responses), so the two stay
76/// identical rather than drifting.
77///
78/// Mirrors upstream's `openai/_chat_client.py` / `_responses_client.py`:
79///
80/// ```text
81/// except BadRequestError as ex:
82///     if ex.code == "content_filter":
83///         raise OpenAIContentFilterException(...)
84///     raise ServiceResponseException(...)
85/// ```
86///
87/// for the content-filter case, and extends upstream's exception
88/// *hierarchy* — `ServiceInvalidAuthError` / `ServiceInvalidRequestError`
89/// already exist in `agent_framework.exceptions`, even though today's Python
90/// OpenAI client folds every other status (auth failures included) into that
91/// same generic `ServiceResponseException` — to also classify by HTTP status:
92///
93/// * `401` / `403` -> [`Error::ServiceInvalidAuth`]
94/// * `400` / `404` / `422` -> [`Error::ServiceInvalidRequest`], unless the
95///   body signals a content-filter refusal (`error.code` or `error.type` ==
96///   `"content_filter"`, checked at the top level and inside a nested
97///   `innererror` for Azure OpenAI's shape) -> [`Error::ServiceContentFilter`]
98/// * anything else — notably `408` / `429` / `5xx`, which
99///   [`RetryOn::Default`](agent_framework_core::client::RetryOn::Default)
100///   depends on — -> [`Error::ServiceStatus`], unchanged
101///
102/// `body` is parsed leniently as JSON purely to look for the content-filter
103/// marker; a non-JSON or differently-shaped body just skips that check and
104/// falls through to the plain status-based classification, so this never
105/// panics or itself errors. `message` is used verbatim as the resulting
106/// error's text (callers already format a provider-specific "OpenAI API
107/// error 400: ..." string; this only picks the variant).
108pub fn classify_service_error(
109    status: u16,
110    body: &str,
111    message: impl Into<String>,
112    retry_after: Option<f64>,
113) -> Error {
114    let message = message.into();
115    match status {
116        401 | 403 => Error::service_invalid_auth(message),
117        400 | 404 | 422 => {
118            if body_signals_content_filter(body) {
119                Error::service_content_filter(message)
120            } else {
121                Error::service_invalid_request(message)
122            }
123        }
124        _ => Error::service_status(status, message, retry_after),
125    }
126}
127
128/// Whether an OpenAI/Azure-OpenAI-shaped error body signals a content-filter
129/// refusal.
130///
131/// Checks `error.code` — the field the OpenAI Python SDK's
132/// `BadRequestError.code` reads, and what upstream compares against
133/// `"content_filter"` — tolerantly falling back to `error.type` and to Azure
134/// OpenAI's nested `error.innererror` (`code`/`type`), since a
135/// `"content_filter"` marker has been observed in slightly different spots
136/// across OpenAI-wire-compatible providers. A non-JSON or unrecognized body
137/// shape is treated as "not a content filter" rather than erroring.
138fn body_signals_content_filter(body: &str) -> bool {
139    let Ok(value) = serde_json::from_str::<Value>(body) else {
140        return false;
141    };
142    let error = value.get("error").unwrap_or(&value);
143    is_content_filter_marker(error)
144        || error
145            .get("innererror")
146            .is_some_and(is_content_filter_marker)
147}
148
149fn is_content_filter_marker(v: &Value) -> bool {
150    v.get("code").and_then(Value::as_str) == Some("content_filter")
151        || v.get("type").and_then(Value::as_str) == Some("content_filter")
152}
153
154/// An OpenAI (or OpenAI-compatible) chat client.
155#[derive(Clone)]
156pub struct OpenAIChatCompletionClient {
157    inner: Arc<Inner>,
158}
159
160#[derive(Clone)]
161struct Inner {
162    http: reqwest::Client,
163    api_key: String,
164    base_url: String,
165    model: String,
166    organization: Option<String>,
167}
168
169impl OpenAIChatCompletionClient {
170    /// Create a client for the given API key and default model.
171    pub fn new(api_key: impl Into<String>, model: impl Into<String>) -> Self {
172        Self {
173            inner: Arc::new(Inner {
174                http: reqwest::Client::new(),
175                api_key: api_key.into(),
176                base_url: DEFAULT_BASE_URL.to_string(),
177                model: model.into(),
178                organization: None,
179            }),
180        }
181    }
182
183    /// Build a client from the `OPENAI_API_KEY` (and optional
184    /// `OPENAI_BASE_URL`) environment variables.
185    pub fn from_env(model: impl Into<String>) -> Result<Self> {
186        let key = std::env::var("OPENAI_API_KEY")
187            .map_err(|_| Error::Configuration("OPENAI_API_KEY is not set".into()))?;
188        let mut client = Self::new(key, model);
189        if let Ok(base) = std::env::var("OPENAI_BASE_URL") {
190            client = client.with_base_url(base);
191        }
192        Ok(client)
193    }
194
195    /// Override the base URL (for Azure OpenAI or compatible servers).
196    pub fn with_base_url(mut self, base_url: impl Into<String>) -> Self {
197        Arc::make_mut(&mut self.inner).base_url = base_url.into();
198        self
199    }
200
201    /// Set the organization header.
202    pub fn with_organization(mut self, org: impl Into<String>) -> Self {
203        Arc::make_mut(&mut self.inner).organization = Some(org.into());
204        self
205    }
206
207    fn build_body(&self, messages: &[Message], options: &ChatOptions, stream: bool) -> Value {
208        let mut body = Map::new();
209        let model = options
210            .model
211            .clone()
212            .unwrap_or_else(|| self.inner.model.clone());
213        body.insert("model".into(), json!(model));
214        body.insert(
215            "messages".into(),
216            json!(convert::messages_to_openai(messages)),
217        );
218        convert::apply_options(&mut body, options);
219        let (tools, tool_choice) = convert::tools_to_openai(options);
220        if let Some(tools) = tools {
221            body.insert("tools".into(), tools);
222        }
223        if let Some(choice) = tool_choice {
224            body.insert("tool_choice".into(), choice);
225        }
226        if stream {
227            body.insert("stream".into(), json!(true));
228            body.insert("stream_options".into(), json!({ "include_usage": true }));
229        }
230        Value::Object(body)
231    }
232
233    async fn post(&self, body: &Value) -> Result<reqwest::Response> {
234        let url = format!(
235            "{}/chat/completions",
236            self.inner.base_url.trim_end_matches('/')
237        );
238        let mut req = self
239            .inner
240            .http
241            .post(&url)
242            .bearer_auth(&self.inner.api_key)
243            .json(body);
244        if let Some(org) = &self.inner.organization {
245            req = req.header("OpenAI-Organization", org);
246        }
247        let resp = req
248            .send()
249            .await
250            .map_err(|e| Error::service(format!("request failed: {e}")))?;
251        if !resp.status().is_success() {
252            let status = resp.status();
253            let retry_after = parse_retry_after(resp.headers());
254            let text = resp.text().await.unwrap_or_default();
255            return Err(classify_service_error(
256                status.as_u16(),
257                &text,
258                format!("OpenAI API error {status}: {text}"),
259                retry_after,
260            ));
261        }
262        Ok(resp)
263    }
264
265    /// The default model id.
266    pub fn model(&self) -> &str {
267        &self.inner.model
268    }
269}
270
271#[async_trait::async_trait]
272impl ChatClient for OpenAIChatCompletionClient {
273    async fn get_response(
274        &self,
275        messages: Vec<Message>,
276        options: ChatOptions,
277    ) -> Result<ChatResponse> {
278        let body = self.build_body(&messages, &options, false);
279        let resp = self.post(&body).await?;
280        let value: Value = resp
281            .json()
282            .await
283            .map_err(|e| Error::service(format!("invalid response json: {e}")))?;
284        Ok(convert::parse_response(&value))
285    }
286
287    async fn get_streaming_response(
288        &self,
289        messages: Vec<Message>,
290        options: ChatOptions,
291    ) -> Result<ChatStream> {
292        let body = self.build_body(&messages, &options, true);
293        let resp = self.post(&body).await?;
294        Ok(parse_sse_stream(resp).boxed())
295    }
296
297    fn model(&self) -> Option<&str> {
298        Some(&self.inner.model)
299    }
300}
301
302type ByteStream =
303    std::pin::Pin<Box<dyn futures::Stream<Item = reqwest::Result<bytes::Bytes>> + Send>>;
304
305/// Turn an SSE HTTP response into a stream of [`ChatResponseUpdate`]s.
306///
307/// Public (but hidden) so `agent-framework-azure` can reuse the exact same
308/// chat-completions SSE parsing for Azure OpenAI's wire-compatible stream.
309#[doc(hidden)]
310pub fn parse_sse_stream(
311    resp: reqwest::Response,
312) -> impl futures::Stream<Item = Result<ChatResponseUpdate>> + Send {
313    let byte_stream: ByteStream = Box::pin(resp.bytes_stream());
314    // `tool_ids` maps a streamed tool-call `index` to its `call_id`, so that
315    // continuation chunks (which carry only the index) can be resolved back to
316    // the id assigned in the first chunk.
317    futures::stream::unfold(
318        SseState {
319            byte_stream,
320            buffer: String::new(),
321            utf8: Utf8StreamDecoder::new(),
322            queued: VecDeque::new(),
323            tool_ids: HashMap::new(),
324            done: false,
325        },
326        |mut state| async move {
327            loop {
328                if let Some(update) = state.queued.pop_front() {
329                    return Some((Ok(update), state));
330                }
331                if state.done {
332                    return None;
333                }
334                match state.byte_stream.next().await {
335                    Some(Ok(bytes)) => {
336                        let decoded = state.utf8.push(&bytes);
337                        state.buffer.push_str(&decoded);
338                        while let Some(pos) = state.buffer.find('\n') {
339                            let line = state.buffer[..pos].trim().to_string();
340                            state.buffer.drain(..=pos);
341                            if let Some(data) = line.strip_prefix("data:") {
342                                let data = data.trim();
343                                if data == "[DONE]" {
344                                    return drain_or_end(state);
345                                }
346                                if let Ok(value) = serde_json::from_str::<Value>(data) {
347                                    // Providers signal mid-stream failures (rate
348                                    // limits, content filter, billing) as an
349                                    // object with an `error` field; surface it.
350                                    if let Some(err) = value.get("error") {
351                                        let msg = err
352                                            .get("message")
353                                            .and_then(Value::as_str)
354                                            .unwrap_or("unknown stream error")
355                                            .to_string();
356                                        state.done = true;
357                                        return Some((Err(Error::service(msg)), state));
358                                    }
359                                    if let Some(update) = parse_delta(&value, &mut state.tool_ids) {
360                                        state.queued.push_back(update);
361                                    }
362                                }
363                            }
364                        }
365                    }
366                    Some(Err(e)) => {
367                        state.done = true;
368                        return Some((Err(Error::service(format!("stream error: {e}"))), state));
369                    }
370                    None => return drain_or_end(state),
371                }
372            }
373        },
374    )
375}
376
377/// State carried across `unfold` iterations while parsing the SSE stream.
378struct SseState {
379    byte_stream: ByteStream,
380    buffer: String,
381    utf8: Utf8StreamDecoder,
382    queued: VecDeque<ChatResponseUpdate>,
383    tool_ids: HashMap<i64, String>,
384    done: bool,
385}
386
387fn drain_or_end(mut state: SseState) -> Option<(Result<ChatResponseUpdate>, SseState)> {
388    match state.queued.pop_front() {
389        Some(update) => {
390            state.done = true;
391            Some((Ok(update), state))
392        }
393        None => None,
394    }
395}
396
397/// Parse one streamed chunk (`chat.completion.chunk`) into an update, resolving
398/// tool-call ids from the index map.
399fn parse_delta(value: &Value, tool_ids: &mut HashMap<i64, String>) -> Option<ChatResponseUpdate> {
400    let mut update = ChatResponseUpdate {
401        response_id: value.get("id").and_then(Value::as_str).map(String::from),
402        model: value.get("model").and_then(Value::as_str).map(String::from),
403        ..Default::default()
404    };
405
406    let mut contents: Vec<Content> = Vec::new();
407
408    if let Some(choice) = value
409        .get("choices")
410        .and_then(|c| c.as_array())
411        .and_then(|a| a.first())
412    {
413        if let Some(delta) = choice.get("delta") {
414            if let Some(r) = delta.get("role").and_then(Value::as_str) {
415                update.role = Some(Role::new(r));
416            }
417            if let Some(text) = delta.get("content").and_then(Value::as_str) {
418                if !text.is_empty() {
419                    contents.push(Content::Text(TextContent::new(text)));
420                }
421            }
422            if let Some(calls) = delta.get("tool_calls").and_then(Value::as_array) {
423                for call in calls {
424                    let index = call.get("index").and_then(Value::as_i64).unwrap_or(0);
425                    let chunk_id = call.get("id").and_then(Value::as_str).unwrap_or_default();
426                    // The first chunk of a call carries its id; record it so
427                    // later index-only chunks resolve to the same call_id.
428                    let id = if chunk_id.is_empty() {
429                        tool_ids.get(&index).cloned().unwrap_or_default()
430                    } else {
431                        tool_ids.insert(index, chunk_id.to_string());
432                        chunk_id.to_string()
433                    };
434                    let func = call.get("function");
435                    let name = func
436                        .and_then(|f| f.get("name"))
437                        .and_then(Value::as_str)
438                        .unwrap_or_default()
439                        .to_string();
440                    let args = func
441                        .and_then(|f| f.get("arguments"))
442                        .and_then(Value::as_str)
443                        .unwrap_or_default()
444                        .to_string();
445                    contents.push(Content::FunctionCall(FunctionCallContent::new(
446                        id,
447                        name,
448                        Some(FunctionArguments::Raw(args)),
449                    )));
450                }
451            }
452        }
453        if let Some(fr) = choice.get("finish_reason").and_then(Value::as_str) {
454            update.finish_reason = Some(FinishReason::new(fr));
455        }
456    }
457
458    // The final chunk (with `stream_options.include_usage`) carries top-level
459    // `usage` and no choices; surface it so streamed runs accumulate token usage
460    // just like non-streaming responses.
461    if let Some(usage) = value.get("usage").filter(|u| u.is_object()) {
462        contents.push(Content::Usage(UsageContent {
463            details: convert::parse_usage(usage),
464        }));
465    }
466
467    if update.role.is_none() {
468        update.role = Some(Role::assistant());
469    }
470    update.contents = contents;
471    Some(update)
472}
473
474#[cfg(test)]
475mod tests {
476    use super::*;
477
478    // -- classify_service_error --------------------------------------------
479    //
480    // Canned status+body combinations run through the exact classification
481    // `OpenAIChatCompletionClient::post` and `responses::OpenAIChatClient::post`
482    // both delegate to.
483
484    #[test]
485    fn classifies_401_and_403_as_invalid_auth() {
486        for status in [401, 403] {
487            let err = classify_service_error(status, "", format!("err {status}"), None);
488            assert!(
489                matches!(err, Error::ServiceInvalidAuth { .. }),
490                "status {status}: {err:?}"
491            );
492        }
493    }
494
495    #[test]
496    fn classifies_400_404_422_as_invalid_request_by_default() {
497        for status in [400, 404, 422] {
498            let body = r#"{"error":{"message":"nope","type":"invalid_request_error"}}"#;
499            let err = classify_service_error(status, body, format!("err {status}"), None);
500            assert!(
501                matches!(err, Error::ServiceInvalidRequest { .. }),
502                "status {status}: {err:?}"
503            );
504        }
505    }
506
507    #[test]
508    fn classifies_400_with_content_filter_code_as_content_filter() {
509        // Plain OpenAI shape: `error.code`.
510        let body = r#"{"error":{"message":"flagged","type":"invalid_request_error","code":"content_filter"}}"#;
511        let err = classify_service_error(400, body, "err", None);
512        assert!(matches!(err, Error::ServiceContentFilter { .. }), "{err:?}");
513    }
514
515    #[test]
516    fn classifies_azure_openai_content_filter_shape_with_innererror() {
517        // Azure OpenAI's shape: outer `error.code` is already "content_filter"
518        // (mirrors `openai/_exceptions.py`'s `OpenAIContentFilterException`,
519        // which is constructed once `ex.code == "content_filter"` is already
520        // known — the nested `innererror.code` carries the more specific
521        // `ResponsibleAIPolicyViolation` detail, not the marker itself), but
522        // this also tolerates the marker living only in `innererror`.
523        let body = r#"{"error":{"message":"The response was filtered","code":"content_filter","innererror":{"code":"ResponsibleAIPolicyViolation"}}}"#;
524        let err = classify_service_error(400, body, "err", None);
525        assert!(matches!(err, Error::ServiceContentFilter { .. }), "{err:?}");
526
527        let nested_only =
528            r#"{"error":{"message":"filtered","innererror":{"code":"content_filter"}}}"#;
529        let err = classify_service_error(400, nested_only, "err", None);
530        assert!(matches!(err, Error::ServiceContentFilter { .. }), "{err:?}");
531    }
532
533    #[test]
534    fn non_json_or_unmarked_body_does_not_trigger_content_filter() {
535        let err = classify_service_error(400, "not json", "err", None);
536        assert!(
537            matches!(err, Error::ServiceInvalidRequest { .. }),
538            "{err:?}"
539        );
540
541        let err = classify_service_error(400, r#"{"error":{"code":"other"}}"#, "err", None);
542        assert!(
543            matches!(err, Error::ServiceInvalidRequest { .. }),
544            "{err:?}"
545        );
546    }
547
548    #[test]
549    fn leaves_retryable_statuses_as_service_status() {
550        // 408/429/5xx must stay `ServiceStatus` (carrying status +
551        // retry_after) exactly as before — the retry layer depends on it.
552        for status in [408, 429, 500, 503] {
553            let err = classify_service_error(status, "", format!("err {status}"), Some(2.0));
554            assert_eq!(err.status(), Some(status), "{err:?}");
555            assert_eq!(err.retry_after(), Some(2.0), "{err:?}");
556        }
557    }
558
559    #[test]
560    fn unclassified_4xx_stays_service_status() {
561        // e.g. 409 Conflict, 413 Payload Too Large: not in the
562        // auth/invalid-request/content-filter buckets, so this falls back to
563        // the generic status-carrying variant rather than guessing.
564        let err = classify_service_error(409, "", "err", None);
565        assert_eq!(err.status(), Some(409), "{err:?}");
566    }
567
568    #[test]
569    fn message_text_is_preserved_verbatim() {
570        let err = classify_service_error(401, "", "OpenAI API error 401: unauthorized", None);
571        assert_eq!(
572            err.to_string(),
573            "service error: OpenAI API error 401: unauthorized"
574        );
575    }
576}
agent_framework_openai/lib.rs

agent_framework_openai/
lib.rs
