|
| 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 | +} |
0 commit comments