@azure/cosmos
Version:
Microsoft Azure Cosmos DB Service Node.js SDK for NOSQL API
273 lines • 13 kB
JavaScript
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.
import { bearerTokenAuthenticationPolicy, createEmptyPipeline, createPipelineRequest, } from "@azure/core-rest-pipeline";
import { createClientLogger } from "@azure/logger";
import { Constants } from "../common/constants.js";
import { StatusCodes } from "../common/statusCodes.js";
import { getCachedDefaultHttpClient } from "../utils/cachedClient.js";
import { ErrorResponse } from "../request/ErrorResponse.js";
import { DiagnosticNodeType } from "../diagnostics/DiagnosticNodeInternal.js";
import { addDiagnosticChild, getEmptyCosmosDiagnostics } from "../utils/diagnostics.js";
import { getCurrentTimestampInMs } from "../utils/time.js";
const logger = createClientLogger("InferenceService");
/** Keys that are not part of the inference service payload. */
const NON_PAYLOAD_KEYS = new Set(["abortSignal"]);
/** HTTP redirection lower bound (kept local so it is not exported in the public API). */
const HTTP_MULTIPLE_CHOICES = 300;
/**
* Provides functionality to interact with the Cosmos DB Inference Service for semantic reranking.
* @internal
*/
export class InferenceService {
pipeline;
httpClient;
inferenceEndpointUrl;
inferenceRequestTimeoutMs;
constructor(cosmosClientOptions) {
if (!cosmosClientOptions.aadCredentials) {
throw new Error("Semantic rerank requires AAD authentication. Provide 'aadCredentials' in CosmosClientOptions.");
}
const semanticRerankConfig = this.getSemanticRerankConfig(cosmosClientOptions);
const endpoint = this.resolveInferenceEndpoint(semanticRerankConfig);
this.inferenceEndpointUrl = `${endpoint}${Constants.Inference.BasePath}`;
this.inferenceRequestTimeoutMs = this.resolveRequestTimeout(semanticRerankConfig);
this.pipeline = this.createInferencePipeline(cosmosClientOptions.aadCredentials);
this.httpClient = cosmosClientOptions.httpClient ?? getCachedDefaultHttpClient();
logger.info(`InferenceService initialized with endpoint: ${endpoint}`);
}
/**
* Sends a semantic rerank request to the inference service.
* @param rerankContext - The context (e.g. query string) to use for reranking.
* @param documents - The documents to be reranked.
* @param options - Optional settings for the reranking request.
* @param diagnosticNode - Optional diagnostic node used to record the inference REST call.
* @returns The reranking results including scores, latency, and token usage.
*/
async semanticRerank(rerankContext, documents, options, diagnosticNode) {
const payload = this.buildPayload(rerankContext, documents, options);
const callerSignal = options?.["abortSignal"];
// Enforce a single-attempt, no-retry timeout for the inference request. A dedicated
// AbortController is the authoritative per-request budget (mirrors the .NET linked
// CancellationTokenSource). Caller cancellation is linked in so it still cancels the
// in-flight request, but it is surfaced unchanged rather than as a timeout.
const timeoutController = new AbortController();
const onCallerAbort = () => timeoutController.abort();
if (callerSignal) {
if (callerSignal.aborted) {
timeoutController.abort();
}
else {
callerSignal.addEventListener("abort", onCallerAbort, { once: true });
}
}
const request = createPipelineRequest({
url: this.inferenceEndpointUrl,
method: "POST",
body: JSON.stringify(payload),
abortSignal: timeoutController.signal,
});
this.setHeaders(request);
const sendAndParse = async (node) => {
const startTimeUTCInMs = getCurrentTimestampInMs();
let timedOut = false;
const timeoutHandle = setTimeout(() => {
timedOut = true;
timeoutController.abort();
}, this.inferenceRequestTimeoutMs);
try {
const response = await this.pipeline.sendRequest(this.httpClient, request);
node?.addData({
startTimeUTCInMs,
durationInMs: getCurrentTimestampInMs() - startTimeUTCInMs,
requestPayloadLengthInBytes: request.body ? String(request.body).length : 0,
responsePayloadLengthInBytes: response.bodyAsText?.length ?? 0,
requestData: { url: this.inferenceEndpointUrl },
});
return this.parseResponse(response);
}
catch (error) {
// Surface only our own timeout as 408; caller cancellation propagates unchanged.
if (timedOut && !callerSignal?.aborted) {
throw this.createTimeoutError(startTimeUTCInMs);
}
throw error;
}
finally {
clearTimeout(timeoutHandle);
if (callerSignal) {
callerSignal.removeEventListener("abort", onCallerAbort);
}
}
};
return diagnosticNode
? addDiagnosticChild((childNode) => sendAndParse(childNode), diagnosticNode, DiagnosticNodeType.HTTP_REQUEST)
: sendAndParse();
}
/**
* Reads the `semanticRerank` preview configuration object from `enablePreviewFeatures`, if present.
*/
getSemanticRerankConfig(cosmosClientOptions) {
const config = cosmosClientOptions.enablePreviewFeatures?.["semanticRerank"];
return typeof config === "object" && config !== null
? config
: undefined;
}
/**
* Resolves the inference endpoint from `enablePreviewFeatures.semanticRerank.inferenceEndpoint`.
*/
resolveInferenceEndpoint(semanticRerankConfig) {
const endpointValue = semanticRerankConfig?.inferenceEndpoint;
const endpoint = typeof endpointValue === "string" ? endpointValue : undefined;
if (!endpoint) {
throw new Error(`Inference endpoint is required for semantic reranking. ` +
`Set 'inferenceEndpoint' under the 'semanticRerank' key of 'enablePreviewFeatures' on CosmosClientOptions.`);
}
// Remove trailing slash if present
return endpoint.replace(/\/+$/, "");
}
/**
* Resolves the per-request timeout (ms) from
* `enablePreviewFeatures.semanticRerank.inferenceRequestTimeout`, falling back to the default
* when not provided or invalid. This is a single-attempt budget with no retries.
*/
resolveRequestTimeout(semanticRerankConfig) {
const timeoutValue = semanticRerankConfig?.inferenceRequestTimeout;
return typeof timeoutValue === "number" && timeoutValue > 0
? timeoutValue
: Constants.Inference.DefaultRequestTimeoutMs;
}
/**
* Creates a pipeline configured for inference service authentication.
*/
createInferencePipeline(credential) {
const pipeline = createEmptyPipeline();
pipeline.addPolicy(bearerTokenAuthenticationPolicy({
credential,
scopes: Constants.Inference.DefaultScope,
}));
return pipeline;
}
/**
* Sets the required HTTP headers on an inference service request.
*/
setHeaders(request) {
request.headers.set("Content-Type", "application/json");
request.headers.set("Accept", "application/json");
request.headers.set("Cache-Control", "no-cache");
request.headers.set(Constants.HttpHeaders.Version, Constants.CurrentVersion);
request.headers.set(Constants.HttpHeaders.UserAgent, Constants.Inference.UserAgent);
request.headers.set(Constants.HttpHeaders.CustomUserAgent, Constants.Inference.UserAgent);
}
/**
* Builds the JSON payload for the semantic rerank request.
*/
buildPayload(rerankContext, documents, options) {
const payload = {};
if (options) {
// Forward all option keys except non-payload keys (e.g. abortSignal)
for (const [key, value] of Object.entries(options)) {
if (!NON_PAYLOAD_KEYS.has(key) && value !== undefined) {
payload[key] = value;
}
}
}
// Required fields are set last to prevent options from overriding them
payload["query"] = rerankContext;
payload["documents"] = documents;
return payload;
}
/**
* Parses the HTTP response into a SemanticRerankResult.
*
* Note: The inference API response uses mixed casing conventions:
* - PascalCase: `Scores` (array of rerank results)
* - camelCase: `latency` (timing info), `document`, `score`, `index`
* - snake_case: `token_usage` (token consumption)
* This is the actual service response format, not a bug.
*/
parseResponse(response) {
if (response.status < StatusCodes.Ok || response.status >= HTTP_MULTIPLE_CHOICES) {
const { code, message } = this.parseServiceError(response.bodyAsText);
throw this.createInferenceError(response, code ?? String(response.status), message ?? `Semantic rerank request failed with status ${response.status}`);
}
if (!response.bodyAsText) {
throw this.createInferenceError(response, String(response.status), "Semantic rerank response body was empty.");
}
const body = JSON.parse(response.bodyAsText);
if (!Array.isArray(body.Scores)) {
throw this.createInferenceError(response, String(response.status), "Semantic rerank response did not contain a Scores array.");
}
const rerankScores = body.Scores.map((item) => ({
document: typeof item.document === "string" ? item.document : "",
score: typeof item.score === "number" ? item.score : 0,
index: typeof item.index === "number" ? item.index : -1,
}));
return {
rerankScores,
latency: body.latency ?? undefined,
tokenUsage: body.token_usage ?? undefined,
headers: response.headers.toJSON(),
diagnostics: getEmptyCosmosDiagnostics(),
};
}
/**
* Parses a service error body into `{ code, message }`. `message` is the body's `message` field
* followed by every other field except `code` (kept even when null) so no detail is lost. A
* non-JSON body is returned verbatim in `message`; an empty body yields an empty object.
*/
parseServiceError(text) {
if (!text) {
return {};
}
let parsed;
try {
parsed = JSON.parse(text);
}
catch {
return { message: text };
}
if (typeof parsed !== "object" || parsed === null) {
return { message: text };
}
const { code, message, ...rest } = parsed;
const parts = [];
if (message !== undefined) {
parts.push(typeof message === "string" ? message : String(JSON.stringify(message)));
}
for (const [key, value] of Object.entries(rest)) {
parts.push(`${key}: ${typeof value === "string" ? value : String(JSON.stringify(value))}`);
}
return {
code: code != null ? String(code) : undefined,
message: parts.join(" ") || undefined,
};
}
/**
* Builds an ErrorResponse carrying the HTTP status on `code` and the service error on `body`.
*/
createInferenceError(response, serviceCode, message) {
const errorBody = { code: serviceCode, message };
const errorResponse = new ErrorResponse(message);
errorResponse.code = response.status;
errorResponse.body = errorBody;
// All response headers (including x-correlation-id) are surfaced here.
errorResponse.headers = response.headers.toJSON();
return errorResponse;
}
/**
* Builds an ErrorResponse for a client-side inference request timeout, carrying HTTP status
* 408 (Request Timeout). No retries are attempted; this is a single-attempt budget.
*/
createTimeoutError(startTimeUTCInMs) {
const elapsedMs = getCurrentTimestampInMs() - startTimeUTCInMs;
const message = `Semantic rerank request timed out after ${this.inferenceRequestTimeoutMs} ms ` +
`(elapsed ${elapsedMs} ms). Adjust 'inferenceRequestTimeout' under ` +
`'enablePreviewFeatures.semanticRerank' on CosmosClientOptions to change this budget.`;
const errorBody = { code: "RequestTimeout", message };
const errorResponse = new ErrorResponse(message);
errorResponse.code = StatusCodes.RequestTimeout;
errorResponse.body = errorBody;
return errorResponse;
}
}
//# sourceMappingURL=InferenceService.js.map