@probelabs/probe
Version:
Node.js wrapper for the probe code search tool
251 lines (249 loc) • 7.34 kB
JavaScript
;
var __defProp = Object.defineProperty;
var __getOwnPropDesc = Object.getOwnPropertyDescriptor;
var __getOwnPropNames = Object.getOwnPropertyNames;
var __hasOwnProp = Object.prototype.hasOwnProperty;
var __export = (target, all) => {
for (var name in all)
__defProp(target, name, { get: all[name], enumerable: true });
};
var __copyProps = (to, from, except, desc) => {
if (from && typeof from === "object" || typeof from === "function") {
for (let key of __getOwnPropNames(from))
if (!__hasOwnProp.call(to, key) && key !== except)
__defProp(to, key, { get: () => from[key], enumerable: !(desc = __getOwnPropDesc(from, key)) || desc.enumerable });
}
return to;
};
var __toCommonJS = (mod) => __copyProps(__defProp({}, "__esModule", { value: true }), mod);
// src/agent/simpleTelemetry.js
var simpleTelemetry_exports = {};
__export(simpleTelemetry_exports, {
SimpleAppTracer: () => SimpleAppTracer,
SimpleTelemetry: () => SimpleTelemetry,
initializeSimpleTelemetryFromOptions: () => initializeSimpleTelemetryFromOptions
});
module.exports = __toCommonJS(simpleTelemetry_exports);
var import_fs = require("fs");
var import_path = require("path");
var SimpleTelemetry = class {
constructor(options = {}) {
this.serviceName = options.serviceName || "probe-agent";
this.enableFile = options.enableFile || false;
this.enableConsole = options.enableConsole || false;
this.filePath = options.filePath || "./traces.jsonl";
this.stream = null;
if (this.enableFile) {
this.initializeFileExporter();
}
}
initializeFileExporter() {
try {
const dir = (0, import_path.dirname)(this.filePath);
if (!(0, import_fs.existsSync)(dir)) {
(0, import_fs.mkdirSync)(dir, { recursive: true });
}
this.stream = (0, import_fs.createWriteStream)(this.filePath, { flags: "a" });
this.stream.on("error", (error) => {
console.error(`[SimpleTelemetry] Stream error: ${error.message}`);
});
console.log(`[SimpleTelemetry] File exporter initialized: ${this.filePath}`);
} catch (error) {
console.error(`[SimpleTelemetry] Failed to initialize file exporter: ${error.message}`);
}
}
createSpan(name, attributes = {}) {
const span = {
traceId: this.generateTraceId(),
spanId: this.generateSpanId(),
name,
startTime: Date.now(),
attributes: { ...attributes, service: this.serviceName },
events: [],
status: "OK"
};
return {
...span,
addEvent: (eventName, eventAttributes = {}) => {
span.events.push({
name: eventName,
time: Date.now(),
attributes: eventAttributes
});
},
setAttributes: (attrs) => {
Object.assign(span.attributes, attrs);
},
setStatus: (status) => {
span.status = status;
},
end: () => {
span.endTime = Date.now();
span.duration = span.endTime - span.startTime;
this.exportSpan(span);
}
};
}
exportSpan(span) {
const spanData = {
...span,
timestamp: (/* @__PURE__ */ new Date()).toISOString()
};
if (this.enableConsole) {
console.log("[Trace]", JSON.stringify(spanData, null, 2));
}
if (this.enableFile && this.stream) {
this.stream.write(JSON.stringify(spanData) + "\n");
}
}
generateTraceId() {
return Math.random().toString(36).substring(2, 15) + Math.random().toString(36).substring(2, 15);
}
generateSpanId() {
return Math.random().toString(36).substring(2, 10);
}
async flush() {
if (this.stream) {
return new Promise((resolve) => {
this.stream.once("drain", resolve);
if (!this.stream.writableNeedDrain) {
resolve();
}
});
}
}
async shutdown() {
if (this.stream) {
return new Promise((resolve) => {
this.stream.end(() => {
console.log(`[SimpleTelemetry] File stream closed: ${this.filePath}`);
resolve();
});
});
}
}
};
var SimpleAppTracer = class {
constructor(telemetry, sessionId = null) {
this.telemetry = telemetry;
this.sessionId = sessionId || this.generateSessionId();
}
generateSessionId() {
return Math.random().toString(36).substring(2, 15);
}
isEnabled() {
return this.telemetry !== null;
}
createSessionSpan(attributes = {}) {
if (!this.isEnabled()) return null;
return this.telemetry.createSpan("agent.session", {
"session.id": this.sessionId,
...attributes
});
}
createAISpan(modelName, provider, attributes = {}) {
if (!this.isEnabled()) return null;
return this.telemetry.createSpan("ai.request", {
"ai.model": modelName,
"ai.provider": provider,
"session.id": this.sessionId,
...attributes
});
}
createToolSpan(toolName, attributes = {}) {
if (!this.isEnabled()) return null;
return this.telemetry.createSpan("tool.call", {
"tool.name": toolName,
"session.id": this.sessionId,
...attributes
});
}
addEvent(name, attributes = {}) {
if (this.telemetry && this.telemetry.enableConsole) {
console.log("[Event]", name, attributes);
}
}
/**
* Record delegation events
*/
recordDelegationEvent(eventType, data = {}) {
if (!this.isEnabled()) return;
this.addEvent(`delegation.${eventType}`, {
"session.id": this.sessionId,
...data
});
}
/**
* Record JSON validation events
*/
recordJsonValidationEvent(eventType, data = {}) {
if (!this.isEnabled()) return;
this.addEvent(`json_validation.${eventType}`, {
"session.id": this.sessionId,
...data
});
}
/**
* Record Mermaid validation events
*/
recordMermaidValidationEvent(eventType, data = {}) {
if (!this.isEnabled()) return;
this.addEvent(`mermaid_validation.${eventType}`, {
"session.id": this.sessionId,
...data
});
}
setAttributes(attributes) {
if (this.telemetry && this.telemetry.enableConsole) {
console.log("[Attributes]", attributes);
}
}
async withSpan(spanName, fn, attributes = {}) {
if (!this.isEnabled()) {
return fn();
}
const span = this.telemetry.createSpan(spanName, {
"session.id": this.sessionId,
...attributes
});
try {
const result = await fn();
span.setStatus("OK");
return result;
} catch (error) {
span.setStatus("ERROR");
span.addEvent("exception", {
"exception.message": error.message,
"exception.stack": error.stack
});
throw error;
} finally {
span.end();
}
}
async flush() {
if (this.telemetry) {
await this.telemetry.flush();
}
}
async shutdown() {
if (this.telemetry) {
await this.telemetry.shutdown();
}
}
};
function initializeSimpleTelemetryFromOptions(options) {
const telemetry = new SimpleTelemetry({
serviceName: "probe-agent",
enableFile: options.traceFile !== void 0,
enableConsole: options.traceConsole,
filePath: options.traceFile || "./traces.jsonl"
});
return telemetry;
}
// Annotate the CommonJS export names for ESM import in node:
0 && (module.exports = {
SimpleAppTracer,
SimpleTelemetry,
initializeSimpleTelemetryFromOptions
});