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

import { action, computed, makeObservable, observable, observe, reaction, runInAction, ObservableSet } from 'mobx';
import { observer } from 'mobx-react';
import OpenAI, { ClientOptions } from 'openai';
import * as React from 'react';
import { Doc, DocListCast } from '../../../../fields/Doc';
import { CsvCast, DocCast, PDFCast, StrCast } from '../../../../fields/Types';
import { DocumentType } from '../../../documents/DocumentTypes';
import { Docs } from '../../../documents/Documents';
import { LinkManager } from '../../../util/LinkManager';
import { ViewBoxAnnotatableComponent } from '../../DocComponent';
import { FieldView, FieldViewProps } from '../FieldView';
import './ChatBox.scss';
import MessageComponentBox from './MessageComponent';
import { ASSISTANT_ROLE, AssistantMessage, AI_Document, Citation, CHUNK_TYPE, RAGChunk, getChunkType, TEXT_TYPE, SimplifiedChunk, ProcessingInfo, MessageContent } from './types';
import { Vectorstore } from './vectorstore/Vectorstore';
import { Agent } from './Agent';
import dotenv from 'dotenv';
import { DocData, DocViews } from '../../../../fields/DocSymbols';
import { AnswerParser } from './AnswerParser';
import { DocumentManager } from '../../../util/DocumentManager';
import { v4 as uuidv4 } from 'uuid';
import { chunk } from 'lodash';
import { DocUtils } from '../../../documents/DocUtils';
import { createRef } from 'react';
import { ClientUtils } from '../../../../ClientUtils';

