doc2vec 2.0.0 → 2.3.0

package/README.md

@@ -22,13 +22,16 @@ The primary goal is to prepare documentation content for Retrieval-Augmented Gen
 *   **Local Directory Processing:** Scans local directories for files, converts content to searchable chunks.
     * **PDF Support:** Automatically extracts text from PDF files and converts them to Markdown format using Mozilla's PDF.js.
     * **Word Document Support:** Processes both legacy `.doc` and modern `.docx` files, extracting text and formatting.
+*   **Code Source Processing:** Ingests code from local directories or GitHub repositories using Chonkie code chunking.
+    * **AST-aware Chunking:** Uses Chonkie-based code chunking with Tree-sitter to preserve code structure.
+    * **Repository Support:** Clones GitHub repos for code ingestion and maps files to GitHub URLs.
 *   **Content Extraction:** Uses Puppeteer for rendering JavaScript-heavy pages and `@mozilla/readability` to extract the main article content.
     *   **Smart H1 Preservation:** Automatically extracts and preserves page titles (H1 headings) that Readability might strip as "page chrome", ensuring proper heading hierarchy.
     *   **Flexible Content Selectors:** Supports multiple content container patterns (`.docs-content`, `.doc-content`, `.markdown-body`, `article`, etc.) for better compatibility with various documentation sites.
 *   **HTML to Markdown:** Converts extracted HTML to clean Markdown using `turndown`, preserving code blocks and basic formatting.
     *   **Clean Heading Text:** Automatically removes anchor links (like `[](#section-id)`) from heading text for cleaner hierarchy display.
 *   **Intelligent Chunking:** Splits Markdown content into manageable chunks based on headings and token limits, preserving context.
-*   **Vector Embeddings:** Generates embeddings for each chunk using OpenAI's `text-embedding-3-large` model.
+*   **Vector Embeddings:** Generates embeddings for each chunk using OpenAI or Azure OpenAI (configurable).
 *   **Vector Storage:** Supports storing chunks, metadata, and embeddings in:
     *   **SQLite:** Using `better-sqlite3` and the `sqlite-vec` extension for efficient vector search.
     *   **Qdrant:** A dedicated vector database, using the `@qdrant/js-client-rest`.
@@ -95,7 +98,7 @@ This ensures that searches for parent topics (like "Installation") will also mat
 *   **Node.js:** Version 18 or higher recommended (check `.nvmrc` if available).
 *   **npm:** Node Package Manager (usually comes with Node.js).
 *   **TypeScript:** As the project is written in TypeScript (`ts-node` is used for execution via `npm start`).
-*   **OpenAI API Key:** You need an API key from OpenAI to generate embeddings.
+*   **OpenAI API Key or Azure OpenAI Credentials:** You need either an OpenAI API key or Azure OpenAI credentials to generate embeddings.
 *   **GitHub Personal Access Token:** Required for accessing GitHub issues (set as `GITHUB_PERSONAL_ACCESS_TOKEN` in your environment).
 *   **Zendesk API Token:** Required for accessing Zendesk tickets and articles (set as `ZENDESK_API_TOKEN` in your environment).
 *   **(Optional) Qdrant Instance:** If using the `qdrant` database type, you need a running Qdrant instance accessible from where you run the script.
