Improve Workspace agent functions and prompt

fixed #14361

Signed-off-by: Jonas Helming <[email protected]>

packages/ai-workspace-agent/src/browser/frontend-module.ts

@@ -17,12 +17,13 @@ import { ContainerModule } from '@theia/core/shared/inversify';
 import { ChatAgent } from '@theia/ai-chat/lib/common';
 import { Agent, ToolProvider } from '@theia/ai-core/lib/common';
 import { WorkspaceAgent } from './workspace-agent';
-import { FileContentFunction, GetWorkspaceFileList } from './functions';
+import { FileContentFunction, GetWorkspaceDirectoryStructure, GetWorkspaceFileList } from './functions';
 
 export default new ContainerModule(bind => {
     bind(WorkspaceAgent).toSelf().inSingletonScope();
     bind(Agent).toService(WorkspaceAgent);
     bind(ChatAgent).toService(WorkspaceAgent);
     bind(ToolProvider).to(GetWorkspaceFileList);
     bind(ToolProvider).to(FileContentFunction);
+    bind(ToolProvider).to(GetWorkspaceDirectoryStructure);
 });

packages/ai-workspace-agent/src/browser/functions.ts

@@ -19,11 +19,61 @@ import { inject, injectable } from '@theia/core/shared/inversify';
 import { FileService } from '@theia/filesystem/lib/browser/file-service';
 import { FileStat } from '@theia/filesystem/lib/common/files';
 import { WorkspaceService } from '@theia/workspace/lib/browser';
-import { FILE_CONTENT_FUNCTION_ID, GET_WORKSPACE_FILE_LIST_FUNCTION_ID } from '../common/functions';
+import { FILE_CONTENT_FUNCTION_ID, GET_WORKSPACE_DIRECTORY_STRUCTURE_FUNCTION_ID, GET_WORKSPACE_FILE_LIST_FUNCTION_ID } from '../common/functions';
+
+function shouldExclude(stat: FileStat): boolean {
+    const excludedFolders = ['node_modules', 'lib'];
+    return stat.resource.path.base.startsWith('.') || excludedFolders.includes(stat.resource.path.base);
+}
+
+@injectable()
+export class GetWorkspaceDirectoryStructure implements ToolProvider {
+    static ID = GET_WORKSPACE_DIRECTORY_STRUCTURE_FUNCTION_ID;
+
+    getTool(): ToolRequest {
+        return {
+            id: GetWorkspaceDirectoryStructure.ID,
+            name: GetWorkspaceDirectoryStructure.ID,
+            description: 'Retrieve the complete directory structure of the workspace, listing only directories (no file contents). This structure excludes specific directories, such as node_modules and hidden files, ensuring paths are within workspace boundaries.',
+            handler: () => this.getDirectoryStructure()
+        };
+    }
+
+    @inject(WorkspaceService)
+    protected workspaceService: WorkspaceService;
+
+    @inject(FileService)
+    protected readonly fileService: FileService;
+
+    private async getDirectoryStructure(): Promise<string[]> {
+        const wsRoots = await this.workspaceService.roots;
+
+        if (wsRoots.length === 0) {
+            throw new Error('Workspace root not found');
+        }
+
+        const workspaceRootUri = wsRoots[0].resource;
+
+        return this.buildDirectoryStructure(workspaceRootUri);
+    }
+
+    private async buildDirectoryStructure(uri: URI, prefix: string = ''): Promise<string[]> {
+        const stat = await this.fileService.resolve(uri);
+        const result: string[] = [];
+
+        if (stat && stat.isDirectory && stat.children) {
+            for (const child of stat.children) {
+                if (!child.isDirectory || shouldExclude(child)) continue;
+                const path = `${prefix}${child.resource.path.base}/`;
+                result.push(path);
+                result.push(...await this.buildDirectoryStructure(child.resource, `${path}`));
+            }
+        }
+
+        return result;
+    }
+}
 
