wip

sourcegraph · Dec 26, 2023 · 5f89832 · 5f89832
1 parent adca3b2
commit 5f89832
Show file tree

Hide file tree

Showing 6 changed files with 74 additions and 32 deletions.
diff --git a/provider/docs/src/corpus/doc/contentExtractor.test.ts b/provider/docs/src/corpus/doc/contentExtractor.test.ts
@@ -0,0 +1,16 @@
+import { describe, expect, test } from 'vitest'
+import { Content, extractContentUsingMozillaReadability } from './contentExtractor'
+
+describe('extractContentUsingMozillaReadability', () => {
+    test('extracts content', () =>
+        expect(
+            extractContentUsingMozillaReadability({
+                id: 1,
+                text: '<html><head><title>Bar - MySite</title></head><body><aside><nav><h1><a href="/">MySite</a></h1> <a href="/foo">foo</a></nav></aside><main><h1>Bar</h1> <p>Baz</p></main></body>',
+            })
+        ).toEqual<Content>({
+            title: 'Bar - MySite',
+            content: '<main><h2>Bar</h2> <p>Baz</p></main>',
+            textContent: 'Bar Baz',
+        }))
+})
diff --git a/provider/docs/src/corpus/doc/contentExtractor.ts b/provider/docs/src/corpus/doc/contentExtractor.ts
@@ -0,0 +1,34 @@
+import { Readability } from '@mozilla/readability'
+import { JSDOM } from 'jsdom'
+import { type Doc } from './doc'
+
+export interface Content {
+    /**
+     * Title of the document.
+     */
+    title: string
+
+    /**
+     * Content of the document, including some markup. Omits non-content-related elements (header,
+     * footer, navigation, etc.).
+     */
+    content: string
+
+    /**
+     * Text content of the document, with all markup removed. Omits all non-content-related elements.
+     */
+    textContent: string
+}
+
+export type ContentExtractor = (doc: Doc) => Content | null
+
+export const extractContentUsingMozillaReadability: ContentExtractor = doc => {
+    const info = new Readability(new JSDOM(doc.text, { url: doc.url }).window.document).parse()
+    return info
+        ? {
+              title: info.title,
+              content: info.content.replace(/^<div id="readability-page-1" class="page">(.*)<\/div>$/, '$1'),
+              textContent: info.textContent,
+          }
+        : null
+}
diff --git a/provider/docs/src/corpus/doc/doc.ts b/provider/docs/src/corpus/doc/doc.ts
@@ -1,5 +1,3 @@
-import { type Chunk } from './chunks'
-
 /**
  * A unique identifier for a document in a corpus.
  */
@@ -14,11 +12,3 @@ export interface Doc {
 
     url?: string
 }
-
-/**
- * An indexed document.
- */
-export interface IndexedDoc {
-    doc: Doc
-    chunks: Chunk[]
-}
diff --git a/provider/docs/src/corpus/doc/summary.ts b/provider/docs/src/corpus/doc/summary.ts
diff --git a/provider/docs/src/corpus/index.ts b/provider/docs/src/corpus/index.ts
@@ -1,19 +1,37 @@
 import { type CorpusData } from './data'
-import { chunk } from './doc/chunks'
-import { type IndexedDoc } from './doc/doc'
+import { chunk, type Chunk } from './doc/chunks'
+import { type Content, type ContentExtractor } from './doc/contentExtractor'
+import { type Doc } from './doc/doc'
 
+/**
+ * An index of a corpus.
+ */
 export interface CorpusIndex {
     data: CorpusData
 
     docs: IndexedDoc[]
 }
 
-export function indexCorpus(data: CorpusData): CorpusIndex {
+/**
+ * An indexed document.
+ */
+export interface IndexedDoc {
+    doc: Doc
+    chunks: Chunk[]
+    content?: Content
+}
+
+interface IndexOptions {
+    contentExtractor?: ContentExtractor
+}
+
+export function indexCorpus(data: CorpusData, { contentExtractor }: IndexOptions = {}): CorpusIndex {
     const index: CorpusIndex = { data, docs: [] }
 
     for (const doc of data.docs) {
         const chunks = chunk(doc.text, { isMarkdown: doc.text.includes('##') })
-        index.docs.push({ doc, chunks })
+        const content = contentExtractor?.(doc) ?? undefined
+        index.docs.push({ doc, chunks, content })
     }
 
     return index

diff --git a/provider/docs/src/corpus/search/tfidf.ts b/provider/docs/src/corpus/search/tfidf.ts
@@ -1,5 +1,6 @@
+import { type IndexedDoc } from '..'
 import { type ChunkIndex } from '../doc/chunks'
-import { type DocID, type IndexedDoc } from '../doc/doc'
+import { type DocID } from '../doc/doc'
 import { terms, type Term } from './terms'
 
 /**