path: root/src/client/views/nodes/ChatBox/ChatBox.tsx

import { MathJaxContext } from 'better-react-mathjax';
import { action, makeObservable, observable, observe, reaction, runInAction } from 'mobx';
import { observer } from 'mobx-react';
import OpenAI, { ClientOptions } from 'openai';
import { ImageFile, Message } from 'openai/resources/beta/threads/messages';
import { RunStep } from 'openai/resources/beta/threads/runs/steps';
import * as React from 'react';
import { Doc } from '../../../../fields/Doc';
import { Id } from '../../../../fields/FieldSymbols';
import { CsvCast, DocCast, PDFCast, StrCast } from '../../../../fields/Types';
import { CsvField } from '../../../../fields/URLField';
import { Networking } from '../../../Network';
import { DocUtils } from '../../../documents/DocUtils';
import { DocumentType } from '../../../documents/DocumentTypes';
import { Docs } from '../../../documents/Documents';
import { DocumentManager } from '../../../util/DocumentManager';
import { LinkManager } from '../../../util/LinkManager';
import { ViewBoxAnnotatableComponent } from '../../DocComponent';
import { FieldView, FieldViewProps } from '../FieldView';
import './ChatBox.scss';
import MessageComponent from './MessageComponent';
import { ANNOTATION_LINK_TYPE, ASSISTANT_ROLE, AssistantMessage, DOWNLOAD_TYPE } from './types';

