picture-information-extraction.tsx

'use client';

import { z } from 'zod';
import {
  Form,
  FormControl,
  FormDescription,
  FormField,
  FormItem,
  FormLabel,
  FormMessage,
} from '@/components/ui/form';
import {
  ResizableHandle,
  ResizablePanel,
  ResizablePanelGroup,
} from '@/components/ui/resizable';
import { ScrollArea } from '@/components/ui/scroll-area';
import { Textarea } from '@/components/ui/textarea';
import { Button } from '@/components/ui/button';
import { Loader2 } from 'lucide-react';
import { useEffect, useMemo, useState } from 'react';
import { useForm } from 'react-hook-form';
import { zodResolver } from '@hookform/resolvers/zod';
import Empty from '@/components/error/empty-data';
import {
  IconInfoCircle,
  IconPresentation,
  IconSettings,
} from '@tabler/icons-react';
import ChunkSettingDrawer from './chunk-setting-drawer';
import useDocApi, { ChunkParams, PreviewResult } from '@/services/doc';
import {
  Dropzone,
  DropZoneArea,
  DropzoneDescription,
  DropzoneFileList,
  DropzoneFileListItem,
  DropzoneMessage,
  DropzoneRemoveFile,
  DropzoneTrigger,
  useDropzone,
} from '@/components/ui/dropzone';
import { CloudUploadIcon, Trash2Icon } from 'lucide-react';
import useConfigStore from '@/store/config';
import useFileApi from '@/services/file';
import { ExtractionApi } from './common';
import { toast } from '@/hooks/use-toast';
import { Chunk } from '@/services/chunk';
import ChunkList from '@/features/chunk/components/chunk-list';
import {
  Tooltip,
  TooltipContent,
  TooltipTrigger,
} from '@/components/ui/tooltip';
import { Separator } from '@/components/ui/separator';
import { KnowledgeSettings } from '@/services/knowledge';
import { Input } from '@/components/ui/input';
import { Switch } from '@/components/ui/switch';
import QAPreview from './qa-preview';

const formSchema = z.object({
  pictures: z.array(z.string()),
  recognition_picture_prompt: z.string(),
  enable_qa: z.boolean().optional(),
  qa_num: z.number().optional(),
  qa_text_prompt: z.string().optional(),
});

type PictureInformationExtractionForm = z.infer<typeof formSchema>;

type Props = {
  settings: KnowledgeSettings;
  getExtractionApi?: (api: ExtractionApi) => void;
};

