🔧 refactor: Improve Agent Context & Minor Fixes (#5349)

* refactor: Improve Context for Agents

* 🔧 fix: Safeguard against undefined properties in OpenAIClient response handling

* refactor: log error before re-throwing for original stack trace

* refactor: remove toolResource state from useFileHandling, allow svg files

* refactor: prevent verbose logs from axios errors when using actions

* refactor: add silent method recordTokenUsage in AgentClient

* refactor: streamline token count assignment in BaseClient

* refactor: enhance safety settings handling for Gemini 2.0 model

* fix: capabilities structure in MCPConnection

* refactor: simplify civic integrity threshold handling in GoogleClient and llm

* refactor: update token count retrieval method in BaseClient tests

* ci: fix test for svg

client/src/components/Chat/Input/Files/AttachFileMenu.tsx

@@ -11,15 +11,15 @@ import { cn } from '~/utils';
 interface AttachFileProps {
   isRTL: boolean;
   disabled?: boolean | null;
-  handleFileChange: (event: React.ChangeEvent<HTMLInputElement>) => void;
-  setToolResource?: React.Dispatch<React.SetStateAction<string | undefined>>;
+  handleFileChange: (event: React.ChangeEvent<HTMLInputElement>, toolResource?: string) => void;
 }
 
-const AttachFile = ({ isRTL, disabled, setToolResource, handleFileChange }: AttachFileProps) => {
+const AttachFile = ({ isRTL, disabled, handleFileChange }: AttachFileProps) => {
   const localize = useLocalize();
   const isUploadDisabled = disabled ?? false;
   const inputRef = useRef<HTMLInputElement>(null);
   const [isPopoverActive, setIsPopoverActive] = useState(false);
+  const [toolResource, setToolResource] = useState<EToolResources | undefined>();
   const { data: endpointsConfig } = useGetEndpointsQuery();
 
   const capabilities = useMemo(
@@ -42,7 +42,7 @@ const AttachFile = ({ isRTL, disabled, setToolResource, handleFileChange }: Atta
       {
         label: localize('com_ui_upload_image_input'),
         onClick: () => {
-          setToolResource?.(undefined);
+          setToolResource(undefined);
           handleUploadClick(true);
         },
         icon: <ImageUpIcon className="icon-md" />,
@@ -53,7 +53,7 @@ const AttachFile = ({ isRTL, disabled, setToolResource, handleFileChange }: Atta
       items.push({
         label: localize('com_ui_upload_file_search'),
         onClick: () => {
-          setToolResource?.(EToolResources.file_search);
+          setToolResource(EToolResources.file_search);
           handleUploadClick();
         },
         icon: <FileSearch className="icon-md" />,
@@ -64,7 +64,7 @@ const AttachFile = ({ isRTL, disabled, setToolResource, handleFileChange }: Atta
       items.push({
         label: localize('com_ui_upload_code_files'),
         onClick: () => {
-          setToolResource?.(EToolResources.execute_code);
+          setToolResource(EToolResources.execute_code);
           handleUploadClick();
         },
         icon: <TerminalSquareIcon className="icon-md" />,
@@ -98,7 +98,12 @@ const AttachFile = ({ isRTL, disabled, setToolResource, handleFileChange }: Atta
   );
 
   return (
-    <FileUpload ref={inputRef} handleFileChange={handleFileChange}>
+    <FileUpload
+      ref={inputRef}
+      handleFileChange={(e) => {
+        handleFileChange(e, toolResource);
+      }}
+    >
       <div className="relative select-none">
         <DropdownPopup
           menuId="attach-file-menu"

client/src/components/Chat/Input/Files/FileFormWrapper.tsx

@@ -27,7 +27,7 @@ function FileFormWrapper({
   const { endpoint: _endpoint, endpointType } = conversation ?? { endpoint: null };
   const isAgents = useMemo(() => isAgentsEndpoint(_endpoint), [_endpoint]);
 
-  const { handleFileChange, abortUpload, setToolResource } = useFileHandling();
+  const { handleFileChange, abortUpload } = useFileHandling();
 
   const { data: fileConfig = defaultFileConfig } = useGetFileConfig({
     select: (data) => mergeFileConfig(data),
@@ -48,7 +48,6 @@ function FileFormWrapper({
         <AttachFileMenu
           isRTL={isRTL}
           disabled={disableInputs}
-          setToolResource={setToolResource}
           handleFileChange={handleFileChange}
         />
       );

client/src/hooks/Files/useFileHandling.ts

@@ -39,7 +39,6 @@ const useFileHandling = (params?: UseFileHandling) => {
   const [errors, setErrors] = useState<string[]>([]);
   const abortControllerRef = useRef<AbortController | null>(null);
   const { startUploadTimer, clearUploadTimer } = useDelayedUploadToast();
-  const [toolResource, setToolResource] = useState<string | undefined>();
   const { files, setFiles, setFilesLoading, conversation } = useChatContext();
   const setError = (error: string) => setErrors((prevErrors) => [...prevErrors, error]);
   const { addFile, replaceFile, updateFileById, deleteFileById } = useUpdateFiles(
@@ -149,9 +148,6 @@ const useFileHandling = (params?: UseFileHandling) => {
             : error?.response?.data?.message ?? 'com_error_files_upload';
         setError(errorMessage);
       },
-      onMutate: () => {
-        setToolResource(undefined);
-      },
     },
     abortControllerRef.current?.signal,
   );
@@ -187,7 +183,7 @@ const useFileHandling = (params?: UseFileHandling) => {
       if (!agent_id) {
         formData.append('message_file', 'true');
       }
-      const tool_resource = extendedFile.tool_resource ?? toolResource;
+      const tool_resource = extendedFile.tool_resource;
       if (tool_resource != null) {
         formData.append('tool_resource', tool_resource);
       }
@@ -365,7 +361,7 @@ const useFileHandling = (params?: UseFileHandling) => {
 
         const isImage = originalFile.type.split('/')[0] === 'image';
         const tool_resource =
-          extendedFile.tool_resource ?? params?.additionalMetadata?.tool_resource ?? toolResource;
+          extendedFile.tool_resource ?? params?.additionalMetadata?.tool_resource;
         if (isAgentsEndpoint(endpoint) && !isImage && tool_resource == null) {
           /** Note: this needs to be removed when we can support files to providers */
           setError('com_error_files_unsupported_capability');
@@ -388,11 +384,11 @@ const useFileHandling = (params?: UseFileHandling) => {
     }
   };
 
-  const handleFileChange = (event: React.ChangeEvent<HTMLInputElement>) => {
+  const handleFileChange = (event: React.ChangeEvent<HTMLInputElement>, _toolResource?: string) => {
     event.stopPropagation();
     if (event.target.files) {
       setFilesLoading(true);
-      handleFiles(event.target.files);
+      handleFiles(event.target.files, _toolResource);
       // reset the input
       event.target.value = '';
     }
@@ -408,7 +404,6 @@ const useFileHandling = (params?: UseFileHandling) => {
 
   return {
     handleFileChange,
-    setToolResource,
     handleFiles,
     abortUpload,
     setFiles,