[WIP] [DO NOT MERGE] Operation Sdk

clients/ai-foundation_models-v1/sdk/imageGenerationSdk.ts

@@ -20,15 +20,16 @@ export class ImageGenerationSdk {
         ClientCallArgs
     >;
 
-    constructor(session: SessionArg) {
+    constructor(session: SessionArg, endpoint = 'llm.api.cloud.yandex.net:443') {
         this.imageGenerationClient = session.client(
             imageGenerationService.ImageGenerationAsyncServiceClient,
+            endpoint,
         );
     }
 
     generateImage(params: GenerateImageProps, args?: ClientCallArgs) {
         const { modelId, folderId, ...restParams } = params;
-        const modelUri = `gpt://${folderId}/${modelId}`;
+        const modelUri = `art://${folderId}/${modelId}`;
 
         return this.imageGenerationClient.generate(
             imageGenerationService.ImageGenerationRequest.fromPartial({ ...restParams, modelUri }),

clients/operation/sdk/index.ts

@@ -0,0 +1 @@
+export * from './operationSdk';

clients/operation/sdk/operationSdk.ts

@@ -0,0 +1,133 @@
+import { isObject, noop } from 'lodash';
+import { operationService } from '..';
+import {
+    CancelOperationRequest,
+    GetOperationRequest,
+    OperationServiceService,
+} from '../generated/yandex/cloud/operation/operation_service';
+
+import { ClientCallArgs, SessionArg, TypeFromProtoc } from './types';
+import { Client } from 'nice-grpc';
+import { Operation } from '../generated/yandex/cloud/operation/operation';
+import { Reader } from 'protobufjs';
+
+export type GetOperationProps = TypeFromProtoc<GetOperationRequest, 'operationId'>;
+
+export type CancelOperationProps = TypeFromProtoc<CancelOperationRequest, 'operationId'>;
+
+type PollArgs<DecoderT> = {
+    operationCallback?: (operation: Operation) => void;
+    decoder?: (input: Reader | Uint8Array, length?: number) => DecoderT;
+};
+
+interface CancellablePromise<T> extends Promise<T> {
+    cancelPolling?: () => void;
+}
+
+export class PollOperationWasCanceled {
+    operation?: Operation;
+    constructor(operation?: Operation) {
+        this.operation = operation;
+    }
+}
+
+export class PollOperationEmptyResponseForDecoder {
+    operation?: Operation;
+    constructor(operation?: Operation) {
+        this.operation = operation;
+    }
+}
+
+export class OperationSdk {
+    private operationClient: Client<typeof OperationServiceService, ClientCallArgs>;
+
+    constructor(session: SessionArg) {
+        this.operationClient = session.client(operationService.OperationServiceClient);
+    }
+
+    static PollOperationWasCanceled = PollOperationWasCanceled;
+    static PollOperationEmptyResponseForDecoder = PollOperationEmptyResponseForDecoder;
+
+    public pollOperation<DecoderT = Operation>(
+        operation: Operation | string,
+        intervalMs: number,
+        args?: PollArgs<DecoderT>,
+    ): CancellablePromise<DecoderT> {
+        let outputReject: (reason?: unknown) => void = noop;
+        let timeoutId: NodeJS.Timeout | undefined;
+
+        const operationStatusHandler = (operation: Operation) => {
+            if (operation.done || operation.error) {
+                if (timeoutId !== undefined) clearTimeout(timeoutId);
+                return true;
+            }
+
+            return false;
+        };
+
+        const operationDecoderHandler = (operation: Operation): DecoderT => {
+            if (args?.decoder) {
+                if (operation.response === undefined) {
+                    throw new PollOperationEmptyResponseForDecoder(operation);
+                }
+
+                return args.decoder(operation.response.value);
+            }
+
+            return operation as DecoderT;
+        };
+
+        if (isObject(operation) && operationStatusHandler(operation)) {
+            return Promise.resolve(operation).then(operationDecoderHandler);
+        }
+
+        const p = new Promise<Operation>((resolver, reject) => {
+            outputReject = reject;
+            const operationId = isObject(operation) ? operation.id : operation;
+
+            const f = () => {
+                this.get({ operationId }).then((operation) => {
+                    args?.operationCallback?.(operation);
+
+                    if (operationStatusHandler(operation)) {
+                        timeoutId = undefined;
+                        resolver(operation);
+                        return;
+                    }
+
+                    timeoutId = setTimeout(() => f(), intervalMs);
+                });
+            };
+
+            f();
+        });
+
+        (p as CancellablePromise<Operation>).cancelPolling = () => {
+            outputReject?.(
+                new PollOperationWasCanceled(isObject(operation) ? operation : undefined),
+            );
+
+            if (timeoutId !== undefined) clearTimeout(timeoutId);
+        };
+
+        return p.then(operationDecoderHandler);
+    }
+
+    public get(params: GetOperationProps, args?: ClientCallArgs) {
+        return this.operationClient.get(
+            operationService.GetOperationRequest.fromPartial(params),
+            args,
+        );
+    }
+
+    public cancel(params: CancelOperationProps, args?: ClientCallArgs) {
+        return this.operationClient.cancel(
+            operationService.CancelOperationRequest.fromPartial(params),
+            args,
+        );
+    }
+}
+
+export const initOperationSdk = (session: SessionArg) => {
+    return new OperationSdk(session);
+};

clients/operation/sdk/types.ts

@@ -0,0 +1,139 @@
+import { ChannelCredentials, ChannelOptions, Client, ServiceDefinition } from '@grpc/grpc-js';
+import { ClientError, RawClient, Status } from 'nice-grpc';
+import { DeadlineOptions } from 'nice-grpc-client-middleware-deadline';
+import { NormalizedServiceDefinition } from 'nice-grpc/lib/service-definitions';
+import { Operation as YandexCloudOperation } from '../generated/yandex/cloud/operation/operation';
+
+import { DeepPartial } from '../generated/typeRegistry';
+
+type RetryOptions = {
+    /**
+     * Boolean indicating whether retries are enabled.
+     *
+     * If the method is marked as idempotent in Protobuf, i.e. has
+     *
+     *     option idempotency_level = IDEMPOTENT;
+     *
+     * then the default is `true`. Otherwise, the default is `false`.
+     *
+     * Method options currently work only when compiling with `ts-proto`.
+     */
+    retry?: boolean;
+    /**
+     * Base delay between retry attempts in milliseconds.
+     *
+     * Defaults to 1000.
+     *
+     * Example: if `retryBaseDelayMs` is 100, then retries will be attempted in
+     * 100ms, 200ms, 400ms etc. (not counting jitter).
+     */
+    retryBaseDelayMs?: number;
+    /**
+     * Maximum delay between attempts in milliseconds.
+     *
+     * Defaults to 15 seconds.
+     *
+     * Example: if `retryBaseDelayMs` is 1000 and `retryMaxDelayMs` is 3000, then
+     * retries will be attempted in 1000ms, 2000ms, 3000ms, 3000ms etc (not
+     * counting jitter).
+     */
+    retryMaxDelayMs?: number;
+    /**
+     * Maximum for the total number of attempts. `Infinity` is supported.
+     *
+     * Defaults to 1, i.e. a single retry will be attempted.
+     */
+    retryMaxAttempts?: number;
+    /**
+     * Array of retryable status codes.
+     *
+     * Default is `[UNKNOWN, RESOURCE_EXHAUSTED, INTERNAL, UNAVAILABLE]`.
+     */
+    retryableStatuses?: Status[];
+    /**
+     * Called after receiving error with retryable status code before setting
+     * backoff delay timer.
+     *
+     * If the error code is not retryable, or the maximum attempts exceeded, this
+     * function will not be called and the error will be thrown from the client
+     * method.
+     */
+    onRetryableError?(error: ClientError, attempt: number, delayMs: number): void;
+};
+
+export interface TokenService {
+    getToken: () => Promise<string>;
+}
+
+export interface GeneratedServiceClientCtor<T extends ServiceDefinition> {
+    service: T;
+
+    new (
+        address: string,
+        credentials: ChannelCredentials,
+        options?: Partial<ChannelOptions>,
+    ): Client;
+}
+
+export interface IIAmCredentials {
+    serviceAccountId: string;
+    accessKeyId: string;
+    privateKey: Buffer | string;
+}
+
+export interface ISslCredentials {
+    rootCertificates?: Buffer;
+    clientPrivateKey?: Buffer;
+    clientCertChain?: Buffer;
+}
+
+export interface ChannelSslOptions {
+    rootCerts?: Buffer;
+    privateKey?: Buffer;
+    certChain?: Buffer;
+}
+
+export interface GenericCredentialsConfig {
+    pollInterval?: number;
+    ssl?: ChannelSslOptions;
+    headers?: Record<string, string>;
+}
+
+export interface OAuthCredentialsConfig extends GenericCredentialsConfig {
+    oauthToken: string;
+}
+
+export interface IamTokenCredentialsConfig extends GenericCredentialsConfig {
+    iamToken: string;
+}
+
+export interface ServiceAccountCredentialsConfig extends GenericCredentialsConfig {
+    serviceAccountJson: IIAmCredentials;
+}
+
+export type SessionConfig =
+    | OAuthCredentialsConfig
+    | IamTokenCredentialsConfig
+    | ServiceAccountCredentialsConfig
+    | GenericCredentialsConfig;
+
+export type ClientCallArgs = DeadlineOptions & RetryOptions;
+
+export type WrappedServiceClientType<S extends ServiceDefinition> = RawClient<
+    NormalizedServiceDefinition<S>,
+    ClientCallArgs
+>;
+
+export type ClientType = <S extends ServiceDefinition>(
+    clientClass: GeneratedServiceClientCtor<S>,
+    customEndpoint?: string,
+) => WrappedServiceClientType<S>;
+
+export type SessionArg = { client: ClientType };
+
+export type TypeFromProtoc<
+    T extends { $type: string },
+    NotPartialKey extends keyof Omit<T, '$type'> = never,
+> = {
+    [Key in NotPartialKey]: T[Key];
+} & DeepPartial<T>;