diff --git a/src/commands/translate/index.spec.ts b/src/commands/translate/index.spec.ts index 1cb647edc..312f4c312 100644 --- a/src/commands/translate/index.spec.ts +++ b/src/commands/translate/index.spec.ts @@ -12,7 +12,7 @@ describe('Translate command', () => { test( 'should fail on unknown provider', '--provider unknown --folder 1', - `error: option '--provider ' argument 'unknown' is invalid. Allowed choices are yandex.`, + `error: option '--provider ' argument 'unknown' is invalid. Allowed choices are yandex, yandexgpt, openai, openrouter, anthropic.`, ); test('should handle default', '--folder 1', { diff --git a/src/commands/translate/index.ts b/src/commands/translate/index.ts index 7ea9d8dd9..59b180c74 100644 --- a/src/commands/translate/index.ts +++ b/src/commands/translate/index.ts @@ -19,6 +19,7 @@ import {DESCRIPTION, NAME, options} from './config'; import {Extract} from './commands/extract'; import {Compose} from './commands/compose'; import {Extension as YandexTranslation} from './providers/yandex'; +import {Extension as AITranslation} from './providers/ai'; import {resolveSource, resolveTargets, resolveVars} from './utils'; import {Run} from './run'; import {configDefaults} from './utils/config'; @@ -89,6 +90,7 @@ export class Translate extends BaseProgram { this.extract, this.compose, new YandexTranslation(), + new AITranslation(), new ExtractOpenapiIncluderFakeExtension(), ]; diff --git a/src/commands/translate/providers/ai/auth.ts b/src/commands/translate/providers/ai/auth.ts new file mode 100644 index 000000000..f26dbe53e --- /dev/null +++ b/src/commands/translate/providers/ai/auth.ts @@ -0,0 +1,35 @@ +import {existsSync, readFileSync} from 'node:fs'; + +const YANDEX_BEARER_PREFIXES = ['y0_', 't1.']; + +/** + * Resolves an auth token from a CLI value. + * The value may be either the raw token or a path to a file containing the token. + */ +export function resolveToken(value: string): string { + if (!value) { + throw new Error('No auth token provided'); + } + + const trimmed = value.trim(); + + if (existsSync(trimmed)) { + return readFileSync(trimmed, 'utf8').trim(); + } + + return trimmed; +} + +/** + * Builds an Authorization header value for Yandex AI Studio. + * Supports IAM/OAuth tokens (Bearer) and service-account API keys (Api-Key). + */ +export function yandexAuthHeader(token: string): string { + for (const prefix of YANDEX_BEARER_PREFIXES) { + if (token.startsWith(prefix)) { + return 'Bearer ' + token; + } + } + + return 'Api-Key ' + token; +} diff --git a/src/commands/translate/providers/ai/clients/anthropic.ts b/src/commands/translate/providers/ai/clients/anthropic.ts new file mode 100644 index 000000000..59cad9422 --- /dev/null +++ b/src/commands/translate/providers/ai/clients/anthropic.ts @@ -0,0 +1,91 @@ +import type {ChatMessage, CompletionOptions, CompletionResult, LLMClient} from './types'; + +import axios from 'axios'; + +import {LLMResponseError, throwLLMError} from '../utils'; + +const DEFAULT_BASE_URL = 'https://api.anthropic.com/v1'; +const ANTHROPIC_VERSION = '2023-06-01'; + +type AnthropicMessagesResponse = { + content: {type: string; text?: string}[]; + stop_reason: string; + usage?: { + input_tokens?: number; + output_tokens?: number; + }; +}; + +export type AnthropicClientOptions = { + token: string; + model: string; + baseUrl?: string; + timeout?: number; +}; + +export class AnthropicClient implements LLMClient { + readonly name = 'anthropic'; + + private readonly token: string; + private readonly model: string; + private readonly baseUrl: string; + private readonly timeout: number; + + constructor(options: AnthropicClientOptions) { + this.token = options.token; + this.model = options.model; + this.baseUrl = (options.baseUrl || DEFAULT_BASE_URL).replace(/\/$/, ''); + this.timeout = options.timeout ?? 60_000; + } + + async complete(messages: ChatMessage[], options: CompletionOptions): Promise { + const system = messages + .filter((m) => m.role === 'system') + .map((m) => m.content) + .join('\n\n'); + const conversation = messages + .filter((m) => m.role !== 'system') + .map((m) => ({role: m.role, content: m.content})); + + try { + const {data} = await axios.post( + `${this.baseUrl}/messages`, + { + model: this.model, + system: system || undefined, + messages: conversation, + temperature: options.temperature, + max_tokens: options.maxTokens, + }, + { + timeout: this.timeout, + headers: { + 'x-api-key': this.token, + 'anthropic-version': ANTHROPIC_VERSION, + 'Content-Type': 'application/json', + 'User-Agent': 'github.com/diplodoc-platform/cli', + }, + }, + ); + + const text = data.content + .filter((block) => block.type === 'text' && block.text) + .map((block) => block.text as string) + .join(''); + + if (!text) { + throw new LLMResponseError('Anthropic returned an empty response'); + } + + return { + text, + usage: { + inputTokens: data.usage?.input_tokens ?? 0, + outputTokens: data.usage?.output_tokens ?? 0, + }, + }; + } catch (error) { + throwLLMError(error, 'Anthropic'); + } + } +} diff --git a/src/commands/translate/providers/ai/clients/openai.ts b/src/commands/translate/providers/ai/clients/openai.ts new file mode 100644 index 000000000..4f126efe9 --- /dev/null +++ b/src/commands/translate/providers/ai/clients/openai.ts @@ -0,0 +1,113 @@ +import type {ChatMessage, CompletionOptions, CompletionResult, LLMClient} from './types'; + +import axios from 'axios'; + +import {LLMResponseError, throwLLMError} from '../utils'; + +const DEFAULT_OPENAI_BASE_URL = 'https://api.openai.com/v1'; +const DEFAULT_OPENROUTER_BASE_URL = 'https://openrouter.ai/api/v1'; + +type OpenAIChatResponse = { + choices: { + message: {role: string; content: string}; + finish_reason: string; + }[]; + usage?: { + prompt_tokens?: number; + completion_tokens?: number; + total_tokens?: number; + }; +}; + +export type OpenAICompatibleClientOptions = { + token: string; + model: string; + baseUrl?: string; + timeout?: number; + extraHeaders?: Record; + name?: string; +}; + +export class OpenAICompatibleClient implements LLMClient { + readonly name: string; + + private readonly token: string; + private readonly model: string; + private readonly baseUrl: string; + private readonly timeout: number; + private readonly extraHeaders: Record; + + constructor(options: OpenAICompatibleClientOptions) { + this.name = options.name || 'openai'; + this.token = options.token; + this.model = options.model; + this.baseUrl = (options.baseUrl || DEFAULT_OPENAI_BASE_URL).replace(/\/$/, ''); + this.timeout = options.timeout ?? 60_000; + this.extraHeaders = options.extraHeaders || {}; + } + + async complete(messages: ChatMessage[], options: CompletionOptions): Promise { + try { + const {data} = await axios.post( + `${this.baseUrl}/chat/completions`, + { + model: this.model, + messages: messages.map((m) => ({role: m.role, content: m.content})), + temperature: options.temperature, + max_tokens: options.maxTokens, + }, + { + timeout: this.timeout, + headers: { + Authorization: `Bearer ${this.token}`, + 'Content-Type': 'application/json', + 'User-Agent': 'github.com/diplodoc-platform/cli', + ...this.extraHeaders, + }, + }, + ); + + const text = data.choices?.[0]?.message?.content; + if (!text) { + throw new LLMResponseError(`${this.name} returned an empty response`); + } + + return { + text, + usage: { + inputTokens: data.usage?.prompt_tokens ?? 0, + outputTokens: data.usage?.completion_tokens ?? 0, + }, + }; + } catch (error) { + throwLLMError(error, this.name); + } + } +} + +export function createOpenAIClient( + opts: Omit & { + baseUrl?: string; + }, +) { + return new OpenAICompatibleClient({ + ...opts, + name: 'openai', + baseUrl: opts.baseUrl || DEFAULT_OPENAI_BASE_URL, + }); +} + +export function createOpenRouterClient( + opts: Omit & {baseUrl?: string}, +) { + return new OpenAICompatibleClient({ + ...opts, + name: 'openrouter', + baseUrl: opts.baseUrl || DEFAULT_OPENROUTER_BASE_URL, + extraHeaders: { + 'HTTP-Referer': 'https://github.com/diplodoc-platform/cli', + 'X-Title': 'Diplodoc CLI', + ...(opts.extraHeaders || {}), + }, + }); +} diff --git a/src/commands/translate/providers/ai/clients/types.ts b/src/commands/translate/providers/ai/clients/types.ts new file mode 100644 index 000000000..61c13382d --- /dev/null +++ b/src/commands/translate/providers/ai/clients/types.ts @@ -0,0 +1,26 @@ +export type ChatRole = 'system' | 'user' | 'assistant'; + +export type ChatMessage = { + role: ChatRole; + content: string; +}; + +export type CompletionOptions = { + temperature: number; + maxTokens: number; +}; + +export type CompletionUsage = { + inputTokens: number; + outputTokens: number; +}; + +export type CompletionResult = { + text: string; + usage?: CompletionUsage; +}; + +export interface LLMClient { + readonly name: string; + complete(messages: ChatMessage[], options: CompletionOptions): Promise; +} diff --git a/src/commands/translate/providers/ai/clients/yandexgpt.ts b/src/commands/translate/providers/ai/clients/yandexgpt.ts new file mode 100644 index 000000000..21a6a55fa --- /dev/null +++ b/src/commands/translate/providers/ai/clients/yandexgpt.ts @@ -0,0 +1,130 @@ +import type {ChatMessage, CompletionOptions, CompletionResult, LLMClient} from './types'; + +import axios from 'axios'; + +import {LLMResponseError, throwLLMError} from '../utils'; +import {yandexAuthHeader} from '../auth'; + +const DEFAULT_ENDPOINT = 'https://llm.api.cloud.yandex.net/foundationModels/v1/completion'; + +type YandexMessage = { + role: 'system' | 'user' | 'assistant'; + text: string; +}; + +type YandexCompletionResponse = { + result: { + alternatives: { + message: {role: string; text: string}; + status: string; + }[]; + usage?: { + inputTextTokens?: string | number; + completionTokens?: string | number; + totalTokens?: string | number; + }; + modelVersion?: string; + }; +}; + +export type YandexGptClientOptions = { + token: string; + model: string; + folder?: string; + endpoint?: string; + timeout?: number; +}; + +/** + * Resolves a `gpt://` URI for Yandex AI Studio. + * Accepts either a short model name (e.g. `yandexgpt-lite`, folder taken from options) + * or a fully qualified URI (e.g. `gpt://b1g.../yandexgpt-lite/latest`). + */ +function resolveModelUri(model: string, folder: string): string { + if (model.startsWith('gpt://') || model.startsWith('ds://')) { + return model; + } + + if (!folder) { + throw new Error( + `Yandex AI Studio: --folder is required when --model is a short name (got "${model}")`, + ); + } + + const parts = model.split('/'); + if (parts.length === 1) { + return `gpt://${folder}/${model}/latest`; + } + + return `gpt://${folder}/${model}`; +} + +export class YandexGptClient implements LLMClient { + readonly name = 'yandexgpt'; + + private readonly token: string; + private readonly folder: string; + private readonly model: string; + private readonly endpoint: string; + private readonly timeout: number; + + constructor(options: YandexGptClientOptions) { + this.token = options.token; + this.folder = options.folder || ''; + this.model = options.model; + this.endpoint = options.endpoint || DEFAULT_ENDPOINT; + this.timeout = options.timeout ?? 60_000; + } + + async complete(messages: ChatMessage[], options: CompletionOptions): Promise { + const yandexMessages: YandexMessage[] = messages.map((m) => ({ + role: m.role, + text: m.content, + })); + + try { + const {data} = await axios.post( + this.endpoint, + { + modelUri: resolveModelUri(this.model, this.folder), + completionOptions: { + stream: false, + temperature: options.temperature, + maxTokens: String(options.maxTokens), + }, + messages: yandexMessages, + }, + { + timeout: this.timeout, + headers: { + Authorization: yandexAuthHeader(this.token), + 'Content-Type': 'application/json', + 'User-Agent': 'github.com/diplodoc-platform/cli', + }, + }, + ); + + const alternative = data.result.alternatives?.[0]; + if (!alternative) { + throw new LLMResponseError('Yandex AI Studio returned no alternatives'); + } + + if (alternative.status === 'ALTERNATIVE_STATUS_CONTENT_FILTER') { + throw new LLMResponseError( + 'Yandex AI Studio rejected the request (content filter)', + false, + ); + } + + return { + text: alternative.message.text, + usage: { + inputTokens: Number(data.result.usage?.inputTextTokens ?? 0), + outputTokens: Number(data.result.usage?.completionTokens ?? 0), + }, + }; + } catch (error) { + throwLLMError(error, 'Yandex AI Studio'); + } + } +} diff --git a/src/commands/translate/providers/ai/config.ts b/src/commands/translate/providers/ai/config.ts new file mode 100644 index 000000000..8003aec5b --- /dev/null +++ b/src/commands/translate/providers/ai/config.ts @@ -0,0 +1,145 @@ +import {cyan, gray} from 'chalk'; +import {dedent} from 'ts-dedent'; + +import {option} from '~/core/config'; + +const auth = option({ + flags: '--auth ', + desc: ` + Authorization token for the AI provider. + Accepts the raw token value or a path to a file containing the token. + + Yandex AI Studio: IAM token (Bearer) or service-account API key. + OpenAI / OpenRouter: Bearer key (sk-...). + Anthropic: x-api-key (sk-ant-...). + `, +}); + +const folder = option({ + flags: '--folder ', + desc: ` + Yandex AI Studio folder ID. Required when --model is a short model name + (e.g. "yandexgpt-lite") so the full gpt:// URI can be built. + `, +}); + +const model = option({ + flags: '--model ', + desc: ` + Target model identifier. + + Yandex AI Studio: short name ("yandexgpt-lite", "yandexgpt") or full URI ("gpt:///yandexgpt/latest"). + OpenAI: e.g. "gpt-4o-mini". + OpenRouter: e.g. "anthropic/claude-3.5-sonnet". + Anthropic: e.g. "claude-sonnet-4-5". + `, +}); + +const apiBase = option({ + flags: '--api-base ', + desc: ` + Override the API base URL (useful for self-hosted or compatible endpoints). + `, +}); + +const systemPrompt = option({ + flags: '--system-prompt ', + desc: ` + System prompt for the LLM. Accepts a string or a path to a file. + + Supports placeholders: {{source}}, {{target}}, {{glossary}}, {{separator}}, {{fragments}}. + + By default the user prompt is appended to the built-in technical-translator system prompt. + Use --prompt-mode replace to fully replace the default. + `, +}); + +const userPrompt = option({ + flags: '--user-prompt ', + desc: ` + User prompt template. Accepts a string or a path to a file. + Supports placeholders: {{source}}, {{target}}, {{glossary}}, {{separator}}, {{fragments}}, {{text}}. + `, +}); + +const promptMode = option({ + flags: '--prompt-mode ', + desc: ` + How the user-supplied system prompt interacts with the default one. + + ${cyan('append')} (default) — supplied system prompt is appended to the built-in default. + ${cyan('replace')} — supplied system prompt fully replaces the built-in default. + `, + choices: ['append', 'replace'], + default: 'append', +}); + +const glossaryExample = gray(dedent` + glossaryPairs: + - sourceText: string + translatedText: string +`); + +const glossary = option({ + flags: '--glossary ', + desc: ` + Path to a YAML file with required term translations. + + Config example: + ${glossaryExample} + `, +}); + +const temperature = option({ + flags: '--temperature ', + desc: 'Sampling temperature. Defaults to 0 for deterministic translation.', + parser: (value: string) => Number(value), + default: 0, +}); + +const maxOutputTokens = option({ + flags: '--max-output-tokens ', + desc: 'Maximum tokens in a single LLM response. Default 4000.', + parser: (value: string) => parseInt(value, 10), + default: 4000, +}); + +const maxBatchTokens = option({ + flags: '--max-batch-tokens ', + desc: ` + Token budget for a single LLM request. Translation units are batched up to this + limit and sent together. Smaller values are safer but slower. Default 2000. + `, + parser: (value: string) => parseInt(value, 10), + default: 2000, +}); + +const maxConcurrency = option({ + flags: '--max-concurrency ', + desc: 'Maximum concurrent LLM requests. Default 5.', + parser: (value: string) => parseInt(value, 10), + default: 5, +}); + +const retry = option({ + flags: '--retry ', + desc: 'Number of retries on retryable LLM errors. Default 3.', + parser: (value: string) => parseInt(value, 10), + default: 3, +}); + +export const options = { + auth, + folder, + model, + apiBase, + systemPrompt, + userPrompt, + promptMode, + glossary, + temperature, + maxOutputTokens, + maxBatchTokens, + maxConcurrency, + retry, +}; diff --git a/src/commands/translate/providers/ai/index.ts b/src/commands/translate/providers/ai/index.ts new file mode 100644 index 000000000..8ac42bb50 --- /dev/null +++ b/src/commands/translate/providers/ai/index.ts @@ -0,0 +1,254 @@ +import type {BaseProgram} from '~/core/program'; +import type {Translate, TranslateArgs, TranslateConfig} from '~/commands/translate'; +import type {LLMClient} from './clients/types'; +import type {GlossaryPair, PromptMode} from './prompts'; + +import {ok} from 'assert'; +import {join} from 'node:path'; + +import {getHooks as getBaseHooks} from '~/core/program'; +import {getHooks} from '~/commands/translate'; +import {defined, resolveConfig} from '~/core/config'; +import {own} from '~/core/utils'; + +import {Provider} from './provider'; +import {options} from './config'; +import {resolveToken} from './auth'; +import {resolvePromptValue} from './prompts'; +import {YandexGptClient} from './clients/yandexgpt'; +import {AnthropicClient} from './clients/anthropic'; +import {createOpenAIClient, createOpenRouterClient} from './clients/openai'; + +const PROVIDER_NAMES = ['yandexgpt', 'openai', 'openrouter', 'anthropic'] as const; +type ProviderName = (typeof PROVIDER_NAMES)[number]; + +const ExtensionName = 'AITranslation'; + +const DEFAULT_MODELS: Record = { + yandexgpt: 'yandexgpt-lite', + openai: 'gpt-4o-mini', + openrouter: 'openai/gpt-4o-mini', + anthropic: 'claude-sonnet-4-5', +}; + +const ENV_VARS: Record = { + yandexgpt: ['YANDEX_API_KEY', 'YC_IAM_TOKEN'], + openai: ['OPENAI_API_KEY'], + openrouter: ['OPENROUTER_API_KEY'], + anthropic: ['ANTHROPIC_API_KEY'], +}; + +type Args = { + auth?: string; + folder?: string; + model?: string; + apiBase?: string; + systemPrompt?: string; + userPrompt?: string; + promptMode?: PromptMode; + glossary?: string; + temperature?: number; + maxOutputTokens?: number; + maxBatchTokens?: number; + maxConcurrency?: number; + retry?: number; +}; + +type Config = { + auth: string; + folder?: string; + model: string; + apiBase?: string; + systemPrompt?: string; + userPrompt?: string; + promptMode: PromptMode; + glossary?: string; + glossaryPairs: GlossaryPair[]; + temperature: number; + maxOutputTokens: number; + maxBatchTokens: number; + maxConcurrency: number; + retry: number; +}; + +export type AITranslationConfig = TranslateConfig & Config; + +function readEnvAuth(provider: ProviderName): string | undefined { + for (const name of ENV_VARS[provider]) { + const value = process.env[name]; + if (value) { + return value; + } + } + return undefined; +} + +function makeClientFactory(provider: ProviderName) { + return function clientFactory(config: AITranslationConfig): LLMClient { + switch (provider) { + case 'yandexgpt': + return new YandexGptClient({ + token: config.auth, + folder: config.folder, + model: config.model, + endpoint: config.apiBase, + }); + case 'openai': + return createOpenAIClient({ + token: config.auth, + model: config.model, + baseUrl: config.apiBase, + }); + case 'openrouter': + return createOpenRouterClient({ + token: config.auth, + model: config.model, + baseUrl: config.apiBase, + }); + case 'anthropic': + return new AnthropicClient({ + token: config.auth, + model: config.model, + baseUrl: config.apiBase, + }); + } + }; +} + +function numberOr(value: unknown, fallback: number): number { + if (value === null || value === undefined || value === '') { + return fallback; + } + const n = Number(value); + return Number.isFinite(n) ? n : fallback; +} + +function intOr(value: unknown, fallback: number): number { + if (value === null || value === undefined || value === '') { + return fallback; + } + const n = parseInt(String(value), 10); + return Number.isFinite(n) ? n : fallback; +} + +export class Extension { + apply(program: Translate) { + getBaseHooks(program).Command.tap(ExtensionName, (_command, opts) => { + const providerOption = opts.find((option) => option.flags.match('--provider')); + ok(providerOption, 'Unable to configure `--provider` option.'); + + const choices = providerOption.argChoices || []; + for (const name of PROVIDER_NAMES) { + if (!choices.includes(name)) { + choices.push(name); + } + } + providerOption.choices(choices); + }); + + for (const providerName of PROVIDER_NAMES) { + this.registerProvider(program, providerName); + } + } + + private registerProvider(program: Translate, providerName: ProviderName) { + getHooks(program) + .Provider.for(providerName) + .tap(`${ExtensionName}.${providerName}`, (_provider, config) => { + getBaseHooks(program).Command.tap(`${ExtensionName}.${providerName}`, (command) => { + command + .addOption(options.auth) + .addOption(options.model) + .addOption(options.apiBase) + .addOption(options.systemPrompt) + .addOption(options.userPrompt) + .addOption(options.promptMode) + .addOption(options.glossary) + .addOption(options.temperature) + .addOption(options.maxOutputTokens) + .addOption(options.maxBatchTokens) + .addOption(options.maxConcurrency) + .addOption(options.retry); + + if (providerName === 'yandexgpt') { + command.addOption(options.folder); + } + }); + + getBaseHooks( + program as BaseProgram< + TranslateConfig & Partial, + TranslateArgs & Partial + >, + ).Config.tapPromise(`${ExtensionName}.${providerName}`, async (config, args) => { + ok(!config.auth, 'Do not store `authToken` in public config'); + + const rawAuth = args.auth || readEnvAuth(providerName); + ok( + rawAuth, + `Required param --auth is not configured for provider "${providerName}"`, + ); + config.auth = resolveToken(rawAuth); + + const model = + (defined('model', args, config) as string | undefined) || + DEFAULT_MODELS[providerName]; + config.model = model; + + const apiBase = defined('apiBase', args, config); + if (apiBase) { + config.apiBase = apiBase; + } + + if (providerName === 'yandexgpt') { + config.folder = defined('folder', args, config); + + const qualified = model.startsWith('gpt://') || model.startsWith('ds://'); + ok( + config.folder || qualified, + 'Yandex AI Studio: --folder is required when --model is a short name', + ); + } + + config.systemPrompt = resolvePromptValue( + defined('systemPrompt', args, config) || undefined, + ); + config.userPrompt = resolvePromptValue( + defined('userPrompt', args, config) || undefined, + ); + config.promptMode = + (defined('promptMode', args, config) as PromptMode) || 'append'; + + config.temperature = numberOr(defined('temperature', args, config), 0); + config.maxOutputTokens = intOr(defined('maxOutputTokens', args, config), 4000); + config.maxBatchTokens = intOr(defined('maxBatchTokens', args, config), 2000); + config.maxConcurrency = intOr(defined('maxConcurrency', args, config), 5); + config.retry = intOr(defined('retry', args, config), 3); + + let glossary: AbsolutePath | undefined; + if (own(args, 'glossary')) { + glossary = join(args.input, args.glossary); + } else if (own(config, 'glossary')) { + glossary = config.resolve(config.glossary); + } + + if (glossary) { + const glossaryConfig = await resolveConfig(glossary, { + defaults: {glossaryPairs: []}, + }); + config.glossaryPairs = glossaryConfig.glossaryPairs || []; + } else { + config.glossaryPairs = []; + } + + return config; + }); + + const provider = new Provider(makeClientFactory(providerName), config); + + provider.pipe(program.logger); + + return provider; + }); + } +} diff --git a/src/commands/translate/providers/ai/prompts.spec.ts b/src/commands/translate/providers/ai/prompts.spec.ts new file mode 100644 index 000000000..394610d91 --- /dev/null +++ b/src/commands/translate/providers/ai/prompts.spec.ts @@ -0,0 +1,94 @@ +import {describe, expect, it} from 'vitest'; + +import {DEFAULT_SYSTEM_PROMPT, FRAGMENT_SEPARATOR, buildMessages, splitFragments} from './prompts'; + +const config = { + promptMode: 'append' as const, + sourceLanguage: 'ru', + targetLanguage: 'en', + glossaryPairs: [], +}; + +describe('translate ai prompts', () => { + describe('splitFragments', () => { + it('should split response by delimiter', () => { + const text = `One\n${FRAGMENT_SEPARATOR}\nTwo\n${FRAGMENT_SEPARATOR}\nThree`; + + expect(splitFragments(text)).toEqual(['One', 'Two', 'Three']); + }); + + it('should tolerate extra whitespace around delimiter', () => { + const text = `One \n\n ${FRAGMENT_SEPARATOR} \n\n Two`; + + expect(splitFragments(text)).toEqual(['One', 'Two']); + }); + + it('should return single fragment when delimiter is absent', () => { + expect(splitFragments('Just one')).toEqual(['Just one']); + }); + }); + + describe('buildMessages', () => { + it('should build system and user messages with substituted placeholders', () => { + const [system, user] = buildMessages(['Hello'], config); + + expect(system.role).toBe('system'); + expect(system.content).toContain('from ru into en'); + expect(user.role).toBe('user'); + expect(user.content).toContain('Hello'); + expect(user.content).toContain(FRAGMENT_SEPARATOR); + }); + + it('should append custom system prompt to the default one', () => { + const [system] = buildMessages(['Hello'], { + ...config, + systemPrompt: 'Prefer formal tone.', + }); + + expect(system.content).toContain('professional technical documentation translator'); + expect(system.content).toContain('Prefer formal tone.'); + }); + + it('should replace the default system prompt in replace mode', () => { + const [system] = buildMessages(['Hello'], { + ...config, + systemPrompt: 'Custom only.', + promptMode: 'replace', + }); + + expect(system.content).toBe('Custom only.'); + }); + + it('should fall back to the default system prompt in replace mode without custom prompt', () => { + const [system] = buildMessages(['Hello'], { + ...config, + promptMode: 'replace', + }); + + expect(system.content).not.toContain('{{source}}'); + expect(system.content).toContain(DEFAULT_SYSTEM_PROMPT.split('\n')[0]); + }); + + it('should render glossary pairs into the user message', () => { + const [, user] = buildMessages(['Hello'], { + ...config, + glossaryPairs: [{sourceText: 'облако', translatedText: 'cloud'}], + }); + + expect(user.content).toContain('облако'); + expect(user.content).toContain('cloud'); + }); + + it('should join fragments with the delimiter', () => { + const [, user] = buildMessages(['One', 'Two'], config); + + expect(user.content).toContain(`One\n${FRAGMENT_SEPARATOR}\nTwo`); + }); + + it('should not substitute placeholders inside fragment content', () => { + const [, user] = buildMessages(['Value of {{source}} var'], config); + + expect(user.content).toContain('Value of {{source}} var'); + }); + }); +}); diff --git a/src/commands/translate/providers/ai/prompts.ts b/src/commands/translate/providers/ai/prompts.ts new file mode 100644 index 000000000..5b283be36 --- /dev/null +++ b/src/commands/translate/providers/ai/prompts.ts @@ -0,0 +1,128 @@ +import type {ChatMessage} from './clients/types'; + +import {existsSync, readFileSync} from 'node:fs'; +import {dedent} from 'ts-dedent'; + +import {escapeRegExp} from './utils'; + +export type PromptMode = 'append' | 'replace'; + +export type GlossaryPair = {sourceText: string; translatedText: string}; + +export type PromptConfig = { + systemPrompt?: string; + userPrompt?: string; + promptMode: PromptMode; + sourceLanguage: string; + targetLanguage: string; + glossaryPairs: GlossaryPair[]; +}; + +const FRAGMENT_SEPARATOR = '<<<§§§>>>'; + +export const DEFAULT_SYSTEM_PROMPT = dedent` + You are a professional technical documentation translator. + Translate the supplied fragments from {{source}} into {{target}}. + + Strict rules: + - Preserve all Markdown syntax, HTML tags, code blocks, inline code, links, images and Liquid/YFM directives exactly as they appear. + - Do not translate code, identifiers, file paths, URLs, or text inside or fenced code blocks. + - Do not add explanations, prefaces, or trailing notes — return translations only. + - Keep the same number of fragments and their original order. + - Each fragment is delimited by the line "${FRAGMENT_SEPARATOR}". Keep this exact delimiter between fragments in your output. + - If a fragment is empty or contains only formatting, return it unchanged. +`; + +export const DEFAULT_USER_PROMPT = dedent` + Translate the following fragments from {{source}} into {{target}}. + Return the translated fragments in the same order, separated by the exact delimiter line "{{separator}}". + + {{glossary}} + + {{fragments}} +`; + +export {FRAGMENT_SEPARATOR}; + +/** + * Resolves a prompt value: if it is an existing file path, read it; otherwise use as-is. + */ +export function resolvePromptValue(value: string | undefined): string | undefined { + if (!value) { + return undefined; + } + + const trimmed = value.trim(); + if (existsSync(trimmed)) { + return readFileSync(trimmed, 'utf8'); + } + + return value; +} + +function applyVars(template: string, vars: Record): string { + return template.replace(/\{\{(\w+)\}\}/g, (match, key) => { + return key in vars ? vars[key] : match; + }); +} + +function renderGlossary(pairs: GlossaryPair[]): string { + if (!pairs.length) { + return ''; + } + const lines = pairs.map( + ({sourceText, translatedText}) => `- ${sourceText} → ${translatedText}`, + ); + return `Use these required term translations:\n${lines.join('\n')}\n`; +} + +function joinFragments(fragments: string[]): string { + return fragments.join(`\n${FRAGMENT_SEPARATOR}\n`); +} + +/** + * Splits an LLM response back into fragments using the delimiter. + */ +export function splitFragments(text: string): string[] { + const delimiter = new RegExp(`\\s*${escapeRegExp(FRAGMENT_SEPARATOR)}\\s*`, 'g'); + return text.split(delimiter).map((part) => part.replace(/^\n+|\n+$/g, '')); +} + +/** + * Builds chat messages for a batch of fragments. + * + * `promptMode`: + * - `append` (default): combines the default system prompt with the user-provided system prompt. + * - `replace`: the user-provided system prompt fully replaces the default. + */ +export function buildMessages(fragments: string[], config: PromptConfig): ChatMessage[] { + const {systemPrompt, userPrompt, promptMode, sourceLanguage, targetLanguage, glossaryPairs} = + config; + + const joined = joinFragments(fragments); + const vars = { + source: sourceLanguage, + target: targetLanguage, + glossary: renderGlossary(glossaryPairs), + separator: FRAGMENT_SEPARATOR, + fragments: joined, + text: joined, + }; + + let system: string; + if (promptMode === 'replace' && systemPrompt) { + system = applyVars(systemPrompt, vars); + } else if (systemPrompt) { + system = applyVars(DEFAULT_SYSTEM_PROMPT, vars) + '\n\n' + applyVars(systemPrompt, vars); + } else { + system = applyVars(DEFAULT_SYSTEM_PROMPT, vars); + } + + const userTemplate = userPrompt || DEFAULT_USER_PROMPT; + const user = applyVars(userTemplate, vars); + + return [ + {role: 'system', content: system}, + {role: 'user', content: user}, + ]; +} diff --git a/src/commands/translate/providers/ai/provider.ts b/src/commands/translate/providers/ai/provider.ts new file mode 100644 index 000000000..16af5179b --- /dev/null +++ b/src/commands/translate/providers/ai/provider.ts @@ -0,0 +1,369 @@ +import type {Logger} from '~/core/logger'; +import type {TranslateConfig} from '~/commands/translate'; +import type {AITranslationConfig} from './index'; +import type {LLMClient} from './clients/types'; + +import {extname, join, resolve} from 'node:path'; +import {asyncify, eachLimit} from 'async'; +import liquid from '@diplodoc/transform/lib/liquid'; + +import {LogLevel} from '~/core/logger'; + +import {FileLoader, TranslateError, compose, extract, resolveSchemas} from '../../utils'; +import {TranslateLogger} from '../../logger'; + +import {Defer, LLMResponseError, backoff, bytes, estimateTokens} from './utils'; +import {buildMessages, splitFragments} from './prompts'; + +const onFatalError = () => { + process.exit(1); +}; + +type AITranslateConfig = TranslateConfig & AITranslationConfig; + +export type ClientFactory = (config: AITranslateConfig) => LLMClient; + +export class Provider { + readonly logger: TranslateLogger; + + private readonly clientFactory: ClientFactory; + + constructor(clientFactory: ClientFactory, config: TranslateConfig) { + this.clientFactory = clientFactory; + this.logger = new TranslateLogger(config); + } + + pipe(logger: Logger) { + this.logger.pipe(logger); + } + + async skip(skipped: [string, string][]) { + this.logger.skipped(skipped); + } + + async translate(files: string[], config: AITranslateConfig) { + const client = this.clientFactory(config); + const {input, output, source, target: targets, vars, dryRun, maxConcurrency} = config; + + try { + for (const target of targets) { + const cache = new Map(); + const stat = {inputTokens: 0, outputTokens: 0, requests: 0, bytes: 0}; + + const translate = makeTranslator({ + client, + config, + sourceLanguage: source.language, + targetLanguage: target.language, + cache, + stat, + logger: this.logger, + }); + + const process = makeProcessor({ + input, + output, + sourceLanguage: source.language, + targetLanguage: target.language, + vars, + translate, + }); + + await eachLimit( + files, + maxConcurrency, + asyncify(async (file: string) => { + try { + this.logger.translate(file); + await process(file); + if (!dryRun) { + this.logger.translated(file); + } + // eslint-disable-next-line @typescript-eslint/no-explicit-any + } catch (error: any) { + if (error instanceof TranslateError) { + this.logger.error(file, `${error.message}`, error.code); + if (error.fatal) { + onFatalError(); + } + } else { + this.logger.error(file, error.message); + } + } + }), + ); + + this.logger.stat( + `requests: ${stat.requests} input-tokens: ${stat.inputTokens} ` + + `output-tokens: ${stat.outputTokens} bytes: ${stat.bytes}`, + ); + } + // eslint-disable-next-line @typescript-eslint/no-explicit-any + } catch (error: any) { + if (error instanceof TranslateError) { + this.logger.topic(LogLevel.ERROR, error.code)(error.message); + } else { + this.logger.error(error); + } + process.exit(1); + } + } +} + +type ProcessorParams = { + input: string; + output: string; + sourceLanguage: string; + targetLanguage: string; + vars: Hash; + translate: Translate; +}; + +type Translate = (path: string, texts: string[]) => Promise; + +function makeProcessor(params: ProcessorParams) { + const {input, output, sourceLanguage, targetLanguage, vars, translate} = params; + const inputRoot = resolve(input); + const outputRoot = resolve(output); + + return async function (path: string) { + const ext = extname(path); + if (!['.yaml', '.md'].includes(ext)) { + return; + } + + const inputPath = join(inputRoot, path); + const outputPath = (path: string) => + join( + outputRoot, + path + .replace(inputRoot, '') + .replace('/' + sourceLanguage + '/', '/' + targetLanguage + '/'), + ); + + const content = new FileLoader(inputPath); + await content.load(); + + if (Object.keys(vars).length && content.isString) { + content.set( + liquid(content.data as string, vars, inputPath, { + conditions: 'strict', + substitutions: false, + cycles: false, + }), + ); + } + + if (!content.data) { + await content.dump(outputPath); + return; + } + + const {schemas, ajvOptions} = await resolveSchemas({content: content.data, path}); + const {units, skeleton} = extract(content.data, { + compact: true, + source: {language: sourceLanguage, locale: 'RU'}, + target: {language: targetLanguage, locale: 'US'}, + schemas, + ajvOptions, + }); + + if (!units.length) { + await content.dump(outputPath); + return; + } + + const parts = await translate(path, units); + + content.set(compose(skeleton, parts, {useSource: true, schemas, ajvOptions})); + await content.dump(outputPath); + }; +} + +type TranslatorParams = { + client: LLMClient; + config: AITranslateConfig; + sourceLanguage: string; + targetLanguage: string; + cache: Map; + stat: {inputTokens: number; outputTokens: number; requests: number; bytes: number}; + logger: Logger; +}; + +function makeTranslator(params: TranslatorParams): Translate { + const {client, config, sourceLanguage, targetLanguage, cache, stat, logger} = params; + const { + systemPrompt, + userPrompt, + promptMode, + glossaryPairs, + temperature, + maxOutputTokens, + maxBatchTokens, + maxConcurrency, + retry, + dryRun, + } = config; + + const schedule = scheduler(maxConcurrency); + + async function translateBatch(fragments: string[]): Promise { + if (!fragments.length) { + return []; + } + + const messages = buildMessages(fragments, { + systemPrompt, + userPrompt, + promptMode, + sourceLanguage, + targetLanguage, + glossaryPairs, + }); + + if (dryRun) { + const inputTokens = messages.reduce((sum, m) => sum + estimateTokens(m.content), 0); + stat.inputTokens += inputTokens; + stat.outputTokens += fragments.reduce((sum, f) => sum + estimateTokens(f), 0); + stat.requests++; + stat.bytes += bytes(fragments); + return fragments; + } + + const result = await backoff( + () => client.complete(messages, {temperature, maxTokens: maxOutputTokens}), + retry, + ); + + stat.requests++; + stat.bytes += bytes(fragments); + if (result.usage) { + stat.inputTokens += result.usage.inputTokens; + stat.outputTokens += result.usage.outputTokens; + } + + const parts = splitFragments(result.text); + + if (parts.length !== fragments.length) { + throw new LLMResponseError( + `Expected ${fragments.length} fragments in LLM response, got ${parts.length}`, + ); + } + + return parts; + } + + async function translateWithFallback(path: string, fragments: string[]): Promise { + try { + return await translateBatch(fragments); + // eslint-disable-next-line @typescript-eslint/no-explicit-any + } catch (error: any) { + if (error instanceof LLMResponseError && fragments.length > 1) { + logger.warn( + path, + `Batch of ${fragments.length} fragments failed (${error.message}); retrying one-by-one.`, + ); + const result: string[] = []; + for (const fragment of fragments) { + const single = await translateBatch([fragment]); + result.push(single[0]); + } + return result; + } + throw error; + } + } + + return async function translate(path: string, texts: string[]) { + const promises: Promise[] = []; + const requests: Promise[] = []; + let buffer: string[] = []; + let bufferTokens = 0; + + const release = () => { + if (!buffer.length) { + return; + } + const batch = buffer; + requests.push( + schedule(async () => { + try { + const translated = await translateWithFallback(path, batch); + translated.forEach((text, i) => { + cache.get(batch[i])?.resolve(text); + }); + } catch (error) { + // Reject and evict pending defers, otherwise files sharing + // the same units would await them forever. + for (const text of batch) { + const defer = cache.get(text); + if (defer) { + cache.delete(text); + defer.promise.catch(() => {}); + defer.reject(error); + } + } + } + }), + ); + buffer = []; + bufferTokens = 0; + }; + + for (const text of texts) { + const tokens = estimateTokens(text); + + const cached = cache.get(text); + if (cached) { + promises.push(cached.promise); + continue; + } + + const defer = new Defer(); + cache.set(text, defer); + promises.push(defer.promise); + + if (bufferTokens + tokens > maxBatchTokens && buffer.length) { + release(); + } + buffer.push(text); + bufferTokens += tokens; + } + + release(); + + await Promise.all(requests); + + return Promise.all(promises); + }; +} + +function scheduler(limit: number) { + let active = 0; + const queue: (() => void)[] = []; + + // Passes the freed slot directly to the next queued task, + // so `active` never exceeds `limit`. + const next = () => { + const task = queue.shift(); + if (task) { + task(); + } else { + active--; + } + }; + + return async function (action: () => Promise): Promise { + if (active >= limit) { + await new Promise((resolve) => queue.push(resolve)); + } else { + active++; + } + + try { + return await action(); + } finally { + next(); + } + }; +} diff --git a/src/commands/translate/providers/ai/utils/errors.ts b/src/commands/translate/providers/ai/utils/errors.ts new file mode 100644 index 000000000..36a7fefce --- /dev/null +++ b/src/commands/translate/providers/ai/utils/errors.ts @@ -0,0 +1,72 @@ +import {AxiosError} from 'axios'; + +import {TranslateError} from '../../../utils'; + +export class LLMRequestError extends TranslateError { + status: number; + + retryable: boolean; + + constructor( + status: number, + message: string, + info: {fatal?: boolean; retryable?: boolean} = {}, + ) { + super(message, 'LLM_REQUEST_ERROR', info.fatal); + this.status = status; + this.retryable = info.retryable ?? false; + } +} + +export class LLMAuthError extends TranslateError { + constructor(message: string) { + super(message, 'LLM_AUTH_ERROR', true); + } +} + +export class LLMRateLimitError extends TranslateError { + retryable = true; + + constructor(message: string) { + super(message, 'LLM_RATE_LIMIT'); + } +} + +export class LLMResponseError extends TranslateError { + retryable: boolean; + + constructor(message: string, retryable = true) { + super(message, 'LLM_RESPONSE_ERROR'); + this.retryable = retryable; + } +} + +/** + * Maps an axios error to the corresponding LLM error. + * Rethrows unknown errors as is. + */ +export function throwLLMError(error: unknown, provider: string): never { + if (error instanceof AxiosError) { + const {response} = error; + + if (response) { + const {status, statusText, data} = response; + const message = data?.error?.message || data?.message || statusText; + + if (status === 401 || status === 403) { + throw new LLMAuthError(`${provider} auth failed: ${message}`); + } + if (status === 429) { + throw new LLMRateLimitError(message); + } + throw new LLMRequestError(status, message, { + retryable: status >= 500 && status < 600, + }); + } + + // No response received (timeout, connection reset) - worth retrying. + throw new LLMRequestError(0, error.message, {retryable: true}); + } + + throw error; +} diff --git a/src/commands/translate/providers/ai/utils/index.spec.ts b/src/commands/translate/providers/ai/utils/index.spec.ts new file mode 100644 index 000000000..ddb373b08 --- /dev/null +++ b/src/commands/translate/providers/ai/utils/index.spec.ts @@ -0,0 +1,66 @@ +import {afterEach, beforeEach, describe, expect, it, vi} from 'vitest'; + +import {LLMAuthError, LLMRateLimitError, backoff} from './index'; + +describe('translate ai utils', () => { + describe('backoff', () => { + beforeEach(() => { + vi.useFakeTimers(); + }); + + afterEach(() => { + vi.useRealTimers(); + }); + + it('should return the result on success', async () => { + const action = vi.fn().mockResolvedValue('ok'); + + await expect(backoff(action, 3)).resolves.toBe('ok'); + expect(action).toHaveBeenCalledTimes(1); + }); + + it('should not retry non-retryable errors', async () => { + const error = new LLMAuthError('denied'); + const action = vi.fn().mockRejectedValue(error); + + await expect(backoff(action, 3)).rejects.toBe(error); + expect(action).toHaveBeenCalledTimes(1); + }); + + it('should throw the original error when retries are disabled', async () => { + const error = new LLMRateLimitError('slow down'); + const action = vi.fn().mockRejectedValue(error); + + await expect(backoff(action, 0)).rejects.toBe(error); + expect(action).toHaveBeenCalledTimes(1); + }); + + it('should retry retryable errors up to the limit', async () => { + const error = new LLMRateLimitError('slow down'); + const action = vi.fn().mockRejectedValue(error); + + const promise = backoff(action, 2); + const failure = expect(promise).rejects.toBe(error); + + await vi.runAllTimersAsync(); + await failure; + + expect(action).toHaveBeenCalledTimes(3); + }); + + it('should resolve when a retry succeeds', async () => { + const action = vi + .fn() + .mockRejectedValueOnce(new LLMRateLimitError('slow down')) + .mockResolvedValue('ok'); + + const promise = backoff(action, 2); + const success = expect(promise).resolves.toBe('ok'); + + await vi.runAllTimersAsync(); + await success; + + expect(action).toHaveBeenCalledTimes(2); + }); + }); +}); diff --git a/src/commands/translate/providers/ai/utils/index.ts b/src/commands/translate/providers/ai/utils/index.ts new file mode 100644 index 000000000..e86890e94 --- /dev/null +++ b/src/commands/translate/providers/ai/utils/index.ts @@ -0,0 +1,64 @@ +export { + LLMRequestError, + LLMAuthError, + LLMRateLimitError, + LLMResponseError, + throwLLMError, +} from './errors'; + +export class Defer { + resolve!: (text: T) => void; + + reject!: (error: unknown) => void; + + promise: Promise; + + constructor() { + this.promise = new Promise((resolve, reject) => { + this.resolve = resolve; + this.reject = reject; + }); + } +} + +export function bytes(texts: string[]) { + return texts.reduce((sum, text) => sum + text.length, 0); +} + +// Rough heuristic: ~4 chars per token. Good enough for budgeting and dry-run. +export function estimateTokens(text: string) { + return Math.ceil(text.length / 4); +} + +export function escapeRegExp(value: string) { + return value.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'); +} + +export async function wait(interval: number) { + return new Promise((resolve) => setTimeout(resolve, interval)); +} + +export async function backoff(action: () => Promise, retries: number): Promise { + const attempts = Math.max(0, retries) + 1; + + for (let attempt = 0; ; attempt++) { + try { + return await action(); + // eslint-disable-next-line @typescript-eslint/no-explicit-any + } catch (error: any) { + if (!canRetry(error) || attempt >= attempts - 1) { + throw error; + } + await wait(Math.pow(2, attempt) * 1000); + } + } +} + +// eslint-disable-next-line @typescript-eslint/no-explicit-any +function canRetry(error: any) { + if (error?.retryable === true) { + return true; + } + const status = error?.status; + return status === 429 || status === 500 || status === 502 || status === 503 || status === 504; +}