const PictureInformationExtraction: React.FC<Props> = ({
  getExtractionApi,
  settings,
}) => {
  const [loading, setLoading] = useState(false);
  const [openChunkSetting, setOpenChunkSetting] = useState<boolean>(false);
  const config = useConfigStore((state) => state.config);
  const [chunks, setChunks] = useState<Chunk[]>([]);
  const [qaList, setQaList] = useState<PreviewResult['qa_list']>([]);
  const fileApi = useFileApi();
  const docApi = useDocApi();

  const [data, setData] = useState<ChunkParams>({
    chunk_method: 'picture',
    chunk_size: settings.chunk_size || 1024,
    chunk_overlap: settings.chunk_overlap || 100,
    model: settings.model || 'gpt-4o',
    embeddings_model: settings.embeddings_model || 'text-embedding-ada-002',
    splitter: settings.splitter || 'sentence',
    separator: settings.separator || '\n',
    keyword_prompt: config?.prompts?.keyword_extract_prompt,
    pictures: [],
    recognition_picture_prompt: config?.prompts?.recognition_picture_prompt,
    enable_qa: false,
    qa_num: 3,
    qa_text_prompt: config?.prompts?.qa_text_prompt,
  });

  const form = useForm<PictureInformationExtractionForm>({
    resolver: zodResolver(formSchema),
    values: {
      pictures: [],
      recognition_picture_prompt: '',
      enable_qa: data.enable_qa,
      qa_num: data.qa_num,
      qa_text_prompt: config?.prompts?.qa_text_prompt,
    },
  });

  const onSubmit = () => {
    api.getChunkParams().then((chunkParams) => {
      setLoading(true);
      docApi
        .preview(chunkParams)
        .then((res) => {
          if (res.code === 200) {
            toast({
              title: 'Preview loaded successfully',
              variant: 'default',
            });
            const data = res.data;
            setChunks(data.chunks);
            setQaList(data.qa_list);
          } else {
            toast({ variant: 'destructive', title: res.msg });
          }
        })
        .finally(() => setLoading(false));
    });
  };

  const dropzone = useDropzone({
    onDropFile: async (file: File) => {
      const r = await fileApi.upload([file]);
      const status = r.code === 200 ? 'success' : 'error';
      const result = {
        status,
        error: r.msg,
        result: status === 'success' && r.data[0],
        file,
      };

      if (status === 'success') {
        const docs = form.getValues('pictures');
        docs.push(result.result!);
        form.setValue('pictures', docs);
      }
      return result;
    },
    onRemoveFile(id) {
      const fileStatuses = dropzone.fileStatuses;
      const findStatus = fileStatuses.find((file) => file.id === id);
      if (findStatus && findStatus.status === 'success') {
        const docs = form.getValues('pictures');
        const newDocs = docs.filter((doc) => doc !== findStatus.result);
        form.setValue('pictures', newDocs);
      }
    },
    validation: {
      accept: {
        'image/*': ['.png', '.jpg', '.jpeg'],
      },
      maxSize: 100 * 1024 * 1024,
      maxFiles: 10,
    },
  });

  const api = useMemo<ExtractionApi>(() => {
    return {
      getChunkParams: async () => {
        const validate = await form.trigger();
        if (validate) {
          const files = dropzone.fileStatuses.values();
          const pictures: string[] = files
            .filter((file) => file.status === 'success')
            .map((file) => file.result)
            .toArray();
          const values = form.getValues();
          return Promise.resolve({ ...data, ...values, pictures });
        } else {
          throw new Error('failed validation');
        }
      },
    };
  }, [dropzone.fileStatuses]);

  getExtractionApi?.(api);

  return (
    <ResizablePanelGroup direction='horizontal'>
      <ResizablePanel defaultSize={40}>
        <div className='flex flex-col gap-2 px-2 h-full w-full'>
          <ScrollArea className='-mx-1 px-3 scroll-smooth'>
            <Form {...form}>
              <form
                id='picture-form'
                onSubmit={form.handleSubmit(onSubmit)}
                className='space-y-6 pb-8'
              >
                <FormField
                  control={form.control}
                  name='pictures'
                  render={({ field }) => (
                    <FormItem className='space-y-1'>
                      <FormLabel>Pictures</FormLabel>
                      <FormControl>
                        <div className='not-prose flex flex-col gap-4'>
                          <Dropzone {...dropzone}>
                            <div>
                              <div className='flex justify-between'>
                                <DropzoneDescription>
                                  Please select up to 10 images
                                </DropzoneDescription>
                                <DropzoneMessage />
                              </div>
                              <DropZoneArea>
                                <DropzoneTrigger className='flex flex-col items-center gap-4 bg-transparent p-10 text-center text-sm'>
                                  <CloudUploadIcon className='size-8' />
                                  <div className='p-2'>
                                    <p className='font-semibold'>
                                      Upload images
                                    </p>
                                    <p className='text-sm text-muted-foreground'>
                                      Click here or drag and drop to upload
                                    </p>
                                  </div>
                                </DropzoneTrigger>
                              </DropZoneArea>
                            </div>

                            <DropzoneFileList className='grid gap-3 p-0 md:grid-cols-2 lg:grid-cols-3'>
                              {dropzone.fileStatuses.map((file) => (
                                <DropzoneFileListItem
                                  className='overflow-hidden rounded-md bg-secondary p-0 shadow-sm'
                                  key={file.id}
                                  file={file}
                                >
                                  {file.status === 'pending' && (
                                    <div className='aspect-video animate-pulse bg-black/20' />
                                  )}
                                  {file.status === 'success' && (
                                    <PreviewImage
                                      filePath={file.result}
                                      file={file.file}
                                    />
                                  )}
                                  <div className='flex items-center justify-between p-2 pl-4'>
                                    <div className='min-w-0'>
                                      <p className='truncate text-sm'>
                                        {file.fileName}
                                      </p>
                                      <p className='text-xs text-muted-foreground'>
                                        {(
                                          file.file.size /
                                          (1024 * 1024)
                                        ).toFixed(2)}{' '}
                                        MB
                                      </p>
                                      {file.status === 'error' && (
                                        <div
                                          className='text-xs flex flex-row items-center gap-1'
                                          style={{
                                            color: 'hsl(var(--destructive))',
                                          }}
                                        >
                                          <Tooltip>
                                            <TooltipTrigger>
                                              <IconInfoCircle size='14px' />
                                            </TooltipTrigger>
                                            <TooltipContent>
                                              {file.error}
                                            </TooltipContent>
                                          </Tooltip>
                                          upload failed
                                          <span
                                            className='cursor-pointer'
                                            style={{
                                              color: '#2563eb',
                                            }}
                                            onClick={(e) => {
                                              e.stopPropagation();
                                              dropzone.onRetry(file.id);
                                            }}
                                          >
                                            retry
                                          </span>
                                        </div>
                                      )}
                                    </div>
                                    <DropzoneRemoveFile
                                      variant='ghost'
                                      className='shrink-0 hover:outline'
                                    >
                                      <Trash2Icon className='size-4' />
                                    </DropzoneRemoveFile>
                                  </div>
                                </DropzoneFileListItem>
                              ))}
                            </DropzoneFileList>
                          </Dropzone>
                        </div>
                      </FormControl>
                      <FormDescription>
                        Maximum length of each text chunk after splitting.
                      </FormDescription>
                      <FormMessage />
                    </FormItem>
                  )}
                />
                <FormField
                  control={form.control}
                  name='recognition_picture_prompt'
                  render={({ field }) => (
                    <FormItem className='space-y-1'>
                      <FormLabel>recognition pictures prompt</FormLabel>
                      <FormControl>
                        <Textarea
                          {...field}
                          placeholder='Enter a prompt'
                          onChange={(e) => field.onChange(e.target.value)}
                          value={field.value}
                          style={{ height: '200px' }}
                        />
                      </FormControl>
                      <FormDescription>
                        Prompt used to interpret or retrieve information related
                        to an image or visual input.
                      </FormDescription>
                      <FormMessage />
                    </FormItem>
                  )}
                />
                <FormField
                  control={form.control}
                  name='enable_qa'
                  render={({ field }) => (
                    <FormItem className='flex flex-row items-center justify-between rounded-lg border p-3 shadow-sm'>
                      <div className='space-y-0.5'>
                        <FormLabel>Enable question-answering</FormLabel>
                        <FormDescription>
                          Flag to turn question-answering (QA) feature on or
                          off.
                        </FormDescription>
                      </div>
                      <FormControl>
                        <Switch
                          checked={field.value}
                          onCheckedChange={field.onChange}
                        />
                      </FormControl>
                    </FormItem>
                  )}
                />
                {form.getValues()['enable_qa'] && (
                  <FormField
                    control={form.control}
                    name='qa_num'
                    render={({ field }) => (
                      <FormItem className='space-y-1'>
                        <FormLabel>Q&A numbers</FormLabel>
                        <FormControl>
                          <Input
                            {...field}
                            placeholder='Enter a chunk numbers'
                            type='number'
                            onChange={(e) =>
                              field.onChange(Number(e.target.value))
                            }
                            value={field.value}
                          />
                        </FormControl>
                        <FormDescription>
                          Number of QA pairs or samples to retrieve or generate.
                        </FormDescription>
                        <FormMessage />
                      </FormItem>
                    )}
                  />
                )}
                {form.getValues()['enable_qa'] && (
                  <FormField
                    control={form.control}
                    name='qa_text_prompt'
                    render={({ field }) => (
                      <FormItem>
                        <FormLabel>Q&A Text prompt</FormLabel>
                        <FormControl>
                          <Textarea
                            {...field}
                            placeholder='Enter a prompt'
                            onChange={(e) => field.onChange(e.target.value)}
                            value={field.value}
                            style={{ height: '200px' }}
                          />
                        </FormControl>
                        <FormDescription>
                          Text prompt used to generate or retrieve QA pairs.
                        </FormDescription>
                        <FormMessage />
                      </FormItem>
                    )}
                  />
                )}
              </form>
            </Form>
          </ScrollArea>
          <div className='flex flex-col items-center justify-between w-full gap-2'>
            <Button
              size='sm'
              variant='outline'
              className='w-full'
              onClick={() => setOpenChunkSetting(true)}
            >
              <IconSettings />
              Chunks Setting
            </Button>
            <Button
              className='w-full'
              size='sm'
              form='text-form'
              type='submit'
              disabled={loading}
            >
              {loading && <Loader2 className='mr-2 h-4 w-4 animate-spin' />}
              <IconPresentation />
              Preview
            </Button>
          </div>
        </div>
      </ResizablePanel>
      <ResizableHandle />
      <ResizablePanel defaultSize={60}>
        <div className='flex flex-col justify-center p-2 gap-2'>
          <span className='text-sm font-semibold'>Results</span>
          <Separator />
          {chunks.length == 0 ? (
            <Empty
              content={
                <span className='text-sm'>
                  Click the 'Preview' button on the left to load the preview
                </span>
              }
            />
          ) : (
            <div className='flex flex-col gap-2'>
              <ChunkList chunks={chunks} />
              <Separator />
              <QAPreview qaList={qaList} />
            </div>
          )}
        </div>
      </ResizablePanel>

      <ChunkSettingDrawer
        params={data}
        open={openChunkSetting}
        onOpenChange={setOpenChunkSetting}
        onChange={(params) => {
          setData((pre) => {
            return { ...pre, ...params };
          });
        }}
      />
    </ResizablePanelGroup>
  );
};

const PreviewImage = ({ file, filePath }: { file: File; filePath: string }) => {
  const [base64, setBase64] = useState<string>();
  const fileApi = useFileApi();

  useEffect(() => {
    fileApi.downloadImage(filePath).then((response) => {
      if (response.code === 200) {
        setBase64(response.data);
      }
    });
  }, [filePath]);

  return (
    <img
      src={base64}
      alt={`uploaded-${file.name}`}
      className='aspect-video object-cover'
    />
  );
};

export default PictureInformationExtraction;