Skip to content

Commit 12f4fd6

Browse files
shepherdjerredshepherdjerred
committed
fix: batch
1 parent 0bc133a commit 12f4fd6

2 files changed

Lines changed: 573 additions & 51 deletions

File tree

‎packages/bun-decompile/src/lib/deminify/deminifier.ts‎

Lines changed: 146 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@ import { ClaudeClient, formatCostEstimate } from "./claude-client.ts";
88
import { OpenAIClient } from "./openai-client.ts";
99
import { BatchDeminifyClient } from "./batch-client.ts";
1010
import type { BatchStatus } from "./batch-client.ts";
11+
import { OpenAIBatchClient } from "./openai-batch.ts";
12+
import type { OpenAIBatchStatus } from "./openai-batch.ts";
1113
import {
1214
saveBatchState,
1315
loadBatchState,
@@ -77,6 +79,7 @@ export class Deminifier {
7779
private config: DeminifyConfig;
7880
private client: ClaudeClient | OpenAIClient;
7981
private batchClient: BatchDeminifyClient | null;
82+
private openAIBatchClient: OpenAIBatchClient | null;
8083
private cache: DeminifyCache | null;
8184
private stats: DeminifyStats;
8285

@@ -86,10 +89,12 @@ export class Deminifier {
8689
// Select client based on provider
8790
if (config.provider === "openai") {
8891
this.client = new OpenAIClient(config);
89-
this.batchClient = null; // OpenAI doesn't have equivalent batch API
92+
this.batchClient = null;
93+
this.openAIBatchClient = new OpenAIBatchClient(config);
9094
} else {
9195
this.client = new ClaudeClient(config);
9296
this.batchClient = new BatchDeminifyClient(config);
97+
this.openAIBatchClient = null;
9398
}
9499

95100
this.cache = config.cacheEnabled
@@ -418,7 +423,13 @@ export class Deminifier {
418423

419424
const result = await processor.processAll(source, graph, {
420425
// maxBatchTokens computed from model context limit if not specified
426+
<<<<<<< Updated upstream
421427
...(options?.maxBatchTokens !== undefined ? { maxBatchTokens: options.maxBatchTokens } : {}),
428+
||||||| Stash base
429+
maxBatchTokens: options?.maxBatchTokens,
430+
=======
431+
...(options?.maxBatchTokens !== undefined && { maxBatchTokens: options.maxBatchTokens }),
432+
>>>>>>> Stashed changes
422433
verbose: this.config.verbose,
423434
onProgress: (progress) => {
424435
const progressUpdate: DeminifyProgress = {
@@ -470,11 +481,13 @@ export class Deminifier {
470481
fileContext: FileContext,
471482
options: DeminifyFileOptions
472483
): Promise<string> {
473-
// Batch mode only available for Anthropic
474-
if (!this.batchClient) {
475-
throw new Error(
476-
"Batch mode is only available with Anthropic provider. Use --provider anthropic or remove --batch flag."
477-
);
484+
// Check that we have a batch client for the provider
485+
const isOpenAI = this.config.provider === "openai";
486+
if (isOpenAI && !this.openAIBatchClient) {
487+
throw new Error("OpenAI batch client not initialized");
488+
}
489+
if (!isOpenAI && !this.batchClient) {
490+
throw new Error("Anthropic batch client not initialized");
478491
}
479492

480493
const fileName = options.fileName ?? "unknown.js";
@@ -520,9 +533,16 @@ export class Deminifier {
520533
contexts.set(func.id, context);
521534
}
522535

523-
// Create batch
524-
console.log("Submitting batch to Anthropic API...");
525-
const batchId = await this.batchClient.createBatch(contexts);
536+
// Create batch using the appropriate client
537+
const providerName = isOpenAI ? "OpenAI" : "Anthropic";
538+
console.log(`Submitting batch to ${providerName} API...`);
539+
540+
let batchId: string;
541+
if (isOpenAI) {
542+
batchId = await this.openAIBatchClient!.createBatch(contexts);
543+
} else {
544+
batchId = await this.batchClient!.createBatch(contexts);
545+
}
526546

527547
// Save state for resume (includes projectId for isolation in shared environments)
528548
await saveBatchState(
@@ -542,27 +562,57 @@ export class Deminifier {
542562
console.log(`\nBatch submitted: ${batchId}`);
543563
console.log("Waiting for results (typically 30-60 minutes)...\n");
544564

545-
// Poll for completion
546-
await this.batchClient.waitForCompletion(batchId, {
547-
onStatusUpdate: (status) => {
548-
options.onBatchStatus?.(status);
549-
if (!options.onBatchStatus) {
550-
// Default console output
551-
const pct =
552-
status.total > 0
553-
? Math.round((status.succeeded / status.total) * 100)
554-
: 0;
555-
process.stdout.write(
556-
`\r Progress: ${status.succeeded}/${status.total} (${pct}%) | Errors: ${status.errored} `
557-
);
558-
}
559-
},
560-
});
565+
// Poll for completion using the appropriate client
566+
if (isOpenAI) {
567+
await this.openAIBatchClient!.waitForCompletion(batchId, {
568+
onStatusUpdate: (status: OpenAIBatchStatus) => {
569+
// Convert to common format for callback
570+
const commonStatus: BatchStatus = {
571+
batchId: status.batchId,
572+
status: status.status === "completed" ? "ended" : "in_progress",
573+
total: status.total,
574+
succeeded: status.completed,
575+
errored: status.failed,
576+
processing: status.total - status.completed - status.failed,
577+
};
578+
options.onBatchStatus?.(commonStatus);
579+
if (!options.onBatchStatus) {
580+
const pct =
581+
status.total > 0
582+
? Math.round((status.completed / status.total) * 100)
583+
: 0;
584+
process.stdout.write(
585+
`\r Progress: ${status.completed}/${status.total} (${pct}%) | Errors: ${status.failed} `
586+
);
587+
}
588+
},
589+
});
590+
} else {
591+
await this.batchClient!.waitForCompletion(batchId, {
592+
onStatusUpdate: (status) => {
593+
options.onBatchStatus?.(status);
594+
if (!options.onBatchStatus) {
595+
const pct =
596+
status.total > 0
597+
? Math.round((status.succeeded / status.total) * 100)
598+
: 0;
599+
process.stdout.write(
600+
`\r Progress: ${status.succeeded}/${status.total} (${pct}%) | Errors: ${status.errored} `
601+
);
602+
}
603+
},
604+
});
605+
}
561606

562607
console.log("\n\nBatch complete! Retrieving results...");
563608

564-
// Get results
565-
const results = await this.batchClient.getResults(batchId, contexts);
609+
// Get results using the appropriate client
610+
let results: Map<string, DeminifyResult>;
611+
if (isOpenAI) {
612+
results = await this.openAIBatchClient!.getResults(batchId, contexts);
613+
} else {
614+
results = await this.batchClient!.getResults(batchId, contexts);
615+
}
566616

567617
console.log(`Retrieved ${results.size} results`);
568618

@@ -610,8 +660,12 @@ export class Deminifier {
610660
batchId: string,
611661
options: DeminifyFileOptions
612662
): Promise<string> {
613-
if (!this.batchClient) {
614-
throw new Error("Batch mode is only available with Anthropic provider.");
663+
const isOpenAI = this.config.provider === "openai";
664+
if (isOpenAI && !this.openAIBatchClient) {
665+
throw new Error("OpenAI batch client not initialized");
666+
}
667+
if (!isOpenAI && !this.batchClient) {
668+
throw new Error("Anthropic batch client not initialized");
615669
}
616670

617671
// Rebuild contexts
@@ -626,33 +680,74 @@ export class Deminifier {
626680
contexts.set(func.id, context);
627681
}
628682

629-
// Check batch status
630-
const status = await this.batchClient.getBatchStatus(batchId);
631-
632-
if (status.status === "in_progress") {
633-
console.log(
634-
`Batch still processing: ${status.succeeded}/${status.total} complete`
635-
);
636-
console.log("Waiting for completion...\n");
683+
// Check batch status using appropriate client
684+
if (isOpenAI) {
685+
const status = await this.openAIBatchClient!.getBatchStatus(batchId);
686+
687+
if (
688+
status.status === "in_progress" ||
689+
status.status === "validating" ||
690+
status.status === "finalizing"
691+
) {
692+
console.log(
693+
`Batch still processing: ${status.completed}/${status.total} complete`
694+
);
695+
console.log("Waiting for completion...\n");
696+
697+
await this.openAIBatchClient!.waitForCompletion(batchId, {
698+
onStatusUpdate: (s: OpenAIBatchStatus) => {
699+
const commonStatus: BatchStatus = {
700+
batchId: s.batchId,
701+
status: s.status === "completed" ? "ended" : "in_progress",
702+
total: s.total,
703+
succeeded: s.completed,
704+
errored: s.failed,
705+
processing: s.total - s.completed - s.failed,
706+
};
707+
options.onBatchStatus?.(commonStatus);
708+
if (!options.onBatchStatus) {
709+
const pct =
710+
s.total > 0 ? Math.round((s.completed / s.total) * 100) : 0;
711+
process.stdout.write(
712+
`\r Progress: ${s.completed}/${s.total} (${pct}%) | Errors: ${s.failed} `
713+
);
714+
}
715+
},
716+
});
717+
}
718+
} else {
719+
const status = await this.batchClient!.getBatchStatus(batchId);
637720

638-
await this.batchClient.waitForCompletion(batchId, {
639-
onStatusUpdate: (s) => {
640-
options.onBatchStatus?.(s);
641-
if (!options.onBatchStatus) {
642-
const pct =
643-
s.total > 0 ? Math.round((s.succeeded / s.total) * 100) : 0;
644-
process.stdout.write(
645-
`\r Progress: ${s.succeeded}/${s.total} (${pct}%) | Errors: ${s.errored} `
646-
);
647-
}
648-
},
649-
});
721+
if (status.status === "in_progress") {
722+
console.log(
723+
`Batch still processing: ${status.succeeded}/${status.total} complete`
724+
);
725+
console.log("Waiting for completion...\n");
726+
727+
await this.batchClient!.waitForCompletion(batchId, {
728+
onStatusUpdate: (s) => {
729+
options.onBatchStatus?.(s);
730+
if (!options.onBatchStatus) {
731+
const pct =
732+
s.total > 0 ? Math.round((s.succeeded / s.total) * 100) : 0;
733+
process.stdout.write(
734+
`\r Progress: ${s.succeeded}/${s.total} (${pct}%) | Errors: ${s.errored} `
735+
);
736+
}
737+
},
738+
});
739+
}
650740
}
651741

652742
console.log("\n\nRetrieving results...");
653743

654-
// Get results
655-
const results = await this.batchClient.getResults(batchId, contexts);
744+
// Get results using appropriate client
745+
let results: Map<string, DeminifyResult>;
746+
if (isOpenAI) {
747+
results = await this.openAIBatchClient!.getResults(batchId, contexts);
748+
} else {
749+
results = await this.batchClient!.getResults(batchId, contexts);
750+
}
656751

657752
console.log(`Retrieved ${results.size} results`);
658753

0 commit comments

Comments
 (0)