@thor123141245r/ai-translate 0.0.0

package/src/split.ts

@@ -0,0 +1,111 @@
+import { RecursiveCharacterTextSplitter } from '@langchain/textsplitters'
+import type {
+  RecursiveCharacterTextSplitterParams,
+  SupportedTextSplitterLanguage
+} from '@langchain/textsplitters'
+import { AsyncTransform } from './asyncTransform.js'
+
+export type TextSplitterFormat = { format?: SupportedTextSplitterLanguage }
+export type TextSplitterParams = Partial<RecursiveCharacterTextSplitterParams> &
+  TextSplitterFormat
+export type Separator = string | RegExp
+export type TextSplitterLike = { splitText(text: string): Promise<string[]> }
+
+const extensionsByFormat: Record<string, string[]> = {
+  cpp: ['.h', '.c', '.cpp'],
+  go: ['.go'],
+  java: ['.java'],
+  js: ['.js', '.cjs', '.mjs', '.jsx', '.ts', '.tsx'],
+  php: ['.php', '.inc'],
+  proto: ['.proto'],
+  python: ['.py'],
+  rst: ['.rst'],
+  ruby: ['.rb', '.gem'],
+  rust: ['.rs'],
+  scala: ['.scala', '.sc'],
+  swift: ['.swift'],
+  markdown: ['.md', '.markdown', '.mkdn'],
+  latex: ['.tex', '.latex', '.ltx'],
+  html: ['.html', '.htm', '.xhtml', '.shtml', '.asp', '.aspx', '.jsp']
+}
+
+let formatByExtension: Record<string, SupportedTextSplitterLanguage> | undefined
+
+export const getFormatByExtension = (
+  extname: string
+): SupportedTextSplitterLanguage | undefined => {
+  if (!formatByExtension) {
+    formatByExtension = {}
+    for (const [key, values] of Object.entries(extensionsByFormat)) {
+      for (const ext of values) {
+        formatByExtension[ext] = key as SupportedTextSplitterLanguage
+      }
+    }
+  }
+  return formatByExtension[extname]
+}
+
+export const recursiveChunkTextSplitter = (options?: TextSplitterParams) => {
+  const { format = 'markdown', ...rest } = options || {}
+  const separators = extensionsByFormat[format]
+    ? RecursiveCharacterTextSplitter.getSeparatorsForLanguage(format)
+    : ['\n\n', '\n', '.', '']
+  return new RecursiveChunkTextSplitter({
+    chunkSize: 1000,
+    chunkOverlap: 0,
+    separators,
+    ...rest
+  })
+}
+
+// @ts-expect-error override joinDocs to preserve whitespace; upstream marks it private
+export class RecursiveChunkTextSplitter extends RecursiveCharacterTextSplitter {
+  protected override joinDocs(docs: string[], separator: string): string | null {
+    const text = docs.join(separator)
+    return text === '' ? null : text
+  }
+}
+
+export class TextSplitterStream extends AsyncTransform {
+  private readonly _textSplitter: TextSplitterLike
+  private _documents: string[]
+  private _buffer: string
+
+  constructor(options?: { textSplitter?: TextSplitterLike }) {
+    const { textSplitter } = options || {}
+    super()
+    this._textSplitter = textSplitter || recursiveChunkTextSplitter()
+    this._documents = []
+    this._buffer = ''
+  }
+
+  protected async _asyncTransform(
+    chunk: Buffer | object,
+    encoding: BufferEncoding | 'buffer'
+  ) {
+    const text = AsyncTransform.toString(chunk, encoding)
+    this._buffer += text
+
+    const documents = await this._textSplitter.splitText(this._buffer)
+    this._buffer = documents.pop() ?? ''
+    this._documents = this._documents.concat(documents)
+
+    while (this._documents.length) {
+      const document = this._documents.shift()
+      if (document === undefined) break
+      const wait = !this.push(document)
+      if (wait) {
+        return true
+      }
+    }
+    return false
+  }
+
+  _flush(done: (error?: Error | null) => void) {
+    for (const document of this._documents) {
+      this.push(document)
+    }
+    this.push(this._buffer)
+    done()
+  }
+}

package/src/utils.ts

@@ -0,0 +1,15 @@
+export const isWhiteSpace = (text = '') => text.trim() === ''
+
+export const whiteSpace = (text = '') => {
+  const head = /^\s{0,200}/.exec(text)?.[0] || ''
+  let tail = ''
+  if (head.length !== text.length) {
+    tail = /\s{0,200}$/.exec(text)?.[0] || ''
+  }
+  return { head, tail }
+}
+
+export const preserveWhiteSpace = (inp: string, out: string) => {
+  const { head, tail } = whiteSpace(inp)
+  return head + out.trim() + tail
+}

package/tsconfig.json

@@ -0,0 +1,19 @@
+{
+  "compilerOptions": {
+    "target": "ES2022",
+    "module": "NodeNext",
+    "moduleResolution": "NodeNext",
+    "rootDir": "src",
+    "outDir": "dist",
+    "declaration": true,
+    "declarationMap": true,
+    "sourceMap": true,
+    "esModuleInterop": true,
+    "resolveJsonModule": false,
+    "strict": true,
+    "skipLibCheck": true,
+    "forceConsistentCasingInFileNames": true
+  },
+  "include": ["src/**/*.ts"],
+  "exclude": ["node_modules", "dist"]
+}