All files / api/src/ai-providers/bedrock create-batch.ts

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
import { generateInvalidProviderResponseError } from '@api/utils/ai-provider';
import type {
  AIProviderFunctionConfig,
  ResponseTransformFunction,
} from '@shared/types/ai-providers/config';
import type { ParameterConfigDefaultFunction } from '@shared/types/api/response/body';
import type {
  CreateBatchRequestBody,
  CreateBatchResponseBody,
} from '@shared/types/api/routes/batch-api';
import { BatchStatus } from '@shared/types/api/routes/batch-api';
import { AIProvider } from '@shared/types/constants';
import { bedrockErrorResponseTransform } from './chat-complete';
 
export const BedrockCreateBatchConfig: AIProviderFunctionConfig = {
  model: {
    param: 'modelId',
    required: true,
  },
  input_file_id: {
    param: 'inputDataConfig',
    required: true,
    transform: (saRequestBody: CreateBatchRequestBody) => {
      return {
        s3InputDataConfig: {
          s3Uri: decodeURIComponent(saRequestBody.input_file_id),
        },
      };
    },
  },
  job_name: {
    param: 'jobName',
    required: true,
    default: () => {
      return `sa-batch-job-${crypto.randomUUID()}`;
    },
  },
  output_data_config: {
    param: 'outputDataConfig',
    required: true,
    default: (({ saRequestBody, saTarget }): Record<string, unknown> => {
      if (!('input_file_id' in saRequestBody)) {
        throw new Error('input_file_id is required');
      }
 
      // TODO: Fix this
      const inputFileId = decodeURIComponent(
        saRequestBody.input_file_id as string,
      );
      const s3URLToContainingFolder = `${inputFileId.split('/').slice(0, -1).join('/')}/`;
      return {
        s3OutputDataConfig: {
          s3Uri: s3URLToContainingFolder,
          ...(saTarget.aws_server_side_encryption_kms_key_id && {
            s3EncryptionKeyId: saTarget.aws_server_side_encryption_kms_key_id,
          }),
        } as Record<string, unknown>,
      };
    }) as ParameterConfigDefaultFunction,
  },
  role_arn: {
    param: 'roleArn',
    required: true,
  },
};
 
export const bedrockCreateBatchResponseTransform: ResponseTransformFunction = (
  response,
  responseStatus,
) => {
  if (responseStatus !== 200) {
    const errorResponse = bedrockErrorResponseTransform(
      response as Record<string, unknown>,
    );
    if (errorResponse) return errorResponse;
  }
 
  if ('jobArn' in response) {
    // AWS Bedrock CreateModelInvocationJob returns a simple response with jobArn
    // We need to construct the OpenAI-compatible batch response
    const awsResponse = response as unknown as { jobArn: string };
 
    const batchResponseBody: CreateBatchResponseBody = {
      id: encodeURIComponent(awsResponse.jobArn),
      object: 'batch',
      endpoint: '/v1/chat/completions', // Default endpoint for batch operations
      input_file_id: '', // This will be populated from the original request context
      status: BatchStatus.VALIDATING, // Initial status for newly created batch
      output_file_id: null, // Will be available when batch completes
      error_file_id: null,
      created_at: Date.now(), // Use current time since AWS doesn't provide this in create response
      in_progress_at: null,
      expires_at: null,
      finalizing_at: null,
      completed_at: null,
      failed_at: null,
      expired_at: null,
      cancelling_at: null,
      cancelled_at: null,
      request_counts: {
        total: 0, // Unknown at creation time
        completed: 0,
        failed: 0,
      },
      completion_window: '24h', // Default completion window
      metadata: null,
    };
 
    return batchResponseBody;
  }
 
  return generateInvalidProviderResponseError(response, AIProvider.BEDROCK);
};