update, stop generation and max tokens

twinnydotdev · Jan 21, 2024 · 59724af · 59724af
1 parent 4b6d90c
commit 59724af
Show file tree

Hide file tree

Showing 4 changed files with 42 additions and 9 deletions.
diff --git a/package.json b/package.json
@@ -107,6 +107,10 @@
         "shortTitle": "Enable twinny",
         "title": "Enable twinny"
       },
+      {
+        "command": "twinny.stopGeneration",
+        "title": "Stop generation"
+      },
       {
         "command": "twinny.disable",
         "title": "Disable twinny",
@@ -183,6 +187,14 @@
         "twinny.temperature": {
           "type": "number",
           "default": 1
+        },
+        "twinny.numPredictChat": {
+          "type": "number",
+          "default": 128
+        },
+        "twinny.numPredictFim": {
+          "type": "number",
+          "default": -2
         }
       }
     }

diff --git a/src/chat-service.ts b/src/chat-service.ts
@@ -15,6 +15,8 @@ export class ChatService {
   private _baseUrl = this._config.get('ollamaBaseUrl') as string
   private _bearerToken = this._config.get('ollamaApiBearerToken') as string
   private _chatModel = this._config.get('chatModelName') as string
+  private _temperature = this._config.get('temperature') as number
+  private _numPredictChat = this._config.get('umPredictChat') as number
   private _completion = ''
   private _port = this._config.get('ollamaApiPort') as string
   private _view?: WebviewView
@@ -34,7 +36,8 @@ export class ChatService {
       model: this._chatModel,
       prompt,
       options: {
-        temperature: this._config.get('temperature') as number
+        temperature: this._temperature,
+        num_predict: this._numPredictChat
       }
     }
 

diff --git a/src/extension.ts b/src/extension.ts
@@ -49,6 +49,9 @@ export async function activate(context: ExtensionContext) {
         sidebarProvider.chatService?.streamTemplateCompletion('explain')
       )
     }),
+    commands.registerCommand('twinny.stopGeneration', () => {
+      completionProvider.destroyStream()
+    }),
     commands.registerCommand('twinny.addTypes', () => {
       commands.executeCommand('workbench.view.extension.twinny-sidebar-view')
       delayExecution(() =>

diff --git a/src/providers/completion.ts b/src/providers/completion.ts
@@ -17,9 +17,10 @@ import { getCache, setCache } from '../cache'
 import { languages } from '../languages'
 import { InlineCompletion, StreamBody } from '../types'
 import { RequestOptions } from 'https'
+import { ClientRequest } from 'http'
 
 export class CompletionProvider implements InlineCompletionItemProvider {
-  private _statusBar: StatusBarItem
+  private statusBar: StatusBarItem
   private _debouncer: NodeJS.Timeout | undefined
   private _document: TextDocument | undefined
   private _config = workspace.getConfiguration('twinny')
@@ -29,11 +30,13 @@ export class CompletionProvider implements InlineCompletionItemProvider {
   private _baseUrl = this._config.get('ollamaBaseUrl') as string
   private _port = this._config.get('ollamaApiPort') as number
   private _temperature = this._config.get('temperature') as number
+  private _numPredictFim = this._config.get('numPredictFim') as number
   private _useFileContext = this._config.get('useFileContext') as number
   private _bearerToken = this._config.get('ollamaApiBearerToken') as number
+  private _currentReq : ClientRequest | undefined = undefined
 
   constructor(statusBar: StatusBarItem) {
-    this._statusBar = statusBar
+    this.statusBar = statusBar
   }
 
   private buildStreamRequest(prompt: string) {
@@ -48,7 +51,7 @@ export class CompletionProvider implements InlineCompletionItemProvider {
       prompt,
       options: {
         temperature: this._temperature,
-        num_predict: -2,
+        num_predict: this._numPredictFim || -2
       }
     }
 
@@ -66,6 +69,11 @@ export class CompletionProvider implements InlineCompletionItemProvider {
     return { requestOptions, requestBody }
   }
 
+  public destroyStream = () => {
+    this._currentReq?.destroy()
+    this.statusBar.text = '🤖'
+  }
+
   public async provideInlineCompletionItems(
     document: TextDocument,
     position: Position
@@ -112,21 +120,27 @@ export class CompletionProvider implements InlineCompletionItemProvider {
         try {
           let completion = ''
           let chunkCount = 0
-          this._statusBar.text = '🤖'
-          this._statusBar.text = '$(loading~spin)'
+          this.statusBar.text = '$(loading~spin)'
+          this.statusBar.command = 'twinny.stopGeneration'
 
           const { requestBody, requestOptions } =
             this.buildStreamRequest(prompt)
 
           streamResponse({
             body: requestBody,
             options: requestOptions,
+            onStart: (req) => {
+              this._currentReq = req
+            },
             onData: (chunk, onDestroy) => {
               const json = JSON.parse(chunk)
               completion = completion + json.response
               chunkCount = chunkCount + 1
-              if (json.response.match('<EOT>')) {
-                this._statusBar.text = '🤖'
+              if (
+                (chunkCount !== 1 && json.response === '\n') ||
+                json.response.match('<EOT>')
+              ) {
+                this.statusBar.text = '🤖'
                 completion = completion.replace('<EOT>', '')
                 onDestroy()
                 resolve(
@@ -141,7 +155,7 @@ export class CompletionProvider implements InlineCompletionItemProvider {
             }
           })
         } catch (error) {
-          this._statusBar.text = '$(alert)'
+          this.statusBar.text = '$(alert)'
           return resolve([] as InlineCompletionItem[])
         }
       }, this._debounceWait as number)
@@ -332,5 +346,6 @@ export class CompletionProvider implements InlineCompletionItemProvider {
     this._temperature = this._config.get('temperature') as number
     this._useFileContext = this._config.get('useFileContext') as number
     this._fimModel = this._config.get('fimModelName') as string
+    this._numPredictFim = this._config.get('numPredictFim') as number
   }
 }