backstory/frontend/src/NewApp/Pages/GenerateCandiate.tsx

import React, { useEffect, useState, useRef } from 'react';
import Box from '@mui/material/Box';
import Tooltip from '@mui/material/Tooltip';
import Button from '@mui/material/Button';
import IconButton from '@mui/material/IconButton';
import CancelIcon from '@mui/icons-material/Cancel';
import SendIcon from '@mui/icons-material/Send';
import PropagateLoader from 'react-spinners/PropagateLoader';
import { CandidateInfo } from '../Components/CandidateInfo';
import { Query } from '../../Components/ChatQuery'
import { streamQueryResponse, StreamQueryController } from '../Components/streamQueryResponse';
import { connectionBase } from 'Global';
import { UserInfo } from '../Components/UserContext';
import { BackstoryElementProps } from 'Components/BackstoryTab';
import { BackstoryTextField, BackstoryTextFieldRef } from 'Components/BackstoryTextField';
import { jsonrepair } from 'jsonrepair';
import { StyledMarkdown } from 'Components/StyledMarkdown';
import { Scrollable } from 'Components/Scrollable';
import { useForkRef } from '@mui/material';

const emptyUser : UserInfo = {
    type: 'candidate',
    description: "[blank]",
    rag_content_size: 0,
    username: "[blank]",
    first_name: "[blank]",
    last_name: "[blank]",
    full_name: "[blank] [blank]",
    contact_info: {},
    questions: [],
    isAuthenticated: false,
    has_profile: false
};

