chore: code optimization

Author: 福晋

.gitignore

@@ -11,3 +11,4 @@ package-lock.json
 yarn.lock
 pnpm-lock.yaml
 pnpm-workspace.yaml
+.vscode/

src/context.ts

@@ -1,22 +1,20 @@
 import * as path from 'path';
 import * as fs from 'fs';
 import { glob } from 'glob';
-import { ContextOptions, QueryOptions, QueryResult } from './types';
+import { ContextOptions, QueryOptions, QueryResult, Loader } from './types';
 import { TransformerVectorizer } from './vectorizer/transformer';
 import { ZVecStore } from './storage/zvec-store';
-import { Loader, MarkdownLoader, JsonLoader, TextLoader } from './loaders';
+import { 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.vectorizer = new TransformerVectorizer(options.model);
     this.loaders = [
       new MarkdownLoader(),
       new JsonLoader(),
@@ -77,7 +75,7 @@ export class Context {
   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);
+    const searchResults = await store.search(vector, options.topK || 5);
 
     return searchResults.map((result) => {
       const doc = store.getDoc(result.id)!;

src/loaders/base.ts

@@ -1,6 +1,6 @@
-import { Document } from '../types';
+import * as path from 'path';
 
-export interface Loader {
-  canHandle(filePath: string): boolean;
-  load(filePath: string): Promise<Document>;
+/** Generate a document ID from a file path (uses the filename to avoid special characters in paths). */
+export function generateDocId(filePath: string): string {
+  return path.basename(filePath);
 }

src/loaders/index.ts

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

src/loaders/json.ts

@@ -1,7 +1,6 @@
 import * as fs from 'fs/promises';
-import * as path from 'path';
-import { Loader } from './base';
-import { Document } from '../types';
+import { Loader, Document } from '../types';
+import { generateDocId } from './base';
 
 export class JsonLoader implements Loader {
   canHandle(filePath: string): boolean {
@@ -13,7 +12,7 @@ export class JsonLoader implements Loader {
     const data = JSON.parse(content);
 
     return {
-      id: path.basename(filePath),
+      id: generateDocId(filePath),
       content: typeof data === 'string' ? data : JSON.stringify(data, null, 2),
     };
   }

src/loaders/markdown.ts

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

src/loaders/text.ts

@@ -1,7 +1,6 @@
 import * as fs from 'fs/promises';
-import * as path from 'path';
-import { Loader } from './base';
-import { Document } from '../types';
+import { Loader, Document } from '../types';
+import { generateDocId } from './base';
 
 export class TextLoader implements Loader {
   canHandle(filePath: string): boolean {
@@ -12,7 +11,7 @@ export class TextLoader implements Loader {
     const content = await fs.readFile(filePath, 'utf-8');
 
     return {
-      id: path.basename(filePath),
+      id: generateDocId(filePath),
       content: content.trim(),
     };
   }

src/storage/zvec-store.ts

@@ -15,13 +15,11 @@ interface DocData {
 
 export class ZVecStore {
   private collection: ZVecCollection;
-  private dimension: number;
   private filePath!: string;
   private docs: Map<string, DocData> = new Map();
 
-  constructor(collection: ZVecCollection, dimension: number) {
+  constructor(collection: ZVecCollection) {
     this.collection = collection;
-    this.dimension = dimension;
   }
 
   static async create(filePath: string, dimension: number = 384): Promise<ZVecStore> {
@@ -43,14 +41,10 @@ export class ZVecStore {
     if (fs.existsSync(filePath)) {
       collection = ZVecOpen(filePath);
     } else {
-      const dir = path.dirname(filePath);
-      if (!fs.existsSync(dir)) {
-        fs.mkdirSync(dir, { recursive: true });
-      }
       collection = ZVecCreateAndOpen(filePath, schema);
     }
 
-    const store = new ZVecStore(collection, dimension);
+    const store = new ZVecStore(collection);
     store.filePath = filePath;
     await store.loadMeta();
     return store;
@@ -75,8 +69,8 @@ export class ZVecStore {
     this.docs.set(id, { content, meta });
   }
 
-  search(queryVector: number[], topK: number): Array<{ id: string; score: number }> {
-    const results = this.collection.querySync({
+  async search(queryVector: number[], topK: number): Promise<Array<{ id: string; score: number }>> {
+    const results = await this.collection.query({
       fieldName: 'embedding',
       vector: queryVector,
       topk: topK,
@@ -93,17 +87,4 @@ export class ZVecStore {
     this.saveMeta();
   }
 
-  close(): void {
-    this.saveMeta();
-    this.collection.closeSync();
-  }
-
-  clear(): void {
-    const ids = Array.from(this.docs.keys());
-    if (ids.length > 0) {
-      this.collection.deleteSync(ids);
-    }
-    this.docs.clear();
-    this.saveMeta();
-  }
 }

src/types.ts

@@ -10,6 +10,14 @@ export interface Document {
   meta?: Record<string, unknown>;
 }
 
+/**
+ * Document loader interface
+ */
+export interface Loader {
+  canHandle(filePath: string): boolean;
+  load(filePath: string): Promise<Document>;
+}
+
 /**
  * Context initialization options
  */

src/vectorizer/transformer.ts

@@ -1,39 +1,15 @@
 import { pipeline, env, FeatureExtractionPipeline } from '@huggingface/transformers';
-import * as path from 'path';
-import * as fs from 'fs';
 
 export class TransformerVectorizer {
   private pipeline: FeatureExtractionPipeline | null = null;
 
-  constructor(
-    private readonly model: string = 'sentence-transformers/all-MiniLM-L6-v2',
-    private readonly dtype?: 'fp32' | 'fp16' | 'q8' | 'q4' | 'q4f16'
-  ) {}
+  constructor(private readonly model: string = 'sentence-transformers/all-MiniLM-L6-v2') {}
 
-  /**
-   * 配置模型下载镜像
-   * 通过环境变量设置: HF_ENDPOINT 或 HF_MIRROR
-   * 常用镜像: https://hf-mirror.com
-   */
-  private configureMirror(): void {
+  async initialize(): Promise<void> {
     const mirror = process.env.HF_ENDPOINT || process.env.HF_MIRROR;
     if (mirror) env.remoteHost = mirror;
-  }
-
-  async initialize(): Promise<void> {
-    this.configureMirror();
-
-    // 查找本地缓存模型
-    const cacheDir = path.join(process.env.HOME || '', '.cache/huggingface/models');
-    let modelPath = this.model;
-
-    if (fs.existsSync(cacheDir)) {
-      const modelName = this.model.replace('sentence-transformers--', '').split('-').slice(0, 2).join('-');
-      const localModel = fs.readdirSync(cacheDir).find(d => d.includes(modelName));
-      if (localModel) modelPath = path.join(cacheDir, localModel);
-    }
 
-    this.pipeline = await pipeline('feature-extraction', modelPath, this.dtype ? { dtype: this.dtype } : undefined) as FeatureExtractionPipeline;
+    this.pipeline = await pipeline('feature-extraction', this.model) as FeatureExtractionPipeline;
   }
 
   async embed(text: string): Promise<number[]> {