@@ -126,8 +129,20 @@ Configuration is managed through two files:
     ```dotenv
     # .env
 
-    # Required: Your OpenAI API Key
+    # Embedding Provider Configuration
+    # Optional: Specify which provider to use (defaults to 'openai' if not set)
+    # Can also be configured in config.yaml
+    EMBEDDING_PROVIDER="azure"  # or "openai"
+
+    # Required: Your OpenAI API Key (if using OpenAI provider)
     OPENAI_API_KEY="sk-..."
+    OPENAI_MODEL="text-embedding-3-large"  # Optional, defaults to text-embedding-3-large
+
+    # Required: Your Azure OpenAI credentials (if using Azure provider)
+    AZURE_OPENAI_KEY="your-azure-key"
+    AZURE_OPENAI_ENDPOINT="https://your-resource.openai.azure.com"
+    AZURE_OPENAI_DEPLOYMENT_NAME="text-embedding-3-large"
+    AZURE_OPENAI_API_VERSION="2024-10-21"
 
     # Required for GitHub sources
     GITHUB_PERSONAL_ACCESS_TOKEN="ghp_..."
@@ -145,7 +160,7 @@ Configuration is managed through two files:
     **Structure:**
 
     *   `sources`: An array of source configurations.
-        *   `type`: Either `'website'`, `'github'`, `'local_directory'`, or `'zendesk'`
+        *   `type`: Either `'website'`, `'github'`, `'local_directory'`, `'code'`, or `'zendesk'`
         
         For websites (`type: 'website'`):
         *   `url`: The starting URL for crawling the documentation site.
@@ -162,6 +177,20 @@ Configuration is managed through two files:
         *   `recursive`: (Optional) Whether to traverse subdirectories (defaults to `true`).
         *   `url_rewrite_prefix` (Optional) URL prefix to rewrite `file://` URLs (e.g., `https://mydomain.com`)
         *   `encoding`: (Optional) File encoding to use (defaults to `'utf8'`). Note: PDF files are processed as binary and this setting doesn't apply to them.
+
+        For code sources (`type: 'code'`):
+        *   `source`: Either `'local_directory'` or `'github'`.
+        *   `path`: Path to the local directory (required when `source: 'local_directory'`).
+        *   `repo`: Repository name in the format `'owner/repo'` (required when `source: 'github'`).
+        *   `branch`: (Optional) Branch to clone for GitHub sources.
+        *   `include_extensions`: (Optional) Array of file extensions to include (defaults to common code extensions).
+        *   `exclude_extensions`: (Optional) Array of file extensions to exclude.
+        *   `recursive`: (Optional) Whether to traverse subdirectories (defaults to `true`).
+        *   `url_rewrite_prefix`: (Optional) URL prefix to rewrite `file://` URLs for local sources.
+        *   `encoding`: (Optional) File encoding to use (defaults to `'utf8'`).
+        *   `chunk_size`: (Optional) Chonkie chunk size for code files.
+        *   `version` is optional for code sources; if omitted it defaults to `branch` (or `local` for local directories).
+        *   `branch` is stored in the database and used by `query_code` filtering.
         
         For Zendesk (`type: 'zendesk'`):
         *   `zendesk_subdomain`: Your Zendesk subdomain (e.g., `'mycompany'` for mycompany.zendesk.com).
