feat: initial code

Author: hustcc

.github/workflows/build.yml

@@ -0,0 +1,34 @@
+name: Build
+
+on:
+  push:
+    branches: ['**']
+  pull_request:
+    branches: ['**']
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Setup pnpm
+        uses: pnpm/action-setup@v4
+        with:
+          version: 9
+
+      - name: Setup Node.js
+        uses: actions/setup-node@v4
+        with:
+          node-version: '20'
+          cache: 'pnpm'
+
+      - name: Install dependencies
+        run: pnpm install
+
+      - name: Run build
+        run: pnpm build
+
+      - name: Run tests
+        run: pnpm test

.gitignore

@@ -0,0 +1,13 @@
+node_modules/
+coverage/
+dist/
+
+*.log
+.DS_Store
+.test-tmp
+.cache
+
+package-lock.json
+yarn.lock
+pnpm-lock.yaml
+pnpm-workspace.yaml

package.json

@@ -0,0 +1,36 @@
+{
+  "name": "@antv/context",
+  "version": "0.1.0",
+  "description": "A local context retrieval library that enables semantic search over your documentation. It loads documents (Markdown, JSON, Text), vectorizes them using [Transformers.js](https://huggingface.co/transformers.js), and stores vectors locally in `.zvec` files for fast semantic querying. ",
+  "main": "dist/index.js",
+  "types": "dist/index.d.ts",
+  "type": "module",
+  "scripts": {
+    "build": "tsc",
+    "test": "vitest run --coverage"
+  },
+  "keywords": [
+    "context",
+    "semantic-search",
+    "embedding",
+    "vector",
+    "transformers",
+    "huggingface",
+    "zvec",
+    "antv"
+  ],
+  "license": "MIT",
+  "dependencies": {
+    "@huggingface/transformers": "^4.2.0",
+    "@zvec/zvec": "^0.5.0",
+    "glob": "^13.0.6",
+    "gray-matter": "^4.0.3"
+  },
+  "devDependencies": {
+    "@types/node": "^20.0.0",
+    "@vitest/coverage-v8": "^3.2.6",
+    "typescript": "^5.0.0",
+    "vitest": "^3.2.6"
+  },
+  "repository": "git@github.com:antvis/context.git"
+}

src/context.ts

@@ -0,0 +1,92 @@
+import * as path from 'path';
+import * as fs from 'fs';
+import { glob } from 'glob';
+import { ContextOptions, QueryOptions, QueryResult } from './types';
+import { TransformerVectorizer } from './vectorizer/transformer';
+import { ZVecStore } from './storage/zvec-store';
+import { Loader, MarkdownLoader, JsonLoader, TextLoader } from './loaders';
+
+export class Context {
+  private readonly vectorsDir: string;
+  private readonly model: string;
+  private readonly vectorizer: TransformerVectorizer;
+  private readonly stores: Map<string, ZVecStore> = new Map();
+  private readonly loaders: Loader[];
+
+  private constructor(options: ContextOptions) {
+    this.vectorsDir = options.vectorsDir;
+    this.model = options.model || 'sentence-transformers/all-MiniLM-L6-v2';
+    this.vectorizer = new TransformerVectorizer(this.model);
+    this.loaders = [
+      new MarkdownLoader(),
+      new JsonLoader(),
+      new TextLoader(),
+    ];
+  }
+
+  static async create(options: ContextOptions): Promise<Context> {
+    const ctx = new Context(options);
+    await ctx.vectorizer.initialize();
+
+    // Ensure vectors directory exists
+    if (!fs.existsSync(options.vectorsDir)) {
+      fs.mkdirSync(options.vectorsDir, { recursive: true });
+    }
+
+    return ctx;
+  }
+
+  private getLoader(filePath: string): Loader | undefined {
+    return this.loaders.find((loader) => loader.canHandle(filePath));
+  }
+
+  private getStoreFilePath(library: string): string {
+    return path.join(this.vectorsDir, `${library}.zvec`);
+  }
+
+  private async getOrCreateStore(library: string): Promise<ZVecStore> {
+    if (this.stores.has(library)) {
+      return this.stores.get(library)!;
+    }
+
+    const filePath = this.getStoreFilePath(library);
+    const store = await ZVecStore.create(filePath);
+
+    this.stores.set(library, store);
+    return store;
+  }
+
+  async load(library: string, pattern: string | string[]): Promise<void> {
+    const patterns = Array.isArray(pattern) ? pattern : [pattern];
+    const files = await glob(patterns, { absolute: true });
+    const store = await this.getOrCreateStore(library);
+
+    for (const filePath of files) {
+      const loader = this.getLoader(filePath);
+      if (!loader) continue;
+
+      const doc = await loader.load(filePath);
+      const vector = await this.vectorizer.embed(doc.content);
+      store.add(doc.id, vector, doc.content, doc.meta);
+    }
+
+    // Save to disk
+    store.save();
+  }
+
+  async query(text: string, options: QueryOptions): Promise<QueryResult[]> {
+    const store = await this.getOrCreateStore(options.library);
+    const vector = await this.vectorizer.embed(text);
+    const searchResults = store.search(vector, options.topK || 5);
+
+    return searchResults.map((result) => {
+      const doc = store.getDoc(result.id)!;
+      return {
+        id: result.id,
+        content: doc.content,
+        score: result.score,
+        meta: doc.meta,
+      };
+    });
+  }
+}

