BAML Streaming Swift Client

BAML CLient

// This is a wrapper client made by Codex (can prob be cleaned up?)

import Foundation
import OpenAPIRuntime
import OpenAPIURLSession

// Minimal wrapper around the generated `Client` to expose convenient calls.
final class BamlClient {
    private let client: Client

    init(address: String = "http://localhost:2024") {
        let session = URLSession(configuration: .default)
        let transport = URLSessionTransport(configuration: .init(session: session))
        let serverURL = (try? Servers.Server1.url(address: address)) ?? URL(string: address)!
        self.client = Client(serverURL: serverURL, transport: transport)
    }

    // MARK: - Streaming
    func streamExtractResume(resumeText: String) -> (AsyncThrowingStream<Components.Schemas.ResumePartial, Error>, () -> Void) {
        let holder = TaskHolder()
        let stream = AsyncThrowingStream<Components.Schemas.ResumePartial, Error> { continuation in
            let task = Task {
                do {
                    let input = Operations.ExtractResumeStream.Input(
                        body: .json(.init(resume: resumeText))
                    )
                    let output = try await client.extractResumeStream(input)
                    switch output {
                    case .ok(let ok):
                        switch ok.body {
                        case .textEventStream(let body):
                            // If OpenAPIRuntime has SSE helpers, you can replace this parsing with asServerSentEvents().
                            let parser = SSEParser()
                            for try await chunk in body {
                                let events = parser.feed(Data(chunk))
                                for evt in events {
                                    if let s = evt.data {
                                        if let obj = try? JSONDecoder().decode(Components.Schemas.ResumePartial.self, from: Data(s.utf8)) {
                                            continuation.yield(obj)
                                        }
                                    }
                                }
                            }
                            // Flush any trailing data
                            for evt in parser.finish() {
                                if let s = evt.data, let obj = try? JSONDecoder().decode(Components.Schemas.ResumePartial.self, from: Data(s.utf8)) {
                                    continuation.yield(obj)
                                }
                            }
                            continuation.finish()
                        }
                    case .undocumented(let code, _):
                        throw NSError(domain: "ExtractResumeStream", code: code, userInfo: [NSLocalizedDescriptionKey: "Undocumented status \(code)"])
                    }
                } catch is CancellationError {
                    continuation.finish()
                } catch {
                    continuation.finish(throwing: error)
                }
            }
            holder.task = task
            continuation.onTermination = { _ in task.cancel() }
        }
        return (stream, { holder.task?.cancel() })
    }

    // MARK: - Non-streaming
    func extractResume(resumeText: String) async throws -> Components.Schemas.Resume {
        let input = Operations.ExtractResume.Input(
            body: .json(.init(resume: resumeText))
        )
        let output = try await client.extractResume(input)
        switch output {
        case .ok(let ok):
            switch ok.body {
            case .json(let resume):
                return resume
            }
        case .undocumented(let code, _):
            throw NSError(domain: "ExtractResume", code: code, userInfo: [NSLocalizedDescriptionKey: "Undocumented status \(code)"])
        }
    }
}

private final class TaskHolder { var task: Task<Void, Never>? }

import Foundation

struct ServerSentEvent {
    var event: String?
    var data: String?
    var id: String?
    var retry: Int?
}

final class SSEParser {
    private var buffer = Data()

    func feed(_ newData: Data) -> [ServerSentEvent] {
        buffer.append(newData)
        var events: [ServerSentEvent] = []
        // Try both LF-LF and CRLF-CRLF as separators
        while let range = buffer.range(of: Data([0x0A, 0x0A])) ?? buffer.range(of: Data([0x0D, 0x0A, 0x0D, 0x0A])) {
            let chunk = buffer.subdata(in: buffer.startIndex..<range.lowerBound)
            buffer.removeSubrange(buffer.startIndex..<range.upperBound)
            if let e = parseEvent(from: chunk) { events.append(e) }
        }
        return events
    }

    func finish() -> [ServerSentEvent] {
        defer { buffer.removeAll(keepingCapacity: false) }
        guard !buffer.isEmpty else { return [] }
        if let e = parseEvent(from: buffer) { return [e] }
        return []
    }

    private func parseEvent(from data: Data) -> ServerSentEvent? {
        guard let text = String(data: data, encoding: .utf8) else { return nil }
        var evt = ServerSentEvent()
        var dataLines: [String] = []
        // Split lines on LF; trim a trailing CR if present to support CRLF
        let lines = text.split(separator: "\n", omittingEmptySubsequences: false).map { line -> Substring in
            if line.last == "\r" { return line.dropLast() }
            return line
        }
        for raw in lines {
            if raw.hasPrefix(":") { continue }
            if let idx = raw.firstIndex(of: ":") {
                let field = String(raw[..<idx])
                let value = String(raw[raw.index(after: idx)...]).trimmingCharacters(in: .whitespaces)
                switch field {
                case "event": evt.event = value
                case "data": dataLines.append(value)
                case "id": evt.id = value
                case "retry": evt.retry = Int(value)
                default: break
                }
            }
        }
        if !dataLines.isEmpty { evt.data = dataLines.joined(separator: "\n") }
        return (evt.event != nil || evt.data != nil || evt.id != nil || evt.retry != nil) ? evt : nil
    }
}

Generated.txt

