Add Toggle AI features, set agents as default, faster autoselect #951

Polleps · Polleps · commit 19a7a17061a7 · 2025-05-08T09:34:48.000+02:00
diff --git a/browser/data-browser/src/components/AI/AgentConfig.tsx b/browser/data-browser/src/components/AI/AgentConfig.tsx
@@ -1,10 +1,10 @@
 import { useEffect, useState } from 'react';
-import { styled } from 'styled-components';
+import { styled, useTheme } from 'styled-components';
 import { Row, Column } from '../Row';
-import { FaPencil, FaPlus, FaTrash } from 'react-icons/fa6';
+import { FaPencil, FaPlus, FaTrash, FaStar } from 'react-icons/fa6';
 import { IconButton } from '../IconButton/IconButton';
 import { ModelSelect } from './ModelSelect';
-import type { AIAgent, MCPServer } from './types';
+import type { AIAgent } from './types';
 import {
   Dialog,
   DialogTitle,
@@ -17,9 +17,6 @@ import { Button } from '../Button';
 import { SkeletonButton } from '../SkeletonButton';
 import { useSettings } from '../../helpers/AppSettings';
 import { Checkbox, CheckboxLabel } from '../forms/Checkbox';
-import { generateObject } from 'ai';
-import { createOpenRouter } from '@openrouter/ai-sdk-provider';
-import { z } from 'zod';
 
 // Helper function to generate a unique ID
 const generateId = () => {
@@ -86,6 +83,10 @@ export const useAIAgentConfig = () => {
     'atomic.ai.autoAgentSelect',
     true,
   );
+  const [defaultAgentId, setDefaultAgentId] = useLocalStorage<string>(
+    'atomic.ai.defaultAgentId',
+    agents[0]?.id || '',
+  );
 
   // Save agents to settings
   const saveAgents = (newAgents: AIAgent[]) => {
@@ -97,57 +98,11 @@ export const useAIAgentConfig = () => {
     autoAgentSelectEnabled,
     setAutoAgentSelectEnabled,
     saveAgents,
+    defaultAgentId,
+    setDefaultAgentId,
   };
 };
 
-function agentToText(agent: AIAgent, mcpServers: MCPServer[]) {
-  return `ID: ${agent.id} Name: ${agent.name} Description: ${agent.description} Tools: ${agent.availableTools.map(t => mcpServers.find(s => s.id === t)?.name).join(', ')}`;
-}
-
-export const useAutoAgentSelect = () => {
-  const { mcpServers, openRouterApiKey } = useSettings();
-  const { agents } = useAIAgentConfig();
-
-  const openrouter = createOpenRouter({
-    apiKey: openRouterApiKey,
-    compatibility: 'strict',
-  });
-
-  const basePrompt = `You are a tool that determines what agent to use to answer the users question.
-These are the agents to choose from
-
-${agents.map(agent => agentToText(agent, mcpServers)).join('\n')}
-
-Answer with only the ID of the agent you pick
-
-User question: `;
-
-  const pickAgent = async (question: string): Promise<AIAgent> => {
-    const prompt = basePrompt + question.trim();
-
-    const { object } = await generateObject({
-      // model: openrouter('google/gemma-3-27b-it:free'),
-      model: openrouter('google/gemini-2.0-flash-lite-preview-02-05:free'),
-      schemaName: 'Agent',
-      schemaDescription: 'The agent to use for the question.',
-      schema: z.object({
-        agentId: z.string(),
-      }),
-      prompt,
-    });
-
-    const agent = agents.find(a => a.id === object.agentId);
-
-    if (!agent) {
-      throw new Error('Agent not found');
-    }
-
-    return agent;
-  };
-
-  return pickAgent;
-};
-
 export const AgentConfig = ({
   open,
   onOpenChange,
@@ -159,10 +114,12 @@ export const AgentConfig = ({
     autoAgentSelectEnabled,
     setAutoAgentSelectEnabled,
     saveAgents,
+    defaultAgentId,
+    setDefaultAgentId,
   } = useAIAgentConfig();
   const [editingAgent, setEditingAgent] = useState<AIAgent | null>(null);
   const [isCreating, setIsCreating] = useState(false);
-
+  const theme = useTheme();
   const [dialogProps, show, close, isOpen] = useDialog({
     bindShow: onOpenChange,
   });
@@ -262,7 +219,29 @@ export const AgentConfig = ({
                       onClick={() => onSelectAgent(agent)}
                     >
                       <Column>
-                        <AgentName>{agent.name}</AgentName>
+                        <Row gap='0.2ch' center>
+                          <IconButton
+                            onClick={e => {
+                              e.stopPropagation();
+                              setDefaultAgentId(agent.id);
+                            }}
+                            title={
+                              defaultAgentId === agent.id
+                                ? 'Default agent'
+                                : 'Set as default'
+                            }
+                            edgeAlign='start'
+                          >
+                            <FaStar
+                              color={
+                                defaultAgentId === agent.id
+                                  ? theme.colors.main
+                                  : theme.colors.bg2
+                              }
+                            />
+                          </IconButton>
+                          <AgentName>{agent.name}</AgentName>
+                        </Row>
                         <AgentDescription>{agent.description}</AgentDescription>
                       </Column>
                       <Row gap='0.5rem'>
diff --git a/browser/data-browser/src/components/AI/NoKeyOverlay.tsx b/browser/data-browser/src/components/AI/NoKeyOverlay.tsx
@@ -0,0 +1,61 @@
+import styled from 'styled-components';
+import { useSettings } from '../../helpers/AppSettings';
+import { paths } from '../../routes/paths';
+import { AtomicLink } from '../AtomicLink';
+import { Column } from '../Row';
+import { FaGear } from 'react-icons/fa6';
+
+export const NoKeyOverlay: React.FC = () => {
+  const { openRouterApiKey } = useSettings();
+
+  if (openRouterApiKey) {
+    return null;
+  }
+
+  return (
+    <Overlay>
+      <Column gap='0.5rem' center>
+        <p>
+          No{' '}
+          <a href='https://openrouter.ai/' target='_blank' rel='noreferrer'>
+            OpenRouter
+          </a>{' '}
+          API key configured.
+        </p>
+        <ButtonLink clean path={paths.appSettings}>
+          <FaGear /> Settings
+        </ButtonLink>
+      </Column>
+    </Overlay>
+  );
+};
+
+const Overlay = styled.div`
+  position: absolute;
+  inset: 0;
+  backdrop-filter: blur(4px);
+  background-color: ${p =>
+    p.theme.darkMode ? 'rgba(0, 0, 0, 0.8)' : 'rgba(255, 255, 255, 0.8)'};
+  border-radius: ${p => p.theme.radius};
+  z-index: 1000;
+  display: grid;
+  place-items: center;
+`;
+
+const ButtonLink = styled(AtomicLink)`
+  display: flex;
+  align-items: center;
+  gap: 0.5rem;
+  border-radius: ${p => p.theme.radius};
+  border: 1px solid ${p => p.theme.colors.bg2};
+  padding: 0.3rem 1rem;
+  background-color: ${p => p.theme.colors.bg};
+  color: ${p => p.theme.colors.textLight};
+
+  &:hover,
+  &:focus-visible {
+    color: ${p => p.theme.colors.main};
+    border-color: ${p => p.theme.colors.main};
+    box-shadow: ${p => p.theme.boxShadowSoft};
+  }
+`;
diff --git a/browser/data-browser/src/components/AI/SimpleAIChat.tsx b/browser/data-browser/src/components/AI/SimpleAIChat.tsx
@@ -36,14 +36,12 @@ import {
   type AIChatDisplayMessage,
   type AIMessageContext,
 } from './types';
-import {
-  AgentConfig,
-  useAIAgentConfig,
-  useAutoAgentSelect,
-} from './AgentConfig';
+import { AgentConfig, useAIAgentConfig } from './AgentConfig';
 import { Button } from '../Button';
 import { MessageContextItem } from './MessageContextItem';
 import { useProcessMessages } from './useProcessMessages';
+import { NoKeyOverlay } from './NoKeyOverlay';
+import { useAutoAgentSelect } from './useAgentAutoSelect';
 type OngoingMessagePart = {
   type: 'reasoning' | 'text';
   text: string;
@@ -95,11 +93,13 @@ export const SimpleAIChat: React.FC<
   const abortSignalRef = useRef<AbortController>(null);
   const [aiState, setAiState] = useState<AIState>(AIState.Stopped);
   const [editedResources, setEditedResources] = useState<string[]>([]);
-  const { agents, autoAgentSelectEnabled } = useAIAgentConfig();
+  const { agents, autoAgentSelectEnabled, defaultAgentId } = useAIAgentConfig();
   const fileInputRef = useRef<HTMLInputElement>(null);
   const [attachedFile, setAttachedFile] = useState<FileAttachment | null>(null);
   const store = useStore();
-  const [selectedAgent, setSelectedAgent] = useState<AIAgent>(agents[0]);
+  const [selectedAgent, setSelectedAgent] = useState<AIAgent>(
+    agents.find(a => a.id === defaultAgentId) || agents[0],
+  );
   const [userInput, setUserInput] = useState('');
   const [userSelectedContextItems, setUserSelectedContextItems] = useState<
     AIMessageContext[]
@@ -108,6 +108,9 @@ export const SimpleAIChat: React.FC<
   const openrouter = createOpenRouter({
     apiKey: openRouterApiKey,
     compatibility: 'strict',
+    extraBody: {
+      transforms: ['middle-out'],
+    },
   });
   const [ongoingMessage, setOngoingMessage] = useState<OngoingMessagePart>({
     type: 'text',
@@ -617,6 +620,7 @@ export const SimpleAIChat: React.FC<
                   </IconButton>
                 </Row>
               </Column>
+              <NoKeyOverlay />
             </ChatInputWrapper>
             <TokensUsed>
               Tokens used: {tokensUsed[0]} input, {tokensUsed[1]} output
@@ -736,6 +740,7 @@ const AttachmentPreview = styled.div`
 
 const ChatWindow = styled.div<{ fullView?: boolean }>`
   padding-top: ${p => (p.fullView ? p.theme.size(2) : 0)};
+  position: relative;
   display: grid;
   grid-template-rows: auto 1fr auto;
   height: 90vh;
diff --git a/browser/data-browser/src/components/AI/useAgentAutoSelect.ts b/browser/data-browser/src/components/AI/useAgentAutoSelect.ts
@@ -0,0 +1,53 @@
+import { createOpenRouter } from '@openrouter/ai-sdk-provider';
+import { useAIAgentConfig } from './AgentConfig';
+import type { AIAgent, MCPServer } from './types';
+import { useSettings } from '../../helpers/AppSettings';
+import { generateObject } from 'ai';
+import { z } from 'zod';
+
+function agentToText(agent: AIAgent, mcpServers: MCPServer[]) {
+  return `ID: ${agent.id} Name: ${agent.name} Description: ${agent.description} Tools: ${agent.availableTools.map(t => mcpServers.find(s => s.id === t)?.name).join(', ')}`;
+}
+
+export const useAutoAgentSelect = () => {
+  const { mcpServers, openRouterApiKey } = useSettings();
+  const { agents } = useAIAgentConfig();
+
+  const openrouter = createOpenRouter({
+    apiKey: openRouterApiKey,
+    compatibility: 'strict',
+  });
+
+  const basePrompt = `You are a tool that determines what agent to use to answer the users question.
+These are the agents to choose from
+
+${agents.map(agent => agentToText(agent, mcpServers)).join('\n')}
+
+Answer with only the ID of the agent you pick
+
+User question: `;
+
+  const pickAgent = async (question: string): Promise<AIAgent> => {
+    const prompt = basePrompt + question.trim();
+
+    const { object } = await generateObject({
+      model: openrouter('google/gemma-3-4b-it:free'),
+      schemaName: 'Agent',
+      schemaDescription: 'The agent to use for the question.',
+      schema: z.object({
+        agentId: z.string(),
+      }),
+      prompt,
+    });
+
+    const agent = agents.find(a => a.id === object.agentId);
+
+    if (!agent) {
+      throw new Error('Agent not found');
+    }
+
+    return agent;
+  };
+
+  return pickAgent;
+};
diff --git a/browser/data-browser/src/components/AI/useGenerativeData.ts b/browser/data-browser/src/components/AI/useGenerativeData.ts
@@ -39,7 +39,7 @@ export const useGenerativeData = () => {
       return cleaned;
     }
 
-    return null;
+    return undefined;
   };
 
   return { generateTitleFromConversation };
diff --git a/browser/data-browser/src/components/Parent.tsx b/browser/data-browser/src/components/Parent.tsx
@@ -28,6 +28,7 @@ type ParentProps = {
 /** Breadcrumb list. Recursively renders parents. */
 function Parent({ resource }: ParentProps): JSX.Element {
   const [parent] = useString(resource, core.properties.parent);
+  const { enableAI } = useSettings();
   const { setIsOpen } = useAISidebar();
 
   return (
@@ -41,12 +42,14 @@ function Parent({ resource }: ParentProps): JSX.Element {
         <BreadCrumbCurrent>{resource.title}</BreadCrumbCurrent>
         <Spacer />
         <ButtonArea>
-          <IconButton
-            title='Toggle AI panel'
-            onClick={() => setIsOpen(prev => !prev)}
-          >
-            <AIIcon />
-          </IconButton>
+          {enableAI && (
+            <IconButton
+              title='Toggle AI panel'
+              onClick={() => setIsOpen(prev => !prev)}
+            >
+              <AIIcon />
+            </IconButton>
+          )}
           <ResourceContextMenu
             isMainMenu
             subject={resource.subject}
diff --git a/browser/data-browser/src/helpers/AppSettings.tsx b/browser/data-browser/src/helpers/AppSettings.tsx
@@ -58,6 +58,8 @@ export const AppSettingsContextProvider = (
   const [baseURL, setBaseURL] = useServerURL();
   const [drive, innerSetDrive] = useLocalStorage('drive', baseURL);
 
+  const [enableAI, setEnableAI] = useLocalStorage('atomic.ai.enabled', true);
+
   const setDrive = useCallback(
     (newDrive: string) => {
       const url = new URL(newDrive);
@@ -108,6 +110,8 @@ export const AppSettingsContextProvider = (
       setOpenRouterApiKey,
       mcpServers,
       setMcpServers,
+      enableAI,
+      setEnableAI,
     }),
     [
       drive,
@@ -135,6 +139,8 @@ export const AppSettingsContextProvider = (
       setOpenRouterApiKey,
       mcpServers,
       setMcpServers,
+      enableAI,
+      setEnableAI,
     ],
   );
 
@@ -186,6 +192,9 @@ export interface AppSettings {
   mcpServers: MCPServer[];
   /** Update the list of MCP servers */
   setMcpServers: (servers: MCPServer[]) => void;
+  /** Enable all AI features in the app */
+  enableAI: boolean;
+  setEnableAI: (b: boolean) => void;
 }
 
 const initialState: AppSettings = {
@@ -214,6 +223,8 @@ const initialState: AppSettings = {
   setOpenRouterApiKey: () => undefined,
   mcpServers: [],
   setMcpServers: () => undefined,
+  enableAI: true,
+  setEnableAI: () => undefined,
 };
 
 /** Hook for using App Settings, such as theme and darkmode */
diff --git a/browser/data-browser/src/routes/AppSettings.tsx b/browser/data-browser/src/routes/AppSettings.tsx
diff --git a/browser/data-browser/src/routes/NewResource/BaseButtons.tsx b/browser/data-browser/src/routes/NewResource/BaseButtons.tsx
diff --git a/browser/data-browser/src/views/AIChat/AIChatPage.tsx b/browser/data-browser/src/views/AIChat/AIChatPage.tsx

Original file line number	Diff line number	Diff line change
`@@ -39,7 +39,7 @@ export const useGenerativeData = () => {`
`39`	`39`	`return cleaned;`
`40`	`40`	`}`
`41`	`41`
`42`		`- return null;`
	`42`	`+ return undefined;`
`43`	`43`	`};`
`44`	`44`
`45`	`45`	`return { generateTitleFromConversation };`