@juspay/neurolink
Version:
The pipe layer of an AI nervous system: one interface connecting provider neurons to your application, across three inference types — generate, stream and decide. `decide` returns typed, calibrated judgements from a non-generative model (~400ms, ~$0.00002
464 lines (463 loc) • 21.1 kB
JavaScript
import { BaseProvider } from "../core/baseProvider.js";
import { createStreamChannel } from "../core/streamChannel.js";
import { logger } from "../utils/logger.js";
import { resolveRequestKind } from "../core/resolveRequestKind.js";
import { resolveToolExecutionRecords, toolCallsFromSummaries, } from "../core/toolExecutionRecorder.js";
import { transformToolExecutions } from "../utils/transformationUtils.js";
import { convertZodToJsonSchema } from "../utils/schemaConversion.js";
import { withProviderRetry } from "../utils/providerRetry.js";
import { DEFAULT_MAX_STEPS } from "../core/constants.js";
import { hasNativeDoGenerate, runNativeGenerateLoop, } from "../core/nativeGenerateLoop.js";
import { withSpan } from "../telemetry/withSpan.js";
import { tracers } from "../telemetry/tracers.js";
// SageMaker-specific imports
import { getDefaultSageMakerEndpoint, getSageMakerConfig, getSageMakerModel, getSageMakerModelConfig, } from "./sagemaker/config.js";
import { handleSageMakerError, SageMakerError } from "./sagemaker/errors.js";
import { SageMakerLanguageModel } from "./sagemaker/language-model.js";
/**
* Amazon SageMaker Provider extending BaseProvider
*/
export class AmazonSageMakerProvider extends BaseProvider {
// Kept as the concrete class so SageMaker-specific members
// (testConnectivity, getModelCapabilities) stay typed; getAISDKModel()
// narrows to the AI SDK handle shape at its boundary.
sagemakerModel;
sagemakerConfig;
modelConfig;
constructor(modelName, endpointName, region, neurolink, credentials) {
super(modelName, "sagemaker", neurolink);
try {
// Load and validate configuration, then overlay per-request credentials
// Credentials are passed in rather than overlaid afterwards, so they
// are present when the config is validated.
const baseConfig = getSageMakerConfig(credentials?.region ?? region, {
...(credentials?.region !== undefined && {
region: credentials.region,
}),
...(credentials?.accessKeyId !== undefined && {
accessKeyId: credentials.accessKeyId,
}),
...(credentials?.secretAccessKey !== undefined && {
secretAccessKey: credentials.secretAccessKey,
}),
...(credentials?.sessionToken !== undefined && {
sessionToken: credentials.sessionToken,
}),
...(credentials?.endpoint !== undefined && {
endpoint: credentials.endpoint,
}),
});
this.sagemakerConfig = {
...baseConfig,
...(credentials?.region !== undefined && {
region: credentials.region,
}),
...(credentials?.accessKeyId !== undefined && {
accessKeyId: credentials.accessKeyId,
}),
...(credentials?.secretAccessKey !== undefined && {
secretAccessKey: credentials.secretAccessKey,
}),
...(credentials?.sessionToken !== undefined && {
sessionToken: credentials.sessionToken,
}),
...(credentials?.endpoint !== undefined && {
endpoint: credentials.endpoint,
}),
};
this.modelConfig = getSageMakerModelConfig(endpointName || getDefaultSageMakerEndpoint());
// Create the SageMaker LanguageModel implementation.
// SageMakerLanguageModel implements SageMakerAsLanguageModel which is
// structurally compatible with LanguageModelV2 (specificationVersion "v2",
// modelId, provider, supportedUrls, doGenerate, doStream).
this.sagemakerModel = new SageMakerLanguageModel(this.modelName, this.sagemakerConfig, this.modelConfig);
logger.debug("Amazon SageMaker Provider initialized", {
modelName: this.modelName,
endpointName: this.modelConfig.endpointName,
region: this.sagemakerConfig.region,
provider: this.providerName,
});
}
catch (error) {
logger.error("Failed to initialize SageMaker provider", {
error: error instanceof Error ? error.message : String(error),
modelName,
endpointName,
});
throw handleSageMakerError(error);
}
}
getProviderName() {
return "sagemaker";
}
getDefaultModel() {
return getSageMakerModel();
}
getAISDKModel() {
// Same sanctioned two-step as construction previously used: the class
// satisfies the structural SageMakerAsLanguageModel shape, which is
// single-assertable to the AI SDK LanguageModel handle.
const smModel = this.sagemakerModel;
return smModel;
}
/**
* Native non-streaming generate.
*
* SageMaker's doGenerate makes one invokeEndpoint call and already returns
* toolCalls; no streaming is involved, so the wire hazard that reverted the
* first migration does not apply here. This supplies only the multi-step
* iteration the ai package used to.
*
* NOT EXERCISED LIVE. This machine has no SageMaker endpoint or credentials.
* The single-step shape is identical by construction — with no tool calls the
* loop breaks after exactly one doGenerate carrying the same options the ai
* loop passed. The multi-step branch is the new code and wants a real
* endpoint before it is trusted.
*/
async generate(optionsOrPrompt, analysisSchema) {
await this.ensureModelLimits();
const options = this.normalizeTextOptions(optionsOrPrompt);
if (resolveRequestKind(options, this.modelName) !== "text") {
return super.generate(options, analysisSchema);
}
this.validateOptions(options);
const mergedTools = await this.getToolsForStream(options);
const callerOwnsFallback = "disableInternalFallback" in options &&
options.disableInternalFallback === true;
// The native loop bypasses BaseProvider.executeGeneration, so the turn
// budget has to be composed here or it stops existing for this provider.
return this.runGenerateWithModelFallback(() => this.withTurnTimeout({ ...options, tools: mergedTools }, this.getDescriptorGenerateMs(), (timedOptions) => this.executeNativeGenerate(timedOptions)), callerOwnsFallback, options);
}
async executeNativeGenerate(options) {
const startTime = Date.now();
// Middleware must wrap the model here. The native loop bypasses
// BaseProvider.executeGeneration, and with it the only place middleware was
// ever applied — a probe showed a caller's wrapGenerate running zero times
// on every native provider while their onFinish still fired, because
// onFinish had been special-cased and nothing else had.
const model = await this.getAISDKModelWithMiddleware(options);
if (!hasNativeDoGenerate(model)) {
throw this.handleProviderError(new Error("sagemaker: model handle exposes no doGenerate()"));
}
const doGenerate = async (call) => {
const format = call.responseFormat;
const result = await model.doGenerate({
...call,
...(call.maxOutputTokens !== undefined
? { maxTokens: call.maxOutputTokens }
: {}),
...(format?.type === "json"
? {
responseFormat: format.schema
? {
type: "json_schema",
json_schema: { name: "response", schema: format.schema },
}
: { type: "json_object" },
}
: {}),
});
// SageMaker's low-level model retains its legacy result for streaming
// and direct consumers. The shared loop consumes V3 content and usage.
if (Array.isArray(result.content)) {
return result;
}
const content = [];
if (typeof result.text === "string" && result.text) {
content.push({ type: "text", text: result.text });
}
if (typeof result.reasoning === "string" && result.reasoning) {
content.push({ type: "reasoning", text: result.reasoning });
}
for (const call of Array.isArray(result.toolCalls)
? result.toolCalls
: []) {
if (typeof call !== "object" || call === null) {
continue;
}
const item = call;
const fn = item.function;
if (typeof item.id === "string" && typeof fn?.name === "string") {
content.push({
type: "tool-call",
toolCallId: item.id,
toolName: fn.name,
input: fn.arguments ?? "{}",
});
}
}
const usage = result.usage;
return {
...result,
content,
usage: {
inputTokens: {
total: typeof usage?.inputTokens === "number" ? usage.inputTokens : 0,
},
outputTokens: {
total: typeof usage?.outputTokens === "number" ? usage.outputTokens : 0,
},
},
};
};
const shouldUseTools = !options.disableTools && this.supportsTools();
const toolsRecord = shouldUseTools
? options.tools || {}
: {};
const v3Tools = Object.entries(toolsRecord).map(([name, t]) => {
const tool = t;
return {
type: "function",
name,
description: tool.description ?? "",
inputSchema: (tool.inputSchema
? convertZodToJsonSchema(tool.inputSchema)
: { type: "object", properties: {} }),
};
});
// Structured output was dropped entirely on this path: the schema never
// reached the request, and nothing downstream re-imposed it, so a caller
// asking for an object got whatever JSON coerceJsonToSchema could scrape
// out of prose.
const responseFormat = options.schema
? {
type: "json",
schema: convertZodToJsonSchema(options.schema),
}
: undefined;
const conversation = (await this.buildMessagesForStream(options));
const toolExecutionSummaries = [];
const loop = await runNativeGenerateLoop({
doGenerate,
conversation,
...(responseFormat ? { responseFormat } : {}),
...(v3Tools.length > 0 ? { tools: v3Tools } : {}),
toolsRecord,
maxSteps: options.maxSteps || DEFAULT_MAX_STEPS,
...(options.maxTokens ? { maxOutputTokens: options.maxTokens } : {}),
...(options.temperature !== undefined
? { temperature: options.temperature }
: {}),
...(options.abortSignal ? { abortSignal: options.abortSignal } : {}),
...(options.toolTimeoutMs !== undefined
? { toolTimeoutMs: options.toolTimeoutMs }
: {}),
runStep: (call) => withProviderRetry(call, undefined, "sagemaker generate").catch((err) => {
throw this.handleProviderError(err);
}),
}, toolExecutionSummaries);
const enhanced = {
content: loop.text,
provider: this.providerName,
model: this.modelName,
finishReason: loop.finishReason,
usage: {
input: loop.inputTokens,
output: loop.outputTokens,
total: loop.inputTokens + loop.outputTokens,
},
responseTime: Date.now() - startTime,
...(loop.reasoning ? { reasoning: loop.reasoning } : {}),
toolsUsed: loop.toolsUsed,
toolCalls: toolCallsFromSummaries(toolExecutionSummaries),
toolExecutions: resolveToolExecutionRecords(options, transformToolExecutions(toolExecutionSummaries)),
enhancedWithTools: loop.toolsUsed.length > 0,
};
return this.finalizeNativeGenerate(enhanced, options, startTime);
}
/**
* Streaming was previously an `executeStream` override that unconditionally
* threw "not yet fully implemented" — while `SageMakerLanguageModel.doStream`
* sat one property access away, complete and working, with its own fallback
* to a synthetic stream when the endpoint does not support true streaming.
*
* This adapts that AI-SDK-shaped result to `BaseProvider`'s `doStream` hook,
* and the inherited default supplies `executeStream`. The two shapes differ:
* the language model emits typed parts (`text-delta`, `finish`) on a
* `ReadableStream`, while the hook wants text chunks plus promises for how
* the turn ended. Those promises are resolved from the `finish` part by a
* detached pump, so they settle whether or not the caller reads a chunk.
*/
async doStream(options) {
return withSpan({
name: "neurolink.provider.sagemaker.stream",
tracer: tracers.stream,
attributes: {
"provider.name": "sagemaker",
"model.name": this.modelName,
"sagemaker.endpoint": this.modelConfig.endpointName,
"sagemaker.region": this.sagemakerConfig.region,
},
}, async () => {
try {
const messages = await this.buildMessagesForStream(options);
const result = await this.sagemakerModel.doStream({
prompt: messages,
maxTokens: options.maxTokens,
temperature: options.temperature,
});
// Settled from the `finish` part below. A turn that ends without one
// — a truncated or errored stream — still settles, so a caller
// awaiting either promise cannot hang.
let settleFinishReason;
let settleUsage;
const finishReason = new Promise((resolve) => {
settleFinishReason = resolve;
});
const usage = new Promise((resolve) => {
settleUsage = resolve;
});
// The source is drained by a detached pump rather than lazily by the
// consumer. Resolving `finishReason`/`usage` from inside a generator
// ties them to somebody iterating it, and the inherited
// `executeStream` chains analytics off exactly those promises — so a
// caller that awaits `result.analytics` without consuming the stream
// would wait forever. Pumping here means the turn completes, and
// both promises settle, whether or not anyone reads a chunk.
const channel = createStreamChannel();
void (async () => {
const reader = result.stream.getReader();
let sawFinish = false;
try {
while (true) {
const { done, value } = await reader.read();
if (done) {
break;
}
const part = value;
if (part.type === "text-delta" && part.textDelta) {
channel.push({ content: part.textDelta });
continue;
}
if (part.type === "finish") {
sawFinish = true;
settleFinishReason(part.finishReason ?? "stop");
settleUsage({
inputTokens: part.usage?.inputTokens ?? part.usage?.promptTokens ?? 0,
outputTokens: part.usage?.outputTokens ??
part.usage?.completionTokens ??
0,
});
}
}
channel.close();
}
catch (error) {
channel.error(error);
}
finally {
reader.releaseLock();
// A stream that ended without a finish part — truncated or
// errored — must still settle both promises, or the same hang
// returns by a different route.
if (!sawFinish) {
settleFinishReason("unknown");
settleUsage({ inputTokens: 0, outputTokens: 0 });
}
}
})();
return {
stream: channel.iterable,
finishReason,
usage,
warnings: (result.warnings ?? []).map((warning) => warning.message ?? String(warning)),
};
}
catch (error) {
throw this.handleProviderError(error);
}
});
}
formatProviderError(error) {
if (error instanceof SageMakerError) {
return error;
}
if (error instanceof Error && error.name === "TimeoutError") {
return new SageMakerError(`SageMaker request timed out. Consider increasing timeout.`, {
code: "NETWORK_ERROR",
statusCode: 408,
cause: error,
endpoint: this.modelConfig.endpointName,
});
}
return handleSageMakerError(error, this.modelConfig.endpointName);
}
/**
* Get SageMaker-specific provider information
*/
getSageMakerInfo() {
return {
endpointName: this.modelConfig.endpointName,
modelType: this.modelConfig.modelType || "custom",
region: this.sagemakerConfig.region,
configured: !!(this.sagemakerConfig.accessKeyId && this.sagemakerConfig.secretAccessKey),
};
}
/**
* Test basic configuration
*/
async testConnection() {
try {
// Basic validation test
if (!this.sagemakerConfig.accessKeyId ||
!this.sagemakerConfig.secretAccessKey) {
return {
connected: false,
error: "AWS credentials not configured",
};
}
if (!this.modelConfig.endpointName ||
this.modelConfig.endpointName === "default-endpoint") {
return {
connected: false,
error: "SageMaker endpoint not configured",
};
}
// For now, just return that configuration looks valid
return {
connected: true,
};
}
catch (error) {
return {
connected: false,
error: error instanceof Error ? error.message : String(error),
};
}
}
/**
* Public method to get the AI SDK model for CLI and external usage
*/
async getModel() {
return this.getAISDKModel();
}
/**
* Test connectivity to the SageMaker endpoint
*/
async testConnectivity() {
const model = this.sagemakerModel;
return model.testConnectivity
? await model.testConnectivity()
: { success: false, error: "Test method not available" };
}
/**
* Get model capabilities and information
*/
getModelCapabilities() {
const model = this.sagemakerModel;
return model.getModelCapabilities
? model.getModelCapabilities()
: {
capabilities: {
streaming: true,
toolCalling: true,
structuredOutput: true,
batchInference: true,
supportedResponseFormats: ["text", "json_object"],
supportedToolTypes: ["function"],
maxBatchSize: 10,
},
};
}
}
export default AmazonSageMakerProvider;