@genkit-ai/telemetry-server
Version:
Genkit AI telemetry server
207 lines • 8.35 kB
JavaScript
;
var __importDefault = (this && this.__importDefault) || function (mod) {
return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.stopTelemetryApi = exports.startTelemetryServer = exports.TraceQuerySchema = exports.LocalFileTraceStore = void 0;
const tools_common_1 = require("@genkit-ai/tools-common");
const utils_1 = require("@genkit-ai/tools-common/utils");
const express_1 = __importDefault(require("express"));
const broadcast_manager_js_1 = require("./broadcast-manager.js");
const otlp_1 = require("./utils/otlp");
var file_trace_store_js_1 = require("./file-trace-store.js");
Object.defineProperty(exports, "LocalFileTraceStore", { enumerable: true, get: function () { return file_trace_store_js_1.LocalFileTraceStore; } });
var types_1 = require("./types");
Object.defineProperty(exports, "TraceQuerySchema", { enumerable: true, get: function () { return types_1.TraceQuerySchema; } });
let server;
const broadcastManager = new broadcast_manager_js_1.BroadcastManager();
async function startTelemetryServer(params) {
await params.traceStore.init();
const api = (0, express_1.default)();
api.use(express_1.default.json({ limit: params.maxRequestBodySize ?? '100mb' }));
api.get('/api/__health', async (_, response) => {
response.status(200).send('OK');
});
api.get('/api/traces/:traceId', async (request, response, next) => {
try {
const { traceId } = request.params;
response.json(await params.traceStore.load(traceId));
}
catch (e) {
next(e);
}
});
api.get('/api/traces/:traceId/stream', async (request, response, next) => {
try {
const { traceId } = request.params;
response.setHeader('Content-Type', 'text/event-stream');
response.setHeader('Cache-Control', 'no-cache');
response.setHeader('Connection', 'keep-alive');
response.setHeader('Access-Control-Allow-Origin', '*');
response.setHeader('Access-Control-Allow-Headers', 'Content-Type');
const currentTrace = await params.traceStore.load(traceId);
if (currentTrace) {
const snapshot = JSON.stringify(currentTrace);
response.write(`data: ${snapshot}\n\n`);
}
broadcastManager.subscribe(traceId, response);
response.on('close', () => {
broadcastManager.unsubscribe(traceId, response);
});
}
catch (e) {
next(e);
}
});
api.post('/api/traces', async (request, response, next) => {
try {
const traceData = tools_common_1.TraceDataSchema.parse(request.body);
await params.traceStore.save(traceData.traceId, traceData);
const allSpans = Object.values(traceData.spans);
const events = [];
for (const span of allSpans) {
events.push({
type: 'span_start',
traceId: traceData.traceId,
span,
});
if (span.endTime > 0) {
events.push({
type: 'span_end',
traceId: traceData.traceId,
span,
});
}
}
events.sort((a, b) => {
const aTime = a.type === 'span_start' ? a.span.startTime : a.span.endTime;
const bTime = b.type === 'span_start' ? b.span.startTime : b.span.endTime;
if (aTime !== bTime) {
return aTime - bTime;
}
return a.type === 'span_start' ? -1 : 1;
});
for (const event of events) {
broadcastManager.broadcast(traceData.traceId, event);
}
response.status(200).send('OK');
}
catch (e) {
next(e);
}
});
api.get('/api/traces', async (request, response, next) => {
try {
const { limit, continuationToken, filter } = request.query;
response.json(await params.traceStore.list({
limit: limit ? Number.parseInt(limit.toString()) : 10,
continuationToken: continuationToken
? continuationToken.toString()
: undefined,
filter: filter
? tools_common_1.TraceQueryFilterSchema.parse(JSON.parse(filter))
: undefined,
}));
}
catch (e) {
next(e);
}
});
api.post('/api/otlp/:parentTraceId/:parentSpanId', async (request, response) => {
try {
const { parentTraceId, parentSpanId } = request.params;
if (!request.body.resourceSpans?.length) {
response.status(200).json({});
return;
}
const traces = (0, otlp_1.traceDataFromOtlp)(request.body);
for (const traceData of traces) {
traceData.traceId = parentTraceId;
for (const span of Object.values(traceData.spans)) {
span.attributes['genkit:otlp-traceId'] = span.traceId;
span.traceId = parentTraceId;
if (!span.parentSpanId) {
span.parentSpanId = parentSpanId;
}
}
await params.traceStore.save(parentTraceId, traceData);
}
response.status(200).json({});
}
catch (err) {
utils_1.logger.error(`Error processing OTLP payload: ${err}`);
response.status(500).json({
code: 13,
message: 'An internal error occurred while processing the OTLP payload.',
});
}
});
api.post('/api/otlp', async (request, response) => {
try {
if (!request.body.resourceSpans?.length) {
response.status(200).json({});
return;
}
const traces = (0, otlp_1.traceDataFromOtlp)(request.body);
for (const trace of traces) {
const traceData = tools_common_1.TraceDataSchema.parse(trace);
await params.traceStore.save(traceData.traceId, traceData);
for (const [_, span] of Object.entries(traceData.spans)) {
const event = {
type: span.endTime > 0 ? 'span_end' : 'span_start',
traceId: traceData.traceId,
span,
};
broadcastManager.broadcast(traceData.traceId, event);
}
}
response.status(200).json({});
}
catch (err) {
utils_1.logger.error(`Error processing OTLP payload: ${err}`);
response.status(500).json({
code: 13,
message: 'An internal error occurred while processing the OTLP payload.',
});
}
});
api.use((err, req, res, next) => {
utils_1.logger.error(err.stack);
const error = err;
const { message, stack } = error;
const errorResponse = {
code: 13,
message,
details: {
stack,
traceId: err.traceId,
},
};
res.status(500).json(errorResponse);
});
server = api.listen(params.port, () => {
utils_1.logger.info(`Telemetry API running on http://localhost:${params.port}`);
});
server.on('error', (error) => {
utils_1.logger.error(error);
});
process.on('SIGTERM', async () => await stopTelemetryApi());
}
exports.startTelemetryServer = startTelemetryServer;
async function stopTelemetryApi() {
await Promise.all([
new Promise((resolve) => {
if (server) {
server.close(() => {
utils_1.logger.debug('Telemetry API has succesfully shut down.');
resolve();
});
}
else {
resolve();
}
}),
]);
}
exports.stopTelemetryApi = stopTelemetryApi;
//# sourceMappingURL=index.js.map