TERES_web_frontend/src/pages/setting/components/Dialog/BedrockDialog.tsx

import React, { useEffect } from 'react';
import {
  Dialog,
  DialogTitle,
  DialogContent,
  DialogActions,
  Button,
  TextField,
  Box,
  Typography,
  IconButton,
  InputAdornment,
  FormControl,
  InputLabel,
  Select,
  MenuItem,
  CircularProgress,
  FormHelperText,
  Link,
} from '@mui/material';
import { Visibility, VisibilityOff } from '@mui/icons-material';
import { Controller, useForm } from 'react-hook-form';
import type { IAddLlmRequestBody } from '@/interfaces/request/llm';

// AWS Bedrock 支持的区域列表
export const BEDROCK_REGIONS = [
  'us-east-2',
  'us-east-1',
  'us-west-1',
  'us-west-2',
  'af-south-1',
  'ap-east-1',
  'ap-south-2',
  'ap-southeast-3',
  'ap-southeast-5',
  'ap-southeast-4',
  'ap-south-1',
  'ap-northeast-3',
  'ap-northeast-2',
  'ap-southeast-1',
  'ap-southeast-2',
  'ap-east-2',
  'ap-southeast-7',
  'ap-northeast-1',
  'ca-central-1',
  'ca-west-1',
  'eu-central-1',
  'eu-west-1',
  'eu-west-2',
  'eu-south-1',
  'eu-west-3',
  'eu-south-2',
  'eu-north-1',
  'eu-central-2',
  'il-central-1',
  'mx-central-1',
  'me-south-1',
  'me-central-1',
  'sa-east-1',
  'us-gov-east-1',
  'us-gov-west-1',
];

// 模型类型选项
const MODEL_TYPE_OPTIONS = [
  { value: 'chat', label: 'Chat' },
  { value: 'embedding', label: 'Embedding' },
];

// 表单数据接口
export interface BedrockFormData extends IAddLlmRequestBody {
  bedrock_ak: string;
  bedrock_sk: string;
  bedrock_region: string;
}

// 对话框 Props 接口
export interface BedrockDialogProps {
  open: boolean;
  onClose: () => void;
  onSubmit: (data: BedrockFormData) => void;
  loading: boolean;
  initialData?: BedrockFormData;
  editMode?: boolean;
}

/**
 * AWS Bedrock 配置对话框
 */
