Fix reading context

CMakeLists.txt

@@ -45,4 +45,5 @@ add_qtc_plugin(QodeAssist
     QodeAssistHoverHandler.hpp QodeAssistHoverHandler.cpp
     QodeAssistClient.hpp QodeAssistClient.cpp
     QodeAssistUtils.hpp
+    DocumentContextReader.hpp DocumentContextReader.cpp
 )

DocumentContextReader.cpp

@@ -0,0 +1,126 @@
+/* 
+ * Copyright (C) 2024 Petr Mironychev
+ *
+ * This file is part of QodeAssist.
+ *
+ * QodeAssist is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * QodeAssist is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with QodeAssist. If not, see <https://www.gnu.org/licenses/>.
+ */
+
+#include "DocumentContextReader.hpp"
+
+#include <QTextBlock>
+
+namespace QodeAssist {
+
+QString DocumentContextReader::getLineText(int lineNumber, int cursorPosition) const
+{
+    if (!m_document || lineNumber < 0)
+        return QString();
+
+    QTextBlock block = m_document->begin();
+    int currentLine = 0;
+
+    while (block.isValid()) {
+        if (currentLine == lineNumber) {
+            QString text = block.text();
+            if (cursorPosition >= 0 && cursorPosition <= text.length()) {
+                text = text.left(cursorPosition);
+            }
+            return text;
+        }
+
+        block = block.next();
+        currentLine++;
+    }
+
+    return QString();
+}
+
+QString DocumentContextReader::getContextBefore(int lineNumber,
+                                                int cursorPosition,
+                                                int linesCount) const
+{
+    QString context;
+    for (int i = qMax(0, lineNumber - linesCount); i <= lineNumber; ++i) {
+        QString line = getLineText(i, i == lineNumber ? cursorPosition : -1);
+        context += line;
+        if (i < lineNumber)
+            context += "\n";
+    }
+    return context;
+}
+
+QString DocumentContextReader::getContextAfter(int lineNumber,
+                                               int cursorPosition,
+                                               int linesCount) const
+{
+    QString context;
+    int maxLine = lineNumber + linesCount;
+    for (int i = lineNumber; i <= maxLine; ++i) {
+        QString line = getLineText(i);
+        if (i == lineNumber && cursorPosition >= 0) {
+            line = line.mid(cursorPosition);
+        }
+        context += line;
+        if (i < maxLine && !line.isEmpty())
+            context += "\n";
+    }
+    return context;
+}
+
+QString DocumentContextReader::readWholeFileBefore(int lineNumber, int cursorPosition) const
+{
+    QString content;
+    QTextBlock block = m_document->begin();
+    int currentLine = 0;
+
+    while (block.isValid() && currentLine <= lineNumber) {
+        if (currentLine == lineNumber) {
+            content += block.text().left(cursorPosition);
+            break;
+        } else {
+            content += block.text() + "\n";
+        }
+        block = block.next();
+        currentLine++;
+    }
+
+    return content;
+}
+
+QString DocumentContextReader::readWholeFileAfter(int lineNumber, int cursorPosition) const
+{
+    QString content;
+    QTextBlock block = m_document->begin();
+    int currentLine = 0;
+
+    while (block.isValid() && currentLine < lineNumber) {
+        block = block.next();
+        currentLine++;
+    }
+
+    while (block.isValid()) {
+        if (currentLine == lineNumber) {
+            content += block.text().mid(cursorPosition) + "\n";
+        } else {
+            content += block.text() + "\n";
+        }
+        block = block.next();
+        currentLine++;
+    }
+
+    return content.trimmed();
+}
+
+} // namespace QodeAssist

DocumentContextReader.hpp

@@ -0,0 +1,43 @@
+/* 
+ * Copyright (C) 2024 Petr Mironychev
+ *
+ * This file is part of QodeAssist.
+ *
+ * QodeAssist is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * QodeAssist is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with QodeAssist. If not, see <https://www.gnu.org/licenses/>.
+ */
+
+#pragma once
+
+#include <QTextDocument>
+
+namespace QodeAssist {
+
+class DocumentContextReader
+{
+public:
+    DocumentContextReader(QTextDocument *doc)
+        : m_document(doc)
+    {}
+
+    QString getLineText(int lineNumber, int cursorPosition = -1) const;
+    QString getContextBefore(int lineNumber, int cursorPosition, int linesCount) const;
+    QString getContextAfter(int lineNumber, int cursorPosition, int linesCount) const;
+    QString readWholeFileBefore(int lineNumber, int cursorPosition) const;
+    QString readWholeFileAfter(int lineNumber, int cursorPosition) const;
+
+private:
+    QTextDocument *m_document;
+};
+
+} // namespace QodeAssist