dotenv.config();

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

    @observable isLoading: boolean = false;
    @observable isUploadingDocs: boolean = false;
    @observable expandedScratchpadIndex: number | null = null;
    @observable inputValue: string = '';
    @observable private linked_docs_to_add: ObservableSet<Doc> = observable.set();
    @observable private linked_csv_files: { filename: string; id: string; text: string }[] = [];
    private openai: OpenAI;
    private vectorstore_id: string;
    private vectorstore: Vectorstore;
    private agent: Agent; // Add the ChatBot instance
    private _oldWheel: HTMLDivElement | null = null;
    private messagesRef: React.RefObject<HTMLDivElement>;

    public static LayoutString(fieldKey: string) {
        return FieldView.LayoutString(ChatBox, fieldKey);
    }

    constructor(props: FieldViewProps) {
        super(props);
        makeObservable(this);
        this.openai = this.initializeOpenAI();
        if (StrCast(this.dataDoc.vectorstore_id) == '') {
            console.log('new_id');
            this.vectorstore_id = uuidv4();
            this.dataDoc.vectorstore_id = this.vectorstore_id;
        } else {
            this.vectorstore_id = StrCast(this.dataDoc.vectorstore_id);
        }
        this.vectorstore = new Vectorstore(this.vectorstore_id, this.retrieveDocIds);
        this.agent = new Agent(this.vectorstore, this.retrieveSummaries, this.retrieveFormattedHistory, this.retrieveCSVData, this.addLinkedUrlDoc);
        this.messagesRef = React.createRef<HTMLDivElement>();

        reaction(
            () => this.history.map((msg: AssistantMessage) => ({ role: msg.role, content: msg.content, follow_up_questions: msg.follow_up_questions, citations: msg.citations })),
            serializableHistory => {
                this.dataDoc.data = JSON.stringify(serializableHistory);
            }
        );
    }

    @action
    addDocToVectorstore = async (newLinkedDoc: Doc) => {
        await this.vectorstore.addAIDoc(newLinkedDoc);
    };

    @action
    addCSVForAnalysis = async (newLinkedDoc: Doc) => {
        console.log('adding csv file for analysis');
        if (!newLinkedDoc.chunk_simpl) {
            const csvData: string = StrCast(newLinkedDoc.text);
            console.log('CSV Data:', csvData);
            const completion = await this.openai.chat.completions.create({
                messages: [
                    {
                        role: 'system',
                        content:
                            'You are an AI assistant tasked with summarizing the content of a CSV file. You will be provided with the data from the CSV file and your goal is to generate a concise summary that captures the main themes, trends, and key points represented in the data.',
                    },
                    {
                        role: 'user',
                        content: `Please provide a comprehensive summary of the CSV file based on the provided data. Ensure the summary highlights the most important information, patterns, and insights. Your response should be in paragraph form and be concise.
                
                    CSV Data:
                
                    ${csvData}
                
                    **********
                    Summary:`,
                    },
                ],
                model: 'gpt-3.5-turbo',
            });
            console.log('CSV Data:', csvData);
            const csvId = uuidv4();

            this.linked_csv_files.push({
                filename: CsvCast(newLinkedDoc.data).url.pathname,
                id: csvId,
                text: csvData,
            });

            console.log(this.linked_csv_files);
            const chunkToAdd = {
                chunkId: csvId,
                chunkType: CHUNK_TYPE.CSV,
            };
            newLinkedDoc.chunk_simpl = JSON.stringify({ chunks: [chunkToAdd] });
            newLinkedDoc.summary = completion.choices[0].message.content!;
        }
    };

    @action
    toggleToolLogs = (index: number) => {
        this.expandedScratchpadIndex = this.expandedScratchpadIndex === index ? null : index;
    };

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

    addScrollListener = () => {
        if (this.messagesRef.current) {
            this.messagesRef.current.addEventListener('wheel', this.onPassiveWheel, { passive: false });
        }
    };

    removeScrollListener = () => {
        if (this.messagesRef.current) {
            this.messagesRef.current.removeEventListener('wheel', this.onPassiveWheel);
        }
    };

    scrollToBottom = () => {
        if (this.messagesRef.current) {
            this.messagesRef.current.scrollTop = this.messagesRef.current.scrollHeight;
        }
    };

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

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

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

        if (trimmedText) {
            try {
                textInput.value = '';
                this.history.push({ role: ASSISTANT_ROLE.USER, content: [{ index: 0, type: TEXT_TYPE.NORMAL, text: trimmedText, citation_ids: null }], processing_info: [] });
                this.isLoading = true;
                this.current_message = { role: ASSISTANT_ROLE.ASSISTANT, content: [], citations: [], processing_info: [] };

                const onProcessingUpdate = (processingUpdate: ProcessingInfo[]) => {
                    runInAction(() => {
                        if (this.current_message) {
                            this.current_message = { ...this.current_message, processing_info: processingUpdate };
                        }
                    });
                    this.scrollToBottom();
                };

                const onAnswerUpdate = (answerUpdate: string) => {
                    runInAction(() => {
                        if (this.current_message) {
                            this.current_message = { ...this.current_message, content: [{ text: answerUpdate, type: TEXT_TYPE.NORMAL, index: 0, citation_ids: [] }] };
                        }
                    });
                };

                const finalMessage = await this.agent.askAgent(trimmedText, onProcessingUpdate, onAnswerUpdate);

                runInAction(() => {
                    if (this.current_message) {
                        this.history.push({ ...finalMessage });
                        this.current_message = undefined;
                        this.dataDoc.data = JSON.stringify(this.history);
                    }
                });
            } catch (err) {
                console.error('Error:', err);
                this.history.push({ role: ASSISTANT_ROLE.ASSISTANT, content: [{ index: 0, type: TEXT_TYPE.ERROR, text: 'Sorry, I encountered an error while processing your request.', citation_ids: null }], processing_info: [] });
            } finally {
                this.isLoading = false;
                this.scrollToBottom();
            }
        }
        this.scrollToBottom();
    };

    @action
    updateMessageCitations = (index: number, citations: Citation[]) => {
        if (this.history[index]) {
            this.history[index].citations = citations;
        }
    };

    @action
    addLinkedUrlDoc = async (url: string, id: string) => {
        const doc = Docs.Create.WebDocument(url);

        const linkDoc = Docs.Create.LinkDocument(this.Document, doc);
        LinkManager.Instance.addLink(linkDoc);

        const chunkToAdd = {
            chunkId: id,
            chunkType: CHUNK_TYPE.URL,
        };

        doc.chunk_simpl = JSON.stringify({ chunks: [chunkToAdd] });
    };

    @computed
    get userName() {
        return ClientUtils.CurrentUserEmail;
    }

    @action
    handleCitationClick = (citation: Citation) => {
        console.log('Citation clicked:', citation);
        const currentLinkedDocs: Doc[] = this.linkedDocs;

        const chunkId = citation.chunk_id;

        for (let doc of currentLinkedDocs) {
            if (doc.chunk_simpl) {
                const docChunkSimpl = JSON.parse(StrCast(doc.chunk_simpl)) as { chunks: SimplifiedChunk[] };
                console.log(docChunkSimpl);
                const foundChunk = docChunkSimpl.chunks.find(chunk => chunk.chunkId === chunkId);
                if (foundChunk) {
                    switch (getChunkType(foundChunk.chunkType)) {
                        case CHUNK_TYPE.IMAGE:
                            const values = foundChunk.location?.replace(/[\[\]]/g, '').split(',');

                            if (values?.length !== 4) {
                                console.error('Location string must contain exactly 4 numbers');
                                return;
                            }

                            const x1 = parseFloat(values[0]) * Doc.NativeWidth(doc);
                            const y1 = parseFloat(values[1]) * Doc.NativeHeight(doc);
                            const x2 = parseFloat(values[2]) * Doc.NativeWidth(doc);
                            const y2 = parseFloat(values[3]) * Doc.NativeHeight(doc);

                            const annotationKey = Doc.LayoutFieldKey(doc) + '_annotations';

                            const existingDoc = DocListCast(doc[DocData][annotationKey]).find(d => d.citation_id === citation.citation_id);
                            const highlightDoc = existingDoc ?? this.createImageCitationHighlight(x1, y1, x2, y2, citation, annotationKey, doc);

                            DocumentManager.Instance.showDocument(highlightDoc, { willZoomCentered: true }, () => {});
                            break;
                        case CHUNK_TYPE.TEXT:
                            DocumentManager.Instance.showDocument(doc, { willZoomCentered: true }, () => {
                                const firstView = Array.from(doc[DocViews])[0];
                                firstView.ComponentView?.search?.(citation.direct_text);
                            });
                            break;
                        case CHUNK_TYPE.URL:
                            DocumentManager.Instance.showDocument(doc, { willZoomCentered: true }, () => {
                                const firstView = Array.from(doc[DocViews])[0];
                            });
                            break;
                        case CHUNK_TYPE.CSV:
                            DocumentManager.Instance.showDocument(doc, { willZoomCentered: true }, () => {
                                const firstView = Array.from(doc[DocViews])[0];
                            });
                            break;
                        default:
                            console.log('Chunk type not supported');
                            break;
                    }
                }
            }
        }
    };

    createImageCitationHighlight = (x1: number, y1: number, x2: number, y2: number, citation: Citation, annotationKey: string, pdfDoc: Doc): Doc => {
        const highlight_doc = Docs.Create.FreeformDocument([], {
            x: x1,
            y: y1,
            _width: x2 - x1,
            _height: y2 - y1,
            backgroundColor: 'rgba(255, 255, 0, 0.5)',
        });
        highlight_doc[DocData].citation_id = citation.citation_id;
        Doc.AddDocToList(pdfDoc[DocData], annotationKey, highlight_doc);
        highlight_doc.annotationOn = pdfDoc;
        Doc.SetContainer(highlight_doc, pdfDoc);
        return highlight_doc;
    };

    componentDidUpdate() {
        this.scrollToBottom();
    }

    componentDidMount() {
        this._props.setContentViewBox?.(this);
        if (this.dataDoc.data) {
            try {
                const storedHistory = JSON.parse(StrCast(this.dataDoc.data));
                runInAction(() => {
                    this.history.push(
                        ...storedHistory.map((msg: AssistantMessage) => ({
                            role: msg.role,
                            content: msg.content,
                            follow_up_questions: msg.follow_up_questions,
                            citations: msg.citations,
                        }))
                    );
                });
            } catch (e) {
                console.error('Failed to parse history from dataDoc:', e);
            }
        } else {
            runInAction(() => {
                this.history.push({
                    role: ASSISTANT_ROLE.ASSISTANT,
                    content: [{ index: 0, type: TEXT_TYPE.NORMAL, text: 'Welcome to the Document Analyser Assistant! Link a document or ask questions to get started.', citation_ids: null }],
                    processing_info: [],
                });
            });
        }
        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 => linked.forEach(doc => this.linked_docs_to_add.add(doc))
        );

        observe(this.linked_docs_to_add, change => {
            if (change.type === 'add') {
                runInAction(() => {
                    this.isUploadingDocs = true;
                });
                if (PDFCast(change.newValue.data)) {
                    this.addDocToVectorstore(change.newValue);
                } else if (CsvCast(change.newValue.data)) {
                    this.addCSVForAnalysis(change.newValue);
                }
                runInAction(() => {
                    this.isUploadingDocs = false;
                });
            } else if (change.type === 'delete') {
                console.log('Deleted docs: ', change.oldValue);
            }
        });
        this.addScrollListener();
    }

    componentWillUnmount() {
        this.removeScrollListener();
    }

    @computed
    get linkedDocs() {
        return LinkManager.Instance.getAllRelatedLinks(this.Document)
            .map(d => DocCast(LinkManager.getOppositeAnchor(d, this.Document)))
            .map(d => DocCast(d?.annotationOn, d))
            .filter(d => d);
    }

    @computed
    get docIds() {
        return LinkManager.Instance.getAllRelatedLinks(this.Document)
            .map(d => DocCast(LinkManager.getOppositeAnchor(d, this.Document)))
            .map(d => DocCast(d?.annotationOn, d))
            .filter(d => d)
            .filter(d => d.ai_doc_id)
            .map(d => StrCast(d.ai_doc_id));
    }

    @computed
    get summaries(): string {
        return (
            LinkManager.Instance.getAllRelatedLinks(this.Document)
                .map(d => DocCast(LinkManager.getOppositeAnchor(d, this.Document)))
                .map(d => DocCast(d?.annotationOn, d))
                .filter(d => d)
                .filter(d => d.summary)
                .map((doc, index) => {
                    if (PDFCast(doc.data)) {
                        return `<summary file_name="${PDFCast(doc.data).url.pathname}" applicable_tools=["rag"]>${doc.summary}</summary>`;
                    } else if (CsvCast(doc.data)) {
                        return `<summary file_name="${CsvCast(doc.data).url.pathname}" applicable_tools=["dataAnalysis"]>${doc.summary}</summary>`;
                    } else {
                        return `${index + 1}) ${doc.summary}`;
                    }
                })
                .join('\n') + '\n'
        );
    }

    @computed
    get linkedCSVs(): { filename: string; id: string; text: string }[] {
        return this.linked_csv_files;
    }

    @computed
    get formattedHistory(): string {
        let history = '<chat_history>\n';
        for (const message of this.history) {
            history += `<${message.role}>${message.content.map(content => content.text).join(' ')}</${message.role}>\n`;
        }
        history += '</chat_history>';
        return history;
    }

    retrieveSummaries = () => {
        return this.summaries;
    };

    retrieveCSVData = () => {
        return this.linkedCSVs;
    };

    retrieveFormattedHistory = () => {
        return this.formattedHistory;
    };

    retrieveDocIds = () => {
        return this.docIds;
    };

    @action
    handleFollowUpClick = (question: string) => {
        console.log('Follow-up question clicked:', question);
        this.inputValue = question;
    };
    render() {
        return (
            <div className="chat-box">
                {this.isUploadingDocs && (
                    <div className="uploading-overlay">
                        <div className="spinner"></div>
                    </div>
                )}
                <div className="chat-header">
                    <h2>{this.userName()}'s AI Assistant</h2>
                </div>
                <div className="chat-messages" ref={this.messagesRef}>
                    {this.history.map((message, index) => (
                        <MessageComponentBox key={index} message={message} index={index} onFollowUpClick={this.handleFollowUpClick} onCitationClick={this.handleCitationClick} updateMessageCitations={this.updateMessageCitations} />
                    ))}
                    {this.current_message && (
                        <MessageComponentBox
                            key={this.history.length}
                            message={this.current_message}
                            index={this.history.length}
                            onFollowUpClick={this.handleFollowUpClick}
                            onCitationClick={this.handleCitationClick}
                            updateMessageCitations={this.updateMessageCitations}
                        />
                    )}
                </div>
                <form onSubmit={this.askGPT} className="chat-input">
                    <input type="text" name="messageInput" autoComplete="off" placeholder="Type your message here..." value={this.inputValue} onChange={e => (this.inputValue = e.target.value)} />
                    <button className="submit-button" type="submit" disabled={this.isLoading}>
                        {this.isLoading ? (
                            <div className="spinner"></div>
                        ) : (
                            <svg viewBox="0 0 24 24" width="24" height="24" stroke="currentColor" strokeWidth="2" fill="none" strokeLinecap="round" strokeLinejoin="round">
                                <line x1="22" y1="2" x2="11" y2="13"></line>
                                <polygon points="22 2 15 22 11 13 2 9 22 2"></polygon>
                            </svg>
                        )}
                    </button>
                </form>
            </div>
        );
    }
}

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: '' },
});