Enhance DocsManager with continuous documentation generation

src/managers/BlockManager.ts

@@ -2,15 +2,27 @@ import { EventEmitter } from 'events'
 import { extractCurrentNewBlocks } from '../utils/string_utils'
 import type { CurrentNewBlock } from '../utils/types'
 
-// Passes through string chunks, looking for code blocks, when it finds one, it starts buffering until it finds the end of the block
-// Then it emits the code block and clears the buffer and continues
+/**
+ * BlockManager class
+ *
+ * This class extends EventEmitter and manages the processing of code blocks within a stream of content.
+ * It passes through string chunks, looking for code blocks. When it finds one, it starts buffering
+ * until it finds the end of the block. Then it emits the code block, clears the buffer, and continues.
+ */
 class BlockManager extends EventEmitter {
+  /** Stores non-code content */
   private buffer: string = ''
+  /** Stores code block content */
   private codeBuffer: string = ''
+  /** Indicates whether currently buffering a code block */
   private isBufferingCode: boolean = false
+  /** Indicates whether waiting for more content to complete a block */
   private isWaitingForContent: boolean = false
 
-  // add content to the buffer, watching for code blocks along the way
+  /**
+   * Adds content to the buffer, watching for code blocks along the way.
+   * @param {string} content - The content to be added and processed.
+   */
   addContent(content: string) {
     if (content.includes('`') || content.includes('\n')) {
       this.isWaitingForContent = true
@@ -36,6 +48,11 @@ class BlockManager extends EventEmitter {
     }
   }
 
+  /**
+   * Processes the content for code blocks.
+   * @param {string} content - The content to be processed for code blocks.
+   * @private
+   */
   private processCodeBlocks(content: string) {
     const codeBlockMarker = '\n```'
     let markerIndex = content.indexOf(codeBlockMarker)
@@ -66,6 +83,11 @@ class BlockManager extends EventEmitter {
     }
   }
 
+  /**
+   * Emits the buffered content as an event.
+   * @param {string} additionalContent - Additional content to be emitted with the buffer.
+   * @private
+   */
   private emitBuffer(additionalContent: string = '') {
     const contentToEmit = this.buffer + additionalContent
     if (contentToEmit) {
@@ -74,6 +96,10 @@ class BlockManager extends EventEmitter {
     }
   }
 
+  /**
+   * Emits the buffered code block as an event.
+   * @private
+   */
   private emitCodeBlock() {
     if (this.codeBuffer.length > 0) {
       const currentNewBlock = this.parseCodeBuffer()
@@ -93,15 +119,18 @@ class BlockManager extends EventEmitter {
    * current start index, divider index, and new end index. It then extracts the current code and new code
    * from the buffer based on these indices.
    *
-   * @return {CurrentNewBlock | null}
-   *   An object containing the language, current start, new end, current code, and new code.
+   * @return {CurrentNewBlock | null} An object containing the language, current start, new end, current code, and new code.
+   * @private
    */
   private parseCodeBuffer(): CurrentNewBlock | null {
     const results = extractCurrentNewBlocks(this.codeBuffer)
 
     return results[0]
   }
 
+  /**
+   * Clears all buffers and resets state flags.
+   */
   clearBuffer() {
     this.buffer = ''
     this.codeBuffer = ''

src/managers/ChatManager.ts

@@ -455,13 +455,8 @@ class ChatManager extends StreamManager {
   }
 
   private async handleSave() {
-    const lastResponse = this.chatHistory
-      .slice()
-      .reverse()
-      .find((msg) => msg.role === 'assistant')
-
-    if (lastResponse) {
-      const currentNewBlocks = extractCurrentNewBlocks(lastResponse.content)
+    if (this.fullResponse) {
+      const currentNewBlocks = extractCurrentNewBlocks(this.fullResponse)
       if (Object.keys(currentNewBlocks).length > 0) {
         await applyAndSaveCurrentNewBlocks(currentNewBlocks)
         console.info(`\n${colors.bold('save')} has finished\n`)
@@ -510,6 +505,7 @@ class ChatManager extends StreamManager {
     this.isProcessing = true
 
     try {
+      this.fullResponse = ''
       await this.refreshContext()
       this.addUserPrompt(this.generatePrompt(input))
 
@@ -588,11 +584,26 @@ class ChatManager extends StreamManager {
   }
 
   protected async finalizeResponse(): Promise<void> {
+    console.log('\n')
+    this.fullResponse += `${this.responseString}\n\n`
     this.addAssistantResponse(this.responseString)
     this.isProcessing = false
 
-    if (this.responseString !== '') {
-      const currentNewBlocks = extractCurrentNewBlocks(this.responseString)
+    if (!this.responseString.includes('======= DONE =======')) {
+      this.addUserPrompt(
+        "If there is more, continue, otherwise print the string '======= DONE ======='.",
+      )
+      const responseStream = await this.gpt.streamText({
+        prompt: this.prompt,
+        messages: this.chatHistory,
+      })
+
+      this.handleResponseStream(responseStream)
+      return
+    }
+
+    if (this.fullResponse !== '') {
+      const currentNewBlocks = extractCurrentNewBlocks(this.fullResponse)
       const [canApply, summary] = await checkBlocksApplicability(currentNewBlocks)
       if (!canApply) {
         if (this.retryCount < this.maxRetries) {

src/managers/CodeManager.ts

@@ -190,11 +190,27 @@ Please briefly explain how the code works in this.`
    * @protected
    */
   protected async finalizeResponse(): Promise<void> {
+    console.log('\n')
+    this.fullResponse += `${this.responseString}\n\n`
     this.addAssistantResponse(this.responseString)
     this.isProcessing = false
 
-    if (this.responseString !== '') {
-      const currentNewBlocks = extractCurrentNewBlocks(this.responseString)
+    if (!this.responseString.includes('======= DONE =======')) {
+      console.log(colors.yellow('Checking if there is more...'))
+      this.addUserPrompt(
+        "If there is more, continue, otherwise print the string '======= DONE ======='.",
+      )
+      const responseStream = await this.gpt.streamText({
+        prompt: this.prompt,
+        messages: this.chatHistory,
+      })
+
+      this.handleResponseStream(responseStream)
+      return
+    }
+
+    if (this.fullResponse !== '') {
+      const currentNewBlocks = extractCurrentNewBlocks(this.fullResponse)
       const [canApply, summary] = await checkBlocksApplicability(currentNewBlocks)
       if (!canApply) {
         if (this.retryCount < this.maxRetries) {
@@ -216,8 +232,10 @@ Please briefly explain how the code works in this.`
         this.retryCount = 0 // Reset the retry count on successful application
       }
     }
+    const extraResponse = this.fullResponse
+    this.fullResponse = '' // Reset the full response
 
-    this.handleUserAction(this.responseString)
+    this.handleUserAction(extraResponse)
   }
 }
 

src/managers/DocsManager.ts

@@ -58,11 +58,27 @@ class DocsManager extends StreamManager {
    * @protected
    */
   protected async finalizeResponse(): Promise<void> {
+    console.log('\n')
+    this.fullResponse += `${this.responseString}\n\n`
     this.addAssistantResponse(this.responseString)
     this.isProcessing = false
 
-    if (this.responseString !== '') {
-      const currentNewBlocks = extractCurrentNewBlocks(this.responseString)
+    if (!this.responseString.includes('======= DONE =======')) {
+      console.log(colors.yellow('Checking if there is more...'))
+      this.addUserPrompt(
+        "If there is more, continue, otherwise print the string '======= DONE ======='.",
+      )
+      const responseStream = await this.gpt.streamText({
+        prompt: this.prompt,
+        messages: this.chatHistory,
+      })
+
+      this.handleResponseStream(responseStream)
+      return
+    }
+
+    if (this.fullResponse !== '') {
+      const currentNewBlocks = extractCurrentNewBlocks(this.fullResponse)
       const [canApply, summary] = await checkBlocksApplicability(currentNewBlocks)
       if (canApply && currentNewBlocks.length > 0) {
         const shouldSave = await confirm({

src/managers/StreamManager.ts

@@ -24,6 +24,7 @@ class StreamManager extends EventEmitter {
   protected gpt: ClovingGPT
   protected prompt: string = ''
   protected responseString: string = ''
+  protected fullResponse: string = ''
   protected contextFiles: Record<string, string> = {}
   protected chatHistory: ChatMessage[] = []
   protected chunkManager: ChunkManager