@mtharrison/loupe 1.2.0 → 1.4.0

package/dist/index.js

@@ -3,21 +3,21 @@ Object.defineProperty(exports, "__esModule", { value: true });
 exports.isTraceEnabled = isTraceEnabled;
 exports.getLocalLLMTracer = getLocalLLMTracer;
 exports.startTraceServer = startTraceServer;
-exports.recordInvokeStart = recordInvokeStart;
-exports.recordInvokeFinish = recordInvokeFinish;
-exports.recordStreamStart = recordStreamStart;
-exports.recordStreamChunk = recordStreamChunk;
-exports.recordStreamFinish = recordStreamFinish;
-exports.recordError = recordError;
+exports.startSpan = startSpan;
+exports.endSpan = endSpan;
+exports.addSpanEvent = addSpanEvent;
+exports.recordException = recordException;
 exports.__resetLocalLLMTracerForTests = __resetLocalLLMTracerForTests;
 exports.wrapChatModel = wrapChatModel;
 exports.wrapOpenAIClient = wrapOpenAIClient;
+const node_async_hooks_1 = require("node:async_hooks");
 const server_1 = require("./server");
 const store_1 = require("./store");
 const ui_build_1 = require("./ui-build");
 const utils_1 = require("./utils");
 let singleton = null;
 const DEFAULT_TRACE_PORT = 4319;
+const activeSpanStorage = new node_async_hooks_1.AsyncLocalStorage();
 function isTraceEnabled() {
     return (0, utils_1.envFlag)('LLM_TRACE_ENABLED');
 }