function BedrockDialog ({
  open,
  onClose,
  onSubmit,
  loading,
  initialData,
  editMode = false,
}: BedrockDialogProps) {
  const [showAccessKey, setShowAccessKey] = React.useState(false);
  const [showSecretKey, setShowSecretKey] = React.useState(false);

  const {
    control,
    handleSubmit,
    reset,
    formState: { errors },
  } = useForm<BedrockFormData>({
    defaultValues: {
      model_type: 'chat',
      llm_name: '',
      bedrock_ak: '',
      bedrock_sk: '',
      bedrock_region: 'us-east-1',
      max_tokens: 4096,
      llm_factory: 'Bedrock',
    },
  });

  // 当对话框打开或初始数据变化时重置表单
  useEffect(() => {
    if (open) {
      reset({
        model_type: 'chat',
        llm_name: '',
        bedrock_ak: '',
        bedrock_sk: '',
        bedrock_region: 'us-east-1',
        max_tokens: 4096,
        llm_factory: initialData?.llm_factory || 'Bedrock',
      });
    }
  }, [open, initialData, reset]);

  const handleFormSubmit = (data: BedrockFormData) => {
    onSubmit(data);
  };

  const toggleShowAccessKey = () => {
    setShowAccessKey(!showAccessKey);
  };

  const toggleShowSecretKey = () => {
    setShowSecretKey(!showSecretKey);
  };

  const docInfo = {
    url: 'https://console.aws.amazon.com/',
    text: '如何集成 Bedrock',
  };

  return (
    <Dialog open={open} onClose={onClose} maxWidth="sm" fullWidth>
      <DialogTitle>
        {editMode ? '编辑' : '添加'} LLM
      </DialogTitle>
      <DialogContent>
        <Box component="form" sx={{ mt: 2 }}>
          {/* 模型类型 */}
          <Controller
            name="model_type"
            control={control}
            rules={{ required: '模型类型是必填项' }}
            render={({ field }) => (
              <FormControl fullWidth margin="normal" error={!!errors.model_type}>
                <InputLabel>* 模型类型</InputLabel>
                <Select {...field} label="* 模型类型">
                  {MODEL_TYPE_OPTIONS.map((option) => (
                    <MenuItem key={option.value} value={option.value}>
                      {option.label}
                    </MenuItem>
                  ))}
                </Select>
                {errors.model_type && (
                  <FormHelperText>{errors.model_type.message}</FormHelperText>
                )}
              </FormControl>
            )}
          />

          {/* 模型名称 */}
          <Controller
            name="llm_name"
            control={control}
            rules={{ required: '模型名称是必填项' }}
            render={({ field }) => (
              <TextField
                {...field}
                fullWidth
                label="* 模型名称"
                margin="normal"
                placeholder="请输入模型名称"
                error={!!errors.llm_name}
                helperText={errors.llm_name?.message}
              />
            )}
          />

          {/* ACCESS KEY */}
          <Controller
            name="bedrock_ak"
            control={control}
            rules={{ required: 'ACCESS KEY 是必填项' }}
            render={({ field }) => (
              <TextField
                {...field}
                fullWidth
                label="* ACCESS KEY"
                type={showAccessKey ? 'text' : 'password'}
                margin="normal"
                placeholder="请输入 ACCESS KEY"
                error={!!errors.bedrock_ak}
                helperText={errors.bedrock_ak?.message}
                InputProps={{
                  endAdornment: (
                    <InputAdornment position="end">
                      <IconButton
                        aria-label="toggle access key visibility"
                        onClick={toggleShowAccessKey}
                        edge="end"
                      >
                        {showAccessKey ? <VisibilityOff /> : <Visibility />}
                      </IconButton>
                    </InputAdornment>
                  ),
                }}
              />
            )}
          />

          {/* SECRET KEY */}
          <Controller
            name="bedrock_sk"
            control={control}
            rules={{ required: 'SECRET KEY 是必填项' }}
            render={({ field }) => (
              <TextField
                {...field}
                fullWidth
                label="* SECRET KEY"
                type={showSecretKey ? 'text' : 'password'}
                margin="normal"
                placeholder="请输入 SECRET KEY"
                error={!!errors.bedrock_sk}
                helperText={errors.bedrock_sk?.message}
                InputProps={{
                  endAdornment: (
                    <InputAdornment position="end">
                      <IconButton
                        aria-label="toggle secret key visibility"
                        onClick={toggleShowSecretKey}
                        edge="end"
                      >
                        {showSecretKey ? <VisibilityOff /> : <Visibility />}
                      </IconButton>
                    </InputAdornment>
                  ),
                }}
              />
            )}
          />

          {/* AWS Region */}
          <Controller
            name="bedrock_region"
            control={control}
            rules={{ required: 'AWS Region 是必填项' }}
            render={({ field }) => (
              <FormControl fullWidth margin="normal" error={!!errors.bedrock_region}>
                <InputLabel>* AWS Region</InputLabel>
                <Select {...field} label="* AWS Region">
                  {BEDROCK_REGIONS.map((region) => (
                    <MenuItem key={region} value={region}>
                      {region}
                    </MenuItem>
                  ))}
                </Select>
                {errors.bedrock_region && (
                  <FormHelperText>{errors.bedrock_region.message}</FormHelperText>
                )}
              </FormControl>
            )}
          />

          {/* 最大token数 */}
          <Controller
            name="max_tokens"
            control={control}
            rules={{
              required: '最大token数是必填项',
              min: { value: 1, message: '最大token数必须大于0' },
            }}
            render={({ field }) => (
              <TextField
                {...field}
                fullWidth
                label="* 最大token数"
                type="number"
                margin="normal"
                placeholder="这设置了模型输出的最大长度，以token（单词或词片段）的数量来衡量"
                error={!!errors.max_tokens}
                helperText={errors.max_tokens?.message}
                onChange={(e) => field.onChange(Number(e.target.value))}
              />
            )}
          />
        </Box>
      </DialogContent>
      <DialogActions>
        <Box sx={{ display: 'flex', justifyContent: 'space-between', width: '100%' }}>
          <Link
            href={docInfo.url}
            target="_blank"
            rel="noopener noreferrer"
            sx={{ alignSelf: 'center', textDecoration: 'none', ml:2 }}
          >
            {docInfo.text}
          </Link>
          <Box>
            <Button onClick={onClose} disabled={loading} sx={{ mr: 1 }}>
              取消
            </Button>
            <Button
              onClick={handleSubmit(handleFormSubmit)}
              variant="contained"
              disabled={loading}
              startIcon={loading ? <CircularProgress size={20} /> : null}
            >
              确定
            </Button>
          </Box>
        </Box>
      </DialogActions>
    </Dialog>
  );
};

export default BedrockDialog;