feat: batch support #753

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft

LeonRuggiero wants to merge 23 commits into preview from feat-batch-support

+119 −18

llumiverse

Submodule llumiverse updated 21 files

+370 −0		BATCH.md
+32 −1		CLAUDE.md
+175 −0		common/src/batch.ts
+1 −0		common/src/index.ts
+43 −4		core/src/Driver.ts
+2 −1		drivers/package.json
+19 −2		drivers/src/test-driver/index.ts
+69 −0		drivers/src/vertexai/batch/REGISTER.MD
+246 −0		drivers/src/vertexai/batch/claude-batch.ts
+495 −0		drivers/src/vertexai/batch/embeddings-batch.ts
+144 −0		drivers/src/vertexai/batch/formatters.ts
+192 −0		drivers/src/vertexai/batch/gemini-batch.ts
+298 −0		drivers/src/vertexai/batch/gemini-files.ts
+68 −0		drivers/src/vertexai/batch/index.ts
+120 −0		drivers/src/vertexai/batch/register_files_http.ts
+471 −0		drivers/src/vertexai/batch/types.ts
+344 −9		drivers/src/vertexai/index.ts
+2 −2		drivers/src/vertexai/models/gemini.ts
+271 −0		drivers/test/batch-completion.test.ts
+463 −0		drivers/test/batch-instant.test.ts
+188 −6		pnpm-lock.yaml

packages/client/src/BatchApi.ts

-Original file line number
+Diff line change
@@ -0,0 +1,90 @@
+    import type { BatchJob, CreateBatchJobOptions, ListBatchJobsOptions, ListBatchJobsResult } from "@llumiverse/common";
+    import { ApiTopic, ClientBase } from "@vertesia/api-fetch-client";
+    export default class BatchApi extends ApiTopic {
+        constructor(parent: ClientBase) {
+            super(parent, "/api/v1/batch");
+        }
+        /**
+         * List batch jobs for a given environment
+         * @param envId The environment ID
+         * @param options Optional list options (pageSize, pageToken, filter)
+         */
+        list(envId: string, options?: ListBatchJobsOptions): Promise<ListBatchJobsResult<any, any>> {
+            return super.get('/', {
+                query: {
+                    envId,
+                    ...options
+                }
+            });
+        }
+        /**
+         * Create a new batch job
+         * @param envId The environment ID
+         * @param options Batch job creation options
+         */
+        create(envId: string, options: CreateBatchJobOptions<any, any>): Promise<BatchJob<any, any>> {
+            return this.post('/', {
+                payload: {
+                    envId,
+                    ...options
+                }
+            });
+        }
+        /**
+         * Get details of a specific batch job
+         * @param batchId The batch job ID
+         * @param envId The environment ID
+         */
+        retrieve(batchId: string, envId: string): Promise<BatchJob<any, any>> {
+            return super.get('/' + batchId, {
+                query: { envId }
+            });
+        }
+        /**
+         * Cancel a running batch job
+         * @param batchId The batch job ID
+         * @param envId The environment ID
+         */
+        cancel(batchId: string, envId: string): Promise<BatchJob<any, any>> {
+            return this.post('/' + batchId + '/cancel', {
+                query: { envId }
+            });
+        }
+        /**
+         * Delete a batch job
+         * @param batchId The batch job ID
+         * @param envId The environment ID
+         */
+        remove(batchId: string, envId: string): Promise<void> {
+            return super.delete('/' + batchId, {
+                query: { envId }
+            });
+        }
+        /**
+         * Test batch embeddings with a given environment
+         * @param envId The environment ID to test with
+         * @param model Optional model to use (defaults to gemini-embedding-001)
+         * @returns Batch job details
+         */
+        testBatchEmbeddings(envId: string, model?: string): Promise<{
+            success: boolean;
+            job_id: string;
+            job_status: string;
+            job_type: string;
+            model: string;
+        }> {
+            return this.post('/test-batch-embeddings', {
+                payload: {
+                    envId,
+                    model
+                }
+            });
+        }
+    }

packages/client/src/EnvironmentsApi.ts

-Original file line number
+Diff line change
@@ -1,6 +1,6 @@
-    import { ExecutionEnvironment, ExecutionEnvironmentCreatePayload, ExecutionEnvironmentRef, ExecutionEnvironmentUpdatePayload, LoadBalancingEnvConfig, MediatorEnvConfig } from "@vertesia/common";
     import type { AIModel, EmbeddingsOptions, EmbeddingsResult, ModelSearchPayload } from "@llumiverse/common";
     import { ApiTopic, ClientBase } from "@vertesia/api-fetch-client";
