From 5ca6a71852341f84a6752bb986151825896871f4 Mon Sep 17 00:00:00 2001
From: "bobo.yang" <bobo.yang@merico.dev>
Date: Sun, 21 Apr 2024 17:01:07 +0800
Subject: [PATCH] feat: Enhance code completion filtering logic

- Implemented filtering based on comment symbols followed by filenames
- Added filtering for comments indicating omitted code sections
- Adjusted prefix gathering for inline completion to consider previous lines
---
 src/contributes/codecomplete/chunkFilter.ts  | 31 ++++++++++++++++++--
 src/contributes/codecomplete/codecomplete.ts |  5 ++--
 src/contributes/codecomplete/status.ts       |  0
 3 files changed, 31 insertions(+), 5 deletions(-)
 delete mode 100644 src/contributes/codecomplete/status.ts

diff --git a/src/contributes/codecomplete/chunkFilter.ts b/src/contributes/codecomplete/chunkFilter.ts
index 17afc17..e545502 100644
--- a/src/contributes/codecomplete/chunkFilter.ts
+++ b/src/contributes/codecomplete/chunkFilter.ts
@@ -5,6 +5,7 @@ import * as vscode from 'vscode';
 
 import { logger } from '../../util/logger';
 import { CodeCompletionChunk, streamComplete } from './llm';
+import { getCommentPrefix } from './ast/language';
 
 
 // 代码补全返回结果定义
@@ -60,6 +61,28 @@ export class LLMStreamComplete {
         }
     }
 
+    async *stopOnFilenameComment(chunks: AsyncIterable<CodeCompletionChunk>, filePath: string): AsyncIterable<CodeCompletionChunk> {
+        const commentPrefix = await getCommentPrefix(filePath);
+        for await (const chunk of chunks) {
+          // 如果遇到特定的注释标记，则停止生成过程
+          if(chunk.text.includes(`${commentPrefix}<filename>`)) {
+            return; // 直接退出生成器函数
+          }
+          yield chunk;
+        }
+    }
+
+    async *stopOnOmittedCodeComment(chunks: AsyncIterable<CodeCompletionChunk>, filePath: string): AsyncIterable<CodeCompletionChunk> {
+        const commentPrefix = await getCommentPrefix(filePath);
+        for await (const chunk of chunks) {
+          // 如果遇到特定的注释标记，则停止生成过程
+          if(chunk.text.includes(`//Code omitted...`)) {
+            return; // 直接退出生成器函数
+          }
+          yield chunk;
+        }
+    }
+
     // 过滤第一个就是换行符的chunk，避免补全时出现空行
     async * stopWhenFirstCharIsNewLine(chunks: AsyncIterable<CodeCompletionChunk>) {
         let isFirst = true;
@@ -140,7 +163,7 @@ export class LLMStreamComplete {
                 firstChunk = false;
             }
 
-            if (chunk.text.trim() === this.nextLine.trim() && this.nextLine.trim().length > 5) {
+            if (chunk.text.trim() === this.nextLine.trim() && this.nextLine.trim().length > 1) {
                 break;
             }
 
@@ -284,10 +307,12 @@ export class LLMStreamComplete {
         const chunks5 = this.stopFirstLineWhenInMiddleLine(chunks4);
         const chunks6 = this.stopWhenSameWithNext(chunks5);
         const chunks7 = this.stopAtSameBlock(chunks6);
-
+        const chunks8 = this.stopOnFilenameComment(chunks7, vscode.window.activeTextEditor?.document.fileName!);
+        const chunks9 = this.stopOnOmittedCodeComment(chunks8, vscode.window.activeTextEditor?.document.fileName!);
+        
         let id = "";
         let lines: string[] = [];
-        for await (const chunk of chunks7) {
+        for await (const chunk of chunks9) {
             lines.push(chunk.text);
             if (chunk.id) {
                 id = chunk.id;
diff --git a/src/contributes/codecomplete/codecomplete.ts b/src/contributes/codecomplete/codecomplete.ts
index 4658738..987fdb0 100644
--- a/src/contributes/codecomplete/codecomplete.ts
+++ b/src/contributes/codecomplete/codecomplete.ts
@@ -203,8 +203,9 @@ export class InlineCompletionProvider implements vscode.InlineCompletionItemProv
 
         let response: CodeCompleteResult | undefined = undefined;
 
-        // 获取当前行光标之前的文本内容
-        const linePrefix = document.lineAt(position.line).text.slice(0, position.character);
+        // 获取当前光标前三行代码
+        const linePrefix = document.getText(new vscode.Range(position.line - 4, 0, position.line, position.character));
+        
         // 如果this.previousPrefix + this.previousCodeComplete包含“当前行光标之前内容”，且index为0，那么不需要执行代码补全
         if (this.previousPrefix && this.previousCodeComplete && this.previousCodeComplete.code.length > 0) {
             const index = (this.previousPrefix + this.previousCodeComplete.code).indexOf(linePrefix);
diff --git a/src/contributes/codecomplete/status.ts b/src/contributes/codecomplete/status.ts
deleted file mode 100644
index e69de29..0000000