halkyon/packages/halkyon_llm/tool/record_llm.dart

// Records real replies of a provider beside the documented ones, which the
// tests then run on too (ADR 0021, D36):
//
//   cd packages/halkyon_llm
//   dart run tool/record_llm.dart --protocol anthropic
//   dart run tool/record_llm.dart --protocol openai --model <a small model>
//   dart run tool/record_llm.dart --protocol openai-responses --model <a small model>
//
// The key comes from LLM_API_KEY, or is asked for without echo; it is never
// written anywhere. Four requests, each a few hundred tokens, cost a few
// cents at most: a short answer; a tool call; the tool's result sent back
// with the model's whole turn (its thinking blocks too), as the app's loop
// does; and a wrong key. Each reply goes to
// test/fixtures/recorded/<protocol>[@<host>]/<scenario>.http, its status,
// content type and body as they came, no request header. Read them before
// committing them.

import 'dart:async';
import 'dart:convert';
import 'dart:io';

import 'package:args/args.dart';
import 'package:halkyon_llm/halkyon_llm.dart';
import 'package:http/http.dart' as http;

import '../test/support/fixture.dart';

/// The tool the recordings ask the model to call, as the tests expect it.
const _findCards = LlmTool(
  name: 'find_cards',
  description:
      "Finds cards in the person's studyset whose front or back contains "
      'the query.',
  inputSchema: {
    'type': 'object',
    'properties': {
      'query': {'type': 'string', 'description': 'What to look for.'},
      'limit': {'type': 'integer', 'description': 'How many cards at most.'},
    },
    'required': ['query'],
  },
);

const _system =
    'You help someone learn Japanese with flashcards. Answer in one short '
    'sentence.';

Future<void> main(List<String> arguments) async {
  final options = ArgParser()
    ..addOption(
      'protocol',
      allowed: [for (final protocol in LlmProtocol.values) protocol.key],
      mandatory: true,
    )
    ..addOption('model', help: 'Anthropic: claude-haiku-5-5 unless given.')
    ..addOption('base-url', help: "The protocol's own provider unless given.")
    ..addFlag('help', abbr: 'h', negatable: false);
  final ArgResults given;
  try {
    given = options.parse(arguments);
  } on FormatException catch (error) {
    stderr
      ..writeln(error.message)
      ..writeln(options.usage);
    exitCode = 64;
    return;
  }
  if (given.flag('help')) {
    stdout.writeln(options.usage);
    return;
  }
  final protocol = LlmProtocol.named(given.option('protocol')!)!;
  final model =
      given.option('model') ??
      (protocol == LlmProtocol.anthropic ? 'claude-haiku-5-5' : null);
  if (model == null) {
    stderr.writeln('Name the model with --model.');
    exitCode = 64;
    return;
  }
  final key = Platform.environment['LLM_API_KEY'] ?? _askKey();
  if (key.isEmpty) {
    stderr.writeln('No key.');
    exitCode = 64;
    return;
  }
  final endpoint = LlmEndpoint(
    protocol: protocol,
    model: model,
    apiKey: key,
    baseUrl: given.option('base-url'),
  );
  final folder = Directory(
    'test/fixtures/recorded/${protocol.key}'
    '${given.option('base-url') == null ? '' : '@${endpoint.host}'}',
  )..createSync(recursive: true);

  Future<LlmFinished?> record(
    String scenario,
    LlmRequest request, {
    LlmEndpoint? at,
  }) async {
    final recorder = _Recorder(File('${folder.path}/$scenario.http'));
    final client = LlmClient(httpClient: recorder);
    try {
      final finished = await client
          .send(at ?? endpoint, request)
          .where((event) => event is LlmFinished)
          .cast<LlmFinished>()
          .single;
      stdout.writeln('$scenario: ${finished.stop.name}, ${finished.usage}');
      return finished;
    } on LlmFailure catch (failure) {
      stdout.writeln('$scenario: $failure');
      return null;
    } finally {
      if (recorder.answered) {
        await recorder.written;
      }
      client.close();
      recorder.close();
    }
  }

  await record(
    'text',
    LlmRequest(
      system: _system,
      messages: [
        LlmMessage.user([const LlmText('What is "dog" in Japanese?')]),
      ],
      maxOutputTokens: 2048,
      personId: 'halkyon-record-tool',
    ),
  );
  final question = LlmMessage.user([
    const LlmText('Find up to three cards about "dog" in my studyset.'),
  ]);
  final tools = LlmRequest(
    system: '$_system Look things up in the studyset with find_cards.',
    messages: [question],
    tools: const [_findCards],
    maxOutputTokens: 2048,
    personId: 'halkyon-record-tool',
  );
  final called = await record('tool', tools);
  final calls = called?.message.toolCalls ?? const [];
  if (calls.isEmpty) {
    stdout.writeln('tool_result: skipped, as the model called no tool.');
  } else {
    await record(
      'tool_result',
      LlmRequest(
        system: tools.system,
        messages: [
          question,
          called!.message,
          LlmMessage.user([
            for (final call in calls)
              LlmToolResult(
                callId: call.id,
                content: jsonEncode([
                  {'front': '犬', 'back': 'dog (いぬ, inu)'},
                ]),
              ),
          ]),
        ],
        tools: tools.tools,
        maxOutputTokens: 2048,
        personId: 'halkyon-record-tool',
      ),
    );
  }
  await record(
    'bad_key',
    LlmRequest(
      messages: [
        LlmMessage.user([const LlmText('Hello')]),
      ],
      maxOutputTokens: 16,
    ),
    at: LlmEndpoint(
      protocol: protocol,
      model: model,
      apiKey: 'sk-halkyon-wrong-key',
      baseUrl: given.option('base-url'),
    ),
  );
  stdout.writeln('Written to ${folder.path}. Now: dart test');
}

String _askKey() {
  stdout.write('API key (not shown): ');
  final echo = stdin.hasTerminal && stdin.echoMode;
  if (stdin.hasTerminal) {
    stdin.echoMode = false;
  }
  try {
    return stdin.readLineSync()?.trim() ?? '';
  } finally {
    if (stdin.hasTerminal) {
      stdin.echoMode = echo;
      stdout.writeln();
    }
  }
}

/// An HTTP client writing each reply to [file] as it streams through.
final class _Recorder(final File file) extends http.BaseClient {
  final _inner = http.Client();
  final _written = Completer<void>();

  /// Done once the reply is written.
  Future<void> get written => _written.future;

  /// Whether a reply began, which is then written.
  bool answered = false;

  @override
  Future<http.StreamedResponse> send(http.BaseRequest request) async {
    final response = await _inner.send(request);
    answered = true;
    final body = <int>[];
    late final StreamSubscription<List<int>> reading;
    final out = StreamController<List<int>>();
    out
      ..onListen = () {
        reading = response.stream.listen(
          (chunk) {
            body.addAll(chunk);
            out.add(chunk);
          },
          onError: out.addError,
          onDone: () {
            _write(response, body);
            unawaited(out.close());
          },
        );
      }
      // The adapter stops reading once the answer is over.
      ..onCancel = () async {
        await reading.cancel();
        _write(response, body);
      };
    return http.StreamedResponse(
      out.stream,
      response.statusCode,
      headers: response.headers,
      request: request,
    );
  }

  void _write(http.StreamedResponse response, List<int> body) {
    if (_written.isCompleted) {
      return;
    }
    file.writeAsBytesSync(
      Fixture.write(response.statusCode, {
        for (final name in ['content-type', 'retry-after'])
          name: ?response.headers[name],
      }, body),
    );
    _written.complete();
  }

  @override
  void close() => _inner.close();
}