+    import { ExecutionEnvironment, ExecutionEnvironmentCreatePayload, ExecutionEnvironmentRef, ExecutionEnvironmentUpdatePayload, LoadBalancingEnvConfig, MediatorEnvConfig } from "@vertesia/common";
     export default class EnvironmentsApi extends ApiTopic {
         constructor(parent: ClientBase) {
@@ Expand Down @@

packages/client/src/client.ts

            
                      Original file line number
                      Diff line number
                      Diff line change
                  
    @@ -5,19 +5,20 @@ import AccountsApi from "./AccountsApi.js";
  
    import AnalyticsApi from "./AnalyticsApi.js";

    import { ApiKeysApi } from "./ApiKeysApi.js";

    import AppsApi from "./AppsApi.js";

    import BatchApi from "./BatchApi.js";

    import CommandsApi from "./CommandsApi.js";

    import EnvironmentsApi from "./EnvironmentsApi.js";

    import { IamApi } from "./IamApi.js";

    import InteractionsApi from "./InteractionsApi.js";

    import ProjectsApi from "./ProjectsApi.js";

    import SkillsApi from "./SkillsApi.js";

    import PromptsApi from "./PromptsApi.js";

    import { RefsApi } from "./RefsApi.js";

    import { RunsApi } from "./RunsApi.js";

    import SkillsApi from "./SkillsApi.js";

    import { ZenoClient } from "./store/client.js";

    import { VERSION, VERSION_HEADER } from "./store/version.js";

    import TrainingApi from "./TrainingApi.js";

    import UsersApi from "./UsersApi.js";

    import { VERSION, VERSION_HEADER } from "./store/version.js";

    /**

    @@ -39,9 +40,9 @@ export type VertesiaClientProps = {
  
         * @since 0.52.0

         */

        site?:

            | "api.vertesia.io"

            | "api-preview.vertesia.io"

            | "api-staging.vertesia.io";

        | "api.vertesia.io"

        | "api-preview.vertesia.io"

        | "api-staging.vertesia.io";

        serverUrl?: string;

        storeUrl?: string;

        tokenServerUrl?: string;

    @@ -93,7 +94,7 @@ export class VertesiaClient extends AbstractFetchClient<VertesiaClient> {
  
            }).withApiKey(token);

        }

        static decodeEndpoints() {}

        static decodeEndpoints() { }

        constructor(

            opts: VertesiaClientProps = {

    @@ -209,16 +210,16 @@ export class VertesiaClient extends AbstractFetchClient<VertesiaClient> {
  
            return this.withAuthCallback(

                apiKey

                    ? async () => {

                          if (!isApiKey(apiKey)) {

                              return `Bearer ${apiKey}`;

                          }

                          if (isTokenExpired(this._jwt)) {

                              const jwt = await this.getAuthToken(apiKey);

                              this._jwt = jwt.token;

                          }

                          return `Bearer ${this._jwt}`;

                      }

                        if (!isApiKey(apiKey)) {

                            return `Bearer ${apiKey}`;

                        }

                        if (isTokenExpired(this._jwt)) {

                            const jwt = await this.getAuthToken(apiKey);

                            this._jwt = jwt.token;

                        }

                        return `Bearer ${this._jwt}`;

                    }

                    : undefined,

            );

        }

    @@ -312,6 +313,7 @@ export class VertesiaClient extends AbstractFetchClient<VertesiaClient> {
  
        projects = new ProjectsApi(this);

        environments = new EnvironmentsApi(this);

        batch = new BatchApi(this);

        interactions = new InteractionsApi(this);

        skills = new SkillsApi(this);

        prompts = new PromptsApi(this);

packages/client/src/store/Batch.ts

-Original file line number
+Diff line change
@@ -0,0 +1,8 @@
+    import { ApiTopic, ClientBase } from "@vertesia/api-fetch-client";
+    export class BatchApi extends ApiTopic {
+        constructor(parent: ClientBase) {
+            super(parent, "/api/v1/batch");
+        }
+    }

packages/client/src/store/index.ts

-Original file line number
+Diff line change
@@ Expand Up / @@ -4,4 +4,5 @@ export * from "./FilesApi.js"; @@
     export * from "./ObjectsApi.js";
     export * from "./TypesApi.js";
     export * from "./WorkflowsApi.js";
+    export * from "./Batch.js";

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: batch support #753

Diff view

Diff view

There are no files selected for viewing

feat: batch support #753

Are you sure you want to change the base?

feat: batch support #753

Uh oh!

Uh oh!

Diff view

Diff view

There are no files selected for viewing