<commit_before>package step.plugins.measurements;

import io.prometheus.client.Collector;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import step.controller.grid.GridPlugin;
import step.core.GlobalContext;
import step.core.plugins.AbstractControllerPlugin;
import step.core.plugins.Plugin;
import step.engine.plugins.ExecutionEnginePlugin;
import step.grid.Grid;
import step.grid.GridReportBuilder;
import step.grid.TokenWrapperState;
import step.grid.reports.TokenGroupCapacity;

import java.util.Arrays;
import java.util.List;

import static step.plugins.measurements.MeasurementPlugin.*;

@Plugin(dependencies = GridPlugin.class)
public class MeasurementControllerPlugin extends AbstractControllerPlugin {

	private static Logger logger = LoggerFactory.getLogger(MeasurementControllerPlugin.class);

	public static String GridGaugeName = "step_grid_tokens";
	public static String ThreadgroupGaugeName = "step_threadgroup";

	GaugeCollectorRegistry gaugeCollectorRegistry;

	@Override
	public void executionControllerStart(GlobalContext context) throws Exception {
		super.executionControllerStart(context);

		initGaugeCollectorRegistry(context);
	}

	protected void initGaugeCollectorRegistry(GlobalContext context) {
		gaugeCollectorRegistry = GaugeCollectorRegistry.getInstance();
		//Register grid gauge metrics
		final String[] labels = {"agent_type","url","name"};
		if (context.get(Grid.class)!=null) {
			gaugeCollectorRegistry.registerCollector(GridGaugeName, new GaugeCollector(GridGaugeName,
					"step grid token usage and capacity", labels) {
				final List<String> groupBy = Arrays.asList(new String[]{"$agenttype", "url"});
				final GridReportBuilder gridReportBuilder = new GridReportBuilder((Grid) context.get(Grid.class));

				@Override
				public List<Collector.MetricFamilySamples> collect() {
					List<TokenGroupCapacity> usageByIdentity = gridReportBuilder.getUsageByIdentity(groupBy);
					for (TokenGroupCapacity tokenGroupCapacity : usageByIdentity) {
						Integer free = tokenGroupCapacity.getCountByState().get(TokenWrapperState.FREE);
						free = (free == null) ? 0 : free;
						int capacity = tokenGroupCapacity.getCapacity();
						getGauge().labels(tokenGroupCapacity.getKey().get("$agenttype"),
								tokenGroupCapacity.getKey().get("url"), "free").set(Double.valueOf(free));
						getGauge().labels(tokenGroupCapacity.getKey().get("$agenttype"),
								tokenGroupCapacity.getKey().get("url"), "capacity").set(Double.valueOf(capacity));
					}
					return getGauge().collect();
				}
			});
		} else {
			logger.warn("No grid instance found in context, the measurements of the grid token usage will be disabled");
		}<commit_msg>SED-711 Agents timeseries metrics are not updated<commit_after>package step.plugins.measurements;

import io.prometheus.client.Collector;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import step.controller.grid.GridPlugin;
import step.core.GlobalContext;
import step.core.plugins.AbstractControllerPlugin;
import step.core.plugins.Plugin;
import step.engine.plugins.ExecutionEnginePlugin;
import step.grid.Grid;
import step.grid.GridReportBuilder;
import step.grid.TokenWrapperState;
import step.grid.reports.TokenGroupCapacity;

import java.util.*;

import static step.plugins.measurements.MeasurementPlugin.*;

@Plugin(dependencies = GridPlugin.class)
public class MeasurementControllerPlugin extends AbstractControllerPlugin {

	private static Logger logger = LoggerFactory.getLogger(MeasurementControllerPlugin.class);

	public static String GridGaugeName = "step_grid_tokens";
	public static String ThreadgroupGaugeName = "step_threadgroup";

	GaugeCollectorRegistry gaugeCollectorRegistry;

	@Override
	public void executionControllerStart(GlobalContext context) throws Exception {
		super.executionControllerStart(context);

		initGaugeCollectorRegistry(context);
	}

	protected void initGaugeCollectorRegistry(GlobalContext context) {
		gaugeCollectorRegistry = GaugeCollectorRegistry.getInstance();
		//Register grid gauge metrics
		final String[] labels = {"agent_type","url","name"};
		if (context.get(Grid.class)!=null) {
			gaugeCollectorRegistry.registerCollector(GridGaugeName, new GaugeCollector(GridGaugeName,
					"step grid token usage and capacity", labels) {
				final List<String> groupBy = Arrays.asList(new String[]{"$agenttype", "url"});
				final GridReportBuilder gridReportBuilder = new GridReportBuilder((Grid) context.get(Grid.class));
				final Map<String, String> urlToType = new HashMap<>();

				@Override
				public List<Collector.MetricFamilySamples> collect() {
					List<TokenGroupCapacity> usageByIdentity = gridReportBuilder.getUsageByIdentity(groupBy);
					Map<String, String> newUrlToType = new HashMap<>();
					for (TokenGroupCapacity tokenGroupCapacity : usageByIdentity) {
						Integer free = tokenGroupCapacity.getCountByState().get(TokenWrapperState.FREE);
						free = (free == null) ? 0 : free;
						int capacity = tokenGroupCapacity.getCapacity();
						String agentUrl = tokenGroupCapacity.getKey().get("url");
						String agentType = tokenGroupCapacity.getKey().get("$agenttype");
						newUrlToType.put(agentUrl, agentType);
						urlToType.remove(agentUrl);
						getGauge().labels(agentType,
								agentUrl, "free").set(Double.valueOf(free));
						getGauge().labels(agentType,
								agentUrl, "capacity").set(Double.valueOf(capacity));
					}
					for (String url : urlToType.keySet()) {
						String type = urlToType.remove(url);
						getGauge().remove(type, url,"free");
						getGauge().remove(type,url, "capacity");
					}
					urlToType.putAll(newUrlToType);
					return getGauge().collect();
				}
			});
		} else {
			logger.warn("No grid instance found in context, the measurements of the grid token usage will be disabled");
		}
{'fraction_non_alphanumeric': 0.10860196517841751, 'fraction_numerical': 0.0015514566454059645, 'mean_word_length': 5.497200447928331, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3879292', 'n_tokens_mistral': 1835, 'n_tokens_neox': 1570, 'n_words': 338}