-/**
- * A Function that can read the contents of a File from the Workspace.
- */
 @injectable()
 export class FileContentFunction implements ToolProvider {
     static ID = FILE_CONTENT_FUNCTION_ID;
@@ -32,13 +82,13 @@ export class FileContentFunction implements ToolProvider {
         return {
             id: FileContentFunction.ID,
             name: FileContentFunction.ID,
-            description: 'Get the content of the file',
+            description: 'The relative path to the target file within the workspace. This path is resolved from the workspace root, and only files within the workspace boundaries are accessible. Attempting to access paths outside the workspace will result in an error.',
             parameters: {
                 type: 'object',
                 properties: {
                     file: {
                         type: 'string',
-                        description: 'The path of the file to retrieve content for',
+                        description: 'Return the content of a specified file within the workspace. The file path must be provided relative to the workspace root. Only files within workspace boundaries are accessible; attempting to access files outside the workspace will return an error.',
                     }
                 }
             },
@@ -61,15 +111,36 @@ export class FileContentFunction implements ToolProvider {
     }
 
     private async getFileContent(file: string): Promise<string> {
-        const uri = new URI(file);
-        const fileContent = await this.fileService.read(uri);
-        return fileContent.value;
+        const wsRoots = await this.workspaceService.roots;
+
+        if (wsRoots.length === 0) {
+            throw new Error('Workspace root not found');
+        }
+
+        const workspaceRootUri = wsRoots[0].resource;
+
+        const targetUri = workspaceRootUri.resolve(file);
+
+        if (!targetUri.toString().startsWith(workspaceRootUri.toString())) {
+            throw new Error('Access outside of the workspace is not allowed');
+        }
+
+        try {
+            const fileStat = await this.fileService.resolve(targetUri);
+
+            if (!fileStat || fileStat.isDirectory) {
+                return JSON.stringify({ error: "File not found" });
+            }
+
+            const fileContent = await this.fileService.read(targetUri);
+            return fileContent.value;
+
+        } catch (error) {
+            return JSON.stringify({ error: "File not found" });
+        }
     }
 }
 
-/**
- * A Function that lists all files in the workspace.
- */
 @injectable()
 export class GetWorkspaceFileList implements ToolProvider {
     static ID = GET_WORKSPACE_FILE_LIST_FUNCTION_ID;
@@ -78,9 +149,20 @@ export class GetWorkspaceFileList implements ToolProvider {
         return {
             id: GetWorkspaceFileList.ID,
             name: GetWorkspaceFileList.ID,
-            description: 'List all files in the workspace',
-
-            handler: () => this.getProjectFileList()
+            parameters: {
+                type: 'object',
+                properties: {
+                    path: {
+                        type: 'string',
+                        description: 'Optional relative path to a directory within the workspace. If no path is specified, the function lists contents directly in the workspace root. Paths are resolved within workspace boundaries only; paths outside the workspace or unvalidated paths will result in an error.'
+                    }
+                }
+            },
+            description: 'List files and directories within a specified workspace directory. Paths are relative to the workspace root, and only workspace-contained paths are allowed. If no path is provided, the root contents are listed. Paths outside the workspace will result in an error.',
+            handler: (arg_string: string) => {
+                const args = JSON.parse(arg_string);
+                return this.getProjectFileList(args.path);
+            }
         };
     }
 
@@ -90,45 +172,51 @@ export class GetWorkspaceFileList implements ToolProvider {
     @inject(FileService)
     protected readonly fileService: FileService;
 
-    async getProjectFileList(): Promise<string[]> {
-        // Get all files from the workspace service as a flat list of qualified file names
+    async getProjectFileList(path?: string): Promise<string[]> {
         const wsRoots = await this.workspaceService.roots;
-        const result: string[] = [];
-        for (const root of wsRoots) {
-            result.push(...await this.listFilesRecursively(root.resource));
+
+        if (wsRoots.length === 0) {
+            throw new Error('Workspace root not found');
+        }
+
+        const workspaceRootUri = wsRoots[0].resource;
+        const targetUri = path ? workspaceRootUri.resolve(path) : workspaceRootUri;
+
+        if (!targetUri.toString().startsWith(workspaceRootUri.toString())) {
+            throw new Error('Access outside of the workspace is not allowed');
+        }
+
+        try {
+            const stat = await this.fileService.resolve(targetUri);
+            if (!stat || !stat.isDirectory) {
+                return ["Error: Directory not found"];
+            }
+            return await this.listFilesDirectly(targetUri, workspaceRootUri);
+
+        } catch (error) {
+            return ["Error: Directory not found"];
         }
-        return result;
     }
 
-    private async listFilesRecursively(uri: URI): Promise<string[]> {
+    private async listFilesDirectly(uri: URI, workspaceRootUri: URI): Promise<string[]> {
         const stat = await this.fileService.resolve(uri);
         const result: string[] = [];
+
         if (stat && stat.isDirectory) {
-            if (this.exclude(stat)) {
+            if (shouldExclude(stat)) {
                 return result;
             }
             const children = await this.fileService.resolve(uri);
             if (children.children) {
                 for (const child of children.children) {
-                    result.push(child.resource.toString());
-                    result.push(...await this.listFilesRecursively(child.resource));
+                    const relativePath = workspaceRootUri.relative(child.resource);
+                    if (relativePath) {
+                        result.push(relativePath.toString());
+                    }
                 }
             }
         }
-        return result;
-    }
 
-    // Exclude folders which are not relevant to the AI Agent
-    private exclude(stat: FileStat): boolean {
-        if (stat.resource.path.base.startsWith('.')) {
-            return true;
-        }
-        if (stat.resource.path.base === 'node_modules') {
-            return true;
-        }
-        if (stat.resource.path.base === 'lib') {
-            return true;
-        }
-        return false;
+        return result;
     }
 }

packages/ai-workspace-agent/src/common/functions.ts

@@ -15,3 +15,4 @@
 // *****************************************************************************
 export const FILE_CONTENT_FUNCTION_ID = 'getFileContent';
 export const GET_WORKSPACE_FILE_LIST_FUNCTION_ID = 'getWorkspaceFileList';
+export const GET_WORKSPACE_DIRECTORY_STRUCTURE_FUNCTION_ID = 'getWorkspaceDirectoryStructure';

packages/ai-workspace-agent/src/common/template.ts

@@ -14,50 +14,30 @@
 // SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-only WITH Classpath-exception-2.0
 // *****************************************************************************
 import { PromptTemplate } from '@theia/ai-core/lib/common';
-import { GET_WORKSPACE_FILE_LIST_FUNCTION_ID, FILE_CONTENT_FUNCTION_ID } from './functions';
+import { GET_WORKSPACE_FILE_LIST_FUNCTION_ID, FILE_CONTENT_FUNCTION_ID, GET_WORKSPACE_DIRECTORY_STRUCTURE_FUNCTION_ID } from './functions';
 
 export const workspaceTemplate = <PromptTemplate>{
    id: 'workspace-system',
    template: `# Instructions
 
-    You are an AI assistant integrated into the Theia IDE, specifically designed to help software developers by
-providing concise and accurate answers to programming-related questions. Your role is to enhance the
-developer's productivity by offering quick solutions, explanations, and best practices.
-Keep responses short and to the point, focusing on delivering valuable insights, best practices and
-simple solutions.
-You are specialized in providing insights based on the Theia IDE's workspace and its files.
-Use the following functions to access the workspace:
-- ~{${GET_WORKSPACE_FILE_LIST_FUNCTION_ID}}
-- ~{${FILE_CONTENT_FUNCTION_ID}}. Never shorten the file paths when using this function.
+You are an AI assistant integrated into Theia IDE, designed to assist software developers with concise answers to programming-related questions. Your goal is to enhance productivity with quick, relevant solutions, explanations, and best practices. Keep responses short, delivering valuable insights and direct solutions.
 
-## Guidelines
+Use the following functions to interact with the workspace files as needed:
+- **~{${GET_WORKSPACE_DIRECTORY_STRUCTURE_FUNCTION_ID}}**: Returns the complete directory structure.
+- **~{${GET_WORKSPACE_FILE_LIST_FUNCTION_ID}}**: Lists files and directories in a specific directory.
+- **~{${FILE_CONTENT_FUNCTION_ID}}**: Retrieves the content of a specific file.
 
-1. **Understand Context:**
-   - **Always answer in context of the workspace and its files. Avoid general answers**.
-   - Use the provided functions to access the workspace files. **Never assume the workspace structure or file contents.**
-   - Tailor responses to be relevant to the programming language, framework, or tools like Eclipse Theia used in the workspace.
-   - Ask clarifying questions if necessary to provide accurate assistance. Always assume it is okay to read additional files from the workspace.
+### Workspace Navigation Guidelines
 
-2. **Provide Clear Solutions:**
-   - Offer direct answers or code snippets that solve the problem or clarify the concept.
-   - Avoid lengthy explanations unless necessary for understanding.
-   - Provide links to official documentation for further reading when applicable.
+1. **Confirm Paths**: Always verify paths by listing directories or files as you navigate. Avoid assumptions based on user input alone.
+2. **Start from Root**: Begin at the root and navigate subdirectories step-by-step.
 
-3. **Support Multiple Languages and Tools:**
-   - Be familiar with popular programming languages, frameworks, IDEs like Eclipse Theia, and command-line tools.
-   - Adapt advice based on the language, environment, or tools specified by the developer.
+### Response Guidelines
 
-4. **Facilitate Learning:**
-   - Encourage learning by explaining why a solution works or why a particular approach is recommended.
-   - Keep explanations concise and educational.
-
-5. **Maintain Professional Tone:**
-   - Communicate in a friendly, professional manner.
-   - Use technical jargon appropriately, ensuring clarity for the target audience.
-
-6. **Stay on Topic:**
-   - Limit responses strictly to topics related to software development, frameworks, Eclipse Theia, terminal usage, and relevant technologies.
-   - Politely decline to answer questions unrelated to these areas by saying, "I'm here to assist with programming-related questions.
-     For other topics, please refer to a specialized source."
+1. **Contextual Focus**: Provide answers relevant to the workspace, avoiding general advice. Use provided functions without assuming file structure or content.
+2. **Clear Solutions**: Offer direct answers and concise explanations. Link to official documentation as needed.
+3. **Tool & Language Adaptability**: Adjust guidance based on the programming language, framework, or tool specified by the developer.
+4. **Supportive Tone**: Maintain a friendly, professional tone with clear, accurate technical language.
+5. **Stay Relevant**: Limit responses to software development, frameworks, Theia, terminal usage, and related technologies. Decline unrelated questions politely.
 `
 };