Add maxContextDepth setting and make it configurable in FileContextTree

allenhutchison · Nov 3, 2024 · 9b72e54 · 9b72e54
1 parent 4090737
commit 9b72e54
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 10 deletions.
diff --git a/main.ts b/main.ts
@@ -11,6 +11,7 @@ interface ObsidianGeminiSettings {
 	modelName: string;
     summaryFrontmatterKey: string;
     userName: string;
+    maxContextDepth: number;
     systemPrompt: string;
     summaryPrompt: string;
     rewritePrompt: string;
@@ -22,6 +23,7 @@ const DEFAULT_SETTINGS: ObsidianGeminiSettings = {
 	modelName: 'gemini-1.5-flash-002',
     summaryFrontmatterKey: 'summary',
     userName: 'User',
+    maxContextDepth: 2,
     rewriteFiles: false,
     systemPrompt: `
         You are a note taking assistant. 

diff --git a/src/file-context.ts b/src/file-context.ts
@@ -12,14 +12,12 @@ interface FileContextNode {
 export class FileContextTree {
     private root: FileContextNode | null = null;
     private plugin: ObsidianGemini;
-    private maxDepth: number = 2; // TODO(adh): Make this configurable
+    private maxDepth: number;
     private readonly MAX_TOTAL_CHARS = 500000; // TODO(adh): Make this configurable
 
-    constructor(plugin: ObsidianGemini, depth: number = this.maxDepth) {
+    constructor(plugin: ObsidianGemini, depth?: number) {
         this.plugin = plugin;
-        if (depth) {
-            this.maxDepth = depth;
-        }
+        this.maxDepth = depth ?? this.plugin.settings.maxContextDepth;
     }
 
     async buildStructure(file: TFile, currentDepth: number = 0): Promise<FileContextNode | null> {

diff --git a/src/settings.ts b/src/settings.ts
@@ -21,7 +21,6 @@ export default class ObsidianGeminiSettingTab extends PluginSettingTab {
                 .setPlaceholder('Enter your API Key')
                 .setValue(this.plugin.settings.apiKey)
                 .onChange(async (value) => {
-                    console.log("Key changed");
                     this.plugin.settings.apiKey = value;
                     await this.plugin.saveSettings();
                 }));
@@ -37,7 +36,6 @@ export default class ObsidianGeminiSettingTab extends PluginSettingTab {
                 .addOption('gemini-1.5-flash-8b', 'gemini-1.5-flash-8b')
             .setValue(this.plugin.settings.modelName)
             .onChange(async (value) => {
-                console.log("Model changed: ", value);
                 this.plugin.settings.modelName = value;
                 await this.plugin.saveSettings();
             }));
@@ -49,7 +47,6 @@ export default class ObsidianGeminiSettingTab extends PluginSettingTab {
                 .setPlaceholder('Enter your key')
                 .setValue(this.plugin.settings.summaryFrontmatterKey)
                 .onChange(async (value) => {
-                    console.log("Frontmatter Key changed");
                     this.plugin.settings.summaryFrontmatterKey = value;
                     await this.plugin.saveSettings();
                 }));
@@ -61,11 +58,25 @@ export default class ObsidianGeminiSettingTab extends PluginSettingTab {
                 .setPlaceholder('Enter your name')
                 .setValue(this.plugin.settings.userName)
                 .onChange(async (value) => {
-                    console.log("User name changed");
                     this.plugin.settings.userName = value;
                     await this.plugin.saveSettings();
                 }));
 
+        new Setting(containerEl)
+            .setName('Context Depth')
+            .setDesc('The number of linked pages to include in the context for the model.')
+            .addDropdown(dropdown => dropdown
+                .addOption('0', '0')
+                .addOption('1', '1')
+                .addOption('2', '2')
+                .addOption('3', '3')
+                .addOption('4', '4')
+            .setValue(this.plugin.settings.maxContextDepth.toString())
+            .onChange(async (value) => {
+                this.plugin.settings.maxContextDepth = parseInt(value);
+                await this.plugin.saveSettings();
+            }));
+
         new Setting(containerEl)
             .setName('Rewrite Files')
             .setDesc('Whether to allow the model to rewrite files during chat.')
@@ -83,7 +94,6 @@ export default class ObsidianGeminiSettingTab extends PluginSettingTab {
                 .setPlaceholder('Enter your system prompt')
                 .setValue(this.plugin.settings.systemPrompt)
                 .onChange(async (value) => {
-                    console.log("System prompt changed");
                     this.plugin.settings.systemPrompt = value;
                     await this.plugin.saveSettings();
                 })