doc2vec 1.3.0 → 2.2.0
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/README.md +216 -8
- package/dist/code-chunker.js +179 -0
- package/dist/content-processor copy.js +819 -0
- package/dist/content-processor.js +620 -72
- package/dist/database.js +205 -19
- package/dist/doc2vec.js +520 -8
- package/dist/electron-app/src/database.js +218 -0
- package/dist/electron-app/src/embeddings.js +80 -0
- package/dist/electron-app/src/main.js +627 -0
- package/dist/electron-app/src/mcp-server.js +496 -0
- package/dist/electron-app/src/preload.js +24 -0
- package/dist/electron-app/src/processor.js +248 -0
- package/dist/electron-app/src/syncer.js +146 -0
- package/dist/macos-app/migrate-to-vec0.js +156 -0
- package/dist/vitest.config.js +16 -0
- package/package.json +13 -3
- package/dist/embedding-providers.js +0 -157
package/README.md
CHANGED
|
@@ -6,6 +6,8 @@ This project provides a configurable tool (`doc2vec`) to crawl specified website
|
|
|
6
6
|
|
|
7
7
|
The primary goal is to prepare documentation content for Retrieval-Augmented Generation (RAG) systems or semantic search applications.
|
|
8
8
|
|
|
9
|
+
> **⚠️ Version 2.0.0 Breaking Change:** Version 2.0.0 introduced enhanced chunking with new metadata fields (`chunk_index` and `total_chunks`) that enable page reconstruction and improved chunk ordering. The database schema has changed, and databases created with versions prior to 2.0.0 use a different format. **If you're upgrading to version 2.0.0 or later, you should start with fresh databases** to take advantage of the new features. While the MCP server maintains backward compatibility for querying old databases, doc2vec itself will create databases in the new format. If you need to migrate existing data, consider re-running doc2vec on your sources to regenerate the databases with the enhanced chunking format.
|
|
10
|
+
|
|
9
11
|
## Key Features
|
|
10
12
|
|
|
11
13
|
* **Website Crawling:** Recursively crawls websites starting from a given base URL.
|
|
@@ -19,10 +21,17 @@ The primary goal is to prepare documentation content for Retrieval-Augmented Gen
|
|
|
19
21
|
* **Flexible Filtering:** Filter tickets by status and priority.
|
|
20
22
|
* **Local Directory Processing:** Scans local directories for files, converts content to searchable chunks.
|
|
21
23
|
* **PDF Support:** Automatically extracts text from PDF files and converts them to Markdown format using Mozilla's PDF.js.
|
|
24
|
+
* **Word Document Support:** Processes both legacy `.doc` and modern `.docx` files, extracting text and formatting.
|
|
25
|
+
* **Code Source Processing:** Ingests code from local directories or GitHub repositories using Chonkie code chunking.
|
|
26
|
+
* **AST-aware Chunking:** Uses Chonkie-based code chunking with Tree-sitter to preserve code structure.
|
|
27
|
+
* **Repository Support:** Clones GitHub repos for code ingestion and maps files to GitHub URLs.
|
|
22
28
|
* **Content Extraction:** Uses Puppeteer for rendering JavaScript-heavy pages and `@mozilla/readability` to extract the main article content.
|
|
29
|
+
* **Smart H1 Preservation:** Automatically extracts and preserves page titles (H1 headings) that Readability might strip as "page chrome", ensuring proper heading hierarchy.
|
|
30
|
+
* **Flexible Content Selectors:** Supports multiple content container patterns (`.docs-content`, `.doc-content`, `.markdown-body`, `article`, etc.) for better compatibility with various documentation sites.
|
|
23
31
|
* **HTML to Markdown:** Converts extracted HTML to clean Markdown using `turndown`, preserving code blocks and basic formatting.
|
|
32
|
+
* **Clean Heading Text:** Automatically removes anchor links (like `[](#section-id)`) from heading text for cleaner hierarchy display.
|
|
24
33
|
* **Intelligent Chunking:** Splits Markdown content into manageable chunks based on headings and token limits, preserving context.
|
|
25
|
-
* **Vector Embeddings:** Generates embeddings for each chunk using OpenAI
|
|
34
|
+
* **Vector Embeddings:** Generates embeddings for each chunk using OpenAI or Azure OpenAI (configurable).
|
|
26
35
|
* **Vector Storage:** Supports storing chunks, metadata, and embeddings in:
|
|
27
36
|
* **SQLite:** Using `better-sqlite3` and the `sqlite-vec` extension for efficient vector search.
|
|
28
37
|
* **Qdrant:** A dedicated vector database, using the `@qdrant/js-client-rest`.
|
|
@@ -32,12 +41,64 @@ The primary goal is to prepare documentation content for Retrieval-Augmented Gen
|
|
|
32
41
|
* **Configuration:** Driven by a YAML configuration file (`config.yaml`) specifying sites, repositories, local directories, Zendesk instances, database types, metadata, and other parameters.
|
|
33
42
|
* **Structured Logging:** Uses a custom logger (`logger.ts`) with levels, timestamps, colors, progress bars, and child loggers for clear execution monitoring.
|
|
34
43
|
|
|
44
|
+
## Chunk Metadata & Page Reconstruction
|
|
45
|
+
|
|
46
|
+
Each chunk stored in the database includes rich metadata that enables powerful retrieval and page reconstruction capabilities.
|
|
47
|
+
|
|
48
|
+
### Metadata Fields
|
|
49
|
+
|
|
50
|
+
| Field | Type | Description |
|
|
51
|
+
|-------|------|-------------|
|
|
52
|
+
| `product_name` | string | Product identifier from config |
|
|
53
|
+
| `version` | string | Version identifier from config |
|
|
54
|
+
| `heading_hierarchy` | string[] | Hierarchical breadcrumb trail (e.g., `["Installation", "Prerequisites", "Docker"]`) |
|
|
55
|
+
| `section` | string | Current section heading |
|
|
56
|
+
| `chunk_id` | string | Unique hash identifier for the chunk |
|
|
57
|
+
| `url` | string | Source URL/path of the original document |
|
|
58
|
+
| `hash` | string | Content hash for change detection |
|
|
59
|
+
| `chunk_index` | number | Position of this chunk within the page (0-based) |
|
|
60
|
+
| `total_chunks` | number | Total number of chunks for this page |
|
|
61
|
+
|
|
62
|
+
### Page Reconstruction
|
|
63
|
+
|
|
64
|
+
The `chunk_index` and `total_chunks` fields enable you to reconstruct full pages from chunks:
|
|
65
|
+
|
|
66
|
+
```typescript
|
|
67
|
+
// Example: Retrieve all chunks for a URL and reconstruct the page
|
|
68
|
+
const chunks = await db.query({
|
|
69
|
+
filter: { url: "https://docs.example.com/guide" },
|
|
70
|
+
sort: { chunk_index: "asc" }
|
|
71
|
+
});
|
|
72
|
+
|
|
73
|
+
// Check if there are more chunks after the current one
|
|
74
|
+
if (currentChunk.chunk_index < currentChunk.total_chunks - 1) {
|
|
75
|
+
// More chunks available - fetch the next one
|
|
76
|
+
const nextChunkIndex = currentChunk.chunk_index + 1;
|
|
77
|
+
}
|
|
78
|
+
|
|
79
|
+
// Reconstruct full page content
|
|
80
|
+
const fullPageContent = chunks
|
|
81
|
+
.sort((a, b) => a.chunk_index - b.chunk_index)
|
|
82
|
+
.map(c => c.content)
|
|
83
|
+
.join("\n\n");
|
|
84
|
+
```
|
|
85
|
+
|
|
86
|
+
### Heading Hierarchy (Breadcrumbs)
|
|
87
|
+
|
|
88
|
+
Each chunk includes a `heading_hierarchy` array that provides context about where the content appears in the document structure. This is injected as a `[Topic: ...]` prefix in the chunk content to improve vector search relevance.
|
|
89
|
+
|
|
90
|
+
For example, a chunk under "Installation > Prerequisites > Docker" will have:
|
|
91
|
+
- `heading_hierarchy`: `["Installation", "Prerequisites", "Docker"]`
|
|
92
|
+
- Content prefix: `[Topic: Installation > Prerequisites > Docker]`
|
|
93
|
+
|
|
94
|
+
This ensures that searches for parent topics (like "Installation") will also match relevant child content.
|
|
95
|
+
|
|
35
96
|
## Prerequisites
|
|
36
97
|
|
|
37
98
|
* **Node.js:** Version 18 or higher recommended (check `.nvmrc` if available).
|
|
38
99
|
* **npm:** Node Package Manager (usually comes with Node.js).
|
|
39
100
|
* **TypeScript:** As the project is written in TypeScript (`ts-node` is used for execution via `npm start`).
|
|
40
|
-
* **OpenAI API Key:** You need an API key
|
|
101
|
+
* **OpenAI API Key or Azure OpenAI Credentials:** You need either an OpenAI API key or Azure OpenAI credentials to generate embeddings.
|
|
41
102
|
* **GitHub Personal Access Token:** Required for accessing GitHub issues (set as `GITHUB_PERSONAL_ACCESS_TOKEN` in your environment).
|
|
42
103
|
* **Zendesk API Token:** Required for accessing Zendesk tickets and articles (set as `ZENDESK_API_TOKEN` in your environment).
|
|
43
104
|
* **(Optional) Qdrant Instance:** If using the `qdrant` database type, you need a running Qdrant instance accessible from where you run the script.
|
|
@@ -68,8 +129,20 @@ Configuration is managed through two files:
|
|
|
68
129
|
```dotenv
|
|
69
130
|
# .env
|
|
70
131
|
|
|
71
|
-
#
|
|
132
|
+
# Embedding Provider Configuration
|
|
133
|
+
# Optional: Specify which provider to use (defaults to 'openai' if not set)
|
|
134
|
+
# Can also be configured in config.yaml
|
|
135
|
+
EMBEDDING_PROVIDER="azure" # or "openai"
|
|
136
|
+
|
|
137
|
+
# Required: Your OpenAI API Key (if using OpenAI provider)
|
|
72
138
|
OPENAI_API_KEY="sk-..."
|
|
139
|
+
OPENAI_MODEL="text-embedding-3-large" # Optional, defaults to text-embedding-3-large
|
|
140
|
+
|
|
141
|
+
# Required: Your Azure OpenAI credentials (if using Azure provider)
|
|
142
|
+
AZURE_OPENAI_KEY="your-azure-key"
|
|
143
|
+
AZURE_OPENAI_ENDPOINT="https://your-resource.openai.azure.com"
|
|
144
|
+
AZURE_OPENAI_DEPLOYMENT_NAME="text-embedding-3-large"
|
|
145
|
+
AZURE_OPENAI_API_VERSION="2024-10-21"
|
|
73
146
|
|
|
74
147
|
# Required for GitHub sources
|
|
75
148
|
GITHUB_PERSONAL_ACCESS_TOKEN="ghp_..."
|
|
@@ -87,7 +160,7 @@ Configuration is managed through two files:
|
|
|
87
160
|
**Structure:**
|
|
88
161
|
|
|
89
162
|
* `sources`: An array of source configurations.
|
|
90
|
-
* `type`: Either `'website'`, `'github'`, `'local_directory'`, or `'zendesk'`
|
|
163
|
+
* `type`: Either `'website'`, `'github'`, `'local_directory'`, `'code'`, or `'zendesk'`
|
|
91
164
|
|
|
92
165
|
For websites (`type: 'website'`):
|
|
93
166
|
* `url`: The starting URL for crawling the documentation site.
|
|
@@ -99,11 +172,25 @@ Configuration is managed through two files:
|
|
|
99
172
|
|
|
100
173
|
For local directories (`type: 'local_directory'`):
|
|
101
174
|
* `path`: Path to the local directory to process.
|
|
102
|
-
* `include_extensions`: (Optional) Array of file extensions to include (e.g., `['.md', '.txt', '.pdf']`). Defaults to `['.md', '.txt', '.html', '.htm', '.pdf']`.
|
|
175
|
+
* `include_extensions`: (Optional) Array of file extensions to include (e.g., `['.md', '.txt', '.pdf', '.doc', '.docx']`). Defaults to `['.md', '.txt', '.html', '.htm', '.pdf']`.
|
|
103
176
|
* `exclude_extensions`: (Optional) Array of file extensions to exclude.
|
|
104
177
|
* `recursive`: (Optional) Whether to traverse subdirectories (defaults to `true`).
|
|
105
178
|
* `url_rewrite_prefix` (Optional) URL prefix to rewrite `file://` URLs (e.g., `https://mydomain.com`)
|
|
106
179
|
* `encoding`: (Optional) File encoding to use (defaults to `'utf8'`). Note: PDF files are processed as binary and this setting doesn't apply to them.
|
|
180
|
+
|
|
181
|
+
For code sources (`type: 'code'`):
|
|
182
|
+
* `source`: Either `'local_directory'` or `'github'`.
|
|
183
|
+
* `path`: Path to the local directory (required when `source: 'local_directory'`).
|
|
184
|
+
* `repo`: Repository name in the format `'owner/repo'` (required when `source: 'github'`).
|
|
185
|
+
* `branch`: (Optional) Branch to clone for GitHub sources.
|
|
186
|
+
* `include_extensions`: (Optional) Array of file extensions to include (defaults to common code extensions).
|
|
187
|
+
* `exclude_extensions`: (Optional) Array of file extensions to exclude.
|
|
188
|
+
* `recursive`: (Optional) Whether to traverse subdirectories (defaults to `true`).
|
|
189
|
+
* `url_rewrite_prefix`: (Optional) URL prefix to rewrite `file://` URLs for local sources.
|
|
190
|
+
* `encoding`: (Optional) File encoding to use (defaults to `'utf8'`).
|
|
191
|
+
* `chunk_size`: (Optional) Chonkie chunk size for code files.
|
|
192
|
+
* `version` is optional for code sources; if omitted it defaults to `branch` (or `local` for local directories).
|
|
193
|
+
* `branch` is stored in the database and used by `query_code` filtering.
|
|
107
194
|
|
|
108
195
|
For Zendesk (`type: 'zendesk'`):
|
|
109
196
|
* `zendesk_subdomain`: Your Zendesk subdomain (e.g., `'mycompany'` for mycompany.zendesk.com).
|
|
@@ -131,6 +218,21 @@ Configuration is managed through two files:
|
|
|
131
218
|
|
|
132
219
|
**Example (`config.yaml`):**
|
|
133
220
|
```yaml
|
|
221
|
+
# Optional: Configure embedding provider
|
|
222
|
+
# Can also be set via EMBEDDING_PROVIDER environment variable
|
|
223
|
+
# Defaults to OpenAI if not specified
|
|
224
|
+
embedding:
|
|
225
|
+
provider: 'openai' # or 'azure'
|
|
226
|
+
openai:
|
|
227
|
+
api_key: '${OPENAI_API_KEY}' # Optional, uses env var by default
|
|
228
|
+
model: 'text-embedding-3-large' # Optional, defaults to text-embedding-3-large
|
|
229
|
+
# For Azure OpenAI, use this instead:
|
|
230
|
+
# azure:
|
|
231
|
+
# api_key: '${AZURE_OPENAI_KEY}'
|
|
232
|
+
# endpoint: '${AZURE_OPENAI_ENDPOINT}'
|
|
233
|
+
# deployment_name: 'text-embedding-3-large'
|
|
234
|
+
# api_version: '2024-10-21' # Optional
|
|
235
|
+
|
|
134
236
|
sources:
|
|
135
237
|
# Website source example
|
|
136
238
|
- type: 'website'
|
|
@@ -161,13 +263,28 @@ Configuration is managed through two files:
|
|
|
161
263
|
product_name: 'project-docs'
|
|
162
264
|
version: 'current'
|
|
163
265
|
path: './docs'
|
|
164
|
-
include_extensions: ['.md', '.txt', '.pdf']
|
|
266
|
+
include_extensions: ['.md', '.txt', '.pdf', '.doc', '.docx']
|
|
165
267
|
recursive: true
|
|
166
|
-
max_size: 10485760 # 10MB recommended for PDF files
|
|
268
|
+
max_size: 10485760 # 10MB recommended for PDF/Word files
|
|
167
269
|
database_config:
|
|
168
270
|
type: 'sqlite'
|
|
169
271
|
params:
|
|
170
272
|
db_path: './project-docs.db'
|
|
273
|
+
|
|
274
|
+
# Code source example (GitHub)
|
|
275
|
+
- type: 'code'
|
|
276
|
+
source: 'github'
|
|
277
|
+
product_name: 'doc2vec'
|
|
278
|
+
version: 'main'
|
|
279
|
+
repo: 'kagent-dev/doc2vec'
|
|
280
|
+
branch: 'main'
|
|
281
|
+
include_extensions: ['.ts', '.tsx', '.md']
|
|
282
|
+
max_size: 1048576
|
|
283
|
+
chunk_size: 2048
|
|
284
|
+
database_config:
|
|
285
|
+
type: 'sqlite'
|
|
286
|
+
params:
|
|
287
|
+
db_path: './doc2vec-code.db'
|
|
171
288
|
|
|
172
289
|
# Zendesk example
|
|
173
290
|
- type: 'zendesk'
|
|
@@ -291,6 +408,67 @@ A PDF file named "user-guide.pdf" will be converted to Markdown format like:
|
|
|
291
408
|
|
|
292
409
|
The resulting Markdown is then chunked and embedded using the same process as other text content.
|
|
293
410
|
|
|
411
|
+
## Word Document Processing
|
|
412
|
+
|
|
413
|
+
Doc2Vec supports processing Microsoft Word documents in both legacy `.doc` format and modern `.docx` format.
|
|
414
|
+
|
|
415
|
+
### Supported Formats
|
|
416
|
+
|
|
417
|
+
| Extension | Format | Library Used |
|
|
418
|
+
|-----------|--------|--------------|
|
|
419
|
+
| `.doc` | Legacy Word (97-2003) | [word-extractor](https://github.com/morungos/node-word-extractor) |
|
|
420
|
+
| `.docx` | Modern Word (2007+) | [mammoth](https://github.com/mwilliamson/mammoth.js) |
|
|
421
|
+
|
|
422
|
+
### Features
|
|
423
|
+
|
|
424
|
+
* **Legacy .doc Support:** Extracts plain text from older Word documents using binary parsing
|
|
425
|
+
* **Modern .docx Support:** Converts DOCX files to HTML first (preserving formatting), then to clean Markdown
|
|
426
|
+
* **Formatting Preservation:** For `.docx` files, headings, lists, bold, italic, and links are preserved
|
|
427
|
+
* **Automatic Title:** Uses the filename as an H1 heading for proper document structure
|
|
428
|
+
* **Local File Support:** Processes Word files found in local directories alongside other documents
|
|
429
|
+
|
|
430
|
+
### Configuration
|
|
431
|
+
|
|
432
|
+
Include `.doc` and/or `.docx` in your `include_extensions` array:
|
|
433
|
+
|
|
434
|
+
```yaml
|
|
435
|
+
- type: 'local_directory'
|
|
436
|
+
product_name: 'company-docs'
|
|
437
|
+
version: 'current'
|
|
438
|
+
path: './documents'
|
|
439
|
+
include_extensions: ['.doc', '.docx', '.pdf', '.md']
|
|
440
|
+
recursive: true
|
|
441
|
+
max_size: 10485760 # 10MB recommended
|
|
442
|
+
database_config:
|
|
443
|
+
type: 'sqlite'
|
|
444
|
+
params:
|
|
445
|
+
db_path: './company-docs.db'
|
|
446
|
+
```
|
|
447
|
+
|
|
448
|
+
### Example Output
|
|
449
|
+
|
|
450
|
+
A Word document named "meeting-notes.docx" will be converted to Markdown like:
|
|
451
|
+
|
|
452
|
+
```markdown
|
|
453
|
+
# meeting-notes
|
|
454
|
+
|
|
455
|
+
## Agenda
|
|
456
|
+
|
|
457
|
+
1. Review Q4 results
|
|
458
|
+
2. Discuss roadmap
|
|
459
|
+
|
|
460
|
+
## Action Items
|
|
461
|
+
|
|
462
|
+
- **John:** Prepare budget report
|
|
463
|
+
- **Sarah:** Schedule follow-up meeting
|
|
464
|
+
```
|
|
465
|
+
|
|
466
|
+
### Notes
|
|
467
|
+
|
|
468
|
+
* **`.doc` files:** Only plain text is extracted. Formatting like bold/italic is not preserved in legacy Word format.
|
|
469
|
+
* **`.docx` files:** Full formatting is preserved including headings, lists, bold, italic, links, and tables.
|
|
470
|
+
* **Embedded Images:** Images embedded in Word documents are not extracted (text-only).
|
|
471
|
+
|
|
294
472
|
## Now Available via npx
|
|
295
473
|
|
|
296
474
|
You can run `doc2vec` without cloning the repo or installing it globally. Just use:
|
|
@@ -333,6 +511,7 @@ If you don't specify a config path, it will look for config.yaml in the current
|
|
|
333
511
|
* Recursively scan directories for files matching the configured extensions.
|
|
334
512
|
* Read file content, converting HTML to Markdown if needed.
|
|
335
513
|
* For PDF files, extract text using Mozilla's PDF.js and convert to Markdown format with proper page structure.
|
|
514
|
+
* For Word documents, extract text from `.doc` files or convert `.docx` files to Markdown with formatting.
|
|
336
515
|
* Process each file's content.
|
|
337
516
|
- **For Zendesk:**
|
|
338
517
|
* Fetch tickets and articles using the Zendesk API.
|
|
@@ -345,4 +524,33 @@ If you don't specify a config path, it will look for config.yaml in the current
|
|
|
345
524
|
* **Embed (if needed):** If the chunk is new or changed, call the OpenAI API (`createEmbeddings`) to get the vector embedding.
|
|
346
525
|
* **Store:** Insert or update the chunk, metadata, hash, and embedding in the database (SQLite `vec_items` table or Qdrant collection).
|
|
347
526
|
4. **Cleanup:** After processing, remove any obsolete chunks from the database.
|
|
348
|
-
4. **Complete:** Log completion status.
|
|
527
|
+
4. **Complete:** Log completion status.
|
|
528
|
+
|
|
529
|
+
## Recent Changes
|
|
530
|
+
|
|
531
|
+
### Word Document Support
|
|
532
|
+
- Added support for legacy `.doc` files using the `word-extractor` library
|
|
533
|
+
- Added support for modern `.docx` files using the `mammoth` library
|
|
534
|
+
- DOCX files preserve formatting (headings, lists, bold, italic, links)
|
|
535
|
+
- Both formats are converted to clean Markdown for embedding
|
|
536
|
+
|
|
537
|
+
### Page Reconstruction Support
|
|
538
|
+
- Added `chunk_index` field to track each chunk's position within a page (0-based)
|
|
539
|
+
- Added `total_chunks` field to indicate the total number of chunks per page
|
|
540
|
+
- Enables AI agents and applications to fetch additional context or reconstruct full pages
|
|
541
|
+
- Works consistently across all content types: websites, GitHub, Zendesk, and local directories
|
|
542
|
+
|
|
543
|
+
### Improved H1/Title Handling
|
|
544
|
+
- Smart H1 preservation ensures page titles aren't stripped by Readability
|
|
545
|
+
- Falls back to `article.title` when H1 extraction fails
|
|
546
|
+
- Proper heading hierarchy starting from H1 through the document structure
|
|
547
|
+
|
|
548
|
+
### Enhanced Content Extraction
|
|
549
|
+
- Added support for multiple content container selectors (`.docs-content`, `.doc-content`, `.markdown-body`, `article`)
|
|
550
|
+
- Cleaner heading text by removing anchor links like `[](#section-id)`
|
|
551
|
+
- Better handling of pages where H1 is outside the main content container
|
|
552
|
+
|
|
553
|
+
### Heading Hierarchy Improvements
|
|
554
|
+
- Fixed sparse array issues that caused `NULL` values in heading hierarchy
|
|
555
|
+
- Proper breadcrumb generation for nested sections
|
|
556
|
+
- Hierarchical context preserved across chunk boundaries
|
|
@@ -0,0 +1,179 @@
|
|
|
1
|
+
"use strict";
|
|
2
|
+
var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
|
|
3
|
+
if (k2 === undefined) k2 = k;
|
|
4
|
+
var desc = Object.getOwnPropertyDescriptor(m, k);
|
|
5
|
+
if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
|
|
6
|
+
desc = { enumerable: true, get: function() { return m[k]; } };
|
|
7
|
+
}
|
|
8
|
+
Object.defineProperty(o, k2, desc);
|
|
9
|
+
}) : (function(o, m, k, k2) {
|
|
10
|
+
if (k2 === undefined) k2 = k;
|
|
11
|
+
o[k2] = m[k];
|
|
12
|
+
}));
|
|
13
|
+
var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
|
|
14
|
+
Object.defineProperty(o, "default", { enumerable: true, value: v });
|
|
15
|
+
}) : function(o, v) {
|
|
16
|
+
o["default"] = v;
|
|
17
|
+
});
|
|
18
|
+
var __importStar = (this && this.__importStar) || (function () {
|
|
19
|
+
var ownKeys = function(o) {
|
|
20
|
+
ownKeys = Object.getOwnPropertyNames || function (o) {
|
|
21
|
+
var ar = [];
|
|
22
|
+
for (var k in o) if (Object.prototype.hasOwnProperty.call(o, k)) ar[ar.length] = k;
|
|
23
|
+
return ar;
|
|
24
|
+
};
|
|
25
|
+
return ownKeys(o);
|
|
26
|
+
};
|
|
27
|
+
return function (mod) {
|
|
28
|
+
if (mod && mod.__esModule) return mod;
|
|
29
|
+
var result = {};
|
|
30
|
+
if (mod != null) for (var k = ownKeys(mod), i = 0; i < k.length; i++) if (k[i] !== "default") __createBinding(result, mod, k[i]);
|
|
31
|
+
__setModuleDefault(result, mod);
|
|
32
|
+
return result;
|
|
33
|
+
};
|
|
34
|
+
})();
|
|
35
|
+
Object.defineProperty(exports, "__esModule", { value: true });
|
|
36
|
+
exports.CodeChunker = void 0;
|
|
37
|
+
const fs = __importStar(require("fs"));
|
|
38
|
+
const path = __importStar(require("path"));
|
|
39
|
+
const web_tree_sitter_1 = require("web-tree-sitter");
|
|
40
|
+
class CodeChunker {
|
|
41
|
+
constructor(lang, chunkSize, tokenCounter) {
|
|
42
|
+
this.lang = lang;
|
|
43
|
+
this.chunkSize = chunkSize;
|
|
44
|
+
this.tokenCounter = tokenCounter;
|
|
45
|
+
}
|
|
46
|
+
static async create(options) {
|
|
47
|
+
if (!CodeChunker.treeSitterInitialized && web_tree_sitter_1.Parser.init) {
|
|
48
|
+
try {
|
|
49
|
+
await web_tree_sitter_1.Parser.init();
|
|
50
|
+
CodeChunker.treeSitterInitialized = true;
|
|
51
|
+
}
|
|
52
|
+
catch (error) {
|
|
53
|
+
console.warn('Failed to initialize tree-sitter parser:', error);
|
|
54
|
+
}
|
|
55
|
+
}
|
|
56
|
+
const chunkSize = options.chunkSize ?? 512;
|
|
57
|
+
if (chunkSize <= 0) {
|
|
58
|
+
throw new Error('chunkSize must be greater than 0');
|
|
59
|
+
}
|
|
60
|
+
const tokenCounter = options.tokenCounter ?? (async (text) => text.length);
|
|
61
|
+
const chunker = new CodeChunker(options.lang, chunkSize, tokenCounter);
|
|
62
|
+
return chunker;
|
|
63
|
+
}
|
|
64
|
+
async chunk(text) {
|
|
65
|
+
if (!text.trim()) {
|
|
66
|
+
return [];
|
|
67
|
+
}
|
|
68
|
+
const parser = await CodeChunker.getParser(this.lang);
|
|
69
|
+
const originalTextBytes = Buffer.from(text, 'utf-8');
|
|
70
|
+
const tree = parser.parse(originalTextBytes.toString());
|
|
71
|
+
if (!tree) {
|
|
72
|
+
throw new Error('Failed to parse code');
|
|
73
|
+
}
|
|
74
|
+
const chunks = [];
|
|
75
|
+
await this.recursiveChunk(tree.rootNode, originalTextBytes.toString(), chunks);
|
|
76
|
+
return this.mergeChunks(chunks);
|
|
77
|
+
}
|
|
78
|
+
static async getParser(lang) {
|
|
79
|
+
const formattedLang = lang.toLowerCase().replace(/-/g, '_');
|
|
80
|
+
const cached = this.parserCache.get(formattedLang);
|
|
81
|
+
if (cached) {
|
|
82
|
+
return cached;
|
|
83
|
+
}
|
|
84
|
+
const parserPromise = (async () => {
|
|
85
|
+
if (!CodeChunker.treeSitterInitialized && web_tree_sitter_1.Parser.init) {
|
|
86
|
+
try {
|
|
87
|
+
await web_tree_sitter_1.Parser.init();
|
|
88
|
+
CodeChunker.treeSitterInitialized = true;
|
|
89
|
+
}
|
|
90
|
+
catch (error) {
|
|
91
|
+
console.warn('Failed to initialize tree-sitter parser:', error);
|
|
92
|
+
}
|
|
93
|
+
}
|
|
94
|
+
const wasmPath = CodeChunker.resolveWasmPath(formattedLang);
|
|
95
|
+
const wasmBuffer = fs.readFileSync(wasmPath);
|
|
96
|
+
const language = await web_tree_sitter_1.Language.load(wasmBuffer);
|
|
97
|
+
const parser = new web_tree_sitter_1.Parser();
|
|
98
|
+
parser.setLanguage(language);
|
|
99
|
+
return parser;
|
|
100
|
+
})();
|
|
101
|
+
this.parserCache.set(formattedLang, parserPromise);
|
|
102
|
+
return parserPromise;
|
|
103
|
+
}
|
|
104
|
+
static resolveWasmPath(formattedLang) {
|
|
105
|
+
const nodeModulesPath = CodeChunker.findNearestNodeModules(__dirname);
|
|
106
|
+
if (!nodeModulesPath) {
|
|
107
|
+
throw new Error('node_modules directory not found.');
|
|
108
|
+
}
|
|
109
|
+
const wasmPath = path.join(nodeModulesPath, `tree-sitter-wasms/out/tree-sitter-${formattedLang}.wasm`);
|
|
110
|
+
if (!fs.existsSync(wasmPath)) {
|
|
111
|
+
throw new Error(`Tree-sitter WASM file for language "${formattedLang}" not found at ${wasmPath}.`);
|
|
112
|
+
}
|
|
113
|
+
return wasmPath;
|
|
114
|
+
}
|
|
115
|
+
static findNearestNodeModules(startDir) {
|
|
116
|
+
let dir = path.resolve(startDir);
|
|
117
|
+
while (true) {
|
|
118
|
+
const candidate = path.join(dir, 'node_modules');
|
|
119
|
+
if (fs.existsSync(candidate) && fs.statSync(candidate).isDirectory()) {
|
|
120
|
+
return candidate;
|
|
121
|
+
}
|
|
122
|
+
const parent = path.dirname(dir);
|
|
123
|
+
if (parent === dir)
|
|
124
|
+
break;
|
|
125
|
+
dir = parent;
|
|
126
|
+
}
|
|
127
|
+
return null;
|
|
128
|
+
}
|
|
129
|
+
async recursiveChunk(node, source, chunks) {
|
|
130
|
+
const nodeText = source.substring(node.startIndex, node.endIndex);
|
|
131
|
+
const tokenCount = await this.tokenCounter(nodeText);
|
|
132
|
+
const children = (node.children || []).filter((child) => Boolean(child));
|
|
133
|
+
if (tokenCount <= this.chunkSize || children.length === 0) {
|
|
134
|
+
if (nodeText.trim()) {
|
|
135
|
+
chunks.push({ text: nodeText, tokenCount });
|
|
136
|
+
}
|
|
137
|
+
return;
|
|
138
|
+
}
|
|
139
|
+
const beforeCount = chunks.length;
|
|
140
|
+
for (const child of children) {
|
|
141
|
+
await this.recursiveChunk(child, source, chunks);
|
|
142
|
+
}
|
|
143
|
+
if (chunks.length === beforeCount && nodeText.trim()) {
|
|
144
|
+
chunks.push({ text: nodeText, tokenCount });
|
|
145
|
+
}
|
|
146
|
+
}
|
|
147
|
+
mergeChunks(chunks) {
|
|
148
|
+
const merged = [];
|
|
149
|
+
let currentText = '';
|
|
150
|
+
let currentTokens = 0;
|
|
151
|
+
const separatorTokens = 1; // Account for the '\n' separator between merged chunks
|
|
152
|
+
for (const chunk of chunks) {
|
|
153
|
+
if (!chunk.text.trim()) {
|
|
154
|
+
continue;
|
|
155
|
+
}
|
|
156
|
+
const nextTokens = currentTokens + separatorTokens + chunk.tokenCount;
|
|
157
|
+
if (currentTokens === 0) {
|
|
158
|
+
currentText = chunk.text;
|
|
159
|
+
currentTokens = chunk.tokenCount;
|
|
160
|
+
continue;
|
|
161
|
+
}
|
|
162
|
+
if (nextTokens <= this.chunkSize) {
|
|
163
|
+
currentText = `${currentText}\n${chunk.text}`;
|
|
164
|
+
currentTokens = nextTokens;
|
|
165
|
+
continue;
|
|
166
|
+
}
|
|
167
|
+
merged.push({ text: currentText, tokenCount: currentTokens });
|
|
168
|
+
currentText = chunk.text;
|
|
169
|
+
currentTokens = chunk.tokenCount;
|
|
170
|
+
}
|
|
171
|
+
if (currentTokens > 0) {
|
|
172
|
+
merged.push({ text: currentText, tokenCount: currentTokens });
|
|
173
|
+
}
|
|
174
|
+
return merged;
|
|
175
|
+
}
|
|
176
|
+
}
|
|
177
|
+
exports.CodeChunker = CodeChunker;
|
|
178
|
+
CodeChunker.treeSitterInitialized = false;
|
|
179
|
+
CodeChunker.parserCache = new Map();
|