MemoryBear/web/src/views/ModelManagement/components/CustomModelModal.tsx

/*
 * @Author: ZhaoYing 
 * @Date: 2026-02-03 16:49:28 
 * @Last Modified by: ZhaoYing
 * @Last Modified time: 2026-02-28 17:24:05
 */
/**
 * Custom Model Modal
 * Modal for creating and editing custom models in the model square
 * Supports logo upload, type/provider selection, and tagging
 */

import { forwardRef, useImperativeHandle, useState } from 'react';
import { Form, Input, App } from 'antd';
import { useTranslation } from 'react-i18next';

import type { CustomModelForm, ModelListItem, CustomModelModalRef, CustomModelModalProps } from '../types';
import RbModal from '@/components/RbModal'
import CustomSelect from '@/components/CustomSelect'
import UploadImages from '@/components/Upload/UploadImages'
import { updateCustomModel, addCustomModel, modelTypeUrl, modelProviderUrl } from '@/api/models'
import { getFileLink } from '@/api/fileStorage'

/**
 * Custom model modal component
 */
const CustomModelModal = forwardRef<CustomModelModalRef, CustomModelModalProps>(({
  refresh
}, ref) => {
  const { t } = useTranslation();
  const { message } = App.useApp();
  const [visible, setVisible] = useState(false);
  const [model, setModel] = useState<ModelListItem>({} as ModelListItem);
  const [isEdit, setIsEdit] = useState(false);
  const [form] = Form.useForm<CustomModelForm>();
  const [loading, setLoading] = useState(false)

  /** Close modal and reset state */
  const handleClose = () => {
    setModel({} as ModelListItem);
    form.resetFields();
    setLoading(false)
    setVisible(false);
  };

  /** Open modal with optional model data for editing */
  const handleOpen = (model?: ModelListItem) => {
    if (model) {
      setIsEdit(true);
      setModel(model);
      form.setFieldsValue({
        ...model,
        logo: model.logo && model.logo.startsWith('http') ? { url: model.logo, uid: model.logo, status: 'done', name: 'logo' } : undefined
      });
    } else {
      setIsEdit(false);
      form.resetFields();
    }
    setVisible(true);
  };
  /** Update or create custom model */
  const handleUpdate = (data: CustomModelForm) => {
    setLoading(true)
    const { type, provider, ...rest} = data
    const res = isEdit ? updateCustomModel(model.id, rest) : addCustomModel(data)

    res.then(() => {
      refresh && refresh(isEdit)
      handleClose()
      message.success(isEdit ? t('common.updateSuccess') : t('common.createSuccess'))
    })
      .catch(() => {
        setLoading(false)
      });
  }
  /** Validate and save custom model */
  const handleSave = () => {
    form
      .validateFields()
      .then((values) => {
        const { logo, ...rest } = values;
        let formData: CustomModelForm = {
          ...rest
        }

        if (typeof logo === 'object' && logo?.response?.data.file_id) {
          getFileLink(logo?.response?.data.file_id)
            .then(res => {
              const logoRes = res as { url: string }
              formData.logo = logoRes.url
              handleUpdate(formData)
            })
            .catch(() => {
              handleUpdate(formData)
            })
        } else {
          formData.logo = typeof logo === 'string' ? logo : logo.url
          handleUpdate(formData)
        }
      })
      .catch((err) => {
        console.log('err', err)
      });
  }

  /** Expose methods to parent component */
  useImperativeHandle(ref, () => ({
    handleOpen,
  }));

  return (
    <RbModal
      title={isEdit ? `${model.name} - ${t('modelNew.modelConfiguration')}` : t('modelNew.createCustomModel')}
      open={visible}
      onCancel={handleClose}
      okText={t(`common.${isEdit ? 'save' : 'create'}`)}
      onOk={handleSave}
      confirmLoading={loading}
    >
      <Form
        form={form}
        layout="vertical"
      >
        <Form.Item
          name="logo"
          label={t('modelNew.logo')}
          valuePropName="fileList"
          rules={[{ required: true, message: t('common.pleaseSelect') }]}
        >
          <UploadImages />
        </Form.Item>
        <Form.Item
          name="name"
          label={t('modelNew.name')}
          rules={[{ required: true, message: t('common.inputPlaceholder', { title: t('modelNew.name') }) }]}
        >
          <Input placeholder={t('common.pleaseEnter')} />
        </Form.Item>
        
        <Form.Item
          name="type"
          label={t('modelNew.type')}
          rules={[{ required: true, message: t('common.selectPlaceholder', { title: t('modelNew.type') }) }]}
        >
          <CustomSelect
            url={modelTypeUrl}
            hasAll={false}
            disabled={isEdit}
            format={(items) => items.map((item) => ({ label: t(`modelNew.${item}`), value: String(item) }))}
          />
        </Form.Item>

        <Form.Item
          name="provider"
          label={t('modelNew.provider')}
          rules={[{ required: true, message: t('common.selectPlaceholder', { title: t('modelNew.provider') }) }]}
        >
          <CustomSelect
            url={modelProviderUrl}
            hasAll={false}
            disabled={isEdit}
            format={(items) => items.map((item) => ({ label: t(`modelNew.${item}`), value: String(item) }))}
          />
        </Form.Item>

        <Form.Item
          name="description"
          label={t('modelNew.description')}
        >
          <Input.TextArea placeholder={t('common.pleaseEnter')} />
        </Form.Item>


        <Form.Item
          name={["api_keys", 0, "api_key"]}
          label={t('modelNew.api_key')}
          rules={[{ required: true, message: t('common.inputPlaceholder', { title: t('modelNew.api_key') }) }]}
        >
          <Input.Password placeholder={t('common.pleaseEnter')} />
        </Form.Item>

        <Form.Item
          name={["api_keys", 0, "api_base"]}
          label={t('modelNew.api_base')}
          rules={[{ required: true, message: t('common.inputPlaceholder', { title: t('modelNew.api_base') }) }]}
        >
          <Input placeholder="https://api.example.com/v1" />
        </Form.Item>
      </Form>
    </RbModal>
  );
});

export default CustomModelModal;