diff --git a/packages/frontend/src/assets/plumber-logo.svg b/packages/frontend/src/assets/plumber-logo.svg
new file mode 100644
index 0000000000..23a7aab094
--- /dev/null
+++ b/packages/frontend/src/assets/plumber-logo.svg
@@ -0,0 +1,5 @@
+
diff --git a/packages/frontend/src/hooks/useChatStream.ts b/packages/frontend/src/hooks/useChatStream.ts
new file mode 100644
index 0000000000..3daa2cbac5
--- /dev/null
+++ b/packages/frontend/src/hooks/useChatStream.ts
@@ -0,0 +1,127 @@
+import { useCallback, useMemo } from 'react'
+import type { UIMessage } from '@ai-sdk/react'
+import { useChat } from '@ai-sdk/react'
+import { useToast } from '@opengovsg/design-system-react'
+import { DefaultChatTransport } from 'ai'
+
+export interface Message {
+ text: string
+ traceId?: string
+ generationId?: string
+ isUser: boolean
+}
+
+// Custom message type with metadata
+type CustomUIMessage = UIMessage<{
+ traceId?: string
+}>
+
+export function useChatStream() {
+ const toast = useToast()
+
+ const {
+ messages: aiMessages,
+ sendMessage,
+ status,
+ error: aiError,
+ stop,
+ } = useChat({
+ transport: new DefaultChatTransport({
+ api: '/api/chat',
+ credentials: 'include',
+ prepareSendMessagesRequest: ({ messages }) => {
+ // Send all messages to maintain conversation context
+ const body = {
+ messages: messages,
+ sessionId: '',
+ }
+ return { body }
+ },
+ }),
+ onError: (error: Error) => {
+ toast({
+ title: 'Error: ' + error.message,
+ status: 'error',
+ duration: 3000,
+ isClosable: true,
+ position: 'top',
+ })
+ },
+ })
+
+ // Helper function to extract text content from UIMessage
+ const extractTextContent = useCallback((msg: CustomUIMessage): string => {
+ return msg.parts
+ .filter((part) => part.type === 'text')
+ .map((part) => (part as any).text)
+ .join('')
+ }, [])
+
+ // Transform AI SDK messages to our Message format
+ const messages = useMemo(() => {
+ const isActivelyStreaming = status === 'streaming' || status === 'submitted'
+
+ // Filter user and assistant messages
+ let messagesToTransform = aiMessages.filter(
+ (msg) => msg.role === 'user' || msg.role === 'assistant',
+ )
+
+ // If streaming, exclude the last assistant message (shown via currentResponse)
+ if (isActivelyStreaming && messagesToTransform.length > 0) {
+ const lastMsg = messagesToTransform[messagesToTransform.length - 1]
+ if (lastMsg.role === 'assistant') {
+ messagesToTransform = messagesToTransform.slice(0, -1)
+ }
+ }
+
+ return messagesToTransform.map((msg) => {
+ // Extract traceId from message metadata
+ const traceId = msg.metadata?.traceId
+
+ return {
+ text: extractTextContent(msg),
+ isUser: msg.role === 'user',
+ traceId: traceId,
+ }
+ })
+ }, [aiMessages, extractTextContent, status])
+
+ // Get the current streaming response (last assistant message that's still being streamed)
+ const currentResponse = useMemo(() => {
+ const isActivelyStreaming = status === 'streaming' || status === 'submitted'
+
+ if (!isActivelyStreaming) {
+ return ''
+ }
+
+ const lastMessage = aiMessages[aiMessages.length - 1]
+ if (lastMessage && lastMessage.role === 'assistant') {
+ return extractTextContent(lastMessage)
+ }
+ return ''
+ }, [aiMessages, status, extractTextContent])
+
+ // Wrapper for sendMessage that matches the expected signature
+ const sendMessageWrapper = useCallback(
+ (userPrompt: string) => {
+ sendMessage({
+ role: 'user',
+ parts: [{ type: 'text', text: userPrompt }],
+ })
+ },
+ [sendMessage],
+ )
+
+ const cancelStream = useCallback(() => {
+ stop()
+ }, [stop])
+
+ return {
+ messages,
+ currentResponse,
+ isStreaming: status === 'submitted' || status === 'streaming',
+ error: aiError?.message || null,
+ sendMessage: sendMessageWrapper,
+ cancelStream,
+ }
+}
diff --git a/packages/frontend/src/pages/AiBuilder/AiBuilderContext.tsx b/packages/frontend/src/pages/AiBuilder/AiBuilderContext.tsx
index 910d9882c4..74c4d2a8a1 100644
--- a/packages/frontend/src/pages/AiBuilder/AiBuilderContext.tsx
+++ b/packages/frontend/src/pages/AiBuilder/AiBuilderContext.tsx
@@ -7,6 +7,7 @@ import { datadogRum } from '@datadog/browser-rum'
import { useIsMobile } from '@opengovsg/design-system-react'
import PrimarySpinner from '@/components/PrimarySpinner'
+import { LaunchDarklyContext } from '@/contexts/LaunchDarkly'
import { GET_APPS } from '@/graphql/queries/get-apps'
import { getStepGroupTypeAndCaption, getStepStructure } from '@/helpers/toolbox'
@@ -30,9 +31,9 @@ interface AiBuilderStep extends IStep {
interface AIBuilderContextValue extends AIBuilderSharedProps {
allApps: IApp[]
- isMobile: boolean
triggerStep: IStep | null
steps: AiBuilderStep[]
+ isMobile: boolean
actionSteps: IStep[]
stepsBeforeGroup: IStep[]
groupedSteps: IStep[][]
@@ -40,6 +41,8 @@ interface AIBuilderContextValue extends AIBuilderSharedProps {
stepGroupCaption: string | null
// DataDog RUM Session ID so we can associate the trace with the RUM
ddSessionId: string
+ // TODO(kevinkim-ogp): remove this once A/B test is complete
+ aiBuilderType: string
}
const AiBuilderContext = createContext(
@@ -69,6 +72,9 @@ export const AiBuilderContextProvider = ({
}: AiBuilderContextProviderProps) => {
const isMobile = useIsMobile()
const ddSessionId = datadogRum.getInternalContext()?.session_id ?? ''
+ // TODO(kevinkim-ogp): remove this once A/B test is complete
+ const { getFlagValue } = useContext(LaunchDarklyContext)
+ const aiBuilderType = getFlagValue('ai-builder-type', 'none')
const { data: getAppsData, loading: isLoadingAllApps } = useQuery(GET_APPS)
@@ -126,6 +132,8 @@ export const AiBuilderContextProvider = ({
stepGroupType,
stepGroupCaption,
ddSessionId,
+ // TODO(kevinkim-ogp): remove this once A/B test is complete
+ aiBuilderType,
}}
>
{children}
diff --git a/packages/frontend/src/pages/AiBuilder/components/ChatInterface/PromptInput.tsx b/packages/frontend/src/pages/AiBuilder/components/ChatInterface/PromptInput.tsx
new file mode 100644
index 0000000000..0e3da32d22
--- /dev/null
+++ b/packages/frontend/src/pages/AiBuilder/components/ChatInterface/PromptInput.tsx
@@ -0,0 +1,168 @@
+import {
+ type FormEvent,
+ type KeyboardEvent,
+ type SyntheticEvent,
+ useRef,
+ useState,
+} from 'react'
+import { FaArrowCircleRight } from 'react-icons/fa'
+import { FaCircleStop } from 'react-icons/fa6'
+import { Box, Flex, Icon, Text, Textarea } from '@chakra-ui/react'
+
+import pairLogo from '@/assets/pair-logo.svg'
+import { ImageBox } from '@/components/FlowStepConfigurationModal/ChooseAndAddConnection/ConfigureExcelConnection'
+import { AI_CHAT_IDEAS, AiChatIdea, AiFormIdea } from '@/pages/Flows/constants'
+
+import IdeaButtons from '../IdeaButtons'
+
+interface PromptInputProps {
+ isStreaming: boolean
+ showIdeas?: boolean
+ placeholder?: string
+ sendMessage: (message: string) => void
+ cancelStream: () => void
+}
+
+export default function PromptInput({
+ isStreaming,
+ showIdeas = false,
+ placeholder = 'Send a message',
+ sendMessage,
+ cancelStream,
+}: PromptInputProps) {
+ const [input, setInput] = useState('')
+ const textareaRef = useRef(null)
+
+ const handleSubmit = async (e: SyntheticEvent) => {
+ e.preventDefault()
+ if (input?.trim()) {
+ sendMessage(input)
+ setInput('')
+ if (textareaRef.current) {
+ textareaRef.current.style.height = 'auto'
+ }
+ }
+ }
+
+ const handleKeyPress = (e: KeyboardEvent) => {
+ if (e.key === 'Enter' && !e.shiftKey) {
+ e.preventDefault()
+ handleSubmit(e)
+ }
+ }
+
+ const handleResize = (e?: FormEvent) => {
+ const target = e?.currentTarget || textareaRef.current
+ if (!target) {
+ return
+ }
+ const maxHeight = window.innerHeight * 0.4 - 100 // 40vh minus padding/margins
+ target.style.height = 'auto'
+ target.style.height = Math.min(target.scrollHeight, maxHeight) + 'px'
+ }
+
+ return (
+
+
+