feat: add map-reduce summary

docs/faq.md

@@ -98,3 +98,50 @@ Example:
 ```bash
 darwin search papers "flash attention" --log-level DEBUG --output ./darwin-data --count 3 --summary --llm ollama
 ```
+
+## What are the different methods of summarization?
+
+When summarizing long documents with LLMs (Large Language Models), the document is divided into chunks due to token limits. Each chunk is processed individually, and then these methods combine these processed chunks into a final summary.
+
+Darwin supports two methods of summarization: `map_reduce` and `refine`.
+
+### Map Reduce Method
+
+The `map_reduce` method first summarizes each document independently in a **"map"** step. Afterward, it combines these individual summaries into a final summary in a **"reduce"** step.
+
+![alt text](./images/map-reduce-summary.png)
+
+**Pros:**
+
+- Each document is summarized independently, ensuring a focused summary per document.
+- Effective for documents with diverse content or where each section merits individual attention.
+
+**Cons:**
+
+- May not refine the summary iteratively, potentially missing nuances that could be captured through iterative refinement.
+- Can be less effective for documents with a single central theme.
+
+**When to Use:**
+
+- Ideal for documents with varied topics or sections that require individual summarization before integration.
+
+These methods enable Darwin to efficiently summarize large documents within the token constraints imposed by LLMs. However, the choice between `refine` and `map_reduce` depends on the nature of the document content and the desired focus of the summary.
+
+### Refine Method
+
+The `refine` method constructs the summary by iteratively updating its answer. It starts by creating a summary for the first chunk of the document. Then, it iteratively updates this summary by adding details from each subsequent chunk one by one, until the entire document is processed.
+
+**Pros:**
+
+- Iterative approach allows for continual improvement of the summary.
+- Well-suited for documents with coherent themes and related content.
+
+**Cons:**
+
+- May get distracted by unrelated content within documents, potentially affecting the accuracy of the final summary.
+
+  > For example, in a webpage document that discusses various topics including cookie preferences, the LLM may focus excessively on the cookie preferences rather than the main content of the document.
+
+**When to Use:**
+
+- Best suited for documents where the content is cohesive and focused, minimizing the risk of getting sidetracked.

src/commands/search/accession.ts

@@ -15,6 +15,7 @@ import outputFlag from '../../inputs/flags/output.flag.js'
 import questionFlag from '../../inputs/flags/question.flag.js'
 import skipCaptchaFlag from '../../inputs/flags/skip-captcha.flag.js'
 import summaryFlag from '../../inputs/flags/summary.flag.js'
+import summaryMethodFlag from '../../inputs/flags/summary-method.flag.js'
 import { PaperSearchService } from '../../services/search/paper-search.service.js'
 
 export default class SearchAccession extends BaseCommand<typeof SearchAccession> {
@@ -47,6 +48,7 @@ export default class SearchAccession extends BaseCommand<typeof SearchAccession>
     legacy: legacyFlag,
     headless: headlessFlag,
     summary: summaryFlag,
+    'summary-method': summaryMethodFlag,
     llm: llmProviderFlag,
     question: questionFlag,
   }
@@ -88,7 +90,14 @@ export default class SearchAccession extends BaseCommand<typeof SearchAccession>
   }
 
   public async run(): Promise<void> {
-    const { count, output, 'accession-number-regex': filterPattern, summary, question } = this.flags
+    const {
+      count,
+      output,
+      'accession-number-regex': filterPattern,
+      summary,
+      question,
+      'summary-method': summaryMethod,
+    } = this.flags
     const { keywords } = this.args
 
     this.logger.info(`Searching papers with Accession Numbers (${filterPattern}) for: ${keywords}`)
@@ -98,6 +107,7 @@ export default class SearchAccession extends BaseCommand<typeof SearchAccession>
       minItemCount: count,
       filterPattern,
       summarize: summary,
+      summaryMethod,
       question,
     })
 

src/commands/search/papers.ts

@@ -14,6 +14,7 @@ import outputFlag from '../../inputs/flags/output.flag.js'
 import questionFlag from '../../inputs/flags/question.flag.js'
 import skipCaptchaFlag from '../../inputs/flags/skip-captcha.flag.js'
 import summaryFlag from '../../inputs/flags/summary.flag.js'
