shelly-shell 0.5.4

# Responses client. The model remains a caller choice; no live requests occur on import.
import std::io::{ Bytes }
import std::json::{ Codec, Document, StructOptions, MappingOptions, UnknownFields,
    enum_codec, struct_codec, convert, wrap, encode, set }
import std::stream::{ Reader }
import std::net::http::{ StatusError, json_body, bearer }
import std::net::sse::{ read }

pub enum Role
{
    Developer, User, Assistant
}

pub fn Role::json_codec(): Codec
{
    enum_codec Role [
        Role::Developer: "developer",
        Role::User: "user",
        Role::Assistant: "assistant",
    ]
}

pub struct Message
{
    pub role: Role,
    pub content: String,
}

pub struct Request
{
    pub model: String,
    pub input: String | [Message],
    pub instructions: String | (),
    pub max_output_tokens: Integer | (),
    pub store: Boolean | (),
}

pub fn Request::json_codec(): Codec
{
    struct_codec Request StructOptions(
        omit_null_fields: ["instructions", "max_output_tokens", "store"],
    )
}

pub struct Usage
{
    pub input_tokens: Integer,
    pub output_tokens: Integer,
    pub total_tokens: Integer,
}

struct Envelope
{
    pub id: String,
    pub status: String,
    pub usage: Usage | (),
}

pub struct Response
{
    pub id: String,
    pub status: String,
    pub usage: Usage | (),
    pub request_id: String | (),
    pub raw: Document,
}

pub fn Response::from_json($document: Document): Response | std::json::Error
{
    let $options = MappingOptions(unknown_fields: UnknownFields::Ignore)
    let $envelope: Envelope = convert $document Envelope $options

    Response(
        id: $envelope.id,
        status: $envelope.status,
        usage: $envelope.usage,
        raw: $document,
    )
}

pub fn Response::to_json(): Document
{
    $self.raw
}

pub enum ResponseFailure
{
    Incomplete, Failed, Refused, NoText, Protocol, Truncated
}

pub error struct ApiError
{
    pub status: Integer,
    pub message: String,
    pub request_id: String | (),
    pub response: Document,
}

pub error struct ResponseError
{
    pub what: ResponseFailure,
    pub message: String,
    pub request_id: String | (),
    pub response: Document | (),
}

type NetworkError = std::net::http::Error | std::net::sse::Error
pub type Error = ApiError | ResponseError | NetworkError | std::json::Error | std::stream::Error

pub fn String::from_api_error($error: ApiError): String
{
    "OpenAI HTTP ${error.status}: ${error.message}"
}

pub fn String::from_response_error($error: ResponseError): String
{
    "OpenAI ${error.what}: ${error.message}"
}

fn Response::failure($what: ResponseFailure, $message: String): ResponseError
{
    ResponseError(what: $what, message: $message,
        request_id: $self.request_id, response: $self.raw)
}

pub fn Response::require_text(): String | ResponseError | std::json::Error
{
    if $self.status != "completed"
    {
        let $what = if $self.status == "incomplete"
            {
                ResponseFailure::Incomplete
            }
            else
            {
                ResponseFailure::Failed
            }

        return ($self.failure $what "Response did not complete")
    }

    let $text = ""
    for $item in $self.raw.array_items "/output"
    {
        if ($item.require_string "/type") != "message"
        {
            continue
        }

        for $part in $item.array_items "/content"
        {
            match ($part.require_string "/type")
            {
                "output_text" => { $text = $text + ($part.require_string "/text") },
                "refusal" =>
                {
                    let $message = $part.require_string "/refusal"
                    return ($self.failure ResponseFailure::Refused $message)
                },
                _ => {},
            }
        }
    }

    if $text == ""
    {
        return ($self.failure ResponseFailure::NoText "Response contains no text output")
    }

    $text
}

fn ApiError::from_status_error(
    $error: StatusError): ApiError | std::json::Error | std::net::http::Error
{
    let $reply = $error.response
    let $document = $reply.read_document

    ApiError(
        status: $reply.status,
        message: ($document.require_string "/error/message"),
        request_id: ($reply.headers.get "x-request-id"),
        response: $document,
    )
}

pub struct Client
{
    transport: std::net::http::Client,
    model: String,
}

pub struct Model
{
    pub id: String,
    pub created: Integer | (),
    pub owned_by: String | (),
}

struct ModelList
{
    pub data: [Model]
}

