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 } 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';

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();
    private openai: OpenAI;
    private vectorstore_id: string;
    private documents: AI_Document[] = [];
    private _oldWheel: any;
    private vectorstore: Vectorstore;
    private agent: Agent; // Add the ChatBot instance

    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);

        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
    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);
    }

    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 = '';
                runInAction(() => {
                    this.history.push({ role: ASSISTANT_ROLE.USER, content: [{ index: 0, type: TEXT_TYPE.NORMAL, text: trimmedText, citation_ids: null }] });
                    this.isLoading = true;
                });

                const response = await this.agent.askAgent(trimmedText); // Use the chatbot to get the response
                runInAction(() => {
                    this.history.push(AnswerParser.parse(response));
                });
                this.dataDoc.data = JSON.stringify(this.history);
            } catch (err) {
                console.error('Error:', err);
                runInAction(() => {
                    this.history.push({ role: ASSISTANT_ROLE.ASSISTANT, content: [{ index: 0, type: TEXT_TYPE.NORMAL, text: 'Sorry, I encountered an error while processing your request.', citation_ids: null }] });
                });
            } finally {
                runInAction(() => {
                    this.isLoading = false;
                });
            }
        }
    };

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

    @action
    handleCitationClick = (citation: Citation) => {
        console.log('Citation clicked:', citation);
        const currentLinkedDocs: Doc[] = this.linkedDocs;
        const chunk_id = citation.chunk_id;
        for (let doc of currentLinkedDocs) {
            if (doc.chunk_simpl) {
                //console.log(JSON.parse(StrCast(doc.chunk_simpl)));
                const doc_chunk_simpl = JSON.parse(StrCast(doc.chunk_simpl));
                console.log(doc_chunk_simpl);
                const text_chunks = doc_chunk_simpl.text_chunks as [{ chunk_id: string; start_page: number; end_page: number }] | [];
                const image_chunks = doc_chunk_simpl.image_chunks as [{ chunk_id: string; location: string; page: number }] | [];

                const found_text_chunk = text_chunks.find(chunk => chunk.chunk_id === chunk_id);
                if (found_text_chunk) {
                    const doc_url = CsvCast(doc.data, PDFCast(doc.data)).url.pathname;
                    console.log('URL: ' + doc_url);

                    //const ai_field_id = doc[this.Document[Id] + '_ai_field_id'];
                    DocumentManager.Instance.showDocument(doc, { willZoomCentered: true }, () => {
                        console.log(doc.data);
                        //look at context path for each docview and choose the doc view that has as
                        //its parent the same collection view the chatbox is in
                        const first_view = Array.from(doc[DocViews])[0];
                        first_view.ComponentView?.search?.(citation.direct_text);
                    });
                }

                const found_image_chunk = image_chunks.find(chunk => chunk.chunk_id === chunk_id);
                if (found_image_chunk) {
                    const location_string: string = found_image_chunk.location;

                    // Extract variables from location_string
                    const values = location_string.replace(/[\[\]]/g, '').split(',');

                    // Ensure we have exactly 4 values
                    if (values.length !== 4) {
                        console.error('Location string must contain exactly 4 numbers');
                        return; // or handle this error as appropriate
                    }

                    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 highlight_doc = existingDoc ?? this.createImageCitationHighlight(x1, y1, x2, y2, citation, annotationKey, doc);

                    DocumentManager.Instance.showDocument(highlight_doc, { willZoomCentered: true }, () => {});
                }
            }
        }
        // You can implement additional functionality here, such as showing a modal with the full citation content
    };

    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;
    };

    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 }],
                });
            });
        }
        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;
                });
                this.addDocToVectorstore(change.newValue);
                runInAction(() => {
                    this.isUploadingDocs = false;
                });
            } else if (change.type === 'delete') {
                console.log('Deleted docs: ', change.oldValue);
            }
        });
    }

    @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) => `${index + 1}) ${doc.summary}`)
                .join('\n') + '\n'
        );
    }

    @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;
    };

    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="chatBox">
                {this.isUploadingDocs && <div className="uploading-overlay"></div>}
                <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) => (
                            <MessageComponentBox key={index} message={message} index={index} onFollowUpClick={this.handleFollowUpClick} onCitationClick={this.handleCitationClick} updateMessageCitations={this.updateMessageCitations} />
                        ))}
                    </div>
                </div>
                <form onSubmit={this.askGPT} className="chat-form">
                    <input type="text" name="messageInput" autoComplete="off" placeholder="Type a message..." value={this.inputValue} onChange={e => (this.inputValue = e.target.value)} />
                    <button type="submit" disabled={this.isLoading}>
                        {this.isLoading ? 'Thinking...' : 'Send'}
                    </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: '' },
});