@maximai/maxim-js
Version:
Maxim AI JS SDK. Visit https://getmaxim.ai for more info.
413 lines • 20.7 kB
JavaScript
"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.MaximAISDKWrapper = void 0;
const uuid_1 = require("uuid");
const utils_1 = require("./../utils");
const utils_2 = require("./utils");
/**
* A wrapper class that adds Maxim logging and tracing to a Vercel AI SDK language model.
*
* This class decorates a LanguageModelV1 instance, intercepting calls to provide
* advanced observability, tracing, and logging via the MaximLogger. It is intended
* for internal use by the wrapMaximAISDKModel function.
*
* @class
* @template T - The type of the language model (must extend LanguageModelV1).
* @param model - The Vercel AI SDK language model instance to wrap.
* @param logger - The MaximLogger instance to use for tracing and logging.
*/
class MaximAISDKWrapper {
/**
* @constructor
* Creates a new MaximAISDKWrapper instance.
*
* @param model - The Vercel AI SDK language model instance to wrap.
* @param logger - The MaximLogger instance to use for tracing and logging.
*/
constructor(model, logger) {
this.model = model;
this.logger = logger;
// Internal state to track trace across multiple doGenerate calls in a tool-call sequence
this.currentTraceId = null;
this.currentTrace = null;
this.currentSession = null;
this.isInToolCallSequence = false;
}
/**
* Sets up Maxim logging and tracing for a model call.
*
* Extracts Maxim metadata, parses prompt messages, and initializes session, trace, and span objects.
* Also logs user input to the trace if appropriate.
*
* @private
* @param options - The call options for the model invocation.
* @param promptMessages - The parsed prompt messages (used to detect tool calls).
* @returns An object containing maximMetadata, trace, session, span, and promptMessages.
*/
setupLogging(options, promptMessages) {
var _a, _b, _c, _d, _e, _f;
// Extracting the maxim object from `providerOptions`
const maximMetadata = (0, utils_1.extractMaximMetadataFromOptions)(options.providerMetadata);
// Check if this is a continuation of a tool-call sequence (has tool results in prompt)
const hasToolResults = promptMessages.some((msg) => msg.role === "tool");
// Determine if we should reuse the existing trace or create a new one
const shouldReuseTrace = !(maximMetadata === null || maximMetadata === void 0 ? void 0 : maximMetadata.traceId) && // User hasn't explicitly provided a traceId
this.isInToolCallSequence && // We're in a tool-call sequence
hasToolResults && // This call has tool results (continuation)
this.currentTraceId !== null; // We have an existing trace
let session = undefined;
let trace;
// If sessionId is passed, then create a session on Maxim. If not passed, do not create a session
if (maximMetadata === null || maximMetadata === void 0 ? void 0 : maximMetadata.sessionId) {
// Reuse existing session if it's the same sessionId, otherwise create new
if (((_a = this.currentSession) === null || _a === void 0 ? void 0 : _a.id) === maximMetadata.sessionId) {
session = this.currentSession;
}
else {
session = this.logger.session({
id: maximMetadata.sessionId,
name: (_b = maximMetadata.sessionName) !== null && _b !== void 0 ? _b : "default-session",
tags: maximMetadata.sessionTags,
});
this.currentSession = session;
}
}
// Determine trace ID: use user-provided, reuse existing, or create new
const traceId = (_c = maximMetadata === null || maximMetadata === void 0 ? void 0 : maximMetadata.traceId) !== null && _c !== void 0 ? _c : (shouldReuseTrace ? this.currentTraceId : (0, uuid_1.v4)());
const traceName = (_d = maximMetadata === null || maximMetadata === void 0 ? void 0 : maximMetadata.traceName) !== null && _d !== void 0 ? _d : "default-trace";
// Reuse existing trace if we're continuing a tool-call sequence
if (shouldReuseTrace && this.currentTrace) {
trace = this.currentTrace;
}
else {
// Create new trace
trace = session
? session.trace({
id: traceId,
name: traceName,
tags: maximMetadata === null || maximMetadata === void 0 ? void 0 : maximMetadata.traceTags,
})
: this.logger.trace({
id: traceId,
name: traceName,
tags: maximMetadata === null || maximMetadata === void 0 ? void 0 : maximMetadata.traceTags,
});
// Store trace state for reuse
this.currentTraceId = traceId;
this.currentTrace = trace;
}
// If this is the start of a new sequence (no tool results), reset the tool-call sequence flag
if (!hasToolResults && !(maximMetadata === null || maximMetadata === void 0 ? void 0 : maximMetadata.traceId)) {
this.isInToolCallSequence = false;
}
const userMessage = promptMessages.findLast((msg) => msg.role === "user");
if (userMessage && userMessage.content) {
const userInput = userMessage.content;
if (typeof userInput === "string") {
trace.input(userInput);
}
else {
const userMessageContent = userInput[0];
switch (userMessageContent.type) {
case "text":
trace.input(userMessageContent.text);
break;
case "image_url":
trace.input(userMessageContent.image_url.url);
trace.addAttachment({
id: (0, uuid_1.v4)(),
type: "url",
url: userMessageContent.image_url.url,
});
break;
default:
break;
}
}
}
const span = trace.span({
id: (_e = maximMetadata === null || maximMetadata === void 0 ? void 0 : maximMetadata.spanId) !== null && _e !== void 0 ? _e : (0, uuid_1.v4)(),
name: (_f = maximMetadata === null || maximMetadata === void 0 ? void 0 : maximMetadata.spanName) !== null && _f !== void 0 ? _f : "default-span",
tags: maximMetadata === null || maximMetadata === void 0 ? void 0 : maximMetadata.spanTags,
});
return { maximMetadata, trace, session, span, promptMessages };
}
/**
* Executes a text or object generation call with Maxim tracing and logging.
*
* This method is called internally by generateText and generateObject, and logs the generation
* result, errors, and relevant metadata to Maxim.
*
* @param options - The call options for the model invocation.
* @returns The result of the underlying model's doGenerate call.
*/
async doGenerate(options) {
var _a, _b, _c, _d;
// Parse prompt messages first to detect tool calls
const promptMessages = (0, utils_2.parsePromptMessages)(options.prompt);
const { maximMetadata, trace, span } = this.setupLogging(options, promptMessages);
let generation = undefined;
let response = undefined;
let hasToolCallsInResponse = false;
try {
generation = span.generation({
id: (0, uuid_1.v4)(),
name: (_a = maximMetadata === null || maximMetadata === void 0 ? void 0 : maximMetadata.generationName) !== null && _a !== void 0 ? _a : "default-generation",
provider: (0, utils_1.determineProvider)(this.model.provider),
model: this.modelId,
messages: promptMessages,
modelParameters: (0, utils_1.extractModelParameters)(options),
tags: maximMetadata === null || maximMetadata === void 0 ? void 0 : maximMetadata.generationTags,
});
(0, utils_2.processToolResultsFromPromptV1)(options.prompt, this.logger);
// Calling the original doGenerate function
response = await this.model.doGenerate(options);
// Check if response has tool calls - in v1, tool calls are in rawResponse.body.choices
const choices = (_d = (_c = (_b = response.rawResponse) === null || _b === void 0 ? void 0 : _b.body) === null || _c === void 0 ? void 0 : _c.choices) !== null && _d !== void 0 ? _d : [];
if (Array.isArray(choices)) {
hasToolCallsInResponse = choices.some((choice) => { var _a; return ((_a = choice.message) === null || _a === void 0 ? void 0 : _a.tool_calls) && Array.isArray(choice.message.tool_calls) && choice.message.tool_calls.length > 0; });
}
if (hasToolCallsInResponse) {
// Mark that we're in a tool-call sequence
this.isInToolCallSequence = true;
}
const res = (0, utils_2.convertDoGenerateResultToChatCompletionResult)(response);
generation.result(res);
generation.end();
return response;
}
catch (error) {
if (generation) {
generation.error({
message: error.message,
});
generation.end();
}
// Log error details
console.error("[MaximSDK] doGenerate failed:", error);
throw error;
}
finally {
span.end();
// End trace if:
// 1. User explicitly provided traceId (they manage it) - but don't reset state
// 2. OR response has no tool calls (sequence is complete or single call)
const shouldEndTrace = (maximMetadata === null || maximMetadata === void 0 ? void 0 : maximMetadata.traceId) || !hasToolCallsInResponse;
if (shouldEndTrace) {
// Reset state when ending trace (only if user didn't provide traceId)
if (!(maximMetadata === null || maximMetadata === void 0 ? void 0 : maximMetadata.traceId)) {
this.currentTraceId = null;
this.currentTrace = null;
this.isInToolCallSequence = false;
trace.end();
}
}
await this.logger.flush();
}
}
/**
* Executes a streaming generation call with Maxim tracing and logging.
*
* This method is called internally by streamText and streamObject, and logs the streaming
* result, errors, and relevant metadata to Maxim.
*
* @param options - The call options for the model invocation.
* @returns The result of the underlying model's doStream call, with a wrapped stream.
*/
async doStream(options) {
var _a;
// Parse prompt messages first to detect tool calls
const promptMessages = (0, utils_2.parsePromptMessages)(options.prompt);
const { maximMetadata, trace, span } = this.setupLogging(options, promptMessages);
let generation = undefined;
let hasToolCallsInResponse = false;
// Capture 'this' reference for use inside the stream
const wrapperInstance = this;
try {
// Process tool results from prompt before streaming (same as doGenerate)
// so toolCallResult/toolCallError are logged even if doStream or reader.read throws.
// Must be inside try so any throw flows into catch/flush path.
(0, utils_2.processToolResultsFromPromptV1)(options.prompt, this.logger);
// Calling the original doStream method
const startTime = performance.now();
const response = await this.model.doStream(options);
const modelProvider = (0, utils_1.determineProvider)(this.model.provider);
const modelId = this.modelId;
generation = span.generation({
id: (0, uuid_1.v4)(),
name: (_a = maximMetadata === null || maximMetadata === void 0 ? void 0 : maximMetadata.generationName) !== null && _a !== void 0 ? _a : "default-generation",
provider: modelProvider,
model: modelId,
modelParameters: (0, utils_1.extractModelParameters)(options),
messages: promptMessages,
});
// going through the original stream to collect chunks and pass them without modifications to the stream
const chunks = [];
const firstToken = {
received: false,
time: null,
};
const stream = new ReadableStream({
async start(controller) {
try {
const reader = response.stream.getReader();
while (true) {
const { done, value } = await reader.read();
if (done) {
// Stream is done, now process before closing
try {
// Check if we have tool calls in the stream
hasToolCallsInResponse = chunks.some((chunk) => chunk.type === "tool-call" || chunk.type === "tool-call-delta");
if (hasToolCallsInResponse) {
// Mark that we're in a tool-call sequence
wrapperInstance.isInToolCallSequence = true;
}
if (firstToken.received && firstToken.time) {
trace.addMetric("time_to_first_token (in ms)", firstToken.time - startTime);
firstToken.received = false;
firstToken.time = null;
}
const endTime = performance.now();
const textChunks = chunks.filter((chunk) => chunk.type === "text-delta" || chunk.type === "tool-call-delta");
trace.addMetric("tokens_per_second", textChunks.length / ((endTime - startTime) / 1000));
if (generation)
(0, utils_2.processStream)(chunks, span, trace, generation, modelId, maximMetadata);
// Handle trace ending after stream processing
// End trace if:
// 1. User explicitly provided traceId (they manage it) - but don't reset state
// 2. OR response has no tool calls (sequence is complete or single call)
const shouldEndTrace = (maximMetadata === null || maximMetadata === void 0 ? void 0 : maximMetadata.traceId) || !hasToolCallsInResponse;
if (shouldEndTrace) {
// Reset state when ending trace (only if user didn't provide traceId)
if (!(maximMetadata === null || maximMetadata === void 0 ? void 0 : maximMetadata.traceId)) {
wrapperInstance.currentTraceId = null;
wrapperInstance.currentTrace = null;
wrapperInstance.isInToolCallSequence = false;
trace.end();
}
}
// Flush logs after processStream and tool-result logging complete
await wrapperInstance.logger.flush();
}
catch (error) {
console.error("[MaximSDK] Processing failed:", error);
if (generation) {
generation.error({
message: error.message,
});
generation.end();
}
// Flush logs even on error
await wrapperInstance.logger.flush();
}
// Now close the stream
controller.close();
break;
}
// Only mark first token when we receive an actual text-delta or tool-call-delta chunk
if (!firstToken.received && (value.type === "text-delta" || value.type === "tool-call-delta")) {
firstToken.received = true;
firstToken.time = performance.now();
}
// Collect chunk and pass it through
chunks.push(value);
controller.enqueue(value);
}
}
catch (error) {
controller.error(error);
if (generation) {
generation.error({
message: error.message,
});
generation.end();
}
// Flush logs on stream error
await wrapperInstance.logger.flush();
}
},
});
// Return response with the logging stream - user gets real-time data without additional delay
return {
...response,
stream: stream,
};
}
catch (error) {
if (generation) {
generation.error({
message: error.message,
});
generation.end();
}
// Log error details
console.error("[MaximSDK] doStream failed:", error);
// Flush logs before throwing error
await this.logger.flush();
throw error;
}
finally {
// Note: For streaming, span ending happens in processStream, and trace ending
// is handled in the stream completion handler above (because hasToolCallsInResponse
// is set asynchronously). We don't end the trace here.
}
}
/**
* Returns the default object generation mode of the wrapped model.
*
* @returns The default object generation mode.
*/
get defaultObjectGenerationMode() {
return this.model.defaultObjectGenerationMode;
}
/**
* Returns the model ID of the wrapped model.
*
* @returns The model ID.
*/
get modelId() {
return this.model.modelId;
}
/**
* Returns the provider name of the wrapped model.
*
* @returns The provider name.
*/
get provider() {
return this.model.provider;
}
/**
* Returns the specification version of the wrapped model.
*
* @returns The specification version.
*/
get specificationVersion() {
return this.model.specificationVersion;
}
/**
* Indicates whether the wrapped model supports image URLs.
*
* @returns True if image URLs are supported, false otherwise.
*/
get supportsImageUrls() {
return this.model.supportsImageUrls;
}
/**
* Indicates whether the wrapped model supports structured outputs.
*
* @returns True if structured outputs are supported, false otherwise.
*/
get supportsStructuredOutputs() {
return this.model.supportsStructuredOutputs;
}
/**
* Indicates whether the wrapped model supports URL input.
*
* @returns True if URL input is supported, false otherwise.
*/
get supportsUrl() {
return this.model.supportsUrl;
}
}
exports.MaximAISDKWrapper = MaximAISDKWrapper;
//# sourceMappingURL=wrapper.js.map