pub fn Client::models(): [Model] | Error
{
    let $reply = $self.::transport.get "models"
    let $options = MappingOptions(unknown_fields: UnknownFields::Ignore)
    let $list: ModelList = $reply.require_success.read_as ModelList $options
    $list.data
}

pub fn Client::new($api_key: String, $model: String, $base_url: String | ()): Client | Error
{
    let $url = if $base_url == ()
    {
        "https://api.openai.com/v1/"
    }
    else
    {
        $base_url
    }

    Client(
        transport: std::net::http::Client(
            base_url: $url,
            auth: (bearer $api_key),
            timeout_ms: 120000,
            error_decoder: ApiError,
        ),
        model: $model,
    )
}

pub fn Client::create($request: Request): Response | Error
{
    # Bind this module's codecs before the generic HTTP body helper changes module scope.
    let $body = wrap $request
    let $wire = json_body $body
    let $reply: std::net::http::Response = $self.::transport.post "responses" $wire

    let $response: Response = $reply.require_success.read_as Response
    $response.request_id = $reply.headers.get "x-request-id"
    $response
}

pub fn Client::text($input: String, $instructions: String | ()): String | Error
{
    let $response: Response = $self.create Request(
        model: $self.::model,
        input: $input,
        instructions: $instructions,
        store: false,
    )

    $response.require_text
}

pub fn Client::close(): () | Error
{
    $self.::transport.close
}

pub struct Event
{
    pub kind: String,
    pub sequence: Integer,
    pub raw: Document,
}

pub fn Event::from_json($document: Document): Event | std::json::Error
{
    Event(kind: ($document.require_string "/type"),
        sequence: ($document.require_integer "/sequence_number"), raw: $document)
}

