evalite
Version:
Test your LLM-powered apps with a TypeScript-native, Vitest-based eval runner. No API key required.
130 lines • 4.76 kB
JavaScript
import { wrapLanguageModel } from "ai";
import { reportTrace, shouldReportTrace } from "./traces.js";
const handlePromptContent = (content) => {
if (typeof content === "string") {
return {
type: "text",
text: content,
};
}
if (content.type === "text") {
return {
type: "text",
text: content.text,
};
}
if (content.type === "tool-call") {
return {
type: "tool-call",
toolName: content.toolName,
input: content.input,
toolCallId: content.toolCallId,
};
}
if (content.type === "tool-result") {
const output = content.output;
// Check for unsupported media content
if (output.type === "content" &&
output.value.find((item) => item.type === "media")) {
throw new Error(`Unsupported content type: media in tool-result. Not supported yet.`);
}
return {
type: "tool-result",
toolCallId: content.toolCallId,
toolName: content.toolName,
output: content.output,
};
}
// Unsupported content types are image and file
throw new Error(`Unsupported content type: ${content.type}. Not supported yet.`);
};
const processPromptForTracing = (prompt) => {
return prompt.map((prompt) => {
if (!Array.isArray(prompt.content)) {
return {
role: prompt.role,
content: prompt.content,
};
}
const content = prompt.content.map(handlePromptContent);
return {
role: prompt.role,
content,
};
});
};
export const traceAISDKModel = (model) => {
if (!shouldReportTrace())
return model;
return wrapLanguageModel({
model,
middleware: {
wrapGenerate: async (opts) => {
const start = performance.now();
const generated = await opts.doGenerate();
const end = performance.now();
const textContent = generated.content
.filter((c) => c.type === "text")
.map((c) => c.text)
.join("");
const toolCalls = generated.content
.filter((c) => c.type === "tool-call")
.map((c) => c.type === "tool-call"
? {
toolName: c.toolName,
input: c.input,
toolCallId: c.toolCallId,
}
: null)
.filter(Boolean);
reportTrace({
output: {
text: textContent,
toolCalls,
},
input: processPromptForTracing(opts.params.prompt),
usage: {
inputTokens: generated.usage.inputTokens ?? 0,
outputTokens: generated.usage.outputTokens ?? 0,
totalTokens: generated.usage.totalTokens ?? 0,
},
start,
end,
});
return generated;
},
wrapStream: async ({ doStream, params }) => {
const start = performance.now();
const { stream, ...rest } = await doStream();
const fullResponse = [];
const transformStream = new TransformStream({
transform(chunk, controller) {
fullResponse.push(chunk);
controller.enqueue(chunk);
},
flush() {
const usage = fullResponse.find((part) => part.type === "finish")?.usage;
reportTrace({
start,
end: performance.now(),
input: processPromptForTracing(params.prompt),
output: fullResponse,
usage: usage
? {
inputTokens: usage.inputTokens ?? 0,
outputTokens: usage.outputTokens ?? 0,
totalTokens: usage.totalTokens ?? 0,
}
: undefined,
});
},
});
return {
stream: stream.pipeThrough(transformStream),
...rest,
};
},
},
});
};
//# sourceMappingURL=ai-sdk.js.map