+import summaryMethodFlag from '../../inputs/flags/summary-method.flag.js'
 import { PaperSearchService } from '../../services/search/paper-search.service.js'
 
 export default class SearchPapers extends BaseCommand<typeof SearchPapers> {
@@ -41,6 +42,7 @@ export default class SearchPapers extends BaseCommand<typeof SearchPapers> {
     legacy: legacyFlag,
     headless: headlessFlag,
     summary: summaryFlag,
+    'summary-method': summaryMethodFlag,
     llm: llmProviderFlag,
     question: questionFlag,
   }
@@ -84,7 +86,7 @@ export default class SearchPapers extends BaseCommand<typeof SearchPapers> {
   }
 
   public async run(): Promise<void> {
-    const { count, output, filter, summary, question } = this.flags
+    const { count, output, filter, summary, question, 'summary-method': summaryMethod } = this.flags
     const { keywords } = this.args
 
     this.logger.info(`Searching papers for: ${keywords}`)
@@ -94,6 +96,7 @@ export default class SearchPapers extends BaseCommand<typeof SearchPapers> {
       minItemCount: count,
       filterPattern: filter,
       summarize: summary,
+      summaryMethod,
       question,
     })
 

src/inputs/flags/summary-method.flag.ts

@@ -0,0 +1,20 @@
+import * as oclif from '@oclif/core'
+
+import { SummaryMethod } from '../../services/llm/llm.service.js'
+
+export default oclif.Flags.custom<SummaryMethod>({
+  summary: 'Selects the method used to generate summaries.',
+  description: 'Refer to the FAQ for details on each method.',
+  options: Object.values(SummaryMethod) as string[],
+  helpValue: Object.values(SummaryMethod).join('|'),
+  default: SummaryMethod.MapReduce,
+  parse: async (input: string): Promise<SummaryMethod> => {
+    if (Object.values(SummaryMethod).includes(input as SummaryMethod)) {
+      return input as SummaryMethod
+    } else {
+      throw new Error(
+        `Invalid Summary Method : ${input}. Must be one of ${Object.values(SummaryMethod).join(', ')}`,
+      )
+    }
+  },
+})()

src/services/llm/llm.service.spec.ts

@@ -2,7 +2,7 @@ import { jest } from '@jest/globals'
 import { BaseLanguageModel } from '@langchain/core/language_models/base'
 import { mock } from 'jest-mock-extended'
 
-import { LLMService } from './llm.service'
+import { LLMService, SummaryMethod } from './llm.service'
 
 describe('LLMService', () => {
   const mockBaseLanguageModel = mock<BaseLanguageModel>()
@@ -27,7 +27,9 @@ describe('LLMService', () => {
       const inputText = 'input text'
       mockBaseLanguageModel.invoke.mockResolvedValue('summary')
 
-      await expect(llmService.summarize(inputText)).resolves.toEqual('summary')
+      await expect(llmService.summarize(inputText, SummaryMethod.Refine)).resolves.toEqual(
+        'summary',
+      )
 
       expect(mockBaseLanguageModel.invoke).toHaveBeenCalledTimes(1)
     })
@@ -36,11 +38,35 @@ describe('LLMService', () => {
       const inputText = 'input text'.repeat(10_000)
       mockBaseLanguageModel.invoke.mockResolvedValue('summary')
 
-      await expect(llmService.summarize(inputText)).resolves.toEqual('summary')
+      await expect(llmService.summarize(inputText, SummaryMethod.Refine)).resolves.toEqual(
+        'summary',
+      )
 
       // 2 calls for each chunk and 1 call for final summary
       expect(mockBaseLanguageModel.invoke).toHaveBeenCalledTimes(3)
     })
+
+    it('should call llm once for short text with map_reduce method', async () => {
+      const inputText = 'input text'
+      mockBaseLanguageModel.invoke.mockResolvedValue('summary')
+
+      await expect(llmService.summarize(inputText, SummaryMethod.MapReduce)).resolves.toEqual(
+        'summary',
+      )
+
+      expect(mockBaseLanguageModel.invoke).toHaveBeenCalledTimes(11)
+    })
+
+    it('should call llm n times for longer text with map_reduce method', async () => {
+      const inputText = 'input text'.repeat(10_000)
+      mockBaseLanguageModel.invoke.mockResolvedValue('summary')
+
+      await expect(llmService.summarize(inputText, SummaryMethod.MapReduce)).resolves.toEqual(
+        'summary',
+      )
+
+      expect(mockBaseLanguageModel.invoke).toHaveBeenCalledTimes(31)
+    })
   })
 
   describe('ask', () => {