# Useful for scripted transports and deterministic tests as well as HTTP.
pub fn text_stream($events: Reader, $request_id: String | ()): Reader | Error
{
    let $session = TextSession(events: $events, request_id: $request_id)
    Reader(
        next: `$session.next,
        result: `$session.result,
        close: `$events.close,
    )
}

fn Client::stream_response($request: Request): std::net::http::Response | Error
{
    let $body = set (wrap $request) "/stream" true
    let $reply = $self.::transport.post "responses" (json_body $body)
    $reply.require_success
}

pub fn Client::events($request: Request): Reader | Error
{
    let $reply = $self.stream_response $request
    (read $reply).read_as Event
}

pub fn Client::stream_text($request: Request): Reader | Error
{
    let $reply = $self.stream_response $request
    let $events = (read $reply).read_as Event
    text_stream $events ($reply.headers.get "x-request-id")
}

enum TextPhase
{
    Waiting, Streaming, TextDone, PartDone, ItemDone
}
struct TextChannel
{
    pub item_id: String,
    pub output_index: Integer,
    pub content_index: Integer,
}

# A stream owns its source and validates every transition before emitting more text.
struct TextSession
{
    events: Reader,
    request_id: String | (),
    response_id: String | (),
    sequence: Integer,
    event_count: Integer,
    encoded_bytes: Integer,
    items: [Integer: Document],
    done_items: [Integer: Boolean],
    phase: TextPhase,
    channel: TextChannel | (),
    text: String,
    completed: Response | (),
}

fn TextSession::new($events: Reader, $request_id: String | ()): TextSession
{
    TextSession(events: $events, request_id: $request_id, sequence: -1,
        event_count: 0, encoded_bytes: 0, items: [:], done_items: [:],
        phase: TextPhase::Waiting, text: "")
}

fn TextSession::failure($what: ResponseFailure, $message: String,
                        $evidence: Document | ()): ResponseError
{
    ResponseError(what: $what, message: $message,
        request_id: $self.::request_id, response: $evidence)
}

fn TextSession::protocol_error($message: String, $evidence: Document | ()): ResponseError
{
    $self.failure ResponseFailure::Protocol $message $evidence
}

fn TextSession::check_channel($event: Document): () | Error
{
    if $self.::channel == ()
    {
        return ($self.protocol_error "Text channel does not match" $event)
    }
    let $channel = $self.::channel
    if ($event.require_string "/item_id") != $channel.item_id ||
        ($event.require_integer "/output_index") != $channel.output_index ||
        ($event.require_integer "/content_index") != $channel.content_index
    {
        return ($self.protocol_error "Text channel does not match" $event)
    }
}

fn TextSession::check_event($event: Event): () | Error
{
    let $doc = $event.raw
    $self.::event_count = $self.::event_count + 1
    $self.::encoded_bytes = $self.::encoded_bytes + Bytes(encode $doc).count
    if ($self.::event_count > 32768) || ($self.::encoded_bytes > 67108864)
    {
        return ($self.protocol_error "Stream exceeds event budget" $doc)
    }
    if ($event.sequence < 0) || ($event.sequence <= $self.::sequence)
    {
        return ($self.protocol_error "Event sequence is not increasing" $doc)
    }
    $self.::sequence = $event.sequence
}

fn TextSession::start_response($doc: Document): () | Error
{
    let $id = $doc.require_string "/response/id"
    if ($self.::response_id != ()) || ($id == "")
    {
        return ($self.protocol_error "Invalid response creation" $doc)
    }
    $self.::response_id = $id
}

fn TextSession::register_item($doc: Document): () | Error
{
    let $index = $doc.require_integer "/output_index"
    let $item = $doc.at "/item"
    let $id = $item.require_string "/id"
    if ($index < 0) || ($id == "") || ($self.::items.count >= 256) ||
        ($self.::items[$index] != ())
    {
        return ($self.protocol_error "Invalid output item registration" $doc)
    }
    for (_, $existing) in $self.::items
    {
        if ($existing.require_string "/id") == $id
        {
            return ($self.protocol_error "Duplicate output item ID" $doc)
        }
    }
    $self.::items[$index] = $item
}

fn TextSession::register_content($doc: Document): () | Error
{
    let $index = $doc.require_integer "/output_index"
    let $content = $doc.require_integer "/content_index"
    let $id = $doc.require_string "/item_id"
    let $item = $self.::items[$index]
    if ($index < 0) || ($content < 0) || ($item == ()) || $self.::done_items[$index] == true ||
        ($item.require_string "/id") != $id
    {
        return ($self.protocol_error "Invalid content part identity" $doc)
    }
    if ($doc.require_string "/part/type") != "output_text"
    {
        return ()
    }
    if $self.::phase != TextPhase::Waiting || ($index < 0) || ($content < 0) || ($item == ())
    {
        return ($self.protocol_error "Invalid text channel registration" $doc)
    }
    if ($item.require_string "/id") != $id ||
        ($item.require_string "/type") != "message" ||
        ($doc.require_string "/part/text") != ""
    {
        return ($self.protocol_error "Invalid initial text part" $doc)
    }
    $self.::channel = TextChannel(item_id: $id,
        output_index: $index, content_index: $content)
    $self.::phase = TextPhase::Streaming
}

fn TextSession::accept_delta($doc: Document): String | Error
{
    $self.check_channel $doc
    if $self.::phase != TextPhase::Streaming
    {
        return ($self.protocol_error "Delta follows text completion" $doc)
    }
    let $delta = $doc.require_string "/delta"
    if Bytes($self.::text).count + Bytes($delta).count > 1048576
    {
        return ($self.protocol_error "Text exceeds 1 MiB" $doc)
    }
    $self.::text = $self.::text + $delta
    return $delta
}

fn TextSession::finish_text($doc: Document): () | Error
{
    $self.check_channel $doc
    if $self.::phase != TextPhase::Streaming || ($doc.require_string "/text") != $self.::text
    {
        return ($self.protocol_error "Final text differs from deltas" $doc)
    }
    $self.::phase = TextPhase::TextDone
}

fn TextSession::finish_content($doc: Document): () | Error
{
    if ($doc.require_string "/part/type") != "output_text"
    {
        return ()
    }
    $self.check_channel $doc
    if $self.::phase != TextPhase::TextDone ||
        ($doc.require_string "/part/text") != $self.::text
    {
        return ($self.protocol_error "Text part completion is inconsistent" $doc)
    }
    $self.::phase = TextPhase::PartDone
}

fn TextSession::finish_item($doc: Document): () | Error
{
    let $channel = $self.::channel
    let $index = $doc.require_integer "/output_index"
    let $item = $self.::items[$index]
    if ($index < 0) || ($item == ()) || $self.::done_items[$index] == true
    {
        return ($self.protocol_error "Invalid output item completion" $doc)
    }
    if ($doc.require_string "/item/id") != ($item.require_string "/id") ||
        ($doc.require_string "/item/type") != ($item.require_string "/type")
    {
        return ($self.protocol_error "Output item identity changed" $doc)
    }
    $self.::done_items[$index] = true
    if ($channel != ()) && ($index == $channel.output_index)
    {
        if $self.::phase != TextPhase::PartDone ||
            ($doc.require_string "/item/id") != $channel.item_id ||
            ($doc.require_string "/item/type") != "message"
        {
            return ($self.protocol_error "Output item completion is inconsistent" $doc)
        }
        let $parts = ($doc.at "/item").array_items "/content"
        if $channel.content_index >= $parts.count ||
            ($parts[$channel.content_index].require_string "/type") != "output_text" ||
            ($parts[$channel.content_index].require_string "/text") != $self.::text
        {
            return ($self.protocol_error "Completed item differs from deltas" $doc)
        }
        $self.::phase = TextPhase::ItemDone
    }
}

fn TextSession::complete_response($doc: Document): () | Error
{
    let $channel = $self.::channel
    let $response = $doc.at "/response"
    if ($response.require_string "/status") != "completed" ||
        $self.::phase != TextPhase::ItemDone
    {
        return ($self.protocol_error "Premature or inconsistent completion" $doc)
    }
    let $items = $response.array_items "/output"
    if $channel.output_index >= $items.count
    {
        return ($self.protocol_error "Terminal output index changed" $doc)
    }
    let $item = $items[$channel.output_index]
    if    (($item.require_string "/id") != $channel.item_id)
       || (($item.require_string "/type") != "message")
    {
        return ($self.protocol_error "Terminal item identity changed" $doc)
    }
    let $parts = $item.array_items "/content"
    if $channel.content_index >= $parts.count ||
        ($parts[$channel.content_index].require_string "/type") != "output_text" ||
        ($parts[$channel.content_index].require_string "/text") != $self.::text
    {
        return ($self.protocol_error "Terminal content channel changed" $doc)
    }
    let $completed: Response = convert $response Response
    $completed.request_id = $self.::request_id
    if $completed.require_text != $self.::text
    {
        return ($self.protocol_error "Completed response differs from emitted text" $doc)
    }
    $self.::completed = $completed
}

fn TextSession::failed_response($doc: Document): ResponseError | Error
{
    let $response = $doc.at "/response"
    if ($response.require_string "/status") != "failed"
    {
        return ($self.protocol_error "Failed event has inconsistent status" $doc)
    }
    return ($self.failure ResponseFailure::Failed "Generation failed" $response)
}

fn TextSession::incomplete_response($doc: Document): ResponseError | Error
{
    let $response = $doc.at "/response"
    if ($response.require_string "/status") != "incomplete"
    {
        return ($self.protocol_error "Incomplete event has inconsistent status" $doc)
    }
    return ($self.failure ResponseFailure::Incomplete "Generation is incomplete" $response)
}

fn TextSession::accept($event: Event): String | () | Error
{
    $self.check_event $event
    let $doc = $event.raw
    if $event.kind == "response.created"
    {
        return ($self.start_response $doc)
    }
    if $event.kind == "error"
    {
        return ($self.failure ResponseFailure::Failed ($doc.require_string "/message") $doc)
    }
    if $self.::response_id == ()
    {
        return ($self.protocol_error "Event precedes response creation" $doc)
    }
    let $response = $doc.find "/response"
    if $response != () && ($response.require_string "/id") != $self.::response_id
    {
        return ($self.protocol_error "Response ID changed" $doc)
    }

    match $event.kind
    {
        "response.output_item.added" => { $self.register_item $doc },
        "response.content_part.added" => { $self.register_content $doc },
        "response.output_text.delta" => { $self.accept_delta $doc },
        "response.output_text.done" => { $self.finish_text $doc },
        "response.content_part.done" => { $self.finish_content $doc },
        "response.output_item.done" => { $self.finish_item $doc },
        "response.completed" => { $self.complete_response $doc },
        "response.failed" => { $self.failed_response $doc },
        "response.incomplete" => { $self.incomplete_response $doc },
        _ => { () },
    }
}

fn TextSession::next(): String | () | Error
{
    if $self.::completed != ()
    {
        return ()
    }
    loop
    {
        let $event = $self.::events.next_item
        if $event == ()
        {
            return ($self.failure ResponseFailure::Truncated "Stream ended before completion" ())
        }
        let $delta = $self.accept $event
        if $delta != ()
        {
            return $delta
        }
        if $self.::completed != ()
        {
            return ()
        }
    }
}

fn TextSession::result(): Response | Error
{
    if $self.::completed == ()
    {
        return ($self.protocol_error "Stream has not completed" ())
    }
    $self.::completed
}