Skip to content

Commit

Permalink
feat: add image support in chat (#132)
Browse files Browse the repository at this point in the history
* feat: add image support in chat

- Add support for pasting images directly into chat input
- Create new MentionableImage type for handling image data
- Add image preview in chat input alongside file content preview
- Add visual indicator for focused mentionable badge

* feat: add image support for LLM providers

- Add support for handling image content in messages across LLM providers (Anthropic, Gemini, Groq, OpenAI)
- Modify PromptGenerator to include image data URLs in message content
- Define new types for content parts (TextContent and ImageContentPart) in request.ts

* feat: add image upload button

- Add ImageUploadButton component for direct image uploads
- Refactor image handling to support multiple images at once

* feat: support drag & drop for images

* style: use text overflow in model select
  • Loading branch information
kevin-on authored Nov 26, 2024
1 parent c403922 commit 5dc627f
Show file tree
Hide file tree
Showing 19 changed files with 640 additions and 112 deletions.
187 changes: 130 additions & 57 deletions src/components/chat-view/chat-input/ChatUserInput.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -2,15 +2,22 @@ import { useQuery } from '@tanstack/react-query'
import { $nodesOfType, LexicalEditor, SerializedEditorState } from 'lexical'
import {
forwardRef,
useCallback,
useEffect,
useImperativeHandle,
useMemo,
useRef,
useState,
} from 'react'

import { useApp } from '../../../contexts/app-context'
import { useDarkModeContext } from '../../../contexts/dark-mode-context'
import { Mentionable, SerializedMentionable } from '../../../types/mentionable'
import {
Mentionable,
MentionableImage,
SerializedMentionable,
} from '../../../types/mentionable'
import { fileToMentionableImage } from '../../../utils/image'
import {
deserializeMentionable,
getMentionableKey,
Expand All @@ -19,6 +26,7 @@ import {
import { openMarkdownFile, readTFileContent } from '../../../utils/obsidian'
import { MemoizedSyntaxHighlighterWrapper } from '../SyntaxHighlighterWrapper'

import { ImageUploadButton } from './ImageUploadButton'
import LexicalContentEditable from './LexicalContentEditable'
import MentionableBadge from './MentionableBadge'
import { ModelSelect } from './ModelSelect'
Expand Down Expand Up @@ -57,7 +65,6 @@ const ChatUserInput = forwardRef<ChatUserInputRef, ChatUserInputProps>(
ref,
) => {
const app = useApp()
const { isDarkMode } = useDarkModeContext()

const editorRef = useRef<LexicalEditor | null>(null)
const contentEditableRef = useRef<HTMLDivElement>(null)
Expand Down Expand Up @@ -139,6 +146,29 @@ const ChatUserInput = forwardRef<ChatUserInputRef, ChatUserInputProps>(
}
}

const handleCreateImageMentionables = useCallback(
(mentionableImages: MentionableImage[]) => {
const newMentionableImages = mentionableImages.filter(
(m) =>
!mentionables.some(
(mentionable) =>
getMentionableKey(serializeMentionable(mentionable)) ===
getMentionableKey(serializeMentionable(m)),
),
)
if (newMentionableImages.length === 0) return
setMentionables([...mentionables, ...newMentionableImages])
setDisplayedMentionableKey(
getMentionableKey(
serializeMentionable(
newMentionableImages[newMentionableImages.length - 1],
),
),
)
},
[mentionables, setMentionables],
)

const handleMentionableDelete = (mentionable: Mentionable) => {
const mentionableKey = getMentionableKey(
serializeMentionable(mentionable),
Expand All @@ -158,47 +188,12 @@ const ChatUserInput = forwardRef<ChatUserInputRef, ChatUserInputProps>(
})
}

const { data: fileContent } = useQuery({
queryKey: [
'file',
displayedMentionableKey,
mentionables.map((m) => getMentionableKey(serializeMentionable(m))), // should be updated when mentionables change (especially on delete)
],
queryFn: async () => {
if (!displayedMentionableKey) return null

const displayedMentionable = mentionables.find(
(m) =>
getMentionableKey(serializeMentionable(m)) ===
displayedMentionableKey,
)

if (!displayedMentionable) return null

if (
displayedMentionable.type === 'file' ||
displayedMentionable.type === 'current-file'
) {
if (!displayedMentionable.file) return null
return await readTFileContent(displayedMentionable.file, app.vault)
} else if (displayedMentionable.type === 'block') {
const fileContent = await readTFileContent(
displayedMentionable.file,
app.vault,
)

return fileContent
.split('\n')
.slice(
displayedMentionable.startLine - 1,
displayedMentionable.endLine,
)
.join('\n')
}

return null
},
})
const handleUploadImages = async (images: File[]) => {
const mentionableImages = await Promise.all(
images.map((image) => fileToMentionableImage(image)),
)
handleCreateImageMentionables(mentionableImages)
}

const handleSubmit = (options: { useVaultSearch?: boolean } = {}) => {
const content = editorRef.current?.getEditorState()?.toJSON()
Expand Down Expand Up @@ -235,23 +230,19 @@ const ChatUserInput = forwardRef<ChatUserInputRef, ChatUserInputProps>(
setDisplayedMentionableKey(mentionableKey)
}
}}
isFocused={
getMentionableKey(serializeMentionable(m)) ===
displayedMentionableKey
}
/>
))}
</div>
)}