const GenerateCandidate = (props: BackstoryElementProps) => {
    const {sessionId, setSnack, submitQuery} = props;
    const [streaming, setStreaming] = useState<string>('');
    const [processing, setProcessing] = useState<boolean>(false);
    const [user, setUser] = useState<UserInfo>(emptyUser);
    const controllerRef = useRef<StreamQueryController>(null);
    const backstoryTextRef = useRef<BackstoryTextFieldRef>(null);
    const promptRef = useRef<string>(null);
    const stateRef = useRef<number>(0); /* Generating persona */
    const userRef = useRef<UserInfo>(user);
    const [prompt, setPrompt] = useState<string>('');
    const [resume, setResume] = useState<string>('');

    const processQuery = (query: Query) => {
        if (controllerRef.current) {
            return;
        }
        setPrompt(query.prompt);
        promptRef.current = query.prompt;
        stateRef.current = 0; 
        setUser(emptyUser);
        setStreaming('');
        setResume('');
        setProcessing(true);

        controllerRef.current = streamQueryResponse({
            query,
            type: "persona",
            sessionId,
            connectionBase,
            onComplete: (msg) => {
                console.log({ msg, state: stateRef.current, prompt: promptRef.current || '' });
                switch (msg.status) {
                    case "partial":
                    case "done":
                        switch (stateRef.current) {
                        case 0: /* Generating persona */
                            let partialUser = JSON.parse(jsonrepair((msg.response || '').trim()));
                            if (!partialUser.full_name) {
                                partialUser.full_name = `${partialUser.first_name} ${partialUser.last_name}`;
                            }
                            console.log(partialUser);
                            setUser(partialUser);
                            stateRef.current = 1 /* Generating resume */
                            break;
                        case 1: /* Generating resume */
                            stateRef.current = 2 /* RAG generation */
                            break;
                        case 2: /* RAG generation */
                            stateRef.current = 2 /* Image generation */
                            break;
                        case 3: /* Generating image */
                            let imageGeneration = JSON.parse(jsonrepair((msg.response || '').trim()));
                            console.log(imageGeneration);
                            if (imageGeneration >= 100) {
                                setUser({...userRef.current});
                            } else {
                                setPrompt(imageGeneration.status);
                            }
                            stateRef.current = 3 /* ... */
                        }
                        if (msg.status === "done") {
                            setProcessing(false);
                            controllerRef.current = null;
                            stateRef.current = 0;
                        }
                        break;
                    case "thinking":
                        setPrompt(msg.response || '');
                        break;

                    case "error":
                        console.log(`Error generating persona: ${msg.response}`);
                        setSnack(msg.response || "", "error");
                        setProcessing(false);
                        setUser({...userRef.current});
                        controllerRef.current = null;
                        stateRef.current = 0;
                        break;
                }
            },
            onStreaming: (chunk) => {
                setStreaming(chunk);
            }
        });
    };

    const cancelQuery = () => {
        if (controllerRef.current) {
            controllerRef.current.abort();
            controllerRef.current = null;
            stateRef.current = 0;
            setProcessing(false);
        }
    }

    useEffect(() => {
        promptRef.current = prompt;
    }, [prompt]);

    useEffect(() => {
        userRef.current = user;
    }, [user]);

    useEffect(() => {
        if (streaming.trim().length === 0) {
            return;
        }

        try {
            switch (stateRef.current) {
                case 0: /* Generating persona */
                    const partialUser = {...emptyUser, ...JSON.parse(jsonrepair(`${streaming.trim()}...`))};
                    if (!partialUser.full_name) {
                        partialUser.full_name = `${partialUser.first_name} ${partialUser.last_name}`;
                    }
                    setUser(partialUser);
                    break;
                case 1: /* Generating resume */
                    setResume(streaming);
                    break;
                case 3: /* RAG streaming */
                    break;
                case 4: /* Image streaming */
                    break;
            }
        } catch {
        }
    }, [streaming]);

    if (!sessionId) {
        return <></>;
    }

    const onEnter = (value: string) => {
        if (processing) {
            return;
        }
        const query: Query = {
            prompt: value
        }
        processQuery(query);
    };

    return (<>
        { user && <CandidateInfo sessionId={sessionId} user={user}/> }
        { resume !== '' && <Scrollable sx={{maxHeight: "20vh"}}><StyledMarkdown  {...{content: resume, setSnack, sessionId, submitQuery}}/></Scrollable> }
        {processing && <Box sx={{
            display: "flex",
            flexDirection: "column",
            alignItems: "center",
            justifyContent: "center",
            m: 2,
        }}>
            <Box sx={{flexDirection: "row"}}><Box>Genearating 
                {stateRef.current === 0 && "persona"}
                {stateRef.current === 1 && "resume"}
                {stateRef.current === 2 && "RAG"}
                {stateRef.current === 3 && "profile image"}
                :</Box><Box sx={{fontWeight: "bold"}}>{prompt}</Box></Box>
            <PropagateLoader
                size="10px"
                loading={processing}
                aria-label="Loading Spinner"
                data-testid="loader"
            />
        </Box> }
        <BackstoryTextField
            ref={backstoryTextRef}
            disabled={processing}
            onEnter={onEnter}
            placeholder='Specify any characteristics you would like the persona to have. For example, "This person likes yo-yos."'
        />
        <Box sx={{ display: "flex", justifyContent: "center", flexDirection: "row" }}>
          <Tooltip title={"Send"}>
            <span style={{ display: "flex", flexGrow: 1 }}>
              <Button
                sx={{ m: 1, gap: 1, flexGrow: 1 }}
                variant="contained"
                disabled={sessionId === undefined || processing}
                onClick={() => { processQuery({ prompt: (backstoryTextRef.current && backstoryTextRef.current.getAndResetValue()) || "" }); }}>
                Send<SendIcon />
              </Button>
            </span>
          </Tooltip>
          <Tooltip title="Cancel">
            <span style={{ display: "flex" }}> { /* This span is used to wrap the IconButton to ensure Tooltip works even when disabled */}
              <IconButton
                aria-label="cancel"
                onClick={() => { cancelQuery(); }}
                sx={{ display: "flex", margin: 'auto 0px' }}
                size="large"
                edge="start"
                disabled={controllerRef.current === null || !sessionId || processing === false}
              >
                <CancelIcon />
              </IconButton>
            </span>
          </Tooltip>          
        </Box>
    </>);
};

export {
  GenerateCandidate
};