// Generated by swift-openapi-generator, do not modify.
@_spi(Generated) import OpenAPIRuntime
#if os(Linux)
@preconcurrency import struct Foundation.URL
@preconcurrency import struct Foundation.Data
@preconcurrency import struct Foundation.Date
#else
import struct Foundation.URL
import struct Foundation.Data
import struct Foundation.Date
#endif
import HTTPTypes
/// baml-cli serve
public struct Client: APIProtocol {
    /// The underlying HTTP client.
    private let client: UniversalClient
    /// Creates a new client.
    /// - Parameters:
    ///   - serverURL: The server URL that the client connects to. Any server
    ///   URLs defined in the OpenAPI document are available as static methods
    ///   on the ``Servers`` type.
    ///   - configuration: A set of configuration values for the client.
    ///   - transport: A transport that performs HTTP operations.
    ///   - middlewares: A list of middlewares to call before the transport.
    public init(
        serverURL: Foundation.URL,
        configuration: Configuration = .init(),
        transport: any ClientTransport,
        middlewares: [any ClientMiddleware] = []
    ) {
        self.client = .init(
            serverURL: serverURL,
            configuration: configuration,
            transport: transport,
            middlewares: middlewares
        )
    }
    private var converter: Converter {
        client.converter
    }
    /// - Remark: HTTP `POST /call/ExtractResume`.
    /// - Remark: Generated from `#/paths//call/ExtractResume/post(ExtractResume)`.
    public func extractResume(_ input: Operations.ExtractResume.Input) async throws -> Operations.ExtractResume.Output {
        try await client.send(
            input: input,
            forOperation: Operations.ExtractResume.id,
            serializer: { input in
                let path = try converter.renderedPath(
                    template: "/call/ExtractResume",
                    parameters: []
                )
                var request: HTTPTypes.HTTPRequest = .init(
                    soar_path: path,
                    method: .post
                )
                suppressMutabilityWarning(&request)
                converter.setAcceptHeader(
                    in: &request.headerFields,
                    contentTypes: input.headers.accept
                )
                let body: OpenAPIRuntime.HTTPBody?
                switch input.body {
                case let .json(value):
                    body = try converter.setRequiredRequestBodyAsJSON(
                        value,
                        headerFields: &request.headerFields,
                        contentType: "application/json; charset=utf-8"
                    )
                }
                return (request, body)
            },
            deserializer: { response, responseBody in
                switch response.status.code {
                case 200:
                    let contentType = converter.extractContentTypeIfPresent(in: response.headerFields)
                    let body: Operations.ExtractResume.Output.Ok.Body
                    let chosenContentType = try converter.bestContentType(
                        received: contentType,
                        options: [
                            "application/json"
                        ]
                    )
                    switch chosenContentType {
                    case "application/json":
                        body = try await converter.getResponseBodyAsJSON(
                            Components.Schemas.Resume.self,
                            from: responseBody,
                            transforming: { value in
                                .json(value)
                            }
                        )
                    default:
                        preconditionFailure("bestContentType chose an invalid content type.")
                    }
                    return .ok(.init(body: body))
                default:
                    return .undocumented(
                        statusCode: response.status.code,
                        .init(
                            headerFields: response.headerFields,
                            body: responseBody
                        )
                    )
                }
            }
        )
    }
    /// - Remark: HTTP `POST /stream/ExtractResume`.
    /// - Remark: Generated from `#/paths//stream/ExtractResume/post(ExtractResumeStream)`.
    public func extractResumeStream(_ input: Operations.ExtractResumeStream.Input) async throws -> Operations.ExtractResumeStream.Output {
        try await client.send(
            input: input,
            forOperation: Operations.ExtractResumeStream.id,
            serializer: { input in
                let path = try converter.renderedPath(
                    template: "/stream/ExtractResume",
                    parameters: []
                )
                var request: HTTPTypes.HTTPRequest = .init(
                    soar_path: path,
                    method: .post
                )
                suppressMutabilityWarning(&request)
                converter.setAcceptHeader(
                    in: &request.headerFields,
                    contentTypes: input.headers.accept
                )
                let body: OpenAPIRuntime.HTTPBody?
                switch input.body {
                case let .json(value):
                    body = try converter.setRequiredRequestBodyAsJSON(
                        value,
                        headerFields: &request.headerFields,
                        contentType: "application/json; charset=utf-8"
                    )
                }
                return (request, body)
            },
            deserializer: { response, responseBody in
                switch response.status.code {
                case 200:
                    let contentType = converter.extractContentTypeIfPresent(in: response.headerFields)
                    let body: Operations.ExtractResumeStream.Output.Ok.Body
                    let chosenContentType = try converter.bestContentType(
                        received: contentType,
                        options: [
                            "text/event-stream"
                        ]
                    )
                    switch chosenContentType {
                    case "text/event-stream":
                        body = try converter.getResponseBodyAsBinary(
                            OpenAPIRuntime.HTTPBody.self,
                            from: responseBody,
                            transforming: { value in
                                .textEventStream(value)
                            }
                        )
                    default:
                        preconditionFailure("bestContentType chose an invalid content type.")
                    }
                    return .ok(.init(body: body))
                default:
                    return .undocumented(
                        statusCode: response.status.code,
                        .init(
                            headerFields: response.headerFields,
                            body: responseBody
                        )
                    )
                }
            }
        )
    }
}