@@ -189,6 +218,21 @@ Configuration is managed through two files:
 
     **Example (`config.yaml`):**
     ```yaml
+    # Optional: Configure embedding provider
+    # Can also be set via EMBEDDING_PROVIDER environment variable
+    # Defaults to OpenAI if not specified
+    embedding:
+      provider: 'openai'  # or 'azure'
+      openai:
+        api_key: '${OPENAI_API_KEY}'  # Optional, uses env var by default
+        model: 'text-embedding-3-large'  # Optional, defaults to text-embedding-3-large
+      # For Azure OpenAI, use this instead:
+      # azure:
+      #   api_key: '${AZURE_OPENAI_KEY}'
+      #   endpoint: '${AZURE_OPENAI_ENDPOINT}'
+      #   deployment_name: 'text-embedding-3-large'
+      #   api_version: '2024-10-21'  # Optional
+
     sources:
       # Website source example
       - type: 'website'
@@ -226,6 +270,21 @@ Configuration is managed through two files:
           type: 'sqlite'
           params:
             db_path: './project-docs.db'
+
+      # Code source example (GitHub)
+      - type: 'code'
+        source: 'github'
+        product_name: 'doc2vec'
+        version: 'main'
+        repo: 'kagent-dev/doc2vec'
+        branch: 'main'
+        include_extensions: ['.ts', '.tsx', '.md']
+        max_size: 1048576
+        chunk_size: 2048
+        database_config:
+          type: 'sqlite'
+          params:
+            db_path: './doc2vec-code.db'
       
       # Zendesk example
       - type: 'zendesk'
@@ -494,4 +553,4 @@ If you don't specify a config path, it will look for config.yaml in the current
 ### Heading Hierarchy Improvements
 - Fixed sparse array issues that caused `NULL` values in heading hierarchy
 - Proper breadcrumb generation for nested sections
-- Hierarchical context preserved across chunk boundaries
+- Hierarchical context preserved across chunk boundaries

package/dist/code-chunker.js

@@ -0,0 +1,179 @@
+"use strict";
+var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
+    if (k2 === undefined) k2 = k;
+    var desc = Object.getOwnPropertyDescriptor(m, k);
+    if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
+      desc = { enumerable: true, get: function() { return m[k]; } };
+    }
+    Object.defineProperty(o, k2, desc);
+}) : (function(o, m, k, k2) {
+    if (k2 === undefined) k2 = k;
+    o[k2] = m[k];
+}));
+var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
+    Object.defineProperty(o, "default", { enumerable: true, value: v });
+}) : function(o, v) {
+    o["default"] = v;
+});
+var __importStar = (this && this.__importStar) || (function () {
+    var ownKeys = function(o) {
+        ownKeys = Object.getOwnPropertyNames || function (o) {
+            var ar = [];
+            for (var k in o) if (Object.prototype.hasOwnProperty.call(o, k)) ar[ar.length] = k;
+            return ar;
+        };
+        return ownKeys(o);
+    };
+    return function (mod) {
+        if (mod && mod.__esModule) return mod;
+        var result = {};
+        if (mod != null) for (var k = ownKeys(mod), i = 0; i < k.length; i++) if (k[i] !== "default") __createBinding(result, mod, k[i]);
+        __setModuleDefault(result, mod);
+        return result;
+    };
+})();
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.CodeChunker = void 0;
+const fs = __importStar(require("fs"));
+const path = __importStar(require("path"));
+const web_tree_sitter_1 = require("web-tree-sitter");
+class CodeChunker {
+    constructor(lang, chunkSize, tokenCounter) {
+        this.lang = lang;
+        this.chunkSize = chunkSize;
+        this.tokenCounter = tokenCounter;
+    }
+    static async create(options) {
+        if (!CodeChunker.treeSitterInitialized && web_tree_sitter_1.Parser.init) {
+            try {
+                await web_tree_sitter_1.Parser.init();
+                CodeChunker.treeSitterInitialized = true;
+            }
+            catch (error) {
+                console.warn('Failed to initialize tree-sitter parser:', error);
+            }
+        }
+        const chunkSize = options.chunkSize ?? 512;
+        if (chunkSize <= 0) {
+            throw new Error('chunkSize must be greater than 0');
+        }
+        const tokenCounter = options.tokenCounter ?? (async (text) => text.length);
+        const chunker = new CodeChunker(options.lang, chunkSize, tokenCounter);
+        return chunker;
+    }
+    async chunk(text) {
+        if (!text.trim()) {
+            return [];
+        }
+        const parser = await CodeChunker.getParser(this.lang);
+        const originalTextBytes = Buffer.from(text, 'utf-8');
+        const tree = parser.parse(originalTextBytes.toString());
+        if (!tree) {
+            throw new Error('Failed to parse code');
+        }
+        const chunks = [];
+        await this.recursiveChunk(tree.rootNode, originalTextBytes.toString(), chunks);
+        return this.mergeChunks(chunks);
+    }
+    static async getParser(lang) {
+        const formattedLang = lang.toLowerCase().replace(/-/g, '_');
+        const cached = this.parserCache.get(formattedLang);
+        if (cached) {
+            return cached;
+        }
+        const parserPromise = (async () => {
+            if (!CodeChunker.treeSitterInitialized && web_tree_sitter_1.Parser.init) {
+                try {
+                    await web_tree_sitter_1.Parser.init();
+                    CodeChunker.treeSitterInitialized = true;
+                }
+                catch (error) {
+                    console.warn('Failed to initialize tree-sitter parser:', error);
+                }
+            }
+            const wasmPath = CodeChunker.resolveWasmPath(formattedLang);
+            const wasmBuffer = fs.readFileSync(wasmPath);
+            const language = await web_tree_sitter_1.Language.load(wasmBuffer);
+            const parser = new web_tree_sitter_1.Parser();
+            parser.setLanguage(language);
+            return parser;
+        })();
+        this.parserCache.set(formattedLang, parserPromise);
+        return parserPromise;
+    }
+    static resolveWasmPath(formattedLang) {
+        const nodeModulesPath = CodeChunker.findNearestNodeModules(__dirname);
+        if (!nodeModulesPath) {
+            throw new Error('node_modules directory not found.');
+        }
+        const wasmPath = path.join(nodeModulesPath, `tree-sitter-wasms/out/tree-sitter-${formattedLang}.wasm`);
+        if (!fs.existsSync(wasmPath)) {
+            throw new Error(`Tree-sitter WASM file for language "${formattedLang}" not found at ${wasmPath}.`);
+        }
+        return wasmPath;
+    }
+    static findNearestNodeModules(startDir) {
+        let dir = path.resolve(startDir);
+        while (true) {
+            const candidate = path.join(dir, 'node_modules');
+            if (fs.existsSync(candidate) && fs.statSync(candidate).isDirectory()) {
+                return candidate;
+            }
+            const parent = path.dirname(dir);
+            if (parent === dir)
+                break;
+            dir = parent;
+        }
+        return null;
+    }
+    async recursiveChunk(node, source, chunks) {
+        const nodeText = source.substring(node.startIndex, node.endIndex);
+        const tokenCount = await this.tokenCounter(nodeText);
+        const children = (node.children || []).filter((child) => Boolean(child));
+        if (tokenCount <= this.chunkSize || children.length === 0) {
+            if (nodeText.trim()) {
+                chunks.push({ text: nodeText, tokenCount });
+            }
+            return;
+        }
+        const beforeCount = chunks.length;
+        for (const child of children) {
+            await this.recursiveChunk(child, source, chunks);
+        }
+        if (chunks.length === beforeCount && nodeText.trim()) {
+            chunks.push({ text: nodeText, tokenCount });
+        }
+    }
+    mergeChunks(chunks) {
+        const merged = [];
+        let currentText = '';
+        let currentTokens = 0;
+        const separatorTokens = 1; // Account for the '\n' separator between merged chunks
+        for (const chunk of chunks) {
+            if (!chunk.text.trim()) {
+                continue;
+            }
+            const nextTokens = currentTokens + separatorTokens + chunk.tokenCount;
+            if (currentTokens === 0) {
+                currentText = chunk.text;
+                currentTokens = chunk.tokenCount;
+                continue;
+            }
+            if (nextTokens <= this.chunkSize) {
+                currentText = `${currentText}\n${chunk.text}`;
+                currentTokens = nextTokens;
+                continue;
+            }
+            merged.push({ text: currentText, tokenCount: currentTokens });
+            currentText = chunk.text;
+            currentTokens = chunk.tokenCount;
+        }
+        if (currentTokens > 0) {
+            merged.push({ text: currentText, tokenCount: currentTokens });
+        }
+        return merged;
+    }
+}
+exports.CodeChunker = CodeChunker;
+CodeChunker.treeSitterInitialized = false;
+CodeChunker.parserCache = new Map();