@@ -33,23 +33,17 @@ function getLocalLLMTracer(config = {}) {
 function startTraceServer(config = {}) {
     return getLocalLLMTracer(config).startServer();
 }
-function recordInvokeStart(context, request, config = {}) {
-    return getLocalLLMTracer(config).recordInvokeStart(context, request);
+function startSpan(context, options = {}, config = {}) {
+    return getLocalLLMTracer(config).startSpan(context, options);
 }
-function recordInvokeFinish(traceId, response, config = {}) {
-    getLocalLLMTracer(config).recordInvokeFinish(traceId, response);
+function endSpan(spanId, response, config = {}) {
+    getLocalLLMTracer(config).endSpan(spanId, response);
 }
-function recordStreamStart(context, request, config = {}) {
-    return getLocalLLMTracer(config).recordStreamStart(context, request);
+function addSpanEvent(spanId, event, config = {}) {
+    getLocalLLMTracer(config).addSpanEvent(spanId, event);
 }
-function recordStreamChunk(traceId, chunk, config = {}) {
-    getLocalLLMTracer(config).recordStreamChunk(traceId, chunk);
-}
-function recordStreamFinish(traceId, chunk, config = {}) {
-    getLocalLLMTracer(config).recordStreamFinish(traceId, chunk);
-}
-function recordError(traceId, error, config = {}) {
-    getLocalLLMTracer(config).recordError(traceId, error);
+function recordException(spanId, error, config = {}) {
+    getLocalLLMTracer(config).recordException(spanId, error);
 }
 function __resetLocalLLMTracerForTests() {
     if (singleton?.uiWatcher) {
@@ -70,14 +64,19 @@ function wrapChatModel(model, getContext, config) {
             if (!tracer.isEnabled()) {
                 return model.invoke(input, options);
             }
-            const traceId = tracer.recordInvokeStart(getContext ? getContext() : {}, { input: input, options: options });
+            const traceId = tracer.startSpan(getContext ? getContext() : {}, {
+                attributes: { 'gen_ai.operation.name': 'chat' },
+                mode: 'invoke',
+                name: 'llm.invoke',
+                request: { input: input, options: options },
+            });
             try {
-                const response = await model.invoke(input, options);
-                tracer.recordInvokeFinish(traceId, response);
+                const response = await tracer.runWithActiveSpan(traceId, () => model.invoke(input, options));
+                tracer.endSpan(traceId, response);
                 return response;
             }
             catch (error) {
-                tracer.recordError(traceId, error);
+                tracer.recordException(traceId, error);
                 throw error;
             }
         },
@@ -87,21 +86,26 @@ function wrapChatModel(model, getContext, config) {
                 yield* model.stream(input, options);
                 return;
             }
-            const traceId = tracer.recordStreamStart(getContext ? getContext() : {}, { input: input, options: options });
+            const traceId = tracer.startSpan(getContext ? getContext() : {}, {
+                attributes: { 'gen_ai.operation.name': 'chat' },
+                mode: 'stream',
+                name: 'llm.stream',
+                request: { input: input, options: options },
+            });
             try {
-                const stream = model.stream(input, options);
+                const stream = tracer.runWithActiveSpan(traceId, () => model.stream(input, options));
                 for await (const chunk of stream) {
                     if (chunk?.type === 'finish') {
-                        tracer.recordStreamFinish(traceId, chunk);
+                        tracer.endSpan(traceId, chunk);
                     }
                     else {
-                        tracer.recordStreamChunk(traceId, chunk);
+                        tracer.addSpanEvent(traceId, (0, utils_1.toSpanEventInputFromChunk)(chunk));
                     }
                     yield chunk;
                 }
             }
             catch (error) {
-                tracer.recordError(traceId, error);
+                tracer.recordException(traceId, error);
                 throw error;
             }
         },
@@ -121,24 +125,34 @@ function wrapOpenAIClient(client, getContext, config) {
                     }
                     const context = withOpenAITraceContext(getContext ? getContext() : {}, params);
                     if (params?.stream) {
-                        const traceId = tracer.recordStreamStart(context, { input: params, options: options });
+                        const traceId = tracer.startSpan(context, {
+                            attributes: { 'gen_ai.operation.name': 'chat' },
+                            mode: 'stream',
+                            name: 'openai.chat.completions',
+                            request: { input: params, options: options },
+                        });
                         try {
-                            const stream = await target.create.call(target, params, options);
+                            const stream = await tracer.runWithActiveSpan(traceId, () => target.create.call(target, params, options));
                             return wrapOpenAIChatCompletionsStream(stream, tracer, traceId);
                         }
                         catch (error) {
-                            tracer.recordError(traceId, error);
+                            tracer.recordException(traceId, error);
                             throw error;
                         }
                     }
-                    const traceId = tracer.recordInvokeStart(context, { input: params, options: options });
+                    const traceId = tracer.startSpan(context, {
+                        attributes: { 'gen_ai.operation.name': 'chat' },
+                        mode: 'invoke',
+                        name: 'openai.chat.completions',
+                        request: { input: params, options: options },
+                    });
                     try {
-                        const response = await target.create.call(target, params, options);
-                        tracer.recordInvokeFinish(traceId, normalizeOpenAIChatCompletionResponse(response));
+                        const response = await tracer.runWithActiveSpan(traceId, () => target.create.call(target, params, options));
+                        tracer.endSpan(traceId, normalizeOpenAIChatCompletionResponse(response));
                         return response;
                     }
                     catch (error) {
-                        tracer.recordError(traceId, error);
+                        tracer.recordException(traceId, error);
                         throw error;
                     }
                 };
@@ -213,6 +227,27 @@ class LocalLLMTracerImpl {
     isEnabled() {
         return isTraceEnabled();
     }
+    startSpan(context, options = {}) {
+        void this.startServer();
+        const parentSpanId = options.parentSpanId || activeSpanStorage.getStore() || null;
+        return this.store.startSpan(context, {
+            ...options,
+            parentSpanId,
+            request: normaliseRequest(options.request || {}),
+        });
+    }
+    runWithActiveSpan(spanId, callback) {
+        return activeSpanStorage.run(spanId, callback);
+    }
+    addSpanEvent(spanId, event) {
+        this.store.addSpanEvent(spanId, (0, utils_1.safeClone)(event));
+    }
+    endSpan(spanId, response) {
+        this.store.endSpan(spanId, (0, utils_1.safeClone)(response));
+    }
+    recordException(spanId, error) {
+        this.store.recordException(spanId, error);
+    }
     startServer() {
         if (!this.isEnabled() || this.serverFailed) {
             return Promise.resolve(this.serverInfo);
@@ -234,7 +269,7 @@ class LocalLLMTracerImpl {
                     this.uiWatcher = await (0, ui_build_1.maybeStartUIWatcher)(() => {
                         this.server?.broadcast({
                             timestamp: new Date().toISOString(),
-                            traceId: null,
+                            spanId: null,
                             type: 'ui:reload',
                         });
                     }, this.config.uiHotReload);
@@ -256,26 +291,6 @@ class LocalLLMTracerImpl {
         })();
         return this.serverStartPromise;
     }
-    recordInvokeStart(context, request) {
-        void this.startServer();
-        return this.store.recordInvokeStart(context, normaliseRequest(request));
-    }
-    recordInvokeFinish(traceId, response) {
-        this.store.recordInvokeFinish(traceId, (0, utils_1.safeClone)(response));
-    }
-    recordStreamStart(context, request) {
-        void this.startServer();
-        return this.store.recordStreamStart(context, normaliseRequest(request));
-    }
-    recordStreamChunk(traceId, chunk) {
-        this.store.recordStreamChunk(traceId, (0, utils_1.safeClone)(chunk));
-    }
-    recordStreamFinish(traceId, chunk) {
-        this.store.recordStreamFinish(traceId, (0, utils_1.safeClone)(chunk));
-    }
-    recordError(traceId, error) {
-        this.store.recordError(traceId, error);
-    }
 }
 function normaliseRequest(request) {
     return {
@@ -322,7 +337,7 @@ function wrapOpenAIChatCompletionsStream(stream, tracer, traceId) {
         state.began = true;
         const nextRole = role || state.role || 'assistant';
         state.role = nextRole;
-        tracer.recordStreamChunk(traceId, { type: 'begin', role: nextRole });
+        tracer.addSpanEvent(traceId, (0, utils_1.toSpanEventInputFromChunk)({ type: 'begin', role: nextRole }));
     };
     const emitFinish = () => {
         if (state.finished) {
@@ -330,7 +345,7 @@ function wrapOpenAIChatCompletionsStream(stream, tracer, traceId) {
         }
         emitBegin(state.role);
         state.finished = true;
-        tracer.recordStreamFinish(traceId, {
+        tracer.endSpan(traceId, {
             type: 'finish',
             finish_reasons: state.finishReasons,
             message: {
@@ -382,23 +397,23 @@ function wrapOpenAIChatCompletionsStream(stream, tracer, traceId) {
             emitBegin(state.role);
         }
         if (contentParts.length > 0) {
-            tracer.recordStreamChunk(traceId, {
+            tracer.addSpanEvent(traceId, (0, utils_1.toSpanEventInputFromChunk)({
                 type: 'chunk',
                 content: contentParts.join(''),
                 finish_reasons: [...finishReasons],
                 raw,
                 tool_calls: chunkToolCalls,
                 usage,
-            });
+            }));
             return;
         }
-        tracer.recordStreamChunk(traceId, {
+        tracer.addSpanEvent(traceId, (0, utils_1.toSpanEventInputFromChunk)({
             type: 'event',
             finish_reasons: [...finishReasons],
             raw,
             tool_calls: chunkToolCalls,
             usage,
-        });
+        }));
     };
     const createWrappedIterator = (iterator) => ({
         async next(...args) {
@@ -412,7 +427,7 @@ function wrapOpenAIChatCompletionsStream(stream, tracer, traceId) {
                 return result;
             }
             catch (error) {
-                tracer.recordError(traceId, error);
+                tracer.recordException(traceId, error);
                 throw error;
             }
         },
@@ -428,12 +443,12 @@ function wrapOpenAIChatCompletionsStream(stream, tracer, traceId) {
                 return result;
             }
             catch (error) {
-                tracer.recordError(traceId, error);
+                tracer.recordException(traceId, error);
                 throw error;
             }
         },
         async throw(error) {
-            tracer.recordError(traceId, error);
+            tracer.recordException(traceId, error);
             if (typeof iterator.throw === 'function') {
                 return iterator.throw(error);
             }

package/dist/session-nav.d.ts

@@ -41,4 +41,4 @@ export declare function findSessionNodePath(nodes: SessionNavHierarchyNode[], id
 export declare function findSessionNodeById(nodes: SessionNavHierarchyNode[], id: string): SessionNavHierarchyNode | null;
 export declare function getNewestTraceIdForNode(node: SessionNavHierarchyNode | null | undefined): string | null;
 export declare function resolveSessionTreeSelection(sessionNodes: SessionNavHierarchyNode[], selectedNodeId: string | null, selectedTraceId: string | null): SessionTreeSelection;
-export declare function getDefaultExpandedSessionTreeNodeIds(sessionNodes: SessionNavHierarchyNode[], activeSessionId: string | null, selectedNodeId: string | null): Set<string>;
+export declare function getDefaultExpandedSessionTreeNodeIds(sessionNodes: SessionNavHierarchyNode[], activeSessionId: string | null, selectedNodeId: string | null, selectedTraceId?: string | null): Set<string>;

package/dist/session-nav.js

@@ -66,7 +66,7 @@ function resolveSessionTreeSelection(sessionNodes, selectedNodeId, selectedTrace
         selectedTraceId: nextSelectedTraceId,
     };
 }
-function getDefaultExpandedSessionTreeNodeIds(sessionNodes, activeSessionId, selectedNodeId) {
+function getDefaultExpandedSessionTreeNodeIds(sessionNodes, activeSessionId, selectedNodeId, selectedTraceId = null) {
     const expanded = new Set();
     const activeSession = (activeSessionId
         ? sessionNodes.find((node) => node.id === activeSessionId) ?? null
@@ -89,6 +89,13 @@ function getDefaultExpandedSessionTreeNodeIds(sessionNodes, activeSessionId, sel
             }
         }
     }
+    if (selectedTraceId) {
+        for (const node of findSessionNodePath([activeSession], `trace:${selectedTraceId}`)) {
+            if (node.children.length) {
+                expanded.add(node.id);
+            }
+        }
+    }
     return expanded;
 }
 function deriveSessionNavItem(node, traceById) {

package/dist/store.d.ts

@@ -1,5 +1,5 @@
 import { EventEmitter } from 'node:events';
-import { type HierarchyResponse, type NormalizedTraceContext, type TraceFilters, type TraceListResponse, type TraceRecord, type TraceRequest } from './types';
+import { type HierarchyResponse, type NormalizedTraceContext, type SpanEventInput, type SpanStartOptions, type TraceFilters, type TraceListResponse, type TraceRecord } from './types';
 export declare class TraceStore extends EventEmitter {
     maxTraces: number;
     order: string[];
@@ -7,17 +7,17 @@ export declare class TraceStore extends EventEmitter {
     constructor(options?: {
         maxTraces?: number;
     });
-    recordInvokeStart(context: NormalizedTraceContext | undefined, request: TraceRequest): string;
-    recordInvokeFinish(traceId: string, response: any): void;
-    recordStreamStart(context: NormalizedTraceContext | undefined, request: TraceRequest): string;
-    recordStreamChunk(traceId: string, chunk: any): void;
-    recordStreamFinish(traceId: string, chunk: any): void;
-    recordError(traceId: string, error: unknown): void;
+    startSpan(context: NormalizedTraceContext | undefined, options?: SpanStartOptions): string;
+    addSpanEvent(spanId: string, event: SpanEventInput): void;
+    endSpan(spanId: string, response: any): void;
+    recordException(spanId: string, error: unknown): void;
+    private applyStreamPayload;
     list(filters?: TraceFilters): TraceListResponse;
     get(traceId: string): TraceRecord | null;
     clear(): void;
     hierarchy(filters?: TraceFilters): HierarchyResponse;
     private recordStart;
+    private findTraceBySpanReference;
     private evictIfNeeded;
     private cloneTrace;
     private filteredTraces;