Merge pull request #5 from wrapss/maxSettings

✨feat: Add maxLines, maxTokens and temperature settings

Author: ex3ndr

package.json

@@ -41,6 +41,21 @@
             "default": "http://127.0.0.1:11434/",
             "description": "Ollama Server Endpoint"
           },
+          "inference.maxLines": {
+            "type": "number",
+            "default": 16,
+            "description": "Max number of lines to be keep."
+          },
+          "inference.maxTokens": {
+            "type": "number",
+            "default": 256,
+            "description": "Max number of new tokens to be generated."
+          },
+          "inference.temperature": {
+            "type": "number",
+            "default": 0.2,
+            "description": "Temperature of the model. Increasing the temperature will make the model answer more creatively."
+          },
           "inference.model": {
             "type": "string",
             "enum": [

src/prompts/autocomplete.ts

@@ -8,6 +8,9 @@ export async function autocomplete(args: {
     model: string,
     prefix: string,
     suffix: string,
+    maxLines: number,
+    maxTokens: number,
+    temperature: number,
     canceled?: () => boolean,
 }): Promise<string> {
 
@@ -17,7 +20,8 @@ export async function autocomplete(args: {
         prompt: adaptPrompt({ prefix: args.prefix, suffix: args.suffix, model: args.model }),
         raw: true,
         options: {
-            num_predict: 256
+            num_predict: args.maxTokens,
+            temperature: args.temperature
         }
     };
 
@@ -75,9 +79,8 @@ export async function autocomplete(args: {
 
         // Update total lines
         totalLines += countSymbol(tokens.response, '\n');
-
         // Break if too many lines and on top level
-        if (totalLines > 16 && blockStack.length === 0) {
+        if (totalLines > args.maxLines && blockStack.length === 0) {
             info('Too many lines, breaking.');
             break;
         }

src/prompts/provider.ts

@@ -65,6 +65,9 @@ export class PromptProvider implements vscode.InlineCompletionItemProvider {
                     let config = vscode.workspace.getConfiguration('inference');
                     let endpoint = config.get('endpoint') as string;
                     let model = config.get('model') as string;
+                    let maxLines = config.get('maxLines') as number;
+                    let maxTokens = config.get('maxTokens') as number;
+                    let temperature = config.get('temperature') as number;
                     if (endpoint.endsWith('/')) {
                         endpoint = endpoint.slice(0, endpoint.length - 1);
                     }
@@ -98,6 +101,9 @@ export class PromptProvider implements vscode.InlineCompletionItemProvider {
                             suffix: prepared.suffix,
                             endpoint: endpoint,
                             model: model,
+                            maxLines: maxLines,
+                            maxTokens: maxTokens,
+                            temperature,
                             canceled: () => token.isCancellationRequested,
                         });
                         info(`AI completion completed: ${res}`);

src/test/suite/extension.test.ts

@@ -17,12 +17,18 @@ suite('Extension Test Suite', () => {
 	test('should perform autocomplete', async () => {
 		let endpoint = 'http://127.0.0.1:11434';
 		let model = 'codellama:7b-code-q4_K_S'; // Lightweight llm for tests
+		let maxLines = 16;
+		let maxTokens = 256;
+		let temperature = 0.2;
 		let prompt = 'fun main(): ';
 		let result = await autocomplete({
 			endpoint,
 			model,
 			prefix: prompt,
-			suffix: ''
+			suffix: '',
+			maxLines,
+			maxTokens,
+			temperature
 		});
 		console.warn(result);
 	});