{fileContent && (
<div className="smtcmp-chat-user-input-file-content-preview">
<MemoizedSyntaxHighlighterWrapper
isDarkMode={isDarkMode}
language="markdown"
hasFilename={false}
wrapLines={false}
>
{fileContent}
</MemoizedSyntaxHighlighterWrapper>
</div>
)}
<MentionableContentPreview
displayedMentionableKey={displayedMentionableKey}
mentionables={mentionables}
/>

<LexicalContentEditable
initialEditorState={(editor) => {
Expand All @@ -267,6 +258,7 @@ const ChatUserInput = forwardRef<ChatUserInputRef, ChatUserInputProps>(
onEnter={() => handleSubmit({ useVaultSearch: false })}
onFocus={onFocus}
onMentionNodeMutation={handleMentionNodeMutation}
onCreateImageMentionables={handleCreateImageMentionables}
autoFocus={autoFocus}
plugins={{
onEnter: {
Expand All @@ -281,8 +273,11 @@ const ChatUserInput = forwardRef<ChatUserInputRef, ChatUserInputProps>(
/>

<div className="smtcmp-chat-user-input-controls">
<ModelSelect />
<div className="smtcmp-chat-user-input-controls-buttons">
<div className="smtcmp-chat-user-input-controls__model-select-container">
<ModelSelect />
</div>
<div className="smtcmp-chat-user-input-controls__buttons">
<ImageUploadButton onUpload={handleUploadImages} />
<SubmitButton onClick={() => handleSubmit()} />
<VaultChatButton
onClick={() => {
Expand All @@ -296,6 +291,84 @@ const ChatUserInput = forwardRef<ChatUserInputRef, ChatUserInputProps>(
},
)

function MentionableContentPreview({
displayedMentionableKey,
mentionables,
}: {
displayedMentionableKey: string | null
mentionables: Mentionable[]
}) {
const app = useApp()
const { isDarkMode } = useDarkModeContext()

const displayedMentionable: Mentionable | null = useMemo(() => {
return (
mentionables.find(
(m) =>
getMentionableKey(serializeMentionable(m)) ===
displayedMentionableKey,
) ?? null
)
}, [displayedMentionableKey, mentionables])

const { data: displayFileContent } = useQuery({
enabled:
!!displayedMentionable &&
['file', 'current-file', 'block'].includes(displayedMentionable.type),
queryKey: [
'file',
displayedMentionableKey,
mentionables.map((m) => getMentionableKey(serializeMentionable(m))), // should be updated when mentionables change (especially on delete)
],
queryFn: async () => {
if (!displayedMentionable) return null
if (
displayedMentionable.type === 'file' ||
displayedMentionable.type === 'current-file'
) {
if (!displayedMentionable.file) return null
return await readTFileContent(displayedMentionable.file, app.vault)
} else if (displayedMentionable.type === 'block') {
const fileContent = await readTFileContent(
displayedMentionable.file,
app.vault,
)

return fileContent
.split('\n')
.slice(
displayedMentionable.startLine - 1,
displayedMentionable.endLine,
)
.join('\n')
}

return null
},
})

const displayImage: MentionableImage | null = useMemo(() => {
return displayedMentionable?.type === 'image' ? displayedMentionable : null
}, [displayedMentionable])

return displayFileContent ? (
<div className="smtcmp-chat-user-input-file-content-preview">
<MemoizedSyntaxHighlighterWrapper
isDarkMode={isDarkMode}
language="markdown"
hasFilename={false}
wrapLines={false}
>
{displayFileContent}
</MemoizedSyntaxHighlighterWrapper>
</div>
) : displayImage ? (
<div className="smtcmp-chat-user-input-file-content-preview">
<img src={displayImage.data} alt={displayImage.name} />
</div>
) : null
}

ChatUserInput.displayName = 'ChatUserInput'

export default ChatUserInput
30 changes: 30 additions & 0 deletions src/components/chat-view/chat-input/ImageUploadButton.tsx
Original file line number Diff line number Diff line change
@@ -0,0 +1,30 @@
import { ImageIcon } from 'lucide-react'

export function ImageUploadButton({
onUpload,
}: {
onUpload: (files: File[]) => void
}) {
const handleFileChange = (event: React.ChangeEvent<HTMLInputElement>) => {
const files = Array.from(event.target.files ?? [])
if (files.length > 0) {
onUpload(files)
}
}

return (
<label className="smtcmp-chat-user-input-submit-button">
<input
type="file"
accept="image/*"
multiple
onChange={handleFileChange}
style={{ display: 'none' }}
/>
<div className="smtcmp-chat-user-input-submit-button-icons">
<ImageIcon size={12} />
</div>
<div>Image</div>
</label>
)
}
Original file line number Diff line number Diff line change
Expand Up @@ -13,8 +13,11 @@ import { LexicalEditor, SerializedEditorState } from 'lexical'
import { RefObject, useCallback, useEffect } from 'react'

import { useApp } from '../../../contexts/app-context'
import { MentionableImage } from '../../../types/mentionable'
import { fuzzySearch } from '../../../utils/fuzzy-search'

import DragDropPaste from './plugins/image/DragDropPastePlugin'
import ImagePastePlugin from './plugins/image/ImagePastePlugin'
import AutoLinkMentionPlugin from './plugins/mention/AutoLinkMentionPlugin'
import { MentionNode } from './plugins/mention/MentionNode'
import MentionPlugin from './plugins/mention/MentionPlugin'
Expand All @@ -33,6 +36,7 @@ export type LexicalContentEditableProps = {
onEnter?: (evt: KeyboardEvent) => void
onFocus?: () => void
onMentionNodeMutation?: (mutations: NodeMutations<MentionNode>) => void
onCreateImageMentionables?: (mentionables: MentionableImage[]) => void
initialEditorState?: InitialEditorStateType
autoFocus?: boolean
plugins?: {
Expand All @@ -52,6 +56,7 @@ export default function LexicalContentEditable({
onEnter,
onFocus,
onMentionNodeMutation,
onCreateImageMentionables,
initialEditorState,
autoFocus = false,
plugins,
Expand Down Expand Up @@ -134,6 +139,8 @@ export default function LexicalContentEditable({
<EditorRefPlugin editorRef={editorRef} />
<NoFormatPlugin />
<AutoLinkMentionPlugin />
<ImagePastePlugin onCreateImageMentionables={onCreateImageMentionables} />
<DragDropPaste onCreateImageMentionables={onCreateImageMentionables} />
<TemplatePlugin />
{plugins?.templatePopover && (
<CreateTemplatePopoverPlugin
Expand Down
Loading

0 comments on commit 5dc627f

Please sign in to comment.