LLMClientInterface.cpp

@@ -25,6 +25,7 @@
 
 #include <texteditor/textdocument.h>
 
+#include "DocumentContextReader.hpp"
 #include "LLMProvidersManager.hpp"
 #include "PromptTemplateManager.hpp"
 #include "QodeAssistSettings.hpp"
@@ -98,38 +99,18 @@ QString LLMClientInterface::сontextBefore(TextEditor::TextEditorWidget *widget,
         return QString();
 
     QTextDocument *doc = widget->document();
-    int totalLines = doc->blockCount();
-    QTextBlock currentBlock = doc->findBlockByLineNumber(lineNumber);
+    DocumentContextReader reader(doc);
 
-    QString beforeCursor;
-
-    if (settings().readFullFile() && totalLines < settings().maxFileThreshold()) {
-        // Read all content from the beginning of the file to the cursor
-        QTextBlock block = doc->begin();
-        while (block.isValid() && block.blockNumber() <= lineNumber) {
-            if (block.blockNumber() == lineNumber) {
-                beforeCursor += block.text().left(cursorPosition);
-                break;
-            } else {
-                beforeCursor += block.text() + "\n";
-            }
-            block = block.next();
-        }
+    QString contextBefore;
+    if (settings().readFullFile()) {
+        contextBefore = reader.readWholeFileBefore(lineNumber, cursorPosition);
     } else {
-        // Read only the specified number of lines before the cursor
-        int contextLinesBefore = settings().readStringsBeforeCursor();
-        QTextBlock block = currentBlock;
-        for (int i = 0; i < contextLinesBefore && block.isValid(); ++i) {
-            if (block.blockNumber() == lineNumber) {
-                beforeCursor = block.text().left(cursorPosition) + beforeCursor;
-            } else {
-                beforeCursor = block.text() + "\n" + beforeCursor;
-            }
-            block = block.previous();
-        }
+        contextBefore = reader.getContextBefore(lineNumber,
+                                                cursorPosition,
+                                                settings().readStringsBeforeCursor());
     }
 
-    return beforeCursor;
+    return contextBefore;
 }
 
 QString LLMClientInterface::сontextAfter(TextEditor::TextEditorWidget *widget,
@@ -140,40 +121,18 @@ QString LLMClientInterface::сontextAfter(TextEditor::TextEditorWidget *widget,
         return QString();
 
     QTextDocument *doc = widget->document();
-    int totalLines = doc->blockCount();
-    QTextBlock currentBlock = doc->findBlockByLineNumber(lineNumber);
+    DocumentContextReader reader(doc);
 
-    QString afterCursor;
-
-    if (settings().readFullFile() && totalLines < settings().maxFileThreshold()) {
-        // Read all content from the cursor to the end of the file
-        QTextBlock block = currentBlock;
-        bool isFirstBlock = true;
-        while (block.isValid()) {
-            if (isFirstBlock) {
-                afterCursor += block.text().mid(cursorPosition) + "\n";
-                isFirstBlock = false;
-            } else {
-                afterCursor += block.text() + "\n";
-            }
-            block = block.next();
-        }
+    QString contextAfter;
+    if (settings().readFullFile()) {
+        contextAfter = reader.readWholeFileAfter(lineNumber, cursorPosition);
     } else {
-        // Read only the specified number of lines after the cursor
-        int contextLinesAfter = settings().readStringsAfterCursor();
-        QTextBlock block = currentBlock;
-        for (int i = 0; i < contextLinesAfter && block.isValid(); ++i) {
-            if (block.blockNumber() == lineNumber) {
-                afterCursor += block.text().mid(cursorPosition);
-            } else {
-                afterCursor += block.text();
-            }
-            afterCursor += "\n";
-            block = block.next();
-        }
+        contextAfter = reader.getContextAfter(lineNumber,
+                                              cursorPosition,
+                                              settings().readStringsAfterCursor());
     }
 
-    return afterCursor;
+    return contextAfter;
 }
 
 void LLMClientInterface::handleInitialize(const QJsonObject &request)