@observer
export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
    @observable modalStatus = false;
    @observable currentFile = { url: '' };
    @observable history: AssistantMessage[] = [];
    @observable.deep current_message: AssistantMessage | undefined = undefined;

    @observable isLoading: boolean = false;
    @observable isInitializing: boolean = true;
    @observable expandedLogIndex: number | null = null;
    @observable linked_docs_to_add: Doc[] = [];

    private openai: OpenAI;
    private interim_history: string = '';
    private assistantID: string = '';
    private threadID: string = '';
    private _oldWheel: any;
    private vectorStoreID: string = '';
    private mathJaxConfig: any;
    private linkedCsvIDs: string[] = [];

    public static LayoutString(fieldKey: string) {
        return FieldView.LayoutString(ChatBox, fieldKey);
    }
    constructor(props: FieldViewProps) {
        super(props);
        makeObservable(this);
        this.openai = this.initializeOpenAI();
        this.history = [];
        this.threadID = StrCast(this.dataDoc.thread_id);
        this.assistantID = StrCast(this.dataDoc.assistant_id);
        this.vectorStoreID = StrCast(this.dataDoc.vector_store_id);
        this.openai = this.initializeOpenAI();
        if (this.assistantID === '' || this.threadID === '' || this.vectorStoreID === '') {
            this.createAssistant();
        } else {
            this.retrieveCsvUrls();
            this.isInitializing = false;
        }
        this.mathJaxConfig = {
            loader: { load: ['input/asciimath'] },
            tex: {
                inlineMath: [
                    ['$', '$'],
                    ['\\(', '\\)'],
                ],
                displayMath: [
                    ['$$', '$$'],
                    ['[', ']'],
                ],
            },
        };
        reaction(
            () => this.history.map((msg: AssistantMessage) => ({ role: msg.role, text: msg.text, image: msg.image, tool_logs: msg.tool_logs, links: msg.links })),
            serializableHistory => {
                this.dataDoc.data = JSON.stringify(serializableHistory);
            }
        );
    }

    toggleToolLogs = (index: number) => {
        this.expandedLogIndex = this.expandedLogIndex === index ? null : index;
    };

    retrieveCsvUrls() {
        const linkedDocs = LinkManager.Instance.getAllRelatedLinks(this.Document)
            .map(d => DocCast(LinkManager.getOppositeAnchor(d, this.Document)))
            .map(d => DocCast(d?.annotationOn, d))
            .filter(d => d);

        linkedDocs.forEach(doc => {
            const aiFieldId = StrCast(doc[this.Document[Id] + '_ai_field_id']);
            if (CsvCast(doc.data)) {
                this.linkedCsvIDs.push(StrCast(aiFieldId));
                console.log(this.linkedCsvIDs);
            }
        });
    }

    initializeOpenAI() {
        const configuration: ClientOptions = {
            apiKey: process.env.OPENAI_KEY,
            dangerouslyAllowBrowser: true,
        };
        return new OpenAI(configuration);
    }

    onPassiveWheel = (e: WheelEvent) => {
        if (this._props.isContentActive()) {
            e.stopPropagation();
        }
    };

    createLink = (linkInfo: string, startIndex: number, endIndex: number, linkType: ANNOTATION_LINK_TYPE, annotationIndex: number = 0) => {
        const text = this.interim_history;
        const subString = this.current_message?.text.substring(startIndex, endIndex) ?? '';
        if (!text) return;
        const textToDisplay = `${annotationIndex}`;
        let fileInfo = linkInfo;
        const fileName = subString.split('/')[subString.split('/').length - 1];
        if (linkType === ANNOTATION_LINK_TYPE.DOWNLOAD_FILE) {
            fileInfo = linkInfo + '!!!' + fileName;
        }

        const formattedLink = `[${textToDisplay}](${fileInfo}~~~${linkType})`;
        console.log(formattedLink);
        const newText = text.replace(subString, formattedLink);
        runInAction(() => {
            this.interim_history = newText;
            console.log(newText);
            this.current_message?.links?.push({
                start: startIndex,
                end: endIndex,
                url: linkType === ANNOTATION_LINK_TYPE.DOWNLOAD_FILE ? fileName : linkInfo,
                id: linkType === ANNOTATION_LINK_TYPE.DOWNLOAD_FILE ? linkInfo : undefined,
                link_type: linkType,
            });
        });
    };

    @action
    createAssistant = async () => {
        this.isInitializing = true;
        try {
            const vectorStore = await this.openai.beta.vectorStores.create({
                name: 'Vector Store for Assistant',
            });
            const assistant = await this.openai.beta.assistants.create({
                name: 'Document Analyser Assistant',
                instructions: `
                You will analyse documents with which you are provided. You will answer questions and provide insights based on the information in the documents.
                For writing math formulas:
                You have a MathJax render environment.
                - Write all in-line equations within a single dollar sign, $, to render them as TeX (this means any time you want to use a dollar sign to represent a dollar sign itself, you must escape it with a backslash: "$");
                - Use a double dollar sign, $$, to render equations on a new line;
                Example: $$x^2 + 3x$$ is output for "x² + 3x" to appear as TeX.`,
                model: 'gpt-4-turbo',
                tools: [{ type: 'file_search' }, { type: 'code_interpreter' }],
                tool_resources: {
                    file_search: {
                        vector_store_ids: [vectorStore.id],
                    },
                    code_interpreter: {
                        file_ids: this.linkedCsvIDs,
                    },
                },
            });
            const thread = await this.openai.beta.threads.create();

            runInAction(() => {
                this.dataDoc.assistant_id = assistant.id;
                this.dataDoc.thread_id = thread.id;
                this.dataDoc.vector_store_id = vectorStore.id;
                this.assistantID = assistant.id;
                this.threadID = thread.id;
                this.vectorStoreID = vectorStore.id;
                this.isInitializing = false;
            });
        } catch (error) {
            console.error('Initialization failed:', error);
            this.isInitializing = false;
        }
    };

    @action
    runAssistant = async (inputText: string) => {
        // Ensure an assistant and thread are created
        if (!this.assistantID || !this.threadID || !this.vectorStoreID) {
            await this.createAssistant();
            console.log('Assistant and thread created:', this.assistantID, this.threadID);
        }
        let currentText: string = '';
        let currentToolCallMessage: string = '';

        // Send the user's input to the assistant
        await this.openai.beta.threads.messages.create(this.threadID, {
            role: 'user',
            content: inputText,
        });

        // Listen to the streaming responses
        const stream = this.openai.beta.threads.runs
            .stream(this.threadID, {
                assistant_id: this.assistantID,
            })
            .on('runStepCreated', (runStep: RunStep) => {
                currentText = '';
                runInAction(() => {
                    this.current_message = { role: ASSISTANT_ROLE.ASSISTANT, text: currentText, tool_logs: '', links: [] };
                });
                this.isLoading = true;
            })
            .on('toolCallDelta', (toolCallDelta, snapshot) => {
                this.isLoading = false;
                if (toolCallDelta.type === 'code_interpreter') {
                    if (toolCallDelta.code_interpreter?.input) {
                        currentToolCallMessage += toolCallDelta.code_interpreter.input;
                        runInAction(() => {
                            if (this.current_message) {
                                this.current_message.tool_logs = currentToolCallMessage;
                            }
                        });
                    }
                    if (toolCallDelta.code_interpreter?.outputs) {
                        currentToolCallMessage += '\n Code interpreter output:';
                        toolCallDelta.code_interpreter.outputs.forEach(output => {
                            if (output.type === 'logs') {
                                runInAction(() => {
                                    if (this.current_message) {
                                        this.current_message.tool_logs += '\n|' + output.logs;
                                    }
                                });
                            }
                        });
                    }
                }
            })
            .on('textDelta', (textDelta, snapshot) => {
                this.isLoading = false;
                currentText += textDelta.value;
                runInAction(() => {
                    if (this.current_message) {
                        // this.current_message = {...this.current_message, text: current_text};
                        this.current_message.text = currentText;
                    }
                });
            })
            .on('messageDone', async event => {
                console.log(event);
                const textItem = event.content.find(item => item.type === 'text');
                if (textItem && textItem.type === 'text') {
                    const { text } = textItem;
                    console.log(text.value);
                    try {
                        runInAction(() => {
                            this.interim_history = text.value;
                        });
                    } catch (e) {
                        console.error('Error parsing JSON response:', e);
                    }

                    const { annotations } = text;
                    console.log('Annotations: ' + annotations);
                    let index = 0;
                    annotations.forEach(async annotation => {
                        console.log(' ' + annotation);
                        console.log(' ' + annotation.text);
                        if (annotation.type === 'file_path') {
                            const { file_path: filePath } = annotation;
                            const fileToDownload = filePath.file_id;
                            console.log(fileToDownload);
                            if (filePath) {
                                console.log(filePath);
                                console.log(fileToDownload);
                                this.createLink(fileToDownload, annotation.start_index, annotation.end_index, ANNOTATION_LINK_TYPE.DOWNLOAD_FILE);
                            }
                        } else {
                            const { file_citation: fileCitation } = annotation;
                            if (fileCitation) {
                                const citedFile = await this.openai.files.retrieve(fileCitation.file_id);
                                const citationUrl = citedFile.filename;
                                this.createLink(citationUrl, annotation.start_index, annotation.end_index, ANNOTATION_LINK_TYPE.DASH_DOC, index);
                                index++;
                            }
                        }
                    });
                    runInAction(() => {
                        if (this.current_message) {
                            console.log('current message: ' + this.current_message.text);
                            this.current_message.text = this.interim_history;
                            this.history.push({ ...this.current_message });
                            this.current_message = undefined;
                        }
                    });
                }
            })
            .on('toolCallDone', toolCall => {
                runInAction(() => {
                    if (this.current_message && currentToolCallMessage) {
                        this.current_message.tool_logs = currentToolCallMessage;
                    }
                });
            })
            .on('imageFileDone', (content: ImageFile, snapshot: Message) => {
                console.log('Image file done:', content);
            })
            .on('end', () => {
                console.log('Streaming done');
            });
    };

    @action
    goToLinkedDoc = async (link: string) => {
        const linkedDocs = LinkManager.Instance.getAllRelatedLinks(this.Document)
            .map(d => DocCast(LinkManager.getOppositeAnchor(d, this.Document)))
            .map(d => DocCast(d?.annotationOn, d))
            .filter(d => d);

        const linkedDoc = linkedDocs.find(doc => {
            const docUrl = CsvCast(doc.data, PDFCast(doc.data)).url.pathname.replace('/files/pdfs/', '').replace('/files/csvs/', '');
            console.log('URL: ' + docUrl + ' Citation URL: ' + link);
            return link === docUrl;
        });

        if (linkedDoc) {
            await DocumentManager.Instance.showDocument(DocCast(linkedDoc), { willZoomCentered: true }, () => {});
        }
    };

    @action
    askGPT = async (event: React.FormEvent<HTMLFormElement>): Promise<void> => {
        event.preventDefault();

        const textInput = event.currentTarget.elements.namedItem('messageInput') as HTMLInputElement;
        const trimmedText = textInput.value.trim();

        if (!this.assistantID || !this.threadID) {
            try {
                await this.createAssistant();
            } catch (err) {
                console.error('Error:', err);
            }
        }

        if (trimmedText) {
            try {
                textInput.value = '';
                runInAction(() => {
                    this.history.push({ role: ASSISTANT_ROLE.USER, text: trimmedText });
                });
                await this.runAssistant(trimmedText);
                this.dataDoc.data = this.history.toString();
            } catch (err) {
                console.error('Error:', err);
            }
        }
    };

    @action
    uploadLinks = async (linkedDocs: Doc[]) => {
        if (this.isInitializing) {
            console.log('Initialization in progress, upload aborted.');
            return;
        }
        const urls = linkedDocs.map(doc => CsvCast(doc.data, PDFCast(doc.data)).url.pathname);
        const csvUrls = urls.filter(url => url.endsWith('.csv'));
        console.log(this.assistantID, this.threadID, urls);

        const { openai_file_ids: openaiFileIds } = await Networking.PostToServer('/uploadPDFToVectorStore', { urls, threadID: this.threadID, assistantID: this.assistantID, vector_store_id: this.vectorStoreID });

        linkedDocs.forEach((doc, i) => {
            doc[this.Document[Id] + '_ai_field_id'] = openaiFileIds[i];
            console.log('AI Field ID: ' + openaiFileIds[i]);
        });

        if (csvUrls.length > 0) {
            for (let i = 0; i < csvUrls.length; i++) {
                this.linkedCsvIDs.push(openaiFileIds[urls.indexOf(csvUrls[i])]);
            }
            console.log('linked csvs:' + this.linkedCsvIDs);
            await this.openai.beta.assistants.update(this.assistantID, {
                tools: [{ type: 'file_search' }, { type: 'code_interpreter' }],
                tool_resources: {
                    file_search: {
                        vector_store_ids: [this.vectorStoreID],
                    },
                    code_interpreter: {
                        file_ids: this.linkedCsvIDs,
                    },
                },
            });
        }
    };

    downloadToComputer = (url: string, fileName: string) => {
        fetch(url, { method: 'get', mode: 'no-cors', referrerPolicy: 'no-referrer' })
            .then(res => res.blob())
            .then(res => {
                const aElement = document.createElement('a');
                aElement.setAttribute('download', fileName);
                const href = URL.createObjectURL(res);
                aElement.href = href;
                aElement.setAttribute('target', '_blank');
                aElement.click();
                URL.revokeObjectURL(href);
            });
    };

    createDocumentInDash = async (url: string) => {
        const fileSuffix = url.substring(url.lastIndexOf('.') + 1);
        console.log(fileSuffix);
        let doc: Doc | null = null;
        switch (fileSuffix) {
            case 'pdf':
                doc = DocCast(await DocUtils.DocumentFromType('pdf', url, {}));
                break;
            case 'csv':
                doc = DocCast(await DocUtils.DocumentFromType('csv', url, {}));
                break;
            case 'png':
            case 'jpg':
            case 'jpeg':
                doc = DocCast(await DocUtils.DocumentFromType('image', url, {}));
                break;
            default:
                console.error('Unsupported file type:', fileSuffix);
                break;
        }
        if (doc) {
            doc && this._props.addDocument?.(doc);
            await DocumentManager.Instance.showDocument(doc, { willZoomCentered: true }, () => {});
        }
    };

    downloadFile = async (fileInfo: string, downloadType: DOWNLOAD_TYPE) => {
        try {
            console.log(fileInfo);
            const [fileId, fileName] = fileInfo.split(/!!!/);
            const { file_path: filePath } = await Networking.PostToServer('/downloadFileFromOpenAI', { file_id: fileId, file_name: fileName });
            const fileLink = CsvCast(new CsvField(filePath)).url.href;
            if (downloadType === DOWNLOAD_TYPE.DASH) {
                this.createDocumentInDash(fileLink);
            } else {
                this.downloadToComputer(fileLink, fileName);
            }
        } catch (error) {
            console.error('Error downloading file:', error);
        }
    };

    handleDownloadToDevice = () => {
        this.downloadFile(this.currentFile.url, DOWNLOAD_TYPE.DEVICE);
        this.modalStatus = false; // Close the modal after the action
        this.currentFile = { url: '' }; // Reset the current file
    };

    handleAddToDash = () => {
        // Assuming `downloadFile` is a method that handles adding to Dash
        this.downloadFile(this.currentFile.url, DOWNLOAD_TYPE.DASH);
        this.modalStatus = false; // Close the modal after the action
        this.currentFile = { url: '' }; // Reset the current file
    };

    renderModal = () => {
        if (!this.modalStatus) return null;

        return (
            <div className="modal">
                <div className="modal-content">
                    <h4>File Actions</h4>
                    <p>Choose an action for the file:</p>
                    <button type="button" onClick={this.handleDownloadToDevice}>
                        Download to Device
                    </button>
                    <button type="button" onClick={this.handleAddToDash}>
                        Add to Dash
                    </button>
                    <button
                        type="button"
                        onClick={() => {
                            this.modalStatus = false;
                        }}>
                        Cancel
                    </button>
                </div>
            </div>
        );
    };
    @action
    showModal = () => {
        this.modalStatus = true;
    };

    @action
    setCurrentFile = (file: { url: string }) => {
        this.currentFile = file;
    };

    componentDidMount() {
        this._props.setContentViewBox?.(this);
        if (this.dataDoc.data) {
            try {
                const storedHistory = JSON.parse(StrCast(this.dataDoc.data));
                runInAction(() => {
                    this.history = storedHistory.map((msg: AssistantMessage) => ({
                        role: msg.role,
                        text: msg.text,
                        quote: msg.quote,
                        tool_logs: msg.tool_logs,
                        image: msg.image,
                    }));
                });
            } catch (e) {
                console.error('Failed to parse history from dataDoc:', e);
            }
        }
        reaction(
            () => {
                const linkedDocs = LinkManager.Instance.getAllRelatedLinks(this.Document)
                    .map(d => DocCast(LinkManager.getOppositeAnchor(d, this.Document)))
                    .map(d => DocCast(d?.annotationOn, d))
                    .filter(d => d);
                return linkedDocs;
            },

            linked => this.linked_docs_to_add.push(...linked.filter(linkedDoc => !this.linked_docs_to_add.includes(linkedDoc)))
        );

        observe(
            // right now this skips during initialization which is necessary because it would be blank
            // However, it will upload the same link twice when it is
            this.linked_docs_to_add,
            change => {
                // observe pushes/splices on a user link DB 'data' field (should only happen for local changes)
                switch (change.type as any) {
                    case 'splice':
                        if ((change as any).addedCount > 0) {
                            // maybe check here if its already in the urls datadoc array so doesn't add twice
                            console.log((change as any).added as Doc[]);
                            this.uploadLinks((change as any).added as Doc[]);
                        }
                        // (change as any).removed.forEach((link: any) => remLinkFromDoc(toRealField(link)));
                        break;
                    case 'update': // let oldValue = change.oldValue;
                    default:
                }
            },
            true
        );
    }

    render() {
        return (
            <MathJaxContext config={this.mathJaxConfig}>
                <div className="chatBox">
                    {this.isInitializing && <div className="initializing-overlay">Initializing...</div>}
                    {this.renderModal()}
                    <div
                        className="scroll-box chat-content"
                        ref={r => {
                            this._oldWheel?.removeEventListener('wheel', this.onPassiveWheel);
                            this._oldWheel = r;
                            r?.addEventListener('wheel', this.onPassiveWheel, { passive: false });
                        }}>
                        <div className="messages">
                            {this.history.map((message, index) => (
                                <MessageComponent
                                    key={index}
                                    message={message}
                                    toggleToolLogs={this.toggleToolLogs}
                                    expandedLogIndex={this.expandedLogIndex}
                                    index={index}
                                    showModal={this.showModal}
                                    goToLinkedDoc={this.goToLinkedDoc}
                                    setCurrentFile={this.setCurrentFile}
                                />
                            ))}
                            {!this.current_message ? null : (
                                <MessageComponent
                                    key={this.history.length}
                                    message={this.current_message}
                                    toggleToolLogs={this.toggleToolLogs}
                                    expandedLogIndex={this.expandedLogIndex}
                                    index={this.history.length}
                                    showModal={this.showModal}
                                    goToLinkedDoc={this.goToLinkedDoc}
                                    setCurrentFile={this.setCurrentFile}
                                    isCurrent
                                />
                            )}
                        </div>
                    </div>
                    <form onSubmit={this.askGPT} className="chat-form">
                        <input type="text" name="messageInput" autoComplete="off" placeholder="Type a message..." />
                        <button type="submit">Send</button>
                    </form>
                </div>
            </MathJaxContext>
        );
    }
}

Docs.Prototypes.TemplateMap.set(DocumentType.CHAT, {
    layout: { view: ChatBox, dataField: 'data' },
    options: { acl: '', chat: '', chat_history: '', chat_thread_id: '', chat_assistant_id: '', chat_vector_store_id: '' },
});