Skip to content

Commit 218eee6

Browse files
authored
Support OpenAI Responses API in AI Core LLM provider (#90)
1 parent ee15098 commit 218eee6

4 files changed

Lines changed: 190 additions & 2 deletions

File tree

crates/llm/src/aicore/mod.rs

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,11 +8,13 @@
88
99
mod anthropic;
1010
mod openai;
11+
mod openai_responses;
1112
mod types;
1213
mod vertex;
1314

1415
pub use anthropic::AiCoreAnthropicClient;
1516
pub use openai::AiCoreOpenAIClient;
17+
pub use openai_responses::AiCoreOpenAIResponsesClient;
1618
pub use types::AiCoreApiType;
1719
pub use vertex::AiCoreVertexClient;
1820

@@ -33,6 +35,11 @@ pub fn create_aicore_client(
3335
AiCoreApiType::OpenAI => {
3436
Box::new(AiCoreOpenAIClient::new(token_manager, base_url, model_id))
3537
}
38+
AiCoreApiType::OpenAIResponses => Box::new(AiCoreOpenAIResponsesClient::new(
39+
token_manager,
40+
base_url,
41+
model_id,
42+
)),
3643
AiCoreApiType::Vertex => {
3744
Box::new(AiCoreVertexClient::new(token_manager, base_url, model_id))
3845
}
@@ -59,6 +66,12 @@ pub fn create_aicore_client_with_recorder<P: AsRef<Path>>(
5966
model_id,
6067
recording_path,
6168
)),
69+
AiCoreApiType::OpenAIResponses => Box::new(AiCoreOpenAIResponsesClient::new_with_recorder(
70+
token_manager,
71+
base_url,
72+
model_id,
73+
recording_path,
74+
)),
6275
AiCoreApiType::Vertex => Box::new(AiCoreVertexClient::new_with_recorder(
6376
token_manager,
6477
base_url,
Lines changed: 147 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,147 @@
1+
//! AI Core client for OpenAI Responses API
2+
//!
3+
//! This client wraps the OpenAI Responses API client with AI Core authentication.
4+
//! The Responses API is OpenAI's modern API format supporting:
5+
//! - Stateless mode with encrypted reasoning
6+
//! - Function calling
7+
//! - Streaming with SSE
8+
9+
use crate::{
10+
auth::TokenManager,
11+
openai_responses::{AuthProvider, OpenAIResponsesClient, RequestCustomizer},
12+
types::*,
13+
LLMProvider, StreamingCallback,
14+
};
15+
use anyhow::Result;
16+
use async_trait::async_trait;
17+
use std::sync::Arc;
18+
19+
// ============================================================================
20+
// AI Core Authentication Provider for OpenAI Responses API
21+
// ============================================================================
22+
23+
/// AI Core authentication provider for OpenAI Responses API
24+
struct AiCoreOpenAIResponsesAuthProvider {
25+
token_manager: Arc<TokenManager>,
26+
}
27+
28+
impl AiCoreOpenAIResponsesAuthProvider {
29+
fn new(token_manager: Arc<TokenManager>) -> Self {
30+
Self { token_manager }
31+
}
32+
}
33+
34+
#[async_trait]
35+
impl AuthProvider for AiCoreOpenAIResponsesAuthProvider {
36+
async fn get_auth_headers(&self) -> Result<Vec<(String, String)>> {
37+
let token = self.token_manager.get_valid_token().await?;
38+
Ok(vec![(
39+
"Authorization".to_string(),
40+
format!("Bearer {token}"),
41+
)])
42+
}
43+
}
44+
45+
// ============================================================================
46+
// AI Core Request Customizer for OpenAI Responses API
47+
// ============================================================================
48+
49+
/// AI Core request customizer for OpenAI Responses API
50+
struct AiCoreOpenAIResponsesRequestCustomizer;
51+
52+
impl RequestCustomizer for AiCoreOpenAIResponsesRequestCustomizer {
53+
fn customize_request(&self, _request: &mut serde_json::Value) -> Result<()> {
54+
// No additional customization needed for Responses API requests
55+
Ok(())
56+
}
57+
58+
fn get_additional_headers(&self) -> Vec<(String, String)> {
59+
vec![
60+
("AI-Resource-Group".to_string(), "default".to_string()),
61+
("Content-Type".to_string(), "application/json".to_string()),
62+
]
63+
}
64+
65+
fn customize_url(&self, base_url: &str, _streaming: bool) -> String {
66+
// AI Core uses /responses endpoint for OpenAI Responses API
67+
format!("{base_url}/responses")
68+
}
69+
}
70+
71+
// ============================================================================
72+
// AI Core OpenAI Responses Client
73+
// ============================================================================
74+
75+
/// AI Core client for OpenAI Responses API
76+
///
77+
/// This client provides access to OpenAI's Responses API through AI Core,
78+
/// supporting features like encrypted reasoning for stateless mode,
79+
/// function calling, and streaming.
80+
pub struct AiCoreOpenAIResponsesClient {
81+
responses_client: OpenAIResponsesClient,
82+
custom_config: Option<serde_json::Value>,
83+
}
84+
85+
impl AiCoreOpenAIResponsesClient {
86+
fn create_responses_client(
87+
token_manager: Arc<TokenManager>,
88+
base_url: String,
89+
model_id: String,
90+
) -> OpenAIResponsesClient {
91+
let auth_provider = Box::new(AiCoreOpenAIResponsesAuthProvider::new(token_manager));
92+
let request_customizer = Box::new(AiCoreOpenAIResponsesRequestCustomizer);
93+
94+
OpenAIResponsesClient::with_customization(
95+
model_id,
96+
base_url,
97+
auth_provider,
98+
request_customizer,
99+
)
100+
}
101+
102+
pub fn new(token_manager: Arc<TokenManager>, base_url: String, model_id: String) -> Self {
103+
let responses_client = Self::create_responses_client(token_manager, base_url, model_id);
104+
Self {
105+
responses_client,
106+
custom_config: None,
107+
}
108+
}
109+
110+
/// Create a new client with recording capability
111+
pub fn new_with_recorder<P: AsRef<std::path::Path>>(
112+
token_manager: Arc<TokenManager>,
113+
base_url: String,
114+
model_id: String,
115+
recording_path: P,
116+
) -> Self {
117+
let responses_client = Self::create_responses_client(token_manager, base_url, model_id)
118+
.with_recorder(recording_path);
119+
Self {
120+
responses_client,
121+
custom_config: None,
122+
}
123+
}
124+
125+
/// Set custom model configuration to be merged into API requests
126+
pub fn with_custom_config(mut self, custom_config: serde_json::Value) -> Self {
127+
self.responses_client = self
128+
.responses_client
129+
.with_custom_config(custom_config.clone());
130+
self.custom_config = Some(custom_config);
131+
self
132+
}
133+
}
134+
135+
#[async_trait]
136+
impl LLMProvider for AiCoreOpenAIResponsesClient {
137+
async fn send_message(
138+
&mut self,
139+
request: LLMRequest,
140+
streaming_callback: Option<&StreamingCallback>,
141+
) -> Result<LLMResponse> {
142+
// Delegate to the wrapped OpenAIResponsesClient
143+
self.responses_client
144+
.send_message(request, streaming_callback)
145+
.await
146+
}
147+
}

crates/llm/src/aicore/types.rs

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,9 @@ pub enum AiCoreApiType {
1111
Anthropic,
1212
/// OpenAI Chat Completions API
1313
OpenAI,
14+
/// OpenAI Responses API (modern format with encrypted reasoning support)
15+
#[serde(rename = "openai-responses")]
16+
OpenAIResponses,
1417
/// Google Vertex AI / Gemini API
1518
Vertex,
1619
}
@@ -20,6 +23,7 @@ impl std::fmt::Display for AiCoreApiType {
2023
match self {
2124
AiCoreApiType::Anthropic => write!(f, "anthropic"),
2225
AiCoreApiType::OpenAI => write!(f, "openai"),
26+
AiCoreApiType::OpenAIResponses => write!(f, "openai-responses"),
2327
AiCoreApiType::Vertex => write!(f, "vertex"),
2428
}
2529
}
@@ -32,9 +36,10 @@ impl std::str::FromStr for AiCoreApiType {
3236
match s.to_lowercase().as_str() {
3337
"anthropic" => Ok(AiCoreApiType::Anthropic),
3438
"openai" => Ok(AiCoreApiType::OpenAI),
39+
"openai-responses" => Ok(AiCoreApiType::OpenAIResponses),
3540
"vertex" => Ok(AiCoreApiType::Vertex),
3641
_ => Err(anyhow::anyhow!(
37-
"Unknown AI Core API type: '{}'. Expected one of: anthropic, openai, vertex",
42+
"Unknown AI Core API type: '{}'. Expected one of: anthropic, openai, openai-responses, vertex",
3843
s
3944
)),
4045
}

crates/llm/src/factory.rs

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,7 @@
1-
use crate::aicore::{AiCoreAnthropicClient, AiCoreApiType, AiCoreOpenAIClient, AiCoreVertexClient};
1+
use crate::aicore::{
2+
AiCoreAnthropicClient, AiCoreApiType, AiCoreOpenAIClient, AiCoreOpenAIResponsesClient,
3+
AiCoreVertexClient,
4+
};
25
use crate::auth::TokenManager;
36
use crate::provider_config::{ConfigurationSystem, ModelConfig, ProviderConfig};
47
use crate::{
@@ -149,6 +152,12 @@ impl WithCustomConfig for AiCoreVertexClient {
149152
}
150153
}
151154

155+
impl WithCustomConfig for AiCoreOpenAIResponsesClient {
156+
fn with_custom_config(self, custom_config: Value) -> Self {
157+
self.with_custom_config(custom_config)
158+
}
159+
}
160+
152161
// ============================================================================
153162
// Macro for Simple Provider Factory Functions
154163
// ============================================================================
@@ -487,6 +496,20 @@ async fn create_ai_core_client(
487496
let client = apply_custom_config(client, model_config);
488497
Ok(Box::new(client))
489498
}
499+
AiCoreApiType::OpenAIResponses => {
500+
let client = if let Some(path) = record_path {
501+
AiCoreOpenAIResponsesClient::new_with_recorder(
502+
token_manager,
503+
api_url,
504+
model_config.id.clone(),
505+
path,
506+
)
507+
} else {
508+
AiCoreOpenAIResponsesClient::new(token_manager, api_url, model_config.id.clone())
509+
};
510+
let client = apply_custom_config(client, model_config);
511+
Ok(Box::new(client))
512+
}
490513
AiCoreApiType::Vertex => {
491514
let client = if let Some(path) = record_path {
492515
AiCoreVertexClient::new_with_recorder(

0 commit comments

Comments
 (0)