@azure/cosmos
Version:
Microsoft Azure Cosmos DB Service Node.js SDK for NOSQL API
292 lines (291 loc) • 12.4 kB
JavaScript
var __defProp = Object.defineProperty;
var __getOwnPropDesc = Object.getOwnPropertyDescriptor;
var __getOwnPropNames = Object.getOwnPropertyNames;
var __hasOwnProp = Object.prototype.hasOwnProperty;
var __export = (target, all) => {
for (var name in all)
__defProp(target, name, { get: all[name], enumerable: true });
};
var __copyProps = (to, from, except, desc) => {
if (from && typeof from === "object" || typeof from === "function") {
for (let key of __getOwnPropNames(from))
if (!__hasOwnProp.call(to, key) && key !== except)
__defProp(to, key, { get: () => from[key], enumerable: !(desc = __getOwnPropDesc(from, key)) || desc.enumerable });
}
return to;
};
var __toCommonJS = (mod) => __copyProps(__defProp({}, "__esModule", { value: true }), mod);
var InferenceService_exports = {};
__export(InferenceService_exports, {
InferenceService: () => InferenceService
});
module.exports = __toCommonJS(InferenceService_exports);
var import_core_rest_pipeline = require("@azure/core-rest-pipeline");
var import_logger = require("@azure/logger");
var import_constants = require("../common/constants.js");
var import_statusCodes = require("../common/statusCodes.js");
var import_cachedClient = require("../utils/cachedClient.js");
var import_ErrorResponse = require("../request/ErrorResponse.js");
var import_DiagnosticNodeInternal = require("../diagnostics/DiagnosticNodeInternal.js");
var import_diagnostics = require("../utils/diagnostics.js");
var import_time = require("../utils/time.js");
const logger = (0, import_logger.createClientLogger)("InferenceService");
const NON_PAYLOAD_KEYS = /* @__PURE__ */ new Set(["abortSignal"]);
const HTTP_MULTIPLE_CHOICES = 300;
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}${import_constants.Constants.Inference.BasePath}`;
this.inferenceRequestTimeoutMs = this.resolveRequestTimeout(semanticRerankConfig);
this.pipeline = this.createInferencePipeline(cosmosClientOptions.aadCredentials);
this.httpClient = cosmosClientOptions.httpClient ?? (0, import_cachedClient.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"];
const timeoutController = new AbortController();
const onCallerAbort = () => timeoutController.abort();
if (callerSignal) {
if (callerSignal.aborted) {
timeoutController.abort();
} else {
callerSignal.addEventListener("abort", onCallerAbort, { once: true });
}
}
const request = (0, import_core_rest_pipeline.createPipelineRequest)({
url: this.inferenceEndpointUrl,
method: "POST",
body: JSON.stringify(payload),
abortSignal: timeoutController.signal
});
this.setHeaders(request);
const sendAndParse = async (node) => {
const startTimeUTCInMs = (0, import_time.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: (0, import_time.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) {
if (timedOut && !callerSignal?.aborted) {
throw this.createTimeoutError(startTimeUTCInMs);
}
throw error;
} finally {
clearTimeout(timeoutHandle);
if (callerSignal) {
callerSignal.removeEventListener("abort", onCallerAbort);
}
}
};
return diagnosticNode ? (0, import_diagnostics.addDiagnosticChild)(
(childNode) => sendAndParse(childNode),
diagnosticNode,
import_DiagnosticNodeInternal.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 : void 0;
}
/**
* Resolves the inference endpoint from `enablePreviewFeatures.semanticRerank.inferenceEndpoint`.
*/
resolveInferenceEndpoint(semanticRerankConfig) {
const endpointValue = semanticRerankConfig?.inferenceEndpoint;
const endpoint = typeof endpointValue === "string" ? endpointValue : void 0;
if (!endpoint) {
throw new Error(
`Inference endpoint is required for semantic reranking. Set 'inferenceEndpoint' under the 'semanticRerank' key of 'enablePreviewFeatures' on CosmosClientOptions.`
);
}
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 : import_constants.Constants.Inference.DefaultRequestTimeoutMs;
}
/**
* Creates a pipeline configured for inference service authentication.
*/
createInferencePipeline(credential) {
const pipeline = (0, import_core_rest_pipeline.createEmptyPipeline)();
pipeline.addPolicy(
(0, import_core_rest_pipeline.bearerTokenAuthenticationPolicy)({
credential,
scopes: import_constants.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(import_constants.Constants.HttpHeaders.Version, import_constants.Constants.CurrentVersion);
request.headers.set(import_constants.Constants.HttpHeaders.UserAgent, import_constants.Constants.Inference.UserAgent);
request.headers.set(import_constants.Constants.HttpHeaders.CustomUserAgent, import_constants.Constants.Inference.UserAgent);
}
/**
* Builds the JSON payload for the semantic rerank request.
*/
buildPayload(rerankContext, documents, options) {
const payload = {};
if (options) {
for (const [key, value] of Object.entries(options)) {
if (!NON_PAYLOAD_KEYS.has(key) && value !== void 0) {
payload[key] = value;
}
}
}
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 < import_statusCodes.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 ?? void 0,
tokenUsage: body.token_usage ?? void 0,
headers: response.headers.toJSON(),
diagnostics: (0, import_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 !== void 0) {
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) : void 0,
message: parts.join(" ") || void 0
};
}
/**
* 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 import_ErrorResponse.ErrorResponse(message);
errorResponse.code = response.status;
errorResponse.body = errorBody;
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 = (0, import_time.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 import_ErrorResponse.ErrorResponse(message);
errorResponse.code = import_statusCodes.StatusCodes.RequestTimeout;
errorResponse.body = errorBody;
return errorResponse;
}
}
// Annotate the CommonJS export names for ESM import in node:
0 && (module.exports = {
InferenceService
});
//# sourceMappingURL=InferenceService.js.map