directory_tree has optional depth limit

src/agents/AgentContext.ts

@@ -76,8 +76,14 @@ export type ProgressData =
 
 export type TreeData = { name: string; children?: TreeData[] };
 
-export type RoutingContextValue = string | string[] | { [key: string]: RoutingContextValue };
-export type RoutingContext = Record<string, RoutingContextValue>;
+export type RoutingContextValue =
+  | string
+  | string[]
+  | undefined
+  | {
+      [key: string]: RoutingContextValue;
+    };
+export type RoutingContext = Record<string, RoutingContextValue> & { modules?: RoutingContextValue[] };
 
 export class AgentContext implements ToolContext {
   routing: RoutingContext = {};
@@ -103,8 +109,8 @@ export class AgentContext implements ToolContext {
     return error;
   }
 
-  public getDirectoryTree(): string {
-    return get_directory_tree(this.workspaceFolder);
+  public getDirectoryTree(depth?: number): string {
+    return get_directory_tree(this.workspaceFolder, depth);
   }
 
   mergeRoutingContext(delta: RoutingContext) {

src/llm/message.ts

@@ -3,14 +3,19 @@ export type ChatCompletionRole = 'system' | 'user' | 'assistant' | 'tool' | 'fun
 
 export type LlmMessage =
   | {
-      role: ChatCompletionRole;
+      role: 'system' | 'user' | 'assistant';
       content: string;
     }
   | {
       role: 'assistant';
       content?: string;
       name?: string;
       tool_calls: ChatCompletionMessageToolCall[];
+    }
+  | {
+      role: 'tool';
+      tool_call_id: string;
+      content: string;
     };
 
 export type LlmRequestMessage =

src/llm/openai/openAiChatRequest.ts

@@ -19,10 +19,17 @@ const getLlmClient = (config: ModelConfig) => {
   return llmClients[config.model];
 };
 
-const toOpenAiMessage = ({ role, content }: LlmMessage): ChatCompletionMessageParam =>
-  ({ role, content }) as ChatCompletionMessageParam;
-export const fromOpenAiMessage = (message: LlmResponseMessage): LlmMessage =>
-  message.role === 'assistant' ? message : { role: 'tool', content: 'TODO: tool response' }; // JSON.stringify(message.tool.function) });
+const toOpenAiMessage = (message: LlmMessage): ChatCompletionMessageParam => {
+  return message as ChatCompletionMessageParam;
+  // if (message.role === 'assistant' && message.tool_calls) {
+  //   return { role: message.role, tool_calls: message.tool_calls } as ChatCompletionMessageParam;
+  // } else if (message.role === 'tool') {
+  //   return { role: message.role, content: message.content, tool_call_id: message.tool_call_id };
+  // }
+  // return { role: message.role, content: message.content } as ChatCompletionMessageParam;
+};
+export const fromOpenAiMessage = (message: LlmResponseMessage): LlmMessage => message as LlmMessage;
+// message.role === 'assistant' ? message : { role: 'tool', content: 'TODO: tool response' }; // JSON.stringify(message.tool.function) });
 
 export const openAiChatRequest = async (
   messages: LlmMessage[],

src/tools/impl/get_directory_tree.ts

@@ -9,11 +9,17 @@ export const TOOL_GET_DIRECTORY_TREE = 'get_directory_tree';
  * Renders the directory tree structure in a simplified format.
  *
  * @param dirPath The starting directory path.
+ * @param depth how deep to go into the directory tree
  * @param prefix Prefix for the current item, used for recursion.
  * @param rootPath The root directory path - also used for recursion.
  * @returns A string representation of the directory tree.
  */
-export const get_directory_tree = (dirPath: string, prefix = '', rootPath: string | null = null): string => {
+export const get_directory_tree = (
+  dirPath: string,
+  depth: number = Number.MAX_VALUE,
+  prefix = '',
+  rootPath: string | null = null,
+): string => {
   let output = '';
   const indent = '  ';
 
@@ -51,11 +57,15 @@ export const get_directory_tree = (dirPath: string, prefix = '', rootPath: strin
     files.sort();
 
     if (dirs.length) {
-      output += '\n';
-      for (const dir of dirs) {
-        const itemPath = path.join(dirPath, dir);
+      if (depth === 0) {
+        output += ` [... ${dirs.length} more directories]\n`;
+      } else {
+        output += '\n';
         const newPrefix = prefix + indent;
-        output += get_directory_tree(itemPath, newPrefix, rootPath);
+        for (const dir of dirs) {
+          const itemPath = path.join(dirPath, dir);
+          output += get_directory_tree(itemPath, depth - 1, newPrefix, rootPath);
+        }
       }
 
       if (files.length) {
@@ -98,8 +108,11 @@ function shouldIgnoreFile(filePath: string): boolean {
 }
 
 ToolManager.registerTool(
-  (context: ToolContext, directoryPath?: string) =>
-    get_directory_tree(directoryPath ? path.join(context.workspaceFolder, directoryPath) : context.workspaceFolder),
+  (context: ToolContext, directoryPath?: string, depth?: number) =>
+    get_directory_tree(
+      directoryPath ? path.join(context.workspaceFolder, directoryPath) : context.workspaceFolder,
+      depth,
+    ),
   {
     name: TOOL_GET_DIRECTORY_TREE,
     description: 'Provides a directory listing in a tree format',
@@ -110,6 +123,10 @@ ToolManager.registerTool(
           type: 'string',
           description: 'The root path to start the directory listing. Defaults to the project root.',
         },
+        depth: {
+          type: 'number',
+          description: 'How deep to go into the directory tree',
+        },
       },
       // required: ['path'],
     },

src/utils/promptUtils.ts

@@ -15,7 +15,7 @@ export const getPromptFromFile = async (filePath: string, context: AgentContext)
 };
 
 function replacePlaceholders(templateStr: string, context: AgentContext): string {
-  Handlebars.registerHelper('directory_tree', context.getDirectoryTree);
+  Handlebars.registerHelper('directory_tree', (depth?: number) => context.getDirectoryTree(depth));
   const template = compile(templateStr);
-  return template(context);
+  return template(context.routing);
 }