diff --git a/src/components/BitteAiChat.tsx b/src/components/BitteAiChat.tsx index 3e642f5..2cfd5b0 100644 --- a/src/components/BitteAiChat.tsx +++ b/src/components/BitteAiChat.tsx @@ -13,6 +13,7 @@ export const BitteAiChat = ({ historyApiUrl, agentId, options, + customToolComponents, }: BitteAiChatProps) => { const [loadedData, setLoadedData] = useState({ agentIdLoaded: "", @@ -62,6 +63,7 @@ export const BitteAiChat = ({ agentId={agentId ?? agentIdLoaded} messages={uiMessages} options={optionsProps} + customToolComponents={customToolComponents} /> ); diff --git a/src/components/chat/ChatContent.tsx b/src/components/chat/ChatContent.tsx index 123c60e..e1a2baf 100644 --- a/src/components/chat/ChatContent.tsx +++ b/src/components/chat/ChatContent.tsx @@ -34,6 +34,7 @@ export const ChatContent = ({ apiKey, options, messages: initialMessages, + customToolComponents, }: BitteAiChatProps) => { const chatId = useRef(options?.chatId || generateId()).current; const [isAtBottom, setIsAtBottom] = useState(true); @@ -269,6 +270,7 @@ export const ChatContent = ({ customMessageContainer={ options?.customComponents?.messageContainer } + customToolComponents={customToolComponents} /> ); })} diff --git a/src/components/chat/MessageGroup.tsx b/src/components/chat/MessageGroup.tsx index 47d21a9..4491b5f 100644 --- a/src/components/chat/MessageGroup.tsx +++ b/src/components/chat/MessageGroup.tsx @@ -45,6 +45,10 @@ interface MessageGroupProps { toolCallId: string; result: BitteToolResult; }) => void; + customToolComponents?: { + toolName: string; + component: React.ComponentType<{ data: any }>; // Changed from JSX.Element to ComponentType + }[]; customMessageContainer?: React.ComponentType; customTxContainer?: React.ComponentType; customApproveTxButton?: React.ComponentType; @@ -68,6 +72,7 @@ export const MessageGroup = ({ customTxContainer, customApproveTxButton, customDeclineTxButton, + customToolComponents, }: MessageGroupProps) => { // State to track agentId for each message const [messagesWithAgentId, setMessagesWithAgentId] = useState< @@ -234,6 +239,15 @@ export const MessageGroup = ({ /> ); } + // Check for custom tool component first + const customTool = customToolComponents?.find( + (tool) => tool.toolName === toolName + ); + if (customTool) { + const CustomComponent = customTool.component; + return ; + } + switch (toolName) { case BittePrimitiveName.GENERATE_IMAGE: { return ( diff --git a/src/types/types.ts b/src/types/types.ts index d37f85d..59d58fc 100644 --- a/src/types/types.ts +++ b/src/types/types.ts @@ -238,6 +238,10 @@ export interface BitteAiChatProps { messages?: Message[]; wallet?: WalletOptions; options?: BitteAiChatOptions; + customToolComponents: { + toolName: string; + component: React.ComponentType<{ data: any }>; + }[]; } /**