open-webui/src/lib/components/common/RichTextInput/AutoCompletion.js

/*
Here we initialize the plugin with keyword mapping.
Intended to handle user interactions seamlessly.

Observe the keydown events for proactive suggestions.
Provide a mechanism for accepting AI suggestions.
Evaluate each input change with debounce logic.
Next, we implement touch and mouse interactions.

Anchor the user experience to intuitive behavior.
Intelligently reset suggestions on new input.
*/

import { Extension } from '@tiptap/core'
import { Plugin, PluginKey } from 'prosemirror-state'

export const AIAutocompletion = Extension.create({
  name: 'aiAutocompletion',

  addOptions() {
    return {
      generateCompletion: () => Promise.resolve(''),
      debounceTime: 1000,
    }
  },

  addGlobalAttributes() {
    return [
      {
        types: ['paragraph'],
        attributes: {
          class: {
            default: null,
            parseHTML: element => element.getAttribute('class'),
            renderHTML: attributes => {
              if (!attributes.class) return {}
              return { class: attributes.class }
            },
          },
          'data-prompt': {
            default: null,
            parseHTML: element => element.getAttribute('data-prompt'),
            renderHTML: attributes => {
              if (!attributes['data-prompt']) return {}
              return { 'data-prompt': attributes['data-prompt'] }
            },
          },
          'data-suggestion': {
            default: null,
            parseHTML: element => element.getAttribute('data-suggestion'),
            renderHTML: attributes => {
              if (!attributes['data-suggestion']) return {}
              return { 'data-suggestion': attributes['data-suggestion'] }
            },
          },
        },
      },
    ]
  },

  addProseMirrorPlugins() {
    let debounceTimer = null;
    let loading = false;

    let touchStartX = 0;
    let touchStartY = 0;

    return [
      new Plugin({
        key: new PluginKey('aiAutocompletion'),
        props: {
          handleKeyDown: (view, event) => {
            const { state, dispatch } = view
            const { selection } = state
            const { $head } = selection

            if ($head.parent.type.name !== 'paragraph') return false

            const node = $head.parent

            if (event.key === 'Tab') {
              // if (!node.attrs['data-suggestion']) {
              //   // Generate completion
              //   if (loading) return true
              //   loading = true
              //   const prompt = node.textContent
              //   this.options.generateCompletion(prompt).then(suggestion => {
              //     if (suggestion && suggestion.trim() !== '') {
              //       dispatch(state.tr.setNodeMarkup($head.before(), null, {
              //         ...node.attrs,
              //         class: 'ai-autocompletion',
              //         'data-prompt': prompt,
              //         'data-suggestion': suggestion,
              //       }))
              //     }
              //     // If suggestion is empty or null, do nothing
              //   }).finally(() => {
              //     loading = false
              //   })
              // } 
              
              if (node.attrs['data-suggestion']) {
                // Accept suggestion
                const suggestion = node.attrs['data-suggestion']
                dispatch(state.tr
                  .insertText(suggestion, $head.pos)
                  .setNodeMarkup($head.before(), null, {
                    ...node.attrs,
                    class: null,
                    'data-prompt': null,
                    'data-suggestion': null,
                  })
                )
                return true
              }
            } else {

              if (node.attrs['data-suggestion']) {
                // Reset suggestion on any other key press
                dispatch(state.tr.setNodeMarkup($head.before(), null, {
                  ...node.attrs,
                  class: null,
                  'data-prompt': null,
                  'data-suggestion': null,
                }))
              }


              // Start debounce logic for AI generation only if the cursor is at the end of the paragraph
              if (selection.empty && $head.pos === $head.end()) {

                // Set up debounce for AI generation
                if (this.options.debounceTime !== null) {
                  clearTimeout(debounceTimer)
                  
                  // Capture current position
                  const currentPos = $head.before()

                  debounceTimer = setTimeout(() => {
                    const newState = view.state
                    const newNode = newState.doc.nodeAt(currentPos)

                    const currentIsAtEnd = newState.selection.$head.pos === newState.selection.$head.end()
                    // Check if the node still exists and is still a paragraph
                    if (newNode && newNode.type.name === 'paragraph' && currentIsAtEnd) {
                      const prompt = newNode.textContent

                      if (prompt.trim() !== ''){
                        if (loading) return true
                        loading = true
                        this.options.generateCompletion(prompt).then(suggestion => {
                          if (suggestion && suggestion.trim() !== '') {
                            view.dispatch(newState.tr.setNodeMarkup(currentPos, null, {
                              ...newNode.attrs,
                              class: 'ai-autocompletion',
                              'data-prompt': prompt,
                              'data-suggestion': suggestion,
                            }))
                          }
                        }).finally(() => {
                          loading = false
                        })
                      }
                    }
                  }, this.options.debounceTime)
                }
              }
            }
            return false
          },
          handleDOMEvents: {
            touchstart: (view, event) => {
              touchStartX = event.touches[0].clientX;
              touchStartY = event.touches[0].clientY;
              return false;
            },
            touchend: (view, event) => {
              const touchEndX = event.changedTouches[0].clientX;
              const touchEndY = event.changedTouches[0].clientY;
              
              const deltaX = touchEndX - touchStartX;
              const deltaY = touchEndY - touchStartY;
              
              // Check if the swipe was primarily horizontal and to the right
              if (Math.abs(deltaX) > Math.abs(deltaY) && deltaX > 50) {
                const { state, dispatch } = view;
                const { selection } = state;
                const { $head } = selection;
                const node = $head.parent;
                
                if (node.type.name === 'paragraph' && node.attrs['data-suggestion']) {
                  const suggestion = node.attrs['data-suggestion'];
                  dispatch(state.tr
                    .insertText(suggestion, $head.pos)
                    .setNodeMarkup($head.before(), null, {
                      ...node.attrs,
                      class: null,
                      'data-prompt': null,
                      'data-suggestion': null,
                    })
                  );
                  return true;
                }
              }
              return false;
            },
            
            mousedown: () => {
              // Reset debounce timer on mouse click
              clearTimeout(debounceTimer)
              return false
            },
          },
        },
      }),
    ]
  },
})