@takara-ai/miru-code 1.4.0 → 1.5.0
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/README.md +2 -0
- package/docs/self-hosted-sagemaker.md +116 -0
- package/package.json +7 -2
- package/src/cli.ts +18 -19
- package/src/credentials.ts +77 -7
- package/src/embeddings/openai.ts +194 -84
- package/src/embeddings/sagemaker.ts +268 -0
- package/src/embeddings/validate.ts +2 -1
- package/src/help.ts +35 -2
- package/src/index/bm25.ts +31 -21
- package/src/index/create.ts +114 -44
- package/src/index/incremental.ts +4 -11
- package/src/index/sparse.ts +14 -0
- package/src/literal.ts +1 -1
- package/src/prompt.ts +13 -0
- package/src/setup.ts +194 -14
package/src/embeddings/openai.ts
CHANGED
|
@@ -1,9 +1,10 @@
|
|
|
1
1
|
import { mapPool, resolveWorkerConcurrency } from "../concurrency.ts";
|
|
2
2
|
import { envFirstString, envOptionalInt, resolveEmbeddingApiKey } from "../env.ts";
|
|
3
|
+
import { createSageMakerClient, resolveSageMakerConfig } from "./sagemaker.ts";
|
|
3
4
|
|
|
4
5
|
export const DEFAULT_EMBEDDING_MODEL = "ds1-miru-int8";
|
|
5
6
|
export const DEFAULT_EMBEDDING_BASE_URL = "https://infer.takara.ai/v1";
|
|
6
|
-
const DEFAULT_BATCH_SIZE =
|
|
7
|
+
const DEFAULT_BATCH_SIZE = 360;
|
|
7
8
|
const DEFAULT_MAX_EMBED_CHARS = 1300;
|
|
8
9
|
const WINDOW_OVERLAP_CHARS = 120;
|
|
9
10
|
|
|
@@ -29,20 +30,40 @@ export interface EmbeddingBackend {
|
|
|
29
30
|
readonly dimensions: number;
|
|
30
31
|
embedDocuments(texts: string[]): Promise<Float32Array[]>;
|
|
31
32
|
embedQuery(text: string): Promise<Float32Array>;
|
|
33
|
+
/**
|
|
34
|
+
* Embed already-windowed strings with no further splitting into windows.
|
|
35
|
+
* Callers should pass at most one API batch (`resolveEmbeddingBatchSize()`).
|
|
36
|
+
* Falls back to `embedDocuments` when omitted (test mocks).
|
|
37
|
+
*/
|
|
38
|
+
embedInputs?(texts: string[]): Promise<Float32Array[]>;
|
|
32
39
|
}
|
|
33
40
|
|
|
34
|
-
|
|
41
|
+
/** SageMaker self-hosted mode has no configurable model name — the endpoint owns that choice. */
|
|
42
|
+
export function sageMakerModelId(endpointName: string): string {
|
|
43
|
+
return `sagemaker:${endpointName}`;
|
|
44
|
+
}
|
|
45
|
+
|
|
46
|
+
/** Takara model name, ignoring SageMaker config — for code that talks to Takara specifically. */
|
|
47
|
+
export function resolveTakaraEmbeddingModel(): string {
|
|
35
48
|
return envFirstString(
|
|
36
49
|
["MIRU_OPENAI_EMBEDDING_MODEL", "OPENAI_EMBEDDING_MODEL"],
|
|
37
50
|
DEFAULT_EMBEDDING_MODEL,
|
|
38
51
|
);
|
|
39
52
|
}
|
|
40
53
|
|
|
54
|
+
export function resolveEmbeddingModel(): string {
|
|
55
|
+
const sageMaker = resolveSageMakerConfig();
|
|
56
|
+
if (sageMaker) {
|
|
57
|
+
return sageMakerModelId(sageMaker.endpointName);
|
|
58
|
+
}
|
|
59
|
+
return resolveTakaraEmbeddingModel();
|
|
60
|
+
}
|
|
61
|
+
|
|
41
62
|
export function resolveMaxEmbedChars(): number {
|
|
42
63
|
return envOptionalInt(["MIRU_MAX_EMBED_CHARS"], 256) ?? DEFAULT_MAX_EMBED_CHARS;
|
|
43
64
|
}
|
|
44
65
|
|
|
45
|
-
function splitIntoWindows(text: string, maxChars: number): string[] {
|
|
66
|
+
export function splitIntoWindows(text: string, maxChars: number): string[] {
|
|
46
67
|
if (text.length <= maxChars) {
|
|
47
68
|
return [text];
|
|
48
69
|
}
|
|
@@ -81,6 +102,128 @@ export function sanitizeEmbeddingInput(text: string): string {
|
|
|
81
102
|
return out;
|
|
82
103
|
}
|
|
83
104
|
|
|
105
|
+
/** One embedding API input tied back to its source document/chunk. */
|
|
106
|
+
export interface EmbeddingWindowJob {
|
|
107
|
+
docIndex: number;
|
|
108
|
+
text: string;
|
|
109
|
+
}
|
|
110
|
+
|
|
111
|
+
export function appendEmbeddingWindowJobs(
|
|
112
|
+
jobs: EmbeddingWindowJob[],
|
|
113
|
+
docIndex: number,
|
|
114
|
+
text: string,
|
|
115
|
+
maxChars: number,
|
|
116
|
+
): void {
|
|
117
|
+
for (const windowText of splitIntoWindows(text, maxChars).map(sanitizeEmbeddingInput)) {
|
|
118
|
+
jobs.push({ docIndex, text: windowText });
|
|
119
|
+
}
|
|
120
|
+
}
|
|
121
|
+
|
|
122
|
+
export function buildEmbeddingWindowJobs(
|
|
123
|
+
texts: string[],
|
|
124
|
+
maxChars: number,
|
|
125
|
+
): { jobs: EmbeddingWindowJob[]; buckets: Float32Array[][] } {
|
|
126
|
+
const jobs: EmbeddingWindowJob[] = [];
|
|
127
|
+
const buckets: Float32Array[][] = texts.map(() => []);
|
|
128
|
+
for (let docIndex = 0; docIndex < texts.length; docIndex++) {
|
|
129
|
+
const text = texts[docIndex];
|
|
130
|
+
if (text === undefined) {
|
|
131
|
+
continue;
|
|
132
|
+
}
|
|
133
|
+
appendEmbeddingWindowJobs(jobs, docIndex, text, maxChars);
|
|
134
|
+
}
|
|
135
|
+
return { jobs, buckets };
|
|
136
|
+
}
|
|
137
|
+
|
|
138
|
+
export function batchEmbeddingWindowJobs(
|
|
139
|
+
jobs: EmbeddingWindowJob[],
|
|
140
|
+
batchSize: number,
|
|
141
|
+
): EmbeddingWindowJob[][] {
|
|
142
|
+
const batches: EmbeddingWindowJob[][] = [];
|
|
143
|
+
for (let i = 0; i < jobs.length; i += batchSize) {
|
|
144
|
+
batches.push(jobs.slice(i, i + batchSize));
|
|
145
|
+
}
|
|
146
|
+
return batches;
|
|
147
|
+
}
|
|
148
|
+
|
|
149
|
+
export function assignEmbeddingWindowVectors(
|
|
150
|
+
jobs: EmbeddingWindowJob[],
|
|
151
|
+
vectors: Float32Array[],
|
|
152
|
+
buckets: Float32Array[][],
|
|
153
|
+
): void {
|
|
154
|
+
for (let i = 0; i < jobs.length; i++) {
|
|
155
|
+
const job = jobs[i];
|
|
156
|
+
const vector = vectors[i];
|
|
157
|
+
if (job === undefined || vector === undefined) {
|
|
158
|
+
continue;
|
|
159
|
+
}
|
|
160
|
+
const bucket = buckets[job.docIndex];
|
|
161
|
+
if (bucket) {
|
|
162
|
+
bucket.push(vector);
|
|
163
|
+
}
|
|
164
|
+
}
|
|
165
|
+
}
|
|
166
|
+
|
|
167
|
+
function normalize(vec: Float32Array): Float32Array {
|
|
168
|
+
let norm = 0;
|
|
169
|
+
for (let i = 0; i < vec.length; i++) {
|
|
170
|
+
const value = vec[i] ?? 0;
|
|
171
|
+
norm += value * value;
|
|
172
|
+
}
|
|
173
|
+
norm = Math.sqrt(norm);
|
|
174
|
+
if (norm === 0) {
|
|
175
|
+
return vec;
|
|
176
|
+
}
|
|
177
|
+
const out = new Float32Array(vec.length);
|
|
178
|
+
for (let i = 0; i < vec.length; i++) {
|
|
179
|
+
out[i] = (vec[i] ?? 0) / norm;
|
|
180
|
+
}
|
|
181
|
+
return out;
|
|
182
|
+
}
|
|
183
|
+
|
|
184
|
+
export function poolWindowVectors(vectors: Float32Array[]): Float32Array {
|
|
185
|
+
if (vectors.length === 0) {
|
|
186
|
+
throw new Error("Embedding API returned no vectors");
|
|
187
|
+
}
|
|
188
|
+
const first = vectors[0];
|
|
189
|
+
if (!first) {
|
|
190
|
+
throw new Error("Embedding API returned no vectors");
|
|
191
|
+
}
|
|
192
|
+
if (vectors.length === 1) {
|
|
193
|
+
return normalize(first);
|
|
194
|
+
}
|
|
195
|
+
const pooled = new Float32Array(first.length);
|
|
196
|
+
for (const vec of vectors) {
|
|
197
|
+
for (let i = 0; i < pooled.length; i++) {
|
|
198
|
+
pooled[i] += vec[i] ?? 0;
|
|
199
|
+
}
|
|
200
|
+
}
|
|
201
|
+
for (let i = 0; i < pooled.length; i++) {
|
|
202
|
+
pooled[i] /= vectors.length;
|
|
203
|
+
}
|
|
204
|
+
return normalize(pooled);
|
|
205
|
+
}
|
|
206
|
+
|
|
207
|
+
export function poolEmbeddingWindowBuckets(buckets: Float32Array[][]): Float32Array[] {
|
|
208
|
+
return buckets.map((vectors, index) => {
|
|
209
|
+
if (vectors.length === 0) {
|
|
210
|
+
throw new Error(`Missing embedding vectors for document ${index}`);
|
|
211
|
+
}
|
|
212
|
+
return poolWindowVectors(vectors);
|
|
213
|
+
});
|
|
214
|
+
}
|
|
215
|
+
|
|
216
|
+
/** Prefer embedInputs when available; fall back for test mocks. */
|
|
217
|
+
export async function embedTextsWithBackend(
|
|
218
|
+
embeddings: EmbeddingBackend,
|
|
219
|
+
texts: string[],
|
|
220
|
+
): Promise<Float32Array[]> {
|
|
221
|
+
if (embeddings.embedInputs) {
|
|
222
|
+
return embeddings.embedInputs(texts);
|
|
223
|
+
}
|
|
224
|
+
return embeddings.embedDocuments(texts);
|
|
225
|
+
}
|
|
226
|
+
|
|
84
227
|
const TRANSIENT_EMBED_STATUSES = new Set([429, 500, 502, 503, 504]);
|
|
85
228
|
const MAX_TRANSIENT_EMBED_RETRIES = 4;
|
|
86
229
|
|
|
@@ -127,16 +270,16 @@ export function resolveEmbeddingBaseUrl(): string {
|
|
|
127
270
|
).replace(/\/$/, "");
|
|
128
271
|
}
|
|
129
272
|
|
|
130
|
-
interface Int8Embedding {
|
|
273
|
+
export interface Int8Embedding {
|
|
131
274
|
dtype: "i8";
|
|
132
275
|
values: number[];
|
|
133
276
|
scale: number;
|
|
134
277
|
zero_point: number;
|
|
135
278
|
}
|
|
136
279
|
|
|
137
|
-
type EmbeddingPayload = number[] | Int8Embedding;
|
|
280
|
+
export type EmbeddingPayload = number[] | Int8Embedding;
|
|
138
281
|
|
|
139
|
-
interface EmbeddingResponseItem {
|
|
282
|
+
export interface EmbeddingResponseItem {
|
|
140
283
|
index: number;
|
|
141
284
|
embedding: EmbeddingPayload;
|
|
142
285
|
}
|
|
@@ -170,11 +313,11 @@ export function dequantizeEmbedding(embedding: EmbeddingPayload): Float32Array {
|
|
|
170
313
|
throw new Error("Embedding API returned unsupported embedding format");
|
|
171
314
|
}
|
|
172
315
|
|
|
173
|
-
interface EmbeddingResponse {
|
|
316
|
+
export interface EmbeddingResponse {
|
|
174
317
|
data: EmbeddingResponseItem[];
|
|
175
318
|
}
|
|
176
319
|
|
|
177
|
-
interface EmbeddingClient {
|
|
320
|
+
export interface EmbeddingClient {
|
|
178
321
|
createEmbeddings(
|
|
179
322
|
input: string[] | string,
|
|
180
323
|
model: string,
|
|
@@ -182,12 +325,20 @@ interface EmbeddingClient {
|
|
|
182
325
|
): Promise<EmbeddingResponse>;
|
|
183
326
|
}
|
|
184
327
|
|
|
328
|
+
/** Shown for 401/403 — API often returns "unauthorized" which is unhelpful. */
|
|
329
|
+
export const EMBEDDING_AUTH_ERROR_MESSAGE =
|
|
330
|
+
"Not authorized. Check your API key and/or token balance.";
|
|
331
|
+
|
|
185
332
|
class EmbeddingApiError extends Error {
|
|
186
333
|
readonly status: number;
|
|
187
334
|
readonly body: string;
|
|
188
335
|
|
|
189
336
|
constructor(status: number, body: string) {
|
|
190
|
-
|
|
337
|
+
const message =
|
|
338
|
+
status === 401 || status === 403
|
|
339
|
+
? EMBEDDING_AUTH_ERROR_MESSAGE
|
|
340
|
+
: `Embedding API error ${status}: ${body.slice(0, 500)}`;
|
|
341
|
+
super(message);
|
|
191
342
|
this.name = "EmbeddingApiError";
|
|
192
343
|
this.status = status;
|
|
193
344
|
this.body = body;
|
|
@@ -266,56 +417,38 @@ export class OpenAIEmbeddingBackend implements EmbeddingBackend {
|
|
|
266
417
|
};
|
|
267
418
|
}
|
|
268
419
|
|
|
269
|
-
|
|
420
|
+
/** Embed pre-windowed inputs; issues a single API request (may 413-split). */
|
|
421
|
+
async embedInputs(texts: string[]): Promise<Float32Array[]> {
|
|
270
422
|
if (texts.length === 0) {
|
|
271
423
|
return [];
|
|
272
424
|
}
|
|
273
|
-
|
|
274
|
-
|
|
275
|
-
|
|
276
|
-
|
|
425
|
+
const vectors = await this.embedBatchRawWithRetry(texts.map(sanitizeEmbeddingInput));
|
|
426
|
+
for (const vec of vectors) {
|
|
427
|
+
if (this.dimensions === 0 && vec.length > 0) {
|
|
428
|
+
this.dimensions = vec.length;
|
|
429
|
+
}
|
|
277
430
|
}
|
|
431
|
+
return vectors;
|
|
432
|
+
}
|
|
278
433
|
|
|
279
|
-
|
|
280
|
-
|
|
281
|
-
|
|
282
|
-
for (let docIndex = 0; docIndex < texts.length; docIndex++) {
|
|
283
|
-
const text = texts[docIndex];
|
|
284
|
-
if (text === undefined) {
|
|
285
|
-
continue;
|
|
286
|
-
}
|
|
287
|
-
const windows = splitIntoWindows(text, this.maxEmbedChars).map(sanitizeEmbeddingInput);
|
|
288
|
-
for (const text of windows) {
|
|
289
|
-
jobs.push({ docIndex, text });
|
|
290
|
-
}
|
|
434
|
+
async embedDocuments(texts: string[]): Promise<Float32Array[]> {
|
|
435
|
+
if (texts.length === 0) {
|
|
436
|
+
return [];
|
|
291
437
|
}
|
|
292
438
|
|
|
439
|
+
const { jobs, buckets } = buildEmbeddingWindowJobs(texts, this.maxEmbedChars);
|
|
293
440
|
if (jobs.length === 0) {
|
|
294
441
|
return texts.map(() => new Float32Array(0));
|
|
295
442
|
}
|
|
296
443
|
|
|
297
|
-
const batches
|
|
298
|
-
for (let i = 0; i < jobs.length; i += this.batchSize) {
|
|
299
|
-
batches.push(jobs.slice(i, i + this.batchSize));
|
|
300
|
-
}
|
|
301
|
-
|
|
444
|
+
const batches = batchEmbeddingWindowJobs(jobs, this.batchSize);
|
|
302
445
|
const concurrency = resolveWorkerConcurrency();
|
|
303
446
|
await mapPool(batches, concurrency, async (batch) => {
|
|
304
|
-
const vectors = await this.
|
|
305
|
-
|
|
306
|
-
const job = batch[i];
|
|
307
|
-
const vector = vectors[i];
|
|
308
|
-
if (job === undefined || vector === undefined) {
|
|
309
|
-
continue;
|
|
310
|
-
}
|
|
311
|
-
const bucket = windowVectors[job.docIndex];
|
|
312
|
-
if (bucket) {
|
|
313
|
-
bucket.push(vector);
|
|
314
|
-
}
|
|
315
|
-
}
|
|
447
|
+
const vectors = await this.embedInputs(batch.map((job) => job.text));
|
|
448
|
+
assignEmbeddingWindowVectors(batch, vectors, buckets);
|
|
316
449
|
});
|
|
317
450
|
|
|
318
|
-
const out =
|
|
451
|
+
const out = poolEmbeddingWindowBuckets(buckets);
|
|
319
452
|
for (const vec of out) {
|
|
320
453
|
if (this.dimensions === 0 && vec.length > 0) {
|
|
321
454
|
this.dimensions = vec.length;
|
|
@@ -461,49 +594,26 @@ function vectorsFromResponse(data: EmbeddingResponseItem[], expected: number): F
|
|
|
461
594
|
});
|
|
462
595
|
}
|
|
463
596
|
|
|
464
|
-
function poolWindowVectors(vectors: Float32Array[]): Float32Array {
|
|
465
|
-
if (vectors.length === 0) {
|
|
466
|
-
throw new Error("Embedding API returned no vectors");
|
|
467
|
-
}
|
|
468
|
-
const first = vectors[0];
|
|
469
|
-
if (!first) {
|
|
470
|
-
throw new Error("Embedding API returned no vectors");
|
|
471
|
-
}
|
|
472
|
-
if (vectors.length === 1) {
|
|
473
|
-
return normalize(first);
|
|
474
|
-
}
|
|
475
|
-
const pooled = new Float32Array(first.length);
|
|
476
|
-
for (const vec of vectors) {
|
|
477
|
-
for (let i = 0; i < pooled.length; i++) {
|
|
478
|
-
pooled[i] += vec[i] ?? 0;
|
|
479
|
-
}
|
|
480
|
-
}
|
|
481
|
-
for (let i = 0; i < pooled.length; i++) {
|
|
482
|
-
pooled[i] /= vectors.length;
|
|
483
|
-
}
|
|
484
|
-
return normalize(pooled);
|
|
485
|
-
}
|
|
486
|
-
|
|
487
|
-
function normalize(vec: Float32Array): Float32Array {
|
|
488
|
-
let norm = 0;
|
|
489
|
-
for (let i = 0; i < vec.length; i++) {
|
|
490
|
-
const value = vec[i] ?? 0;
|
|
491
|
-
norm += value * value;
|
|
492
|
-
}
|
|
493
|
-
norm = Math.sqrt(norm);
|
|
494
|
-
if (norm === 0) {
|
|
495
|
-
return vec;
|
|
496
|
-
}
|
|
497
|
-
const out = new Float32Array(vec.length);
|
|
498
|
-
for (let i = 0; i < vec.length; i++) {
|
|
499
|
-
out[i] = (vec[i] ?? 0) / norm;
|
|
500
|
-
}
|
|
501
|
-
return out;
|
|
502
|
-
}
|
|
503
|
-
|
|
504
597
|
let defaultBackend: OpenAIEmbeddingBackend | null = null;
|
|
598
|
+
let sageMakerBackend: OpenAIEmbeddingBackend | null = null;
|
|
505
599
|
|
|
506
600
|
export function getEmbeddingBackend(model?: string): OpenAIEmbeddingBackend {
|
|
601
|
+
// Self-hosted mode: an endpoint ARN/name overrides Takara entirely, ignoring any
|
|
602
|
+
// explicit `model` override since the deployed endpoint is the only model available.
|
|
603
|
+
// Cached separately from the Takara singleton so the two never shadow each other.
|
|
604
|
+
const sageMaker = resolveSageMakerConfig();
|
|
605
|
+
if (sageMaker) {
|
|
606
|
+
if (!sageMakerBackend) {
|
|
607
|
+
const sageMakerModel = sageMakerModelId(sageMaker.endpointName);
|
|
608
|
+
sageMakerBackend = new OpenAIEmbeddingBackend({
|
|
609
|
+
model: sageMakerModel,
|
|
610
|
+
client: createSageMakerClient(sageMaker),
|
|
611
|
+
dimensions: resolveEmbeddingDimensions(sageMakerModel),
|
|
612
|
+
});
|
|
613
|
+
}
|
|
614
|
+
return sageMakerBackend;
|
|
615
|
+
}
|
|
616
|
+
|
|
507
617
|
if (model) {
|
|
508
618
|
return new OpenAIEmbeddingBackend({
|
|
509
619
|
model,
|
|
@@ -0,0 +1,268 @@
|
|
|
1
|
+
import type { InvokeEndpointCommandOutput } from "@aws-sdk/client-sagemaker-runtime";
|
|
2
|
+
import type { EmbeddingClient, EmbeddingPayload, EmbeddingResponse } from "./openai.ts";
|
|
3
|
+
|
|
4
|
+
type SageMakerRuntimeModule = typeof import("@aws-sdk/client-sagemaker-runtime");
|
|
5
|
+
|
|
6
|
+
/** Lazily loaded so Takara-only MCP/CLI processes never pay the AWS SDK startup cost. */
|
|
7
|
+
let sageMakerRuntimePromise: Promise<SageMakerRuntimeModule> | null = null;
|
|
8
|
+
|
|
9
|
+
function loadSageMakerRuntime(): Promise<SageMakerRuntimeModule> {
|
|
10
|
+
sageMakerRuntimePromise ??= import("@aws-sdk/client-sagemaker-runtime");
|
|
11
|
+
return sageMakerRuntimePromise;
|
|
12
|
+
}
|
|
13
|
+
|
|
14
|
+
/** arn:aws:sagemaker:<region>:<account-id>:endpoint/<endpoint-name> (also covers aws-cn/aws-us-gov). */
|
|
15
|
+
const ENDPOINT_ARN_PATTERN = /^arn:aws[a-z0-9-]*:sagemaker:([a-z0-9-]+):(\d{12}):endpoint\/(.+)$/;
|
|
16
|
+
|
|
17
|
+
export interface SageMakerEmbeddingConfig {
|
|
18
|
+
endpointName: string;
|
|
19
|
+
region: string;
|
|
20
|
+
normalize: boolean;
|
|
21
|
+
truncate: boolean;
|
|
22
|
+
truncationDirection: "Left" | "Right";
|
|
23
|
+
promptName?: string;
|
|
24
|
+
}
|
|
25
|
+
|
|
26
|
+
export function parseSageMakerEndpointArn(arn: string): {
|
|
27
|
+
region: string;
|
|
28
|
+
accountId: string;
|
|
29
|
+
endpointName: string;
|
|
30
|
+
} {
|
|
31
|
+
const match = ENDPOINT_ARN_PATTERN.exec(arn.trim());
|
|
32
|
+
if (!match) {
|
|
33
|
+
throw new Error(
|
|
34
|
+
`Invalid SageMaker endpoint ARN: "${arn}". Expected format ` +
|
|
35
|
+
"arn:aws:sagemaker:<region>:<account-id>:endpoint/<endpoint-name>.",
|
|
36
|
+
);
|
|
37
|
+
}
|
|
38
|
+
const [, region, accountId, endpointName] = match;
|
|
39
|
+
if (!region || !accountId || !endpointName) {
|
|
40
|
+
throw new Error(`Invalid SageMaker endpoint ARN: "${arn}".`);
|
|
41
|
+
}
|
|
42
|
+
return { region, accountId, endpointName };
|
|
43
|
+
}
|
|
44
|
+
|
|
45
|
+
function envBool(name: string, fallback: boolean): boolean {
|
|
46
|
+
const raw = process.env[name]?.trim().toLowerCase();
|
|
47
|
+
if (!raw) {
|
|
48
|
+
return fallback;
|
|
49
|
+
}
|
|
50
|
+
return !["false", "0", "no"].includes(raw);
|
|
51
|
+
}
|
|
52
|
+
|
|
53
|
+
function resolveTruncationDirection(): "Left" | "Right" {
|
|
54
|
+
return process.env.MIRU_SAGEMAKER_TRUNCATION_DIRECTION?.trim() === "Left" ? "Left" : "Right";
|
|
55
|
+
}
|
|
56
|
+
|
|
57
|
+
/**
|
|
58
|
+
* Reads MIRU_SAGEMAKER_ENDPOINT_ARN (or MIRU_SAGEMAKER_ENDPOINT_NAME + a region) to enable
|
|
59
|
+
* self-hosted mode. AWS credentials are resolved by the SDK's default provider chain
|
|
60
|
+
* (AWS_ACCESS_KEY_ID/AWS_SECRET_ACCESS_KEY/AWS_SESSION_TOKEN, AWS_PROFILE, or an IAM role) —
|
|
61
|
+
* there is nothing Miru-specific to configure for auth.
|
|
62
|
+
*/
|
|
63
|
+
export function resolveSageMakerConfig(): SageMakerEmbeddingConfig | null {
|
|
64
|
+
const arn = process.env.MIRU_SAGEMAKER_ENDPOINT_ARN?.trim();
|
|
65
|
+
const explicitName = process.env.MIRU_SAGEMAKER_ENDPOINT_NAME?.trim();
|
|
66
|
+
|
|
67
|
+
let endpointName: string;
|
|
68
|
+
let region: string;
|
|
69
|
+
|
|
70
|
+
if (arn) {
|
|
71
|
+
const parsed = parseSageMakerEndpointArn(arn);
|
|
72
|
+
endpointName = parsed.endpointName;
|
|
73
|
+
region = parsed.region;
|
|
74
|
+
} else if (explicitName) {
|
|
75
|
+
endpointName = explicitName;
|
|
76
|
+
const explicitRegion =
|
|
77
|
+
process.env.MIRU_SAGEMAKER_REGION?.trim() ||
|
|
78
|
+
process.env.AWS_REGION?.trim() ||
|
|
79
|
+
process.env.AWS_DEFAULT_REGION?.trim();
|
|
80
|
+
if (!explicitRegion) {
|
|
81
|
+
throw new Error(
|
|
82
|
+
"MIRU_SAGEMAKER_ENDPOINT_NAME is set but no AWS region was found. Set " +
|
|
83
|
+
"MIRU_SAGEMAKER_REGION, AWS_REGION, or AWS_DEFAULT_REGION.",
|
|
84
|
+
);
|
|
85
|
+
}
|
|
86
|
+
region = explicitRegion;
|
|
87
|
+
} else {
|
|
88
|
+
return null;
|
|
89
|
+
}
|
|
90
|
+
|
|
91
|
+
return {
|
|
92
|
+
endpointName,
|
|
93
|
+
region,
|
|
94
|
+
normalize: envBool("MIRU_SAGEMAKER_NORMALIZE", true),
|
|
95
|
+
truncate: envBool("MIRU_SAGEMAKER_TRUNCATE", true),
|
|
96
|
+
truncationDirection: resolveTruncationDirection(),
|
|
97
|
+
promptName: process.env.MIRU_SAGEMAKER_PROMPT_NAME?.trim() || undefined,
|
|
98
|
+
};
|
|
99
|
+
}
|
|
100
|
+
|
|
101
|
+
export function isSageMakerConfigured(): boolean {
|
|
102
|
+
return resolveSageMakerConfig() != null;
|
|
103
|
+
}
|
|
104
|
+
|
|
105
|
+
/** Shown for credential/permission failures — mirrors EMBEDDING_AUTH_ERROR_MESSAGE for the AWS path. */
|
|
106
|
+
export const SAGEMAKER_AUTH_ERROR_MESSAGE =
|
|
107
|
+
"Not authorized to invoke this SageMaker endpoint. Check your AWS credentials " +
|
|
108
|
+
"(AWS_ACCESS_KEY_ID/AWS_SECRET_ACCESS_KEY/AWS_SESSION_TOKEN or AWS_PROFILE) and that the " +
|
|
109
|
+
"identity has sagemaker:InvokeEndpoint permission on the endpoint.";
|
|
110
|
+
|
|
111
|
+
const AUTH_EXCEPTION_NAMES = new Set([
|
|
112
|
+
"AccessDeniedException",
|
|
113
|
+
"UnrecognizedClientException",
|
|
114
|
+
"ExpiredTokenException",
|
|
115
|
+
"CredentialsProviderError",
|
|
116
|
+
]);
|
|
117
|
+
|
|
118
|
+
class SageMakerInvokeError extends Error {
|
|
119
|
+
readonly status: number;
|
|
120
|
+
|
|
121
|
+
constructor(status: number, message: string) {
|
|
122
|
+
super(message);
|
|
123
|
+
this.name = "SageMakerInvokeError";
|
|
124
|
+
this.status = status;
|
|
125
|
+
}
|
|
126
|
+
}
|
|
127
|
+
|
|
128
|
+
function toInvokeError(err: unknown): SageMakerInvokeError {
|
|
129
|
+
const awsErr = err as {
|
|
130
|
+
name?: string;
|
|
131
|
+
message?: string;
|
|
132
|
+
$metadata?: { httpStatusCode?: number };
|
|
133
|
+
OriginalStatusCode?: string | number;
|
|
134
|
+
OriginalMessage?: string;
|
|
135
|
+
};
|
|
136
|
+
|
|
137
|
+
if (awsErr?.name && AUTH_EXCEPTION_NAMES.has(awsErr.name)) {
|
|
138
|
+
return new SageMakerInvokeError(403, SAGEMAKER_AUTH_ERROR_MESSAGE);
|
|
139
|
+
}
|
|
140
|
+
|
|
141
|
+
// ModelError wraps the embedding container's real status/message (e.g. 424 when the
|
|
142
|
+
// deployed model isn't an embedding model) behind SageMaker's own 424 response.
|
|
143
|
+
const originalStatus =
|
|
144
|
+
awsErr?.OriginalStatusCode != null ? Number(awsErr.OriginalStatusCode) : undefined;
|
|
145
|
+
const status =
|
|
146
|
+
originalStatus != null && Number.isFinite(originalStatus)
|
|
147
|
+
? originalStatus
|
|
148
|
+
: (awsErr?.$metadata?.httpStatusCode ?? 500);
|
|
149
|
+
const message = awsErr?.OriginalMessage ?? awsErr?.message ?? String(err);
|
|
150
|
+
|
|
151
|
+
if (status === 424) {
|
|
152
|
+
return new SageMakerInvokeError(
|
|
153
|
+
424,
|
|
154
|
+
`SageMaker endpoint returned 424 — the deployed model is not an embedding model (${message}).`,
|
|
155
|
+
);
|
|
156
|
+
}
|
|
157
|
+
return new SageMakerInvokeError(status, `SageMaker endpoint error ${status}: ${message}`);
|
|
158
|
+
}
|
|
159
|
+
|
|
160
|
+
function extractEmbeddings(parsed: unknown): EmbeddingPayload[] {
|
|
161
|
+
if (Array.isArray(parsed)) {
|
|
162
|
+
return parsed as EmbeddingPayload[];
|
|
163
|
+
}
|
|
164
|
+
if (parsed && typeof parsed === "object") {
|
|
165
|
+
const obj = parsed as {
|
|
166
|
+
embeddings?: EmbeddingPayload[];
|
|
167
|
+
data?: Array<{ embedding: EmbeddingPayload }>;
|
|
168
|
+
};
|
|
169
|
+
if (Array.isArray(obj.embeddings)) {
|
|
170
|
+
return obj.embeddings;
|
|
171
|
+
}
|
|
172
|
+
if (Array.isArray(obj.data)) {
|
|
173
|
+
return obj.data.map((item) => item.embedding);
|
|
174
|
+
}
|
|
175
|
+
}
|
|
176
|
+
throw new Error("SageMaker endpoint returned an unrecognized embedding payload shape");
|
|
177
|
+
}
|
|
178
|
+
|
|
179
|
+
/** Builds the TEI-style /embed request body this schema expects (`inputs`, not `input`). */
|
|
180
|
+
function buildRequestBody(
|
|
181
|
+
config: SageMakerEmbeddingConfig,
|
|
182
|
+
input: string[] | string,
|
|
183
|
+
dimensions?: number,
|
|
184
|
+
): Record<string, unknown> {
|
|
185
|
+
const body: Record<string, unknown> = {
|
|
186
|
+
inputs: input,
|
|
187
|
+
normalize: config.normalize,
|
|
188
|
+
truncate: config.truncate,
|
|
189
|
+
truncation_direction: config.truncationDirection,
|
|
190
|
+
};
|
|
191
|
+
if (dimensions != null) {
|
|
192
|
+
body.dimensions = dimensions;
|
|
193
|
+
}
|
|
194
|
+
if (config.promptName) {
|
|
195
|
+
body.prompt_name = config.promptName;
|
|
196
|
+
}
|
|
197
|
+
return body;
|
|
198
|
+
}
|
|
199
|
+
|
|
200
|
+
export function createSageMakerClient(config: SageMakerEmbeddingConfig): EmbeddingClient {
|
|
201
|
+
let client: InstanceType<SageMakerRuntimeModule["SageMakerRuntimeClient"]> | null = null;
|
|
202
|
+
|
|
203
|
+
return {
|
|
204
|
+
async createEmbeddings(
|
|
205
|
+
input: string[] | string,
|
|
206
|
+
_model: string,
|
|
207
|
+
dimensions?: number,
|
|
208
|
+
): Promise<EmbeddingResponse> {
|
|
209
|
+
const { InvokeEndpointCommand, SageMakerRuntimeClient } = await loadSageMakerRuntime();
|
|
210
|
+
client ??= new SageMakerRuntimeClient({ region: config.region });
|
|
211
|
+
|
|
212
|
+
const body = buildRequestBody(config, input, dimensions);
|
|
213
|
+
|
|
214
|
+
let response: InvokeEndpointCommandOutput;
|
|
215
|
+
try {
|
|
216
|
+
response = await client.send(
|
|
217
|
+
new InvokeEndpointCommand({
|
|
218
|
+
EndpointName: config.endpointName,
|
|
219
|
+
ContentType: "application/json",
|
|
220
|
+
Accept: "application/json",
|
|
221
|
+
Body: new TextEncoder().encode(JSON.stringify(body)),
|
|
222
|
+
}),
|
|
223
|
+
);
|
|
224
|
+
} catch (err) {
|
|
225
|
+
throw toInvokeError(err);
|
|
226
|
+
}
|
|
227
|
+
|
|
228
|
+
if (!response.Body) {
|
|
229
|
+
throw new Error("SageMaker endpoint returned an empty response body");
|
|
230
|
+
}
|
|
231
|
+
|
|
232
|
+
let parsed: unknown;
|
|
233
|
+
try {
|
|
234
|
+
parsed = JSON.parse(new TextDecoder().decode(response.Body));
|
|
235
|
+
} catch {
|
|
236
|
+
throw new Error("SageMaker endpoint returned invalid JSON");
|
|
237
|
+
}
|
|
238
|
+
|
|
239
|
+
const embeddings = extractEmbeddings(parsed);
|
|
240
|
+
return { data: embeddings.map((embedding, index) => ({ index, embedding })) };
|
|
241
|
+
},
|
|
242
|
+
};
|
|
243
|
+
}
|
|
244
|
+
|
|
245
|
+
export interface ValidateSageMakerResult {
|
|
246
|
+
valid: boolean;
|
|
247
|
+
status?: number;
|
|
248
|
+
message: string;
|
|
249
|
+
}
|
|
250
|
+
|
|
251
|
+
/** Real InvokeEndpoint round trip, used by `miru setup --sagemaker` to confirm auth + shape. */
|
|
252
|
+
export async function validateSageMakerConnection(
|
|
253
|
+
config: SageMakerEmbeddingConfig,
|
|
254
|
+
): Promise<ValidateSageMakerResult> {
|
|
255
|
+
try {
|
|
256
|
+
const client = createSageMakerClient(config);
|
|
257
|
+
const response = await client.createEmbeddings("miru setup validation", "sagemaker-setup");
|
|
258
|
+
const embedding = response.data[0]?.embedding;
|
|
259
|
+
if (!embedding) {
|
|
260
|
+
return { valid: false, message: "SageMaker endpoint returned an empty response." };
|
|
261
|
+
}
|
|
262
|
+
return { valid: true, message: "SageMaker endpoint responded successfully." };
|
|
263
|
+
} catch (err) {
|
|
264
|
+
const status = (err as { status?: number }).status;
|
|
265
|
+
const message = err instanceof Error ? err.message : String(err);
|
|
266
|
+
return { valid: false, status, message };
|
|
267
|
+
}
|
|
268
|
+
}
|
|
@@ -1,4 +1,5 @@
|
|
|
1
1
|
import {
|
|
2
|
+
EMBEDDING_AUTH_ERROR_MESSAGE,
|
|
2
3
|
embeddingDimensions,
|
|
3
4
|
resolveEmbeddingBaseUrl,
|
|
4
5
|
resolveEmbeddingDimensions,
|
|
@@ -54,7 +55,7 @@ export async function validateEmbeddingApiKey(options: {
|
|
|
54
55
|
return {
|
|
55
56
|
valid: false,
|
|
56
57
|
status: response.status,
|
|
57
|
-
message:
|
|
58
|
+
message: EMBEDDING_AUTH_ERROR_MESSAGE,
|
|
58
59
|
};
|
|
59
60
|
}
|
|
60
61
|
return {
|