decode method

Stream<ChatResponseUpdate> decode(Stream<String> tokens)

Implementation

Stream<ChatResponseUpdate> decode(Stream<String> tokens) async* {
  // Hold back the longest-marker-minus-one trailing chars so a marker that
  // straddles two pieces is never emitted as content.
  final holdback =
      <int>[
        _channelOpen.length,
        _channelClose.length,
        _callOpen.length,
        _turnOpen.length,
      ].reduce((a, b) => a > b ? a : b) -
      1;

  var buf = '';
  final tail = StringBuffer();
  var thinking = false;
  var buffering = false;
  var labelPending = false;
  var turnLabelPending = false;

  await for (final piece in tokens) {
    if (buffering) {
      tail.write(piece);
      continue;
    }
    buf += piece;

    var progressed = true;
    while (progressed) {
      progressed = false;

      if (thinking) {
        // Drop the leading `thought` label once we have enough to see it
        // whole (or the channel has already closed).
        if (labelPending) {
          if (buf.length > _thoughtLabel.length ||
              buf.contains(_channelClose)) {
            buf = _stripLabel(buf);
            labelPending = false;
            progressed = true;
            continue;
          }
          break;
        }

        final closeAt = buf.indexOf(_channelClose);
        if (closeAt >= 0) {
          final reasoning = buf.substring(0, closeAt);
          if (reasoning.isNotEmpty) yield _reasoning(reasoning);
          buf = buf.substring(closeAt + _channelClose.length);
          thinking = false;
          progressed = true;
          continue;
        }
        if (buf.length > holdback) {
          final emit = buf.substring(0, buf.length - holdback);
          if (emit.isNotEmpty) yield _reasoning(emit);
          buf = buf.substring(buf.length - holdback);
        }
        break;
      }

      // A stray `<|turn>role` header is dropped through its newline.
      if (turnLabelPending) {
        final newline = buf.indexOf('\n');
        if (newline < 0) break;
        buf = buf.substring(newline + 1);
        turnLabelPending = false;
        progressed = true;
        continue;
      }

      // Act on whichever marker appears first; the rest of the buffer is
      // re-scanned after the state change.
      var at = -1;
      var marker = '';
      void consider(int index, String m) {
        if (index >= 0 && (at < 0 || index < at)) {
          at = index;
          marker = m;
        }
      }

      consider(buf.indexOf(_callOpen), _callOpen);
      consider(buf.indexOf(_channelOpen), _channelOpen);
      consider(buf.indexOf(_channelClose), _channelClose);
      consider(buf.indexOf(_turnOpen), _turnOpen);

      if (at >= 0) {
        final prose = buf.substring(0, at);
        if (prose.isNotEmpty) yield _text(prose);
        if (marker == _callOpen) {
          tail.write(buf.substring(at));
          buf = '';
          buffering = true;
          break;
        }
        buf = buf.substring(at + marker.length);
        if (marker == _channelOpen) {
          thinking = true;
          labelPending = true;
        } else if (marker == _turnOpen) {
          turnLabelPending = true;
        }
        // An unmatched `<channel|>` close needs no state: it is dropped.
        progressed = true;
        continue;
      }
      if (buf.length > holdback) {
        final emit = buf.substring(0, buf.length - holdback);
        if (emit.isNotEmpty) yield _text(emit);
        buf = buf.substring(buf.length - holdback);
      }
      break;
    }
  }

  if (buffering) {
    // A truncated or malformed tool call (e.g. the run hit maxTokens
    // mid-call) must not error the whole turn; surface the raw tail as
    // text instead so the user sees what the model produced.
    GemmaTurn turn;
    try {
      turn = template.parse(tail.toString());
    } on FormatException {
      yield _text(tail.toString());
      return;
    }
    if (turn.text.isNotEmpty) yield _text(turn.text);
    if (turn.calls.isNotEmpty) {
      yield ChatResponseUpdate(
        role: ChatRole.assistant,
        contents: List<AIContent>.of(turn.calls),
      );
    }
  } else if (thinking) {
    if (labelPending) buf = _stripLabel(buf);
    if (buf.isNotEmpty) yield _reasoning(buf);
  } else if (turnLabelPending) {
    // The stream ended inside a stray `<|turn>role` header; drop it.
  } else if (buf.isNotEmpty) {
    yield _text(buf);
  }
}