@maximai/maxim-js
Version:
Maxim AI JS SDK. Visit https://getmaxim.ai for more info.
253 lines • 10.4 kB
JavaScript
;
Object.defineProperty(exports, "__esModule", { value: true });
exports.MaximOpenAIChatCompletions = void 0;
const uuid_1 = require("uuid");
const utils_1 = require("./utils");
/**
* Extract tool calls and their results from messages.
* Matches assistant tool_calls with corresponding tool messages.
*/
function extractToolCalls(messages) {
const toolCalls = [];
// Map of tool_call_id to tool call info from assistant messages
const toolCallMap = new Map();
// First pass: collect tool calls from assistant messages
for (const msg of messages) {
if (msg.role === "assistant" && "tool_calls" in msg && msg.tool_calls) {
for (const tc of msg.tool_calls) {
if (tc.type === "function" && "function" in tc && tc.function) {
toolCallMap.set(tc.id, {
name: tc.function.name,
arguments: tc.function.arguments,
});
}
}
}
}
// Second pass: match tool messages with their corresponding tool calls
for (const msg of messages) {
if (msg.role === "tool" && "tool_call_id" in msg) {
const toolCallInfo = toolCallMap.get(msg.tool_call_id);
if (toolCallInfo) {
toolCalls.push({
id: msg.tool_call_id,
name: toolCallInfo.name,
arguments: toolCallInfo.arguments,
result: typeof msg.content === "string" ? msg.content : JSON.stringify(msg.content),
});
}
}
}
return toolCalls;
}
/**
* Async stream wrapper that logs to Maxim when the stream completes.
*/
class AsyncStreamWrapper {
constructor(stream, generation, trace, isLocalTrace) {
this.stream = stream;
this.generation = generation;
this.trace = trace;
this.isLocalTrace = isLocalTrace;
this.chunks = [];
this.consumed = false;
this.firstTokenTime = null;
this.startTime = performance.now();
}
async *[Symbol.asyncIterator]() {
try {
for await (const chunk of this.stream) {
// Track time to first token
if (this.firstTokenTime === null && this.hasContent(chunk)) {
this.firstTokenTime = performance.now();
}
this.chunks.push(chunk);
yield chunk;
}
}
finally {
if (!this.consumed) {
this.consumed = true;
this.finalizeLogging();
}
}
}
/**
* Check if the chunk has actual content (not just metadata).
*/
hasContent(chunk) {
var _a, _b;
for (const choice of chunk.choices) {
if (((_a = choice.delta) === null || _a === void 0 ? void 0 : _a.content) || ((_b = choice.delta) === null || _b === void 0 ? void 0 : _b.tool_calls)) {
return true;
}
}
return false;
}
/**
* Get completion tokens from the usage field of the last chunk.
* The usage field is present when stream_options.include_usage is true.
*/
getCompletionTokens() {
var _a;
for (let i = this.chunks.length - 1; i >= 0; i--) {
const chunk = this.chunks[i];
if (((_a = chunk.usage) === null || _a === void 0 ? void 0 : _a.completion_tokens) !== undefined) {
return chunk.usage.completion_tokens;
}
}
return null;
}
finalizeLogging() {
const endTime = performance.now();
try {
if (this.generation && this.chunks.length > 0) {
const combinedResponse = utils_1.OpenAIUtils.parseCompletionFromChunks(this.chunks);
this.generation.result(combinedResponse);
if (this.firstTokenTime !== null) {
const ttftMs = this.firstTokenTime - this.startTime;
this.generation.addMetric("ttft_ms", ttftMs);
const completionTokens = this.getCompletionTokens();
const generationDurationMs = endTime - this.firstTokenTime;
if (generationDurationMs > 0 && completionTokens !== null && generationDurationMs > 0 && completionTokens > 0) {
const tps = (completionTokens / generationDurationMs) * 1000;
this.generation.addMetric("tokens_per_second", tps);
}
}
}
if (this.isLocalTrace && this.trace) {
const combinedText = utils_1.OpenAIUtils.extractTextFromChunks(this.chunks);
this.trace.output(combinedText);
this.trace.end();
}
}
catch (e) {
if (this.generation) {
this.generation.error({ message: e instanceof Error ? e.message : String(e) });
}
console.warn(`[MaximSDK][MaximOpenAIChatCompletions] Error in logging stream completion: ${e}`);
}
}
}
/**
* Wrapped OpenAI chat completions that automatically logs to Maxim.
*/
class MaximOpenAIChatCompletions {
constructor(client, logger) {
this.client = client;
this.logger = logger;
}
async create(params, options) {
var _a, _b;
// Extract Maxim metadata from extra_headers
const extraHeaders = options === null || options === void 0 ? void 0 : options.headers;
let traceId;
let generationName;
let sessionId;
let traceTags;
if (extraHeaders) {
traceId = extraHeaders["maxim-trace-id"] || undefined;
generationName = extraHeaders["maxim-generation-name"] || undefined;
sessionId = extraHeaders["maxim-session-id"] || undefined;
traceTags = extraHeaders["maxim-trace-tags"] || undefined;
}
const isLocalTrace = traceId === undefined;
const finalTraceId = traceId || (0, uuid_1.v4)();
const isStreaming = params.stream === true;
let generation = null;
let trace = null;
// Add stream_options with include_usage for streaming if not present
const modifiedParams = { ...params };
if (isStreaming) {
if (!modifiedParams.stream_options) {
modifiedParams.stream_options = { include_usage: true };
}
else if (!modifiedParams.stream_options.include_usage) {
modifiedParams.stream_options = { ...modifiedParams.stream_options, include_usage: true };
}
}
try {
// Create trace and generation
trace = this.logger.trace({ id: finalTraceId, sessionId });
if (isStreaming) {
trace.addTag("stream", "true");
}
if (traceTags) {
try {
const parsedTags = JSON.parse(traceTags);
for (const [key, value] of Object.entries(parsedTags)) {
trace.addTag(key, value);
}
}
catch (error) {
console.warn(`[MaximSDK][MaximOpenAIChatCompletions] Error in parsing trace tags: ${error}`);
}
}
// Log any tool calls present in the messages (from previous interactions)
const extractedToolCalls = extractToolCalls(params.messages);
for (const tc of extractedToolCalls) {
const toolCall = trace.toolCall({
id: tc.id,
name: tc.name,
description: `Tool call: ${tc.name}`,
args: tc.arguments,
});
toolCall.result(tc.result);
}
const messages = utils_1.OpenAIUtils.parseMessageParam(params.messages);
const modelParams = utils_1.OpenAIUtils.getModelParams(params);
generation = trace.generation({
id: (0, uuid_1.v4)(),
model: params.model,
provider: "openai",
name: generationName,
modelParameters: modelParams,
messages: messages,
});
}
catch (e) {
console.warn(`[MaximSDK][MaximOpenAIChatCompletions] Error in setting up logging: ${e}`);
}
try {
// Call the original OpenAI API
const response = await this.client.chat.completions.create(modifiedParams, options);
if (isStreaming) {
// For streaming responses, return a wrapped stream that handles logging
// The OpenAI SDK returns an AsyncIterable stream when stream: true
return new AsyncStreamWrapper(response, generation, trace, isLocalTrace);
}
else {
// For non-streaming responses, log immediately
const completion = response;
try {
if (generation) {
const result = utils_1.OpenAIUtils.parseCompletion(completion);
generation.result(result);
}
if (isLocalTrace && trace) {
const output = ((_b = (_a = completion.choices[0]) === null || _a === void 0 ? void 0 : _a.message) === null || _b === void 0 ? void 0 : _b.content) || "";
trace.output(output);
trace.end();
}
}
catch (e) {
if (generation) {
generation.error({ message: e instanceof Error ? e.message : String(e) });
}
console.warn(`[MaximSDK][MaximOpenAIChatCompletions] Error in logging generation: ${e}`);
}
return completion;
}
}
catch (e) {
// Log error if generation was created
if (generation) {
generation.error({ message: e instanceof Error ? e.message : String(e) });
}
console.warn(`[MaximSDK][MaximOpenAIChatCompletions] Error in generating content: ${e}`);
throw e;
}
}
}
exports.MaximOpenAIChatCompletions = MaximOpenAIChatCompletions;
//# sourceMappingURL=completions.js.map