Add support for custom HTTP transport via URLSessionProtocol and URLSessionFactory

Sources/OpenAI/OpenAI.swift

@@ -106,6 +106,62 @@ final public class OpenAI: OpenAIProtocol, @unchecked Sendable {
         )
     }
 
+    /// Creates an OpenAI client with a custom URLSession protocol implementation.
+    ///
+    /// - Important: This initializer only uses the custom session for non-streaming requests.
+    ///   For streaming requests, use the initializer that accepts a `URLSessionFactory`.
+    ///
+    /// - Parameters:
+    ///   - configuration: The client configuration
+    ///   - customSession: Custom URLSession protocol implementation
+    ///   - middlewares: Optional middlewares for request/response interception
+    public convenience init(
+        configuration: Configuration,
+        customSession: any URLSessionProtocol,
+        middlewares: [OpenAIMiddleware] = []
+    ) {
+        let streamingSessionFactory = ImplicitURLSessionStreamingSessionFactory(
+            middlewares: middlewares,
+            parsingOptions: configuration.parsingOptions,
+            sslDelegate: nil
+        )
+
+        self.init(
+            configuration: configuration,
+            session: customSession,
+            streamingSessionFactory: streamingSessionFactory,
+            middlewares: middlewares
+        )
+    }
+
+    /// Creates an OpenAI client with custom session handling for both regular and streaming requests.
+    ///
+    /// - Parameters:
+    ///   - configuration: The client configuration
+    ///   - customSession: Custom URLSession protocol implementation for non-streaming requests
+    ///   - streamingURLSessionFactory: Factory for creating sessions for streaming requests
+    ///   - middlewares: Optional middlewares for request/response interception
+    public convenience init(
+        configuration: Configuration,
+        customSession: any URLSessionProtocol,
+        streamingURLSessionFactory: URLSessionFactory,
+        middlewares: [OpenAIMiddleware] = []
+    ) {
+        let streamingSessionFactory = ImplicitURLSessionStreamingSessionFactory(
+            urlSessionFactory: streamingURLSessionFactory,
+            middlewares: middlewares,
+            parsingOptions: configuration.parsingOptions,
+            sslDelegate: nil
+        )
+
+        self.init(
+            configuration: configuration,
+            session: customSession,
+            streamingSessionFactory: streamingSessionFactory,
+            middlewares: middlewares
+        )
+    }
+
     init(
         configuration: Configuration,
         session: URLSessionProtocol,

Sources/OpenAI/Private/Streaming/InvalidatableSession.swift

@@ -7,7 +7,7 @@
 
 import Foundation
 
-protocol InvalidatableSession: Sendable {
+public protocol InvalidatableSession: Sendable {
     func invalidateAndCancel()
     func finishTasksAndInvalidate()
 }

Sources/OpenAI/Private/Streaming/ServerSentEventsStreamingSessionFactory.swift

@@ -35,17 +35,31 @@ protocol StreamingSessionFactory: Sendable {
 }
 
 struct ImplicitURLSessionStreamingSessionFactory: StreamingSessionFactory {
+    let urlSessionFactory: URLSessionFactory
     let middlewares: [OpenAIMiddleware]
     let parsingOptions: ParsingOptions
     let sslDelegate: SSLDelegateProtocol?
-
+
+    init(
+        urlSessionFactory: URLSessionFactory = FoundationURLSessionFactory(),
+        middlewares: [OpenAIMiddleware],
+        parsingOptions: ParsingOptions,
+        sslDelegate: SSLDelegateProtocol?
+    ) {
+        self.urlSessionFactory = urlSessionFactory
+        self.middlewares = middlewares
+        self.parsingOptions = parsingOptions
+        self.sslDelegate = sslDelegate
+    }
+
     func makeServerSentEventsStreamingSession<ResultType>(
         urlRequest: URLRequest,
         onReceiveContent: @Sendable @escaping (StreamingSession<ServerSentEventsStreamInterpreter<ResultType>>, ResultType) -> Void,
         onProcessingError: @Sendable @escaping (StreamingSession<ServerSentEventsStreamInterpreter<ResultType>>, any Error) -> Void,
         onComplete: @Sendable @escaping (StreamingSession<ServerSentEventsStreamInterpreter<ResultType>>, (any Error)?) -> Void
     ) -> StreamingSession<ServerSentEventsStreamInterpreter<ResultType>> where ResultType : Decodable, ResultType : Encodable, ResultType : Sendable {
         .init(
+            urlSessionFactory: urlSessionFactory,
             urlRequest: urlRequest,
             interpreter: .init(parsingOptions: parsingOptions),
             sslDelegate: sslDelegate,
@@ -55,14 +69,15 @@ struct ImplicitURLSessionStreamingSessionFactory: StreamingSessionFactory {
             onComplete: onComplete
         )
     }
-    
+
     func makeAudioSpeechStreamingSession(
         urlRequest: URLRequest,
         onReceiveContent: @Sendable @escaping (StreamingSession<AudioSpeechStreamInterpreter>, AudioSpeechResult) -> Void,
         onProcessingError: @Sendable @escaping (StreamingSession<AudioSpeechStreamInterpreter>, any Error) -> Void,
         onComplete: @Sendable @escaping (StreamingSession<AudioSpeechStreamInterpreter>, (any Error)?) -> Void
     ) -> StreamingSession<AudioSpeechStreamInterpreter> {
         .init(
+            urlSessionFactory: urlSessionFactory,
             urlRequest: urlRequest,
             interpreter: .init(),
             sslDelegate: sslDelegate,
@@ -72,14 +87,15 @@ struct ImplicitURLSessionStreamingSessionFactory: StreamingSessionFactory {
             onComplete: onComplete
         )
     }
-    
+
     func makeModelResponseStreamingSession(
         urlRequest: URLRequest,
         onReceiveContent: @Sendable @escaping (StreamingSession<ModelResponseEventsStreamInterpreter>, ResponseStreamEvent) -> Void,
         onProcessingError: @Sendable @escaping (StreamingSession<ModelResponseEventsStreamInterpreter>, any Error) -> Void,
         onComplete: @Sendable @escaping (StreamingSession<ModelResponseEventsStreamInterpreter>, (any Error)?) -> Void
     ) -> StreamingSession<ModelResponseEventsStreamInterpreter> {
         .init(
+            urlSessionFactory: urlSessionFactory,
             urlRequest: urlRequest,
             interpreter: .init(),
             sslDelegate: sslDelegate,

Sources/OpenAI/Private/URLSessionCombine.swift

@@ -14,19 +14,19 @@ import FoundationNetworking
 #if canImport(Combine)
 import Combine
 
-protocol URLSessionCombine {
+public protocol URLSessionCombine {
     func dataTaskPublisher(for request: URLRequest) -> AnyPublisher<(data: Data, response: URLResponse), URLError>
 }
 
 extension URLSession: URLSessionCombine {
-    func dataTaskPublisher(for request: URLRequest) -> AnyPublisher<(data: Data, response: URLResponse), URLError> {
+    public func dataTaskPublisher(for request: URLRequest) -> AnyPublisher<(data: Data, response: URLResponse), URLError> {
         let typedPublisher: URLSession.DataTaskPublisher = dataTaskPublisher(for: request)
         return typedPublisher.eraseToAnyPublisher()
     }
 }
 
 #else
-protocol URLSessionCombine {
+public protocol URLSessionCombine {
 }
 
 extension URLSession: URLSessionCombine {}

Sources/OpenAI/Private/URLSessionDataTaskProtocol.swift

@@ -10,14 +10,14 @@ import Foundation
 import FoundationNetworking
 #endif
 
-protocol URLSessionTaskProtocol: Sendable {
+public protocol URLSessionTaskProtocol: Sendable {
     var originalRequest: URLRequest? { get }
     func cancel()
 }
 
 extension URLSessionTask: URLSessionTaskProtocol {}
 
-protocol URLSessionDataTaskProtocol: URLSessionTaskProtocol {
+public protocol URLSessionDataTaskProtocol: URLSessionTaskProtocol {
     func resume()
 }
 

Sources/OpenAI/Private/URLSessionDelegateProtocol.swift

@@ -11,19 +11,24 @@ import Foundation
 import FoundationNetworking
 #endif
 
-protocol URLSessionDelegateProtocol: Sendable { // Sendable to make a better match with URLSessionDelegate, it's sendable too
+/// Protocol for handling URLSession delegate callbacks.
+/// Sendable to match URLSessionDelegate behavior.
+/// AnyObject constraint allows weak references to delegate implementations.
+public protocol URLSessionDelegateProtocol: AnyObject, Sendable {
     func urlSession(_ session: URLSessionProtocol, task: URLSessionTaskProtocol, didCompleteWithError error: Error?)
-    
+
     func urlSession(
         _ session: URLSession,
         didReceive challenge: URLAuthenticationChallenge,
         completionHandler: @escaping @Sendable (URLSession.AuthChallengeDisposition, URLCredential?) -> Void
     )
 }
 
-protocol URLSessionDataDelegateProtocol: URLSessionDelegateProtocol {
+/// Protocol for handling URLSession data delegate callbacks.
+/// Used for streaming data reception.
+public protocol URLSessionDataDelegateProtocol: URLSessionDelegateProtocol {
     func urlSession(_ session: URLSessionProtocol, dataTask: URLSessionDataTaskProtocol, didReceive data: Data)
-    
+
     func urlSession(
         _ session: URLSessionProtocol,
         dataTask: URLSessionDataTaskProtocol,

Sources/OpenAI/Private/URLSessionFactory.swift

@@ -10,12 +10,20 @@ import Foundation
 import FoundationNetworking
 #endif
 
-protocol URLSessionFactory: Sendable {
+/// Factory protocol for creating URLSession instances.
+/// Implement this protocol to provide custom session creation for streaming requests.
+public protocol URLSessionFactory: Sendable {
+    /// Creates a URLSession for streaming requests.
+    /// - Parameter delegate: The delegate to receive streaming data callbacks
+    /// - Returns: A URLSession protocol implementation
     func makeUrlSession(delegate: URLSessionDataDelegateProtocol) -> URLSessionProtocol
 }
 
-struct FoundationURLSessionFactory: URLSessionFactory {
-    func makeUrlSession(delegate: URLSessionDataDelegateProtocol) -> any URLSessionProtocol {
+/// Default factory that creates standard Foundation URLSession instances.
+public struct FoundationURLSessionFactory: URLSessionFactory {
+    public init() {}
+
+    public func makeUrlSession(delegate: URLSessionDataDelegateProtocol) -> any URLSessionProtocol {
         let forwarder = URLSessionDataDelegateForwarder(target: delegate)
         return URLSession(configuration: .default, delegate: forwarder, delegateQueue: nil)
     }

Sources/OpenAI/Private/URLSessionProtocol.swift

@@ -10,20 +10,20 @@ import Foundation
 import FoundationNetworking
 #endif
 
-protocol URLSessionProtocol: InvalidatableSession, URLSessionCombine {
+public protocol URLSessionProtocol: InvalidatableSession, URLSessionCombine {
     func dataTask(with request: URLRequest, completionHandler: @escaping @Sendable (Data?, URLResponse?, Error?) -> Void) -> URLSessionDataTaskProtocol
     func dataTask(with request: URLRequest) -> URLSessionDataTaskProtocol
-    
+
     @available(iOS 15.0, macOS 12.0, tvOS 15.0, watchOS 8.0, *)
     func data(for request: URLRequest, delegate: (any URLSessionTaskDelegate)?) async throws -> (Data, URLResponse)
 }
 
 extension URLSession: URLSessionProtocol {
-    func dataTask(with request: URLRequest) -> URLSessionDataTaskProtocol {
+    public func dataTask(with request: URLRequest) -> URLSessionDataTaskProtocol {
         dataTask(with: request) as URLSessionDataTask
     }
-    
-    func dataTask(with request: URLRequest, completionHandler: @escaping @Sendable (Data?, URLResponse?, Error?) -> Void) -> URLSessionDataTaskProtocol {
+
+    public func dataTask(with request: URLRequest, completionHandler: @escaping @Sendable (Data?, URLResponse?, Error?) -> Void) -> URLSessionDataTaskProtocol {
         dataTask(with: request, completionHandler: completionHandler) as URLSessionDataTask
     }
 }