From 3e683de6203271742eaa455c1943e015616ec5ec Mon Sep 17 00:00:00 2001 From: Saumya Date: Fri, 8 Aug 2025 19:46:46 +0530 Subject: [PATCH] quick connect option in model selector --- .../composio_integration/toolkit_service.py | 6 +- .../thread/chat-input/message-input.tsx | 1 + .../thread/chat-input/model-selector.tsx | 518 ++++++++++++------ 3 files changed, 345 insertions(+), 180 deletions(-) diff --git a/backend/composio_integration/toolkit_service.py b/backend/composio_integration/toolkit_service.py index 91253cd2..54755086 100644 --- a/backend/composio_integration/toolkit_service.py +++ b/backend/composio_integration/toolkit_service.py @@ -45,13 +45,13 @@ class ToolkitService: logger.error(f"Failed to list categories: {e}", exc_info=True) raise - async def list_toolkits(self, limit: int = 100, category: Optional[str] = None) -> List[ToolkitInfo]: + async def list_toolkits(self, limit: int = 500, category: Optional[str] = None) -> List[ToolkitInfo]: try: logger.info(f"Fetching toolkits with limit: {limit}, category: {category}") if category: - toolkits_response = self.client.toolkits.list(limit=limit, category=category) + toolkits_response = self.client.toolkits.list(limit=500, category=category, managed_by="composio") else: - toolkits_response = self.client.toolkits.list(limit=limit) + toolkits_response = self.client.toolkits.list(limit=500, managed_by="composio") items = getattr(toolkits_response, 'items', []) if hasattr(toolkits_response, '__dict__'): diff --git a/frontend/src/components/thread/chat-input/message-input.tsx b/frontend/src/components/thread/chat-input/message-input.tsx index c635b7db..d82deb50 100644 --- a/frontend/src/components/thread/chat-input/message-input.tsx +++ b/frontend/src/components/thread/chat-input/message-input.tsx @@ -177,6 +177,7 @@ export const MessageInput = forwardRef( refreshCustomModels={refreshCustomModels} billingModalOpen={billingModalOpen} setBillingModalOpen={setBillingModalOpen} + selectedAgentId={selectedAgentId} /> ); diff --git a/frontend/src/components/thread/chat-input/model-selector.tsx b/frontend/src/components/thread/chat-input/model-selector.tsx index ef0d21f2..e66c7046 100644 --- a/frontend/src/components/thread/chat-input/model-selector.tsx +++ b/frontend/src/components/thread/chat-input/model-selector.tsx @@ -6,7 +6,18 @@ import { DropdownMenuContent, DropdownMenuItem, DropdownMenuTrigger, + DropdownMenuSub, + DropdownMenuSubTrigger, + DropdownMenuSubContent, + DropdownMenuPortal, + DropdownMenuSeparator, } from '@/components/ui/dropdown-menu'; +import { + Dialog, + DialogContent, + DialogHeader, + DialogTitle, +} from '@/components/ui/dialog'; import { Tooltip, TooltipContent, @@ -14,7 +25,7 @@ import { TooltipTrigger, } from '@/components/ui/tooltip'; import { Button } from '@/components/ui/button'; -import { Check, ChevronDown, Search, AlertTriangle, Crown, ArrowUpRight, Brain, Plus, Edit, Trash, Cpu, Key, KeyRound } from 'lucide-react'; +import { Check, Search, AlertTriangle, Crown, ArrowUpRight, Brain, Plus, Edit, Trash, Cpu, KeyRound, ExternalLink, Settings } from 'lucide-react'; import { ModelOption, SubscriptionStatus, @@ -24,15 +35,47 @@ import { DEFAULT_PREMIUM_MODEL_ID, formatModelName, getCustomModels, - MODELS // Import the centralized MODELS constant + MODELS } from './_use-model-selection'; import { PaywallDialog } from '@/components/payment/paywall-dialog'; -import { BillingModal } from '@/components/billing/billing-modal'; import { cn } from '@/lib/utils'; -import { useRouter } from 'next/navigation'; import { isLocalMode } from '@/lib/config'; import { CustomModelDialog, CustomModelFormData } from './custom-model-dialog'; import Link from 'next/link'; +import { IntegrationsRegistry } from '@/components/agents/integrations-registry'; +import { ComposioConnector } from '@/components/agents/composio/composio-connector'; +import { useComposioToolkitIcon } from '@/hooks/react-query/composio/use-composio'; +import { useComposioProfiles } from '@/hooks/react-query/composio/use-composio-profiles'; +import { useAgent } from '@/hooks/react-query/agents/use-agents'; +import { Skeleton } from '@/components/ui/skeleton'; +import { useFeatureFlag } from '@/lib/feature-flags'; + +const PREDEFINED_APPS = [ + { + id: 'googledrive', + name: 'Google Drive', + slug: 'googledrive', + description: 'Access and manage files in Google Drive' + }, + { + id: 'slack', + name: 'Slack', + slug: 'slack', + description: 'Send messages and manage channels' + }, + { + id: 'gmail', + name: 'Gmail', + slug: 'gmail', + description: 'Send and manage emails' + }, + { + id: 'notion', + name: 'Notion', + slug: 'notion', + description: 'Create and manage Notion pages' + } +]; interface CustomModel { id: string; @@ -50,6 +93,7 @@ interface ModelSelectorProps { billingModalOpen: boolean; setBillingModalOpen: (open: boolean) => void; hasBorder?: boolean; + selectedAgentId?: string; } export const ModelSelector: React.FC = ({ @@ -62,6 +106,7 @@ export const ModelSelector: React.FC = ({ billingModalOpen, setBillingModalOpen, hasBorder = false, + selectedAgentId, }) => { const [paywallOpen, setPaywallOpen] = useState(false); const [lockedModel, setLockedModel] = useState(null); @@ -69,33 +114,65 @@ export const ModelSelector: React.FC = ({ const [searchQuery, setSearchQuery] = useState(''); const [highlightedIndex, setHighlightedIndex] = useState(-1); const searchInputRef = useRef(null); - const router = useRouter(); - // Custom models state const [customModels, setCustomModels] = useState([]); const [isCustomModelDialogOpen, setIsCustomModelDialogOpen] = useState(false); const [dialogInitialData, setDialogInitialData] = useState({ id: '', label: '' }); const [dialogMode, setDialogMode] = useState<'add' | 'edit'>('add'); const [editingModelId, setEditingModelId] = useState(null); - // Load custom models from localStorage on component mount + const [showIntegrationsManager, setShowIntegrationsManager] = useState(false); + const [selectedApp, setSelectedApp] = useState(null); + const [showComposioConnector, setShowComposioConnector] = useState(false); + + const { data: googleDriveIcon } = useComposioToolkitIcon('googledrive', { enabled: true }); + const { data: slackIcon } = useComposioToolkitIcon('slack', { enabled: true }); + const { data: gmailIcon } = useComposioToolkitIcon('gmail', { enabled: true }); + const { data: notionIcon } = useComposioToolkitIcon('notion', { enabled: true }); + + const { data: selectedAppIcon } = useComposioToolkitIcon(selectedApp?.slug || '', { + enabled: !!selectedApp?.slug && showComposioConnector + }); + + const appIconMap = { + 'googledrive': googleDriveIcon?.icon_url, + 'slack': slackIcon?.icon_url, + 'gmail': gmailIcon?.icon_url, + 'notion': notionIcon?.icon_url, + }; + + const { data: agent } = useAgent(selectedAgentId || ''); + const { data: profiles } = useComposioProfiles(); + + // Feature flag for custom agents + const { enabled: customAgentsEnabled } = useFeatureFlag('custom_agents'); + + const isAppConnectedToAgent = (appSlug: string): boolean => { + if (!selectedAgentId || !agent?.custom_mcps || !profiles) return false; + + return agent.custom_mcps.some((mcpConfig: any) => { + if (mcpConfig.config?.profile_id) { + const profile = profiles.find(p => p.profile_id === mcpConfig.config.profile_id); + return profile?.toolkit_slug === appSlug; + } + return false; + }); + }; + useEffect(() => { if (isLocalMode()) { setCustomModels(getCustomModels()); } }, []); - // Save custom models to localStorage whenever they change useEffect(() => { if (isLocalMode() && customModels.length > 0) { localStorage.setItem(STORAGE_KEY_CUSTOM_MODELS, JSON.stringify(customModels)); } }, [customModels]); - // Enhance model options with capabilities - using a Map to ensure uniqueness const modelMap = new Map(); - // First add all standard models to the map modelOptions.forEach(model => { modelMap.set(model.id, { ...model, @@ -103,12 +180,8 @@ export const ModelSelector: React.FC = ({ }); }); - // Then add custom models from the current customModels state (not from props) - // This ensures we're using the most up-to-date list of custom models if (isLocalMode()) { - // Get current custom models from state (not from storage) customModels.forEach(model => { - // Only add if it doesn't exist or mark it as a custom model if it does if (!modelMap.has(model.id)) { modelMap.set(model.id, { id: model.id, @@ -118,7 +191,6 @@ export const ModelSelector: React.FC = ({ isCustom: true }); } else { - // If it already exists (rare case), mark it as a custom model const existingModel = modelMap.get(model.id); modelMap.set(model.id, { ...existingModel, @@ -128,22 +200,16 @@ export const ModelSelector: React.FC = ({ }); } - // Convert map back to array const enhancedModelOptions = Array.from(modelMap.values()); - // Filter models based on search query const filteredOptions = enhancedModelOptions.filter((opt) => opt.label.toLowerCase().includes(searchQuery.toLowerCase()) || opt.id.toLowerCase().includes(searchQuery.toLowerCase()) ); - // Get free models from modelOptions (helper function) const getFreeModels = () => modelOptions.filter(m => !m.requiresSubscription).map(m => m.id); - // No sorting needed - models are already sorted in the hook const sortedModels = filteredOptions; - - // Simplified premium models function - just filter without sorting const getPremiumModels = () => { return modelOptions .filter(m => m.requiresSubscription) @@ -153,12 +219,10 @@ export const ModelSelector: React.FC = ({ })); } - // Make sure model IDs are unique for rendering const getUniqueModelKey = (model: any, index: number): string => { return `model-${model.id}-${index}`; }; - // Map models to ensure unique IDs for React keys const uniqueModels = sortedModels.map((model, index) => ({ ...model, uniqueKey: getUniqueModelKey(model, index) @@ -179,17 +243,12 @@ export const ModelSelector: React.FC = ({ enhancedModelOptions.find((o) => o.id === selectedModel)?.label || 'Select model'; const handleSelect = (id: string) => { - // Check if it's a custom model const isCustomModel = customModels.some(model => model.id === id); - - // Custom models are always accessible in local mode if (isCustomModel && isLocalMode()) { onModelChange(id); setIsOpen(false); return; } - - // Otherwise use the regular canAccessModel check if (canAccessModel(id)) { onModelChange(id); setIsOpen(false); @@ -233,38 +292,30 @@ export const ModelSelector: React.FC = ({ const shouldDisplayAll = (!isLocalMode() && subscriptionStatus === 'no_subscription') && premiumModels.length > 0; - // Handle opening the custom model dialog const openAddCustomModelDialog = (e?: React.MouseEvent) => { e?.stopPropagation(); setDialogInitialData({ id: '', label: '' }); setDialogMode('add'); setIsCustomModelDialogOpen(true); - setIsOpen(false); // Close dropdown when opening modal + setIsOpen(false); }; - // Handle opening the edit model dialog const openEditCustomModelDialog = (model: CustomModel, e?: React.MouseEvent) => { e?.stopPropagation(); setDialogInitialData({ id: model.id, label: model.label }); - setEditingModelId(model.id); // Keep the original ID with prefix for reference + setEditingModelId(model.id); setDialogMode('edit'); setIsCustomModelDialogOpen(true); - setIsOpen(false); // Close dropdown when opening modal + setIsOpen(false); }; - // Handle saving a custom model const handleSaveCustomModel = (formData: CustomModelFormData) => { - // Get model ID without automatically adding prefix const modelId = formData.id.trim(); - - // Generate display name based on model ID (remove prefix if present for display name) const displayId = modelId.startsWith('openrouter/') ? modelId.replace('openrouter/', '') : modelId; const modelLabel = formData.label.trim() || formatModelName(displayId); if (!modelId) return; - - // Check for duplicates - only for new models or if ID changed during edit const checkId = modelId; if (customModels.some(model => model.id === checkId && (dialogMode === 'add' || model.id !== editingModelId))) { @@ -272,44 +323,32 @@ export const ModelSelector: React.FC = ({ return; } - // First close the dialog to prevent UI issues closeCustomModelDialog(); - - // Create the new model object const newModel = { id: modelId, label: modelLabel }; - // Update models array (add new or update existing) const updatedModels = dialogMode === 'add' ? [...customModels, newModel] : customModels.map(model => model.id === editingModelId ? newModel : model); - // Save to localStorage first try { localStorage.setItem(STORAGE_KEY_CUSTOM_MODELS, JSON.stringify(updatedModels)); } catch (error) { console.error('Failed to save custom models to localStorage:', error); } - // Update state with new models setCustomModels(updatedModels); - - // Refresh custom models in the parent hook if the function is available if (refreshCustomModels) { refreshCustomModels(); } - // Handle model selection changes if (dialogMode === 'add') { - // Always select newly added models onModelChange(modelId); - // Also save the selection to localStorage try { localStorage.setItem(STORAGE_KEY_MODEL, modelId); } catch (error) { console.warn('Failed to save selected model to localStorage:', error); } } else if (selectedModel === editingModelId) { - // For edits, only update if the edited model was selected onModelChange(modelId); try { localStorage.setItem(STORAGE_KEY_MODEL, modelId); @@ -317,23 +356,16 @@ export const ModelSelector: React.FC = ({ console.warn('Failed to save selected model to localStorage:', error); } } - - // Force dropdown to close to ensure fresh data on next open setIsOpen(false); - - // Force a UI refresh by delaying the state update setTimeout(() => { setHighlightedIndex(-1); }, 0); }; - // Handle closing the custom model dialog const closeCustomModelDialog = () => { setIsCustomModelDialogOpen(false); setDialogInitialData({ id: '', label: '' }); setEditingModelId(null); - - // Improved fix for pointer-events issue: ensure dialog closes properly document.body.classList.remove('overflow-hidden'); const bodyStyle = document.body.style; setTimeout(() => { @@ -342,15 +374,11 @@ export const ModelSelector: React.FC = ({ }, 150); }; - // Handle deleting a custom model const handleDeleteCustomModel = (modelId: string, e?: React.MouseEvent) => { e?.stopPropagation(); e?.preventDefault(); - // Filter out the model to delete const updatedCustomModels = customModels.filter(model => model.id !== modelId); - - // Update localStorage first to ensure data consistency if (isLocalMode() && typeof window !== 'undefined') { try { localStorage.setItem(STORAGE_KEY_CUSTOM_MODELS, JSON.stringify(updatedCustomModels)); @@ -358,16 +386,10 @@ export const ModelSelector: React.FC = ({ console.error('Failed to update custom models in localStorage:', error); } } - - // Update state with the new list setCustomModels(updatedCustomModels); - - // Refresh custom models in the parent hook if the function is available if (refreshCustomModels) { refreshCustomModels(); } - - // Check if we need to change the selected model if (selectedModel === modelId) { const defaultModel = isLocalMode() ? DEFAULT_PREMIUM_MODEL_ID : DEFAULT_FREE_MODEL_ID; onModelChange(defaultModel); @@ -377,17 +399,9 @@ export const ModelSelector: React.FC = ({ console.warn('Failed to update selected model in localStorage:', error); } } - - // Force dropdown to close setIsOpen(false); - - // Update the modelMap and recreate enhancedModelOptions on next render - // This will force a complete refresh of the model list setTimeout(() => { - // Force React to fully re-evaluate the component with fresh data setHighlightedIndex(-1); - - // Reopen dropdown with fresh data if it was open if (isOpen) { setIsOpen(false); setTimeout(() => setIsOpen(true), 50); @@ -395,15 +409,28 @@ export const ModelSelector: React.FC = ({ }, 10); }; + const handleAppSelect = (app: typeof PREDEFINED_APPS[0]) => { + setSelectedApp(app); + setShowComposioConnector(true); + setIsOpen(false); + }; + + const handleComposioComplete = (profileId: string, appName: string, appSlug: string) => { + console.log('Composio integration complete:', { profileId, appName, appSlug, selectedAgentId }); + setShowComposioConnector(false); + setSelectedApp(null); + }; + + const handleOpenIntegrationsManager = () => { + setShowIntegrationsManager(true); + setIsOpen(false); + }; + const renderModelOption = (opt: any, index: number) => { - // More accurate check for custom models - use the actual customModels array - // from both the opt.isCustom flag and by checking if it exists in customModels const isCustom = Boolean(opt.isCustom) || (isLocalMode() && customModels.some(model => model.id === opt.id)); const accessible = isCustom ? true : canAccessModel(opt.id); - - // Fix the highlighting logic to use the index parameter instead of searching in filteredOptions const isHighlighted = index === highlightedIndex; const isPremium = opt.requiresSubscription; const isLowQuality = MODELS[opt.id]?.lowQuality || false; @@ -427,7 +454,6 @@ export const ModelSelector: React.FC = ({ {opt.label}
- {/* Show capabilities */} {isLowQuality && ( )} @@ -439,7 +465,6 @@ export const ModelSelector: React.FC = ({ {isPremium && !accessible && ( )} - {/* Custom model actions */} {isLocalMode() && isCustom && ( <> - - - Add a custom model - - - -
+
+
+ + + + Models + + + +
+ All Models + {isLocalMode() && ( +
+ + + + + + + + + Local .Env Manager + + + + + + + + + + Add a custom model + + + +
+ )} +
+
+
+ + setSearchQuery(e.target.value)} + onKeyDown={handleSearchInputKeyDown} + className="w-full h-8 px-8 py-1 rounded-lg text-sm focus:outline-none bg-muted" + /> +
+
+ {uniqueModels + .filter(m => + m.label.toLowerCase().includes(searchQuery.toLowerCase()) || + m.id.toLowerCase().includes(searchQuery.toLowerCase()) + ) + .map((model, index) => renderModelOption(model, index))} + {uniqueModels.length === 0 && ( +
+ No models match your search +
+ )} + + + + + {/* Quick Connect Apps Submenu - only show if custom_agents is enabled */} + {customAgentsEnabled && ( + + +
+ + Quick Connect +
+
+ {googleDriveIcon?.icon_url && slackIcon?.icon_url && notionIcon?.icon_url ? ( + <> + Google Drive + Slack + Notion + + ) : ( + <> + + + + + )} +
+ + + +
+ Popular Apps +
+
+ {!selectedAgentId || !agent || !profiles ? ( + <> + {Array.from({ length: 4 }).map((_, index) => ( +
+
+ + +
+ +
+ ))} + + ) : ( + PREDEFINED_APPS.map((app) => { + const isConnected = isAppConnectedToAgent(app.slug); + return ( + + + + handleAppSelect(app)} + disabled={isConnected} + > +
+ {appIconMap[app.slug] ? ( + {app.name} + ) : ( +
+ {app.name.charAt(0)} +
+ )} + {app.name} + {isConnected && ( + + Connected + + )} +
+
+ {isConnected ? ( + + ) : ( + + )} +
+ + + +

{isConnected ? `Manage ${app.name} tools` : app.description}

+ + + + ); + }) + )} +
+ +
+ + + + +
+ + Discover more apps +
+ + + + +

Open full integrations manager

+ + + +
+ + + )}
- {uniqueModels - .filter(m => - m.label.toLowerCase().includes(searchQuery.toLowerCase()) || - m.id.toLowerCase().includes(searchQuery.toLowerCase()) - ) - .map((model, index) => renderModelOption(model, index))} - - {uniqueModels.length === 0 && ( -
- No models match your search -
- )}
)} - {!shouldDisplayAll &&
-
- - setSearchQuery(e.target.value)} - onKeyDown={handleSearchInputKeyDown} - className="w-full h-8 px-8 py-1 rounded-lg text-sm focus:outline-none bg-muted" - /> -
-
} + - - {/* Custom Model Dialog - moved to separate component */} = ({ initialData={dialogInitialData} mode={dialogMode} /> - + + + + Integrations Manager + + setShowIntegrationsManager(false)} + /> + + + {selectedApp && ( + + )} {paywallOpen && ( = ({ )} ); -}; \ No newline at end of file +}; +