📂 refactor: Show File Search and Code File Upload Options Based on Agent Tools (#9532)

This commit is contained in:
Danny Avila 2025-09-09 20:48:29 -04:00 committed by GitHub
parent 957fa7a994
commit 5c0e9d8fbb
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 297 additions and 9 deletions

View file

@ -39,6 +39,7 @@ function AttachFileChat({
<AttachFileMenu
disabled={disableInputs}
conversationId={conversationId}
agentId={conversation?.agent_id}
endpointFileConfig={endpointFileConfig}
/>
);

View file

@ -11,7 +11,13 @@ import {
SharePointIcon,
} from '@librechat/client';
import type { EndpointFileConfig } from 'librechat-data-provider';
import { useLocalize, useGetAgentsConfig, useFileHandling, useAgentCapabilities } from '~/hooks';
import {
useAgentToolPermissions,
useAgentCapabilities,
useGetAgentsConfig,
useFileHandling,
useLocalize,
} from '~/hooks';
import useSharePointFileHandling from '~/hooks/Files/useSharePointFileHandling';
import { SharePointPickerDialog } from '~/components/SharePoint';
import { useGetStartupConfig } from '~/data-provider';
@ -21,11 +27,17 @@ import { cn } from '~/utils';
interface AttachFileMenuProps {
conversationId: string;
agentId?: string | null;
disabled?: boolean | null;
endpointFileConfig?: EndpointFileConfig;
}
const AttachFileMenu = ({ disabled, conversationId, endpointFileConfig }: AttachFileMenuProps) => {
const AttachFileMenu = ({
agentId,
disabled,
conversationId,
endpointFileConfig,
}: AttachFileMenuProps) => {
const localize = useLocalize();
const isUploadDisabled = disabled ?? false;
const inputRef = useRef<HTMLInputElement>(null);
@ -52,6 +64,8 @@ const AttachFileMenu = ({ disabled, conversationId, endpointFileConfig }: Attach
* */
const capabilities = useAgentCapabilities(agentsConfig?.capabilities ?? defaultAgentCapabilities);
const { fileSearchAllowedByAgent, codeAllowedByAgent } = useAgentToolPermissions(agentId);
const handleUploadClick = (isImage?: boolean) => {
if (!inputRef.current) {
return;
@ -86,7 +100,7 @@ const AttachFileMenu = ({ disabled, conversationId, endpointFileConfig }: Attach
});
}
if (capabilities.fileSearchEnabled) {
if (capabilities.fileSearchEnabled && fileSearchAllowedByAgent) {
items.push({
label: localize('com_ui_upload_file_search'),
onClick: () => {
@ -101,7 +115,7 @@ const AttachFileMenu = ({ disabled, conversationId, endpointFileConfig }: Attach
});
}
if (capabilities.codeEnabled) {
if (capabilities.codeEnabled && codeAllowedByAgent) {
items.push({
label: localize('com_ui_upload_code_files'),
onClick: () => {
@ -142,6 +156,8 @@ const AttachFileMenu = ({ disabled, conversationId, endpointFileConfig }: Attach
setToolResource,
setEphemeralAgent,
sharePointEnabled,
codeAllowedByAgent,
fileSearchAllowedByAgent,
setIsSharePointDialogOpen,
]);

View file

@ -2,7 +2,13 @@ import React, { useMemo } from 'react';
import { OGDialog, OGDialogTemplate } from '@librechat/client';
import { ImageUpIcon, FileSearch, TerminalSquareIcon, FileType2Icon } from 'lucide-react';
import { EToolResources, defaultAgentCapabilities } from 'librechat-data-provider';
import { useLocalize, useGetAgentsConfig, useAgentCapabilities } from '~/hooks';
import {
useAgentToolPermissions,
useAgentCapabilities,
useGetAgentsConfig,
useLocalize,
} from '~/hooks';
import { useChatContext } from '~/Providers';
interface DragDropModalProps {
onOptionSelect: (option: EToolResources | undefined) => void;
@ -26,6 +32,11 @@ const DragDropModal = ({ onOptionSelect, setShowModal, files, isVisible }: DragD
* Use definition for agents endpoint for ephemeral agents
* */
const capabilities = useAgentCapabilities(agentsConfig?.capabilities ?? defaultAgentCapabilities);
const { conversation } = useChatContext();
const { fileSearchAllowedByAgent, codeAllowedByAgent } = useAgentToolPermissions(
conversation?.agent_id,
);
const options = useMemo(() => {
const _options: FileOption[] = [
{
@ -35,14 +46,14 @@ const DragDropModal = ({ onOptionSelect, setShowModal, files, isVisible }: DragD
condition: files.every((file) => file.type?.startsWith('image/')),
},
];
if (capabilities.fileSearchEnabled) {
if (capabilities.fileSearchEnabled && fileSearchAllowedByAgent) {
_options.push({
label: localize('com_ui_upload_file_search'),
value: EToolResources.file_search,
icon: <FileSearch className="icon-md" />,
});
}
if (capabilities.codeEnabled) {
if (capabilities.codeEnabled && codeAllowedByAgent) {
_options.push({
label: localize('com_ui_upload_code_files'),
value: EToolResources.execute_code,
@ -58,7 +69,7 @@ const DragDropModal = ({ onOptionSelect, setShowModal, files, isVisible }: DragD
}
return _options;
}, [capabilities, files, localize]);
}, [capabilities, files, localize, fileSearchAllowedByAgent, codeAllowedByAgent]);
if (!isVisible) {
return null;

View file

@ -0,0 +1,180 @@
import { renderHook } from '@testing-library/react';
import { Tools } from 'librechat-data-provider';
import useAgentToolPermissions from '../useAgentToolPermissions';
// Mock the dependencies
jest.mock('~/data-provider', () => ({
useGetAgentByIdQuery: jest.fn(),
}));
jest.mock('~/Providers', () => ({
useAgentsMapContext: jest.fn(),
}));
const mockUseGetAgentByIdQuery = jest.requireMock('~/data-provider').useGetAgentByIdQuery;
const mockUseAgentsMapContext = jest.requireMock('~/Providers').useAgentsMapContext;
describe('useAgentToolPermissions', () => {
beforeEach(() => {
jest.clearAllMocks();
});
describe('when no agentId is provided', () => {
it('should allow all tools for ephemeral agents', () => {
mockUseAgentsMapContext.mockReturnValue({});
mockUseGetAgentByIdQuery.mockReturnValue({ data: undefined });
const { result } = renderHook(() => useAgentToolPermissions(null));
expect(result.current.fileSearchAllowedByAgent).toBe(true);
expect(result.current.codeAllowedByAgent).toBe(true);
expect(result.current.tools).toBeUndefined();
});
it('should allow all tools when agentId is undefined', () => {
mockUseAgentsMapContext.mockReturnValue({});
mockUseGetAgentByIdQuery.mockReturnValue({ data: undefined });
const { result } = renderHook(() => useAgentToolPermissions(undefined));
expect(result.current.fileSearchAllowedByAgent).toBe(true);
expect(result.current.codeAllowedByAgent).toBe(true);
expect(result.current.tools).toBeUndefined();
});
it('should allow all tools when agentId is empty string', () => {
mockUseAgentsMapContext.mockReturnValue({});
mockUseGetAgentByIdQuery.mockReturnValue({ data: undefined });
const { result } = renderHook(() => useAgentToolPermissions(''));
expect(result.current.fileSearchAllowedByAgent).toBe(true);
expect(result.current.codeAllowedByAgent).toBe(true);
expect(result.current.tools).toBeUndefined();
});
});
describe('when agentId is provided but agent not found', () => {
it('should disallow all tools', () => {
mockUseAgentsMapContext.mockReturnValue({});
mockUseGetAgentByIdQuery.mockReturnValue({ data: undefined });
const { result } = renderHook(() => useAgentToolPermissions('non-existent-agent'));
expect(result.current.fileSearchAllowedByAgent).toBe(false);
expect(result.current.codeAllowedByAgent).toBe(false);
expect(result.current.tools).toBeUndefined();
});
});
describe('when agent is found with tools', () => {
it('should allow tools that are included in the agent tools array', () => {
const agentId = 'test-agent';
const agent = {
id: agentId,
tools: [Tools.file_search],
};
mockUseAgentsMapContext.mockReturnValue({ [agentId]: agent });
mockUseGetAgentByIdQuery.mockReturnValue({ data: undefined });
const { result } = renderHook(() => useAgentToolPermissions(agentId));
expect(result.current.fileSearchAllowedByAgent).toBe(true);
expect(result.current.codeAllowedByAgent).toBe(false);
expect(result.current.tools).toEqual([Tools.file_search]);
});
it('should allow both tools when both are included', () => {
const agentId = 'test-agent';
const agent = {
id: agentId,
tools: [Tools.file_search, Tools.execute_code],
};
mockUseAgentsMapContext.mockReturnValue({ [agentId]: agent });
mockUseGetAgentByIdQuery.mockReturnValue({ data: undefined });
const { result } = renderHook(() => useAgentToolPermissions(agentId));
expect(result.current.fileSearchAllowedByAgent).toBe(true);
expect(result.current.codeAllowedByAgent).toBe(true);
expect(result.current.tools).toEqual([Tools.file_search, Tools.execute_code]);
});
it('should use data from API query when available', () => {
const agentId = 'test-agent';
const agentMapData = {
id: agentId,
tools: [Tools.file_search],
};
const agentApiData = {
id: agentId,
tools: [Tools.execute_code, Tools.file_search],
};
mockUseAgentsMapContext.mockReturnValue({ [agentId]: agentMapData });
mockUseGetAgentByIdQuery.mockReturnValue({ data: agentApiData });
const { result } = renderHook(() => useAgentToolPermissions(agentId));
// API data should take precedence
expect(result.current.fileSearchAllowedByAgent).toBe(true);
expect(result.current.codeAllowedByAgent).toBe(true);
expect(result.current.tools).toEqual([Tools.execute_code, Tools.file_search]);
});
it('should fallback to agent map data when API data is not available', () => {
const agentId = 'test-agent';
const agentMapData = {
id: agentId,
tools: [Tools.execute_code],
};
mockUseAgentsMapContext.mockReturnValue({ [agentId]: agentMapData });
mockUseGetAgentByIdQuery.mockReturnValue({ data: undefined });
const { result } = renderHook(() => useAgentToolPermissions(agentId));
expect(result.current.fileSearchAllowedByAgent).toBe(false);
expect(result.current.codeAllowedByAgent).toBe(true);
expect(result.current.tools).toEqual([Tools.execute_code]);
});
});
describe('when agent has no tools', () => {
it('should disallow all tools with empty array', () => {
const agentId = 'test-agent';
const agent = {
id: agentId,
tools: [],
};
mockUseAgentsMapContext.mockReturnValue({ [agentId]: agent });
mockUseGetAgentByIdQuery.mockReturnValue({ data: undefined });
const { result } = renderHook(() => useAgentToolPermissions(agentId));
expect(result.current.fileSearchAllowedByAgent).toBe(false);
expect(result.current.codeAllowedByAgent).toBe(false);
expect(result.current.tools).toEqual([]);
});
it('should disallow all tools with undefined tools', () => {
const agentId = 'test-agent';
const agent = {
id: agentId,
tools: undefined,
};
mockUseAgentsMapContext.mockReturnValue({ [agentId]: agent });
mockUseGetAgentByIdQuery.mockReturnValue({ data: undefined });
const { result } = renderHook(() => useAgentToolPermissions(agentId));
expect(result.current.fileSearchAllowedByAgent).toBe(false);
expect(result.current.codeAllowedByAgent).toBe(false);
expect(result.current.tools).toBeUndefined();
});
});
});

View file

@ -5,3 +5,4 @@ export type { ProcessedAgentCategory } from './useAgentCategories';
export { default as useAgentCapabilities } from './useAgentCapabilities';
export { default as useGetAgentsConfig } from './useGetAgentsConfig';
export { default as useAgentDefaultPermissionLevel } from './useAgentDefaultPermissionLevel';
export { default as useAgentToolPermissions } from './useAgentToolPermissions';

View file

@ -0,0 +1,65 @@
import { useMemo } from 'react';
import { Tools } from 'librechat-data-provider';
import { useGetAgentByIdQuery } from '~/data-provider';
import { useAgentsMapContext } from '~/Providers';
interface AgentToolPermissionsResult {
fileSearchAllowedByAgent: boolean;
codeAllowedByAgent: boolean;
tools: string[] | undefined;
}
/**
* Hook to determine whether specific tools are allowed for a given agent.
*
* @param agentId - The ID of the agent. If null/undefined/empty, returns true for all tools (ephemeral agent behavior)
* @returns Object with boolean flags for file_search and execute_code permissions, plus the tools array
*/
export default function useAgentToolPermissions(
agentId: string | null | undefined,
): AgentToolPermissionsResult {
const agentsMap = useAgentsMapContext();
// Get the agent from the map if available
const selectedAgent = useMemo(() => {
return agentId != null && agentId !== '' ? agentsMap?.[agentId] : undefined;
}, [agentId, agentsMap]);
// Query for agent data from the API
const { data: agentData } = useGetAgentByIdQuery(agentId ?? '', {
enabled: !!agentId,
});
// Get tools from either the API data or the agents map
const tools = useMemo(
() =>
(agentData?.tools as string[] | undefined) || (selectedAgent?.tools as string[] | undefined),
[agentData?.tools, selectedAgent?.tools],
);
// Determine if file_search is allowed
const fileSearchAllowedByAgent = useMemo(() => {
// If no agentId, allow for ephemeral agents
if (!agentId) return true;
// If agentId exists but agent not found, disallow
if (!selectedAgent) return false;
// Check if the agent has the file_search tool
return tools?.includes(Tools.file_search) ?? false;
}, [agentId, selectedAgent, tools]);
// Determine if execute_code is allowed
const codeAllowedByAgent = useMemo(() => {
// If no agentId, allow for ephemeral agents
if (!agentId) return true;
// If agentId exists but agent not found, disallow
if (!selectedAgent) return false;
// Check if the agent has the execute_code tool
return tools?.includes(Tools.execute_code) ?? false;
}, [agentId, selectedAgent, tools]);
return {
fileSearchAllowedByAgent,
codeAllowedByAgent,
tools,
};
}

View file

@ -15,6 +15,7 @@ import {
import type { DropTargetMonitor } from 'react-dnd';
import type * as t from 'librechat-data-provider';
import store, { ephemeralAgentByConvoId } from '~/store';
import { useAgentToolPermissions } from '~/hooks';
import useFileHandling from './useFileHandling';
export default function useDragHelpers() {
@ -22,6 +23,8 @@ export default function useDragHelpers() {
const [showModal, setShowModal] = useState(false);
const [draggedFiles, setDraggedFiles] = useState<File[]>([]);
const conversation = useRecoilValue(store.conversationByIndex(0)) || undefined;
const agentId = conversation?.agent_id ?? '';
const { fileSearchAllowedByAgent, codeAllowedByAgent } = useAgentToolPermissions(agentId);
const setEphemeralAgent = useSetRecoilState(
ephemeralAgentByConvoId(conversation?.conversationId ?? Constants.NEW_CONVO),
);
@ -64,7 +67,18 @@ export default function useDragHelpers() {
const fileSearchEnabled = capabilities.includes(AgentCapabilities.file_search) === true;
const codeEnabled = capabilities.includes(AgentCapabilities.execute_code) === true;
const ocrEnabled = capabilities.includes(AgentCapabilities.ocr) === true;
if (!codeEnabled && !fileSearchEnabled && !ocrEnabled) {
/** Determine if dragged files are all images (enables the base image option) */
const allImages = item.files.every((f) => f.type?.startsWith('image/'));
const shouldShowModal =
allImages ||
(fileSearchEnabled && fileSearchAllowedByAgent) ||
(codeEnabled && codeAllowedByAgent) ||
ocrEnabled;
if (!shouldShowModal) {
// Fallback: directly handle files without showing modal
handleFiles(item.files);
return;
}