halkyon/packages/halkyon_llm/lib/testing.dart

/// An [LlmClient] for the tests of what uses this package: it answers from a
/// script and keeps what it was sent.
library;

import 'dart:async';
import 'dart:collection';

import 'package:halkyon_llm/halkyon_llm.dart';
import 'package:halkyon_llm/src/client.dart' show checkEffort;

/// An [LlmClient] answering each request with the next reply scripted.
final class FakeLlmClient() implements LlmClient {
  final _replies = Queue<_Reply>();

  /// What it was sent, in order.
  final List<(LlmEndpoint, LlmRequest)> sent = [];

  /// How many answers were cancelled before they ended.
  int cancelled = 0;

  /// Whether [close] was called.
  bool closed = false;

  /// Scripts the next reply: [events], then [failure] if given.
  void reply(List<LlmEvent> events, {LlmFailure? failure}) =>
      _replies.add(_Reply(events, failure));

  /// Scripts an answer of [text], in two pieces, counted as [usage].
  void replyText(
    String text, {
    LlmUsage usage = const LlmUsage(inputTokens: 100, outputTokens: 20),
  }) {
    final half = text.length ~/ 2;
    reply([
      LlmTextDelta(text.substring(0, half)),
      LlmTextDelta(text.substring(half)),
      LlmFinished(
        message: LlmMessage.assistant([LlmText(text)]),
        stop: LlmStop.end,
        usage: usage,
      ),
    ]);
  }

  /// Scripts a failure of [kind] before anything is written, with [usage]
  /// counted.
  void fail(LlmFailureKind kind, {LlmUsage usage = LlmUsage.none}) => reply(
    const [],
    failure: LlmFailure(kind: kind, detail: 'scripted', usage: usage),
  );

  /// Whether every reply scripted was given.
  bool get allReplied => _replies.isEmpty;

  @override
  Stream<LlmEvent> send(LlmEndpoint endpoint, LlmRequest request) {
    checkEffort(endpoint, request);
    sent.add((endpoint, request));
    if (_replies.isEmpty) {
      throw StateError('No reply scripted for request ${sent.length}');
    }
    final reply = _replies.removeFirst();
    var ended = false;
    late final StreamController<LlmEvent> events;
    events = StreamController(
      onListen: () async {
        for (final event in reply.events) {
          if (events.isClosed) {
            return;
          }
          events.add(event);
          await Future<void>.delayed(Duration.zero);
        }
        if (reply.failure case final failure?) {
          events.addError(failure);
        }
        ended = true;
        await events.close();
      },
      onCancel: () {
        if (!ended) {
          cancelled++;
        }
      },
    );
    return events.stream;
  }

  @override
  void close() => closed = true;
}

final class _Reply(final List<LlmEvent> events, final LlmFailure? failure);