@genkit-ai/telemetry-server
Version:
Genkit AI telemetry server
281 lines • 11.2 kB
JavaScript
;
var __importDefault = (this && this.__importDefault) || function (mod) {
return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.Index = exports.LocalFileTraceStore = void 0;
const tools_common_1 = require("@genkit-ai/tools-common");
const utils_1 = require("@genkit-ai/tools-common/utils");
const async_mutex_1 = require("async-mutex");
const fs_1 = __importDefault(require("fs"));
const lockfile_1 = __importDefault(require("lockfile"));
const path_1 = __importDefault(require("path"));
const version_1 = require("./utils/version");
const MAX_TRACES = 1000;
const MAX_INDEX_FILES = 10;
class LocalFileTraceStore {
storeRoot;
indexRoot;
mutexes = {};
filters;
index;
static defaultFilters = {
'genkit:metadata:subtype': 'prompt',
};
constructor(options) {
this.storeRoot = path_1.default.resolve(options.storeRoot, `.genkit/traces`);
fs_1.default.mkdirSync(this.storeRoot, { recursive: true });
this.indexRoot = path_1.default.resolve(options.indexRoot, `.genkit/traces_idx`);
fs_1.default.mkdirSync(this.indexRoot, { recursive: true });
utils_1.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 !== version_1.version ||
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);
}
utils_1.logger.info(`Indexed ${list.traces.length} traces in ${Date.now() - time}ms in ${this.indexRoot}`);
}
async load(id) {
const filePath = path_1.default.resolve(this.storeRoot, `${id}`);
if (!fs_1.default.existsSync(filePath)) {
return undefined;
}
const data = fs_1.default.readFileSync(filePath, 'utf8');
const parsed = JSON.parse(data);
if (!parsed.traceId) {
parsed.traceId = id;
}
return tools_common_1.TraceDataSchema.parse(parsed);
}
getMutex(id) {
if (!this.mutexes[id]) {
this.mutexes[id] = new async_mutex_1.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) => (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;
}
});
}
existing.displayName = trace.displayName;
existing.startTime = trace.startTime;
existing.endTime = trace.endTime;
trace = existing;
}
fs_1.default.writeFileSync(path_1.default.resolve(this.storeRoot, `${id}`), JSON.stringify(trace));
const hasRootSpan = !!Object.values(trace.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'])));
return {
traces: loadedTraces.filter((t) => !!t),
continuationToken: searchResult.pageLastIndex
? `${searchResult.pageLastIndex}`
: undefined,
};
}
async listFromFiles(query) {
const files = fs_1.default.readdirSync(this.storeRoot);
files.sort((a, b) => {
return (fs_1.default.statSync(path_1.default.resolve(this.storeRoot, `${b}`)).mtime.getTime() -
fs_1.default.statSync(path_1.default.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_1.default.resolve(this.storeRoot, `${id}`);
const data = fs_1.default.readFileSync(filePath, 'utf8');
const parsed = JSON.parse(data);
if (!parsed.traceId) {
parsed.traceId = id;
}
return tools_common_1.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;
}
}
exports.LocalFileTraceStore = LocalFileTraceStore;
class Index {
indexRoot;
currentIndexFile;
constructor(indexRoot) {
this.indexRoot = indexRoot;
this.currentIndexFile = path_1.default.resolve(this.indexRoot, this.newIndexFileName());
fs_1.default.mkdirSync(this.indexRoot, { recursive: true });
}
clear() {
fs_1.default.rmSync(this.indexRoot, { recursive: true, force: true });
fs_1.default.mkdirSync(this.indexRoot, { recursive: true });
fs_1.default.appendFileSync(this.metadataFileName(), JSON.stringify({ version: version_1.version }));
}
metadataFileName() {
return path_1.default.resolve(this.indexRoot, 'genkit.metadata');
}
getMetadata() {
if (!fs_1.default.existsSync(this.metadataFileName())) {
return undefined;
}
return JSON.parse(fs_1.default.readFileSync(this.metadataFileName(), { encoding: 'utf8' }));
}
newIndexFileName() {
return `idx_${(Date.now() + '').padStart(17, '0')}.json`;
}
listIndexFiles() {
return fs_1.default.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_1.default.lockSync(lockFile(this.currentIndexFile));
fs_1.default.appendFileSync(this.currentIndexFile, JSON.stringify(indexData) + '\n');
}
finally {
lockfile_1.default.unlockSync(lockFile(this.currentIndexFile));
}
}
search(query) {
const startFromIndex = query.startFromIndex ?? 0;
const fullData = [];
for (const idxFile of this.listIndexFiles()) {
const idxTxt = fs_1.default.readFileSync(path_1.default.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 result = {
data: fullData.slice(startFromIndex, startFromIndex + query.limit),
};
if (startFromIndex + query.limit < fullData.length) {
result.pageLastIndex = startFromIndex + query.limit;
}
return result;
}
}
exports.Index = Index;
function lockFile(file) {
return `${file}.lock`;
}
//# sourceMappingURL=localFileTraceStore.js.map