src/index.ts

@@ -0,0 +1,2 @@
+export { Context } from './context';
+export * from './types';

src/loaders/base.ts

@@ -0,0 +1,6 @@
+import { Document } from '../types';
+
+export interface Loader {
+  canHandle(filePath: string): boolean;
+  load(filePath: string): Promise<Document>;
+}

src/loaders/index.ts

@@ -0,0 +1,4 @@
+export { Loader } from './base';
+export { MarkdownLoader } from './markdown';
+export { JsonLoader } from './json';
+export { TextLoader } from './text';

src/loaders/json.ts

@@ -0,0 +1,20 @@
+import * as fs from 'fs/promises';
+import * as path from 'path';
+import { Loader } from './base';
+import { Document } from '../types';
+
+export class JsonLoader implements Loader {
+  canHandle(filePath: string): boolean {
+    return filePath.endsWith('.json');
+  }
+
+  async load(filePath: string): Promise<Document> {
+    const content = await fs.readFile(filePath, 'utf-8');
+    const data = JSON.parse(content);
+
+    return {
+      id: path.basename(filePath),
+      content: typeof data === 'string' ? data : JSON.stringify(data, null, 2),
+    };
+  }
+}

src/loaders/markdown.ts

@@ -0,0 +1,25 @@
+import * as fs from 'fs/promises';
+import * as path from 'path';
+import matter from 'gray-matter';
+import { Loader } from './base';
+import { Document } from '../types';
+
+export class MarkdownLoader implements Loader {
+  canHandle(filePath: string): boolean {
+    return filePath.endsWith('.md') || filePath.endsWith('.markdown');
+  }
+
+  async load(filePath: string): Promise<Document> {
+    const content = await fs.readFile(filePath, 'utf-8');
+    const { data: meta, content: body } = matter(content);
+
+    // 使用文件名作为 ID，避免路径中的特殊字符
+    const id = path.basename(filePath);
+
+    return {
+      id,
+      content: body.trim(),
+      meta,
+    };
+  }
+}

src/loaders/text.ts

@@ -0,0 +1,19 @@
+import * as fs from 'fs/promises';
+import * as path from 'path';
+import { Loader } from './base';
+import { Document } from '../types';
+
+export class TextLoader implements Loader {
+  canHandle(filePath: string): boolean {
+    return filePath.endsWith('.txt');
+  }
+
+  async load(filePath: string): Promise<Document> {
+    const content = await fs.readFile(filePath, 'utf-8');
+
+    return {
+      id: path.basename(filePath),
+      content: content.trim(),
+    };
+  }
+}