private TokenUsage parseUsage()

in server/src/main/java/com/epam/aidial/core/server/token/TokenUsageParser.java [28:66]

32 lines of code
7 McCabe index (conditional complexity)


    private TokenUsage parseUsage(Buffer body) throws Exception {
        int index = findUsage(body);
        if (index < 0) {
            return null;
        }

        ByteBuf slice = body.slice(index, body.length()).getByteBuf();
        JsonFactory factory = new JsonFactory();

        try (InputStream stream = new ByteBufInputStream(slice); JsonParser parser = factory.createParser(stream)) {
            TokenUsage usage = new TokenUsage();
            verify(parser.nextToken(), JsonToken.START_OBJECT);

            while (true) {
                JsonToken token = parser.nextToken();
                if (token == JsonToken.END_OBJECT) {
                    return usage;
                }

                verify(token, JsonToken.FIELD_NAME);
                String name = parser.getCurrentName();

                token = parser.nextValue();
                verify(token,
                        JsonToken.VALUE_NUMBER_INT, JsonToken.VALUE_NUMBER_FLOAT,
                        JsonToken.VALUE_STRING, JsonToken.VALUE_NULL,
                        JsonToken.VALUE_FALSE, JsonToken.VALUE_TRUE);

                switch (name) {
                    case "completion_tokens" -> usage.setCompletionTokens(parser.getLongValue());
                    case "prompt_tokens" -> usage.setPromptTokens(parser.getLongValue());
                    case "total_tokens" -> usage.setTotalTokens(parser.getLongValue());
                    default -> {
                        // ignore
                    }
                }
            }
        }
    }