mirror of
https://github.com/danny-avila/LibreChat.git
synced 2025-12-19 09:50:15 +01:00
📂 refactor: Show File Search and Code File Upload Options Based on Agent Tools (#9532)
This commit is contained in:
parent
957fa7a994
commit
5c0e9d8fbb
7 changed files with 297 additions and 9 deletions
|
|
@ -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();
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -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';
|
||||
|
|
|
|||
65
client/src/hooks/Agents/useAgentToolPermissions.ts
Normal file
65
client/src/hooks/Agents/useAgentToolPermissions.ts
Normal 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,
|
||||
};
|
||||
}
|
||||
|
|
@ -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;
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue