235 lines
9.6 KiB
TypeScript
235 lines
9.6 KiB
TypeScript
/** Token counts the provider reports for one completion. */
|
|
export interface AiCompletionTokenUsage {
|
|
promptTokens?: number;
|
|
completionTokens?: number;
|
|
totalTokens?: number;
|
|
}
|
|
|
|
export interface LiteLlmEnvironment {
|
|
ASSESSOR_AI_MODE?: string;
|
|
ASSESSOR_AI_BASE_URL?: string;
|
|
ASSESSOR_AI_MODEL?: string;
|
|
ASSESSOR_AI_API_KEY?: string;
|
|
ASSESSOR_AI_TIMEOUT_MS?: string;
|
|
ASSESSOR_AI_MAX_OUTPUT_BYTES?: string;
|
|
}
|
|
|
|
export interface AiCompletion {
|
|
value: unknown;
|
|
usage?: AiCompletionTokenUsage;
|
|
}
|
|
|
|
export interface AiCompletionClient {
|
|
readonly model: string;
|
|
complete(systemPrompt: string, input: string): Promise<AiCompletion>;
|
|
}
|
|
|
|
export class LiteLlmClientError extends Error {
|
|
constructor(readonly code: string, message: string) {
|
|
super(message);
|
|
this.name = 'LiteLlmClientError';
|
|
}
|
|
}
|
|
|
|
interface LiteLlmClientOptions {
|
|
endpoint: URL;
|
|
model: string;
|
|
apiKey: string;
|
|
timeoutMs: number;
|
|
maximumOutputBytes: number;
|
|
fetcher?: typeof fetch;
|
|
}
|
|
|
|
export class LiteLlmClient implements AiCompletionClient {
|
|
readonly model: string;
|
|
private readonly endpoint: URL;
|
|
private readonly apiKey: string;
|
|
private readonly timeoutMs: number;
|
|
private readonly maximumOutputBytes: number;
|
|
private readonly fetcher: typeof fetch;
|
|
|
|
constructor(options: LiteLlmClientOptions) {
|
|
this.endpoint = new URL(options.endpoint.href);
|
|
this.model = options.model;
|
|
this.apiKey = options.apiKey;
|
|
this.timeoutMs = options.timeoutMs;
|
|
this.maximumOutputBytes = options.maximumOutputBytes;
|
|
this.fetcher = options.fetcher ?? globalThis.fetch.bind(globalThis);
|
|
}
|
|
|
|
async complete(systemPrompt: string, input: string): Promise<AiCompletion> {
|
|
const controller = new AbortController();
|
|
const timeout = setTimeout(() => controller.abort(), this.timeoutMs);
|
|
try {
|
|
const response = await this.fetcher(this.endpoint, {
|
|
method: 'POST',
|
|
redirect: 'error',
|
|
signal: controller.signal,
|
|
headers: {
|
|
Accept: 'application/json',
|
|
'Content-Type': 'application/json',
|
|
Authorization: `Bearer ${this.apiKey}`
|
|
},
|
|
body: JSON.stringify({
|
|
model: this.model,
|
|
response_format: { type: 'json_object' },
|
|
messages: [
|
|
{ role: 'system', content: systemPrompt },
|
|
{ role: 'user', content: input }
|
|
]
|
|
})
|
|
});
|
|
if (!response.ok) throw new LiteLlmClientError('litellm_http_error', `Work LiteLLM returned HTTP ${response.status}.${await safeErrorHint(response)}`);
|
|
const text = await readBoundedText(response, this.maximumOutputBytes);
|
|
let envelope: unknown;
|
|
try { envelope = JSON.parse(text); } catch { throw new LiteLlmClientError('litellm_invalid_response', 'Work LiteLLM returned invalid JSON.'); }
|
|
const content = completionContent(envelope);
|
|
let value: unknown;
|
|
try { value = JSON.parse(content); } catch { throw new LiteLlmClientError('litellm_invalid_output', 'Work LiteLLM did not return a JSON object.'); }
|
|
return { value, usage: completionUsage(envelope) };
|
|
} catch (error) {
|
|
if (error instanceof LiteLlmClientError) throw error;
|
|
if (controller.signal.aborted) throw new LiteLlmClientError('litellm_timeout', 'Work LiteLLM request timed out.');
|
|
throw new LiteLlmClientError('litellm_unavailable', 'Work LiteLLM could not be reached.');
|
|
} finally {
|
|
clearTimeout(timeout);
|
|
}
|
|
}
|
|
}
|
|
|
|
export function createLiteLlmClientFromEnvironment(
|
|
environment: LiteLlmEnvironment = {
|
|
ASSESSOR_AI_MODE: process.env.ASSESSOR_AI_MODE,
|
|
ASSESSOR_AI_BASE_URL: process.env.ASSESSOR_AI_BASE_URL,
|
|
ASSESSOR_AI_MODEL: process.env.ASSESSOR_AI_MODEL,
|
|
ASSESSOR_AI_API_KEY: process.env.ASSESSOR_AI_API_KEY,
|
|
ASSESSOR_AI_TIMEOUT_MS: process.env.ASSESSOR_AI_TIMEOUT_MS,
|
|
ASSESSOR_AI_MAX_OUTPUT_BYTES: process.env.ASSESSOR_AI_MAX_OUTPUT_BYTES
|
|
},
|
|
fetcher?: typeof fetch
|
|
): LiteLlmClient | undefined {
|
|
const mode = environment.ASSESSOR_AI_MODE?.trim().toLowerCase();
|
|
if (!mode) return undefined;
|
|
if (mode !== 'litellm') throw new Error('ASSESSOR_AI_MODE must be litellm when AI generation is enabled.');
|
|
const baseUrl = environment.ASSESSOR_AI_BASE_URL?.trim();
|
|
const model = environment.ASSESSOR_AI_MODEL?.trim();
|
|
const apiKey = environment.ASSESSOR_AI_API_KEY?.trim();
|
|
if (!baseUrl || !model || !apiKey) throw new Error('LiteLLM mode requires ASSESSOR_AI_BASE_URL, ASSESSOR_AI_MODEL, and ASSESSOR_AI_API_KEY.');
|
|
const endpoint = liteLlmCompletionEndpoint(baseUrl);
|
|
return new LiteLlmClient({
|
|
endpoint,
|
|
model,
|
|
apiKey,
|
|
timeoutMs: boundedInteger(environment.ASSESSOR_AI_TIMEOUT_MS, 60_000, 1_000, 300_000, 'ASSESSOR_AI_TIMEOUT_MS'),
|
|
maximumOutputBytes: boundedInteger(environment.ASSESSOR_AI_MAX_OUTPUT_BYTES, 256 * 1024, 4 * 1024, 2 * 1024 * 1024, 'ASSESSOR_AI_MAX_OUTPUT_BYTES'),
|
|
fetcher
|
|
});
|
|
}
|
|
|
|
export function liteLlmCompletionEndpoint(baseUrl: string): URL {
|
|
let url: URL;
|
|
try { url = new URL(baseUrl); } catch { throw new Error('ASSESSOR_AI_BASE_URL must be an absolute HTTPS URL.'); }
|
|
if (url.protocol !== 'https:') throw new Error('ASSESSOR_AI_BASE_URL must use HTTPS.');
|
|
if (url.username || url.password || url.search || url.hash) throw new Error('ASSESSOR_AI_BASE_URL must not contain credentials, query parameters, or a fragment.');
|
|
if (['localhost', '127.0.0.1', '[::1]'].includes(url.hostname.toLowerCase())) {
|
|
throw new Error('ASSESSOR_AI_BASE_URL must reference the company-managed HTTPS gateway, not a local endpoint.');
|
|
}
|
|
const path = url.pathname.replace(/\/+$/, '');
|
|
url.pathname = path.endsWith('/chat/completions') ? path
|
|
: path ? `${path}/chat/completions` : '/v1/chat/completions';
|
|
return url;
|
|
}
|
|
|
|
async function readBoundedText(response: Response, maximumBytes: number): Promise<string> {
|
|
const declared = Number(response.headers.get('content-length') ?? 0);
|
|
if (declared > maximumBytes) throw new LiteLlmClientError('litellm_response_too_large', 'Work LiteLLM response exceeded the configured size limit.');
|
|
if (!response.body) return '';
|
|
const reader = response.body.getReader();
|
|
const chunks: Uint8Array[] = [];
|
|
let size = 0;
|
|
while (true) {
|
|
const { done, value } = await reader.read();
|
|
if (done) break;
|
|
size += value.byteLength;
|
|
if (size > maximumBytes) {
|
|
await reader.cancel();
|
|
throw new LiteLlmClientError('litellm_response_too_large', 'Work LiteLLM response exceeded the configured size limit.');
|
|
}
|
|
chunks.push(value);
|
|
}
|
|
const merged = new Uint8Array(size);
|
|
let offset = 0;
|
|
chunks.forEach(chunk => { merged.set(chunk, offset); offset += chunk.byteLength; });
|
|
return new TextDecoder().decode(merged);
|
|
}
|
|
|
|
async function safeErrorHint(response: Response): Promise<string> {
|
|
let envelope: unknown;
|
|
try {
|
|
const text = await readTruncatedText(response, 8 * 1024);
|
|
envelope = text ? JSON.parse(text) : undefined;
|
|
} catch {
|
|
return '';
|
|
}
|
|
const error = isObject(envelope) ? envelope.error : undefined;
|
|
const type = safeErrorToken(isObject(error) ? error.type : undefined);
|
|
const code = safeErrorToken(isObject(error) ? error.code : undefined);
|
|
const parts = [type && `type=${type}`, code && `code=${code}`].filter(Boolean);
|
|
return parts.length ? ` (${parts.join(', ')})` : '';
|
|
}
|
|
|
|
// Only short enum-like identifiers pass; free-text messages, params, and raw bodies are never echoed.
|
|
function safeErrorToken(value: unknown): string | undefined {
|
|
const token = typeof value === 'string' ? value
|
|
: typeof value === 'number' && Number.isFinite(value) ? String(value) : undefined;
|
|
return token && /^[A-Za-z0-9_.:-]{1,64}$/.test(token) ? token : undefined;
|
|
}
|
|
|
|
async function readTruncatedText(response: Response, maximumBytes: number): Promise<string> {
|
|
if (!response.body) return '';
|
|
const reader = response.body.getReader();
|
|
const chunks: Uint8Array[] = [];
|
|
let size = 0;
|
|
while (size < maximumBytes) {
|
|
const { done, value } = await reader.read();
|
|
if (done) break;
|
|
chunks.push(value);
|
|
size += value.byteLength;
|
|
}
|
|
await reader.cancel().catch(() => {});
|
|
const merged = new Uint8Array(size);
|
|
let offset = 0;
|
|
chunks.forEach(chunk => { merged.set(chunk, offset); offset += chunk.byteLength; });
|
|
return new TextDecoder().decode(merged.subarray(0, maximumBytes));
|
|
}
|
|
|
|
function completionContent(value: unknown): string {
|
|
if (!isObject(value) || !Array.isArray(value.choices) || !isObject(value.choices[0])
|
|
|| !isObject(value.choices[0].message) || typeof value.choices[0].message.content !== 'string') {
|
|
throw new LiteLlmClientError('litellm_invalid_response', 'Work LiteLLM response did not contain message content.');
|
|
}
|
|
return value.choices[0].message.content;
|
|
}
|
|
|
|
function completionUsage(value: unknown): AiCompletionTokenUsage | undefined {
|
|
if (!isObject(value) || !isObject(value.usage)) return undefined;
|
|
const number = (entry: unknown) => typeof entry === 'number' && Number.isFinite(entry) ? entry : undefined;
|
|
const usage = {
|
|
promptTokens: number(value.usage.prompt_tokens),
|
|
completionTokens: number(value.usage.completion_tokens),
|
|
totalTokens: number(value.usage.total_tokens)
|
|
};
|
|
return Object.values(usage).some(item => item !== undefined) ? usage : undefined;
|
|
}
|
|
|
|
function boundedInteger(value: string | undefined, fallback: number, minimum: number, maximum: number, name: string): number {
|
|
if (!value?.trim()) return fallback;
|
|
const parsed = Number(value);
|
|
if (!Number.isInteger(parsed) || parsed < minimum || parsed > maximum) throw new Error(`${name} must be an integer between ${minimum} and ${maximum}.`);
|
|
return parsed;
|
|
}
|
|
|
|
function isObject(value: unknown): value is Record<string, any> {
|
|
return Boolean(value) && typeof value === 'object' && !Array.isArray(value);
|
|
}
|