@genkit-ai/telemetry-server
Version:
Genkit AI telemetry server
231 lines • 8.42 kB
JavaScript
import { TraceDataSchema, } from '@genkit-ai/tools-common';
import { Mutex } from 'async-mutex';
import fs from 'fs';
import lockfile from 'lockfile';
import os from 'os';
import path from 'path';
const MAX_TRACES = 1000;
export class LocalFileTraceStore {
storeRoot;
indexRoot;
mutexes = {};
filters;
index;
static defaultFilters = {
'genkit:metadata:subtype': 'prompt',
};
constructor(options) {
this.storeRoot =
options?.storeRoot ?? path.resolve(process.cwd(), `.genkit/traces`);
this.indexRoot = path.resolve(os.tmpdir(), `./genkit-trace-index-${Date.now()}-${Math.random()}`);
fs.mkdirSync(this.storeRoot, { recursive: true });
console.info(`[Telemetry Server] initialized local file trace store at root: ${this.storeRoot}`);
this.filters = options?.filters ?? LocalFileTraceStore.defaultFilters;
if (options?.useIndex) {
this.index = new Index(this.indexRoot);
}
}
async init() {
if (!this.index) {
return;
}
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);
}
console.log(`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);
if (Object.keys(trace.spans).length === 0) {
return;
}
const mutex = this.getMutex(id);
await mutex.waitForUnlock();
const release = await mutex.acquire();
try {
const existing = await this.load(id);
if (existing) {
Object.keys(trace.spans).forEach((spanId) => (existing.spans[spanId] = trace.spans[spanId]));
existing.displayName = trace.displayName;
existing.startTime = trace.startTime;
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) {
if (!this.index) {
return this.listFromFiles(query);
}
const searchResult = this.index.search({
limit: query?.limit ?? 10,
startFromIndex: query?.continuationToken
? parseInt(query?.continuationToken)
: undefined,
filter: query?.filter,
});
const loadedTraces = await Promise.all(searchResult.data.map((d) => this.load(d['id'])));
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
? 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;
}
}
export class Index {
indexRoot;
currentIndexFile;
constructor(indexRoot) {
this.indexRoot = indexRoot;
this.currentIndexFile = path.resolve(this.indexRoot, '1.idx');
fs.mkdirSync(this.indexRoot, { recursive: true });
if (!fs.existsSync(this.currentIndexFile)) {
fs.writeFileSync(this.currentIndexFile, '');
}
}
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;
}
try {
lockfile.lockSync(lockFile(this.currentIndexFile));
fs.appendFileSync(this.currentIndexFile, JSON.stringify(indexData) + '\n');
}
finally {
lockfile.unlockSync(lockFile(this.currentIndexFile));
}
}
search(query) {
const idxTxt = fs.readFileSync(this.currentIndexFile, 'utf8');
const fullData = 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();
const startFromIndex = query.startFromIndex ?? 0;
const result = {
data: fullData.slice(startFromIndex, startFromIndex + query.limit),
};
if (startFromIndex + query.limit < fullData.length) {
result.pageLastIndex = startFromIndex + query.limit;
}
return result;
}
}
function lockFile(file) {
return `${file}.lock`;
}
//# sourceMappingURL=localFileTraceStore.js.map