@genkit-ai/telemetry-server
Version:
Genkit AI telemetry server
331 lines • 12.2 kB
JavaScript
import { TraceDataSchema, } from '@genkit-ai/tools-common';
import { logger } from '@genkit-ai/tools-common/utils';
import { Mutex } from 'async-mutex';
import fs from 'fs';
import lockfile from 'lockfile';
import path from 'path';
import { version as currentVersion } from './utils/version';
const MAX_TRACES = 1000;
const MAX_INDEX_FILES = 10;
const MAX_LIST_ATTR_LENGTH = 1000;
export class LocalFileTraceStore {
storeRoot;
indexRoot;
mutexes = {};
filters;
index;
static defaultFilters = {
'genkit:metadata:subtype': 'prompt',
};
constructor(options) {
this.storeRoot = path.resolve(options.storeRoot, `.genkit/traces`);
fs.mkdirSync(this.storeRoot, { recursive: true });
this.indexRoot = path.resolve(options.indexRoot, `.genkit/traces_idx`);
fs.mkdirSync(this.indexRoot, { recursive: true });
logger.debug(`[Telemetry Server] initialized local file trace store at root: ${this.storeRoot}`);
this.filters = options?.filters ?? LocalFileTraceStore.defaultFilters;
this.index = new Index(this.indexRoot);
}
async init() {
const metadata = this.index.getMetadata();
if (!metadata ||
metadata.version !== currentVersion ||
this.index.listIndexFiles().length > MAX_INDEX_FILES) {
await this.reIndex();
}
}
async reIndex() {
this.index.clear();
const time = Date.now();
const list = await this.listFromFiles({ limit: MAX_TRACES });
for (const trace of list.traces.reverse()) {
const hasRootSpan = !!Object.values(trace.spans).find((s) => !s.parentSpanId);
if (!hasRootSpan)
continue;
this.index.add(trace);
}
logger.info(`Indexed ${list.traces.length} traces in ${Date.now() - time}ms in ${this.indexRoot}`);
}
async load(id) {
const filePath = path.resolve(this.storeRoot, `${id}`);
if (!fs.existsSync(filePath)) {
return undefined;
}
const data = fs.readFileSync(filePath, 'utf8');
const parsed = JSON.parse(data);
if (!parsed.traceId) {
parsed.traceId = id;
}
return TraceDataSchema.parse(parsed);
}
getMutex(id) {
if (!this.mutexes[id]) {
this.mutexes[id] = new Mutex();
}
return this.mutexes[id];
}
async save(id, rawTrace) {
let trace = this.filter(rawTrace);
const possibleRoot = Object.keys(trace.spans).length === 0;
const mutex = this.getMutex(id);
await mutex.waitForUnlock();
const release = await mutex.acquire();
try {
const existing = (await this.load(id)) || trace;
if (existing) {
Object.keys(trace.spans).forEach((spanId) => {
if (existing.spans[spanId] &&
existing.spans[spanId].endTime &&
!trace.spans[spanId].endTime) {
return;
}
existing.spans[spanId] = trace.spans[spanId];
});
if (possibleRoot) {
Object.keys(existing.spans).forEach((spanId) => {
const span = existing.spans[spanId];
if (possibleRoot &&
span.parentSpanId &&
!existing.spans[span.parentSpanId]) {
delete span.parentSpanId;
}
});
}
if (trace.displayName !== undefined) {
existing.displayName = trace.displayName;
}
if (trace.startTime !== undefined) {
existing.startTime = trace.startTime;
}
if (trace.endTime !== undefined) {
existing.endTime = trace.endTime;
}
trace = existing;
}
fs.writeFileSync(path.resolve(this.storeRoot, `${id}`), JSON.stringify(trace));
const hasRootSpan = !!Object.values(rawTrace.spans).find((s) => !s.parentSpanId);
if (this.index && hasRootSpan) {
const fullTrace = await this.load(rawTrace.traceId);
if (!fullTrace) {
throw new Error('unable to read the trace that was just written... "this should never happen"');
}
this.index.add(fullTrace);
}
}
finally {
release();
}
}
async list(query) {
const searchResult = this.index.search({
limit: query?.limit ?? 10,
startFromIndex: query?.continuationToken
? Number.parseInt(query?.continuationToken)
: undefined,
filter: query?.filter,
});
const loadedTraces = await Promise.all(searchResult.data.map((d) => this.load(d['id']).then(trucateTraceDetails)));
return {
traces: loadedTraces.filter((t) => !!t),
continuationToken: searchResult.pageLastIndex
? `${searchResult.pageLastIndex}`
: undefined,
};
}
async listFromFiles(query) {
const files = fs.readdirSync(this.storeRoot);
files.sort((a, b) => {
return (fs.statSync(path.resolve(this.storeRoot, `${b}`)).mtime.getTime() -
fs.statSync(path.resolve(this.storeRoot, `${a}`)).mtime.getTime());
});
const startFrom = query?.continuationToken
? Number.parseInt(query?.continuationToken)
: 0;
const stopAt = startFrom + (query?.limit || 10);
const traces = files.slice(startFrom, stopAt).map((id) => {
const filePath = path.resolve(this.storeRoot, `${id}`);
const data = fs.readFileSync(filePath, 'utf8');
const parsed = JSON.parse(data);
if (!parsed.traceId) {
parsed.traceId = id;
}
return TraceDataSchema.parse(parsed);
});
return {
traces,
continuationToken: files.length > stopAt ? stopAt.toString() : undefined,
};
}
filter(trace) {
Object.keys(trace.spans).forEach((spanId) => {
const span = trace.spans[spanId];
Object.keys(this.filters).forEach((f) => {
if (span.attributes[f] === this.filters[f]) {
delete trace.spans[spanId];
}
});
});
if (Object.keys(trace.spans).length === 1) {
Object.keys(trace.spans).forEach((spanId) => {
const span = trace.spans[spanId];
if (span.attributes['genkit:name'] === 'dev-run-action-wrapper') {
delete trace.spans[spanId];
}
});
}
return trace;
}
}
function trucateTraceDetails(t) {
if (!t)
return t;
const { spans: originalSpans, ...restOfTrace } = t;
const spans = {};
for (const spanId of Object.keys(originalSpans)) {
if (!originalSpans[spanId].parentSpanId) {
const { attributes: originalAttributes, ...restOfSpan } = originalSpans[spanId];
spans[spanId] = {
attributes: trucateLargeAttrs(originalAttributes),
...restOfSpan,
};
break;
}
}
return { spans, ...restOfTrace };
}
export function trucateLargeAttrs(input) {
if (input === undefined ||
input === null ||
Array.isArray(input) ||
typeof input !== 'object') {
return input;
}
for (const key in input) {
if (typeof input[key] === 'string' &&
input[key].length > MAX_LIST_ATTR_LENGTH) {
input[key] = (input[key].substring(0, MAX_LIST_ATTR_LENGTH) +
'...');
}
}
return input;
}
export class Index {
indexRoot;
currentIndexFile;
constructor(indexRoot) {
this.indexRoot = indexRoot;
this.currentIndexFile = path.resolve(this.indexRoot, this.newIndexFileName());
fs.mkdirSync(this.indexRoot, { recursive: true });
}
clear() {
fs.rmSync(this.indexRoot, { recursive: true, force: true });
fs.mkdirSync(this.indexRoot, { recursive: true });
fs.appendFileSync(this.metadataFileName(), JSON.stringify({ version: currentVersion }));
}
metadataFileName() {
return path.resolve(this.indexRoot, 'genkit.metadata');
}
getMetadata() {
try {
return JSON.parse(fs.readFileSync(this.metadataFileName(), { encoding: 'utf8' }));
}
catch (e) {
return undefined;
}
}
newIndexFileName() {
return `idx_${(Date.now() + '').padStart(17, '0')}.json`;
}
listIndexFiles() {
return fs.readdirSync(this.indexRoot).filter((f) => f.startsWith('idx_'));
}
add(traceData) {
const rootSpans = Object.values(traceData.spans).filter((s) => !s.parentSpanId);
const rootSpan = rootSpans.length > 0 ? rootSpans[0] : undefined;
const indexData = {
id: traceData.traceId,
};
indexData['type'] =
`${rootSpan?.attributes?.['genkit:metadata:subtype'] || rootSpan?.attributes?.['genkit:type'] || 'UNKNOWN'}`;
if (rootSpan?.startTime) {
indexData['start'] = rootSpan.startTime;
}
if (rootSpan?.displayName) {
indexData['name'] = rootSpan.displayName;
}
if (rootSpan?.endTime) {
indexData['end'] = rootSpan.endTime;
}
if (rootSpan?.displayName) {
indexData['status'] = rootSpan.status?.code ?? 'UNKNOWN';
}
Object.keys(rootSpan?.attributes ?? {})
.filter((k) => k.startsWith('genkitx:'))
.forEach((k) => {
indexData[k] = `${rootSpan.attributes[k]}`;
});
try {
lockfile.lockSync(lockFile(this.currentIndexFile));
fs.appendFileSync(this.currentIndexFile, JSON.stringify(indexData) + '\n');
}
finally {
lockfile.unlockSync(lockFile(this.currentIndexFile));
}
}
search(query) {
const startFromIndex = query.startFromIndex ?? 0;
const fullData = [];
for (const idxFile of this.listIndexFiles()) {
const idxTxt = fs.readFileSync(path.resolve(this.indexRoot, idxFile), 'utf8');
const fileData = idxTxt
.split('\n')
.map((l) => {
try {
return JSON.parse(l);
}
catch {
return undefined;
}
})
.filter((d) => {
if (!d)
return false;
if (!query?.filter)
return true;
if (query.filter.eq &&
Object.keys(query.filter.eq).find((k) => d[k] !== query.filter.eq[k])) {
return false;
}
if (query.filter.neq &&
Object.keys(query.filter.neq).find((k) => d[k] === query.filter.neq[k])) {
return false;
}
return true;
})
.reverse();
fullData.push(...fileData);
}
fullData
.sort((a, b) => b['start'] - a['start']);
const deduped = [];
const seenIds = new Set();
for (const entry of fullData) {
const id = entry['id'];
if (!seenIds.has(id)) {
seenIds.add(id);
deduped.push(entry);
}
}
const result = {
data: deduped.slice(startFromIndex, startFromIndex + query.limit),
};
if (startFromIndex + query.limit < deduped.length) {
result.pageLastIndex = startFromIndex + query.limit;
}
return result;
}
}
function lockFile(file) {
return `${file}.lock`;
}
//# sourceMappingURL=file-trace-store.js.map