Merge pull request #26 from appwrite/feat-embeddings

Author: ChiragAgg5k

package-lock.json

@@ -21,15 +21,16 @@
   },
   "license": "MIT",
   "dependencies": {
+    "@huggingface/transformers": "^3.5.2",
     "archiver": "^7.0.1",
     "chokidar": "^4.0.3",
     "eslint": "^9.27.0",
     "ignore": "^7.0.4",
     "mime-types": "^3.0.1",
+    "node-appwrite": "17.0.0",
     "node-pty": "^1.0.0",
     "prettier": "^3.5.3",
-    "ws": "^8.18.2",
-    "node-appwrite": "17.0.0"
+    "ws": "^8.18.2"
   },
   "devDependencies": {
     "@types/archiver": "^6.0.3",

package.json

@@ -0,0 +1,41 @@
+export interface EmbeddingConfig {
+  [key: string]: any;
+}
+
+export abstract class EmbeddingAdapter {
+  protected config: EmbeddingConfig;
+  public isInitializing: boolean = false;
+
+  constructor(config: EmbeddingConfig = {}) {
+    this.config = config;
+  }
+
+  /**
+   * Initialize the embedding model/service
+   */
+  abstract initialize(): Promise<void>;
+
+  /**
+   * Generate embeddings for the given text
+   * @param text The text to generate embeddings for
+   * @returns Promise<number[]> The embedding vector
+   */
+  abstract generateEmbedding(text: string): Promise<number[]>;
+
+  /**
+   * Get the name/identifier of this adapter
+   */
+  abstract getName(): string;
+
+  /**
+   * Check if the adapter is initialized and ready to use
+   */
+  abstract isInitialized(): boolean;
+
+  /**
+   * Clean up resources when done
+   */
+  async cleanup(): Promise<void> {
+    // Default implementation - can be overridden by specific adapters
+  }
+}

src/adapters/embeddings.ts

@@ -0,0 +1,74 @@
+import {
+  pipeline,
+  type FeatureExtractionPipeline,
+} from "@huggingface/transformers";
+import { EmbeddingAdapter, EmbeddingConfig } from "../embeddings";
+
+export interface HuggingFaceConfig extends EmbeddingConfig {
+  modelName?: string;
+  pooling?: "mean" | "none" | "cls" | undefined;
+  normalize?: boolean;
+}
+
+export class HuggingFaceEmbeddingAdapter extends EmbeddingAdapter {
+  private pipeline: FeatureExtractionPipeline | null = null;
+  private modelName: string;
+  private pooling: "mean" | "none" | "cls" | undefined;
+  private normalize: boolean;
+
+  constructor(config: HuggingFaceConfig = {}) {
+    super(config);
+    this.modelName = config.modelName || "jinaai/jina-embeddings-v2-base-code";
+    this.pooling = config.pooling || "mean";
+    this.normalize = config.normalize !== false; // default to true
+  }
+
+  async initialize(): Promise<void> {
+    if (!this.pipeline && !this.isInitializing) {
+      this.isInitializing = true;
+      console.log(`[HuggingFace] Initializing model: ${this.modelName}...`);
+      try {
+        this.pipeline = await pipeline("feature-extraction", this.modelName);
+        console.log("[HuggingFace] Model initialized successfully");
+      } catch (error) {
+        console.error(`[HuggingFace] Error initializing model: ${error}`);
+        throw error;
+      }
+    }
+    this.isInitializing = false;
+  }
+
+  async generateEmbedding(text: string): Promise<number[]> {
+    if (!this.pipeline) {
+      throw new Error(
+        "HuggingFace adapter not initialized. Call initialize() first.",
+      );
+    }
+
+    try {
+      const output = await this.pipeline(text, {
+        pooling: this.pooling,
+        normalize: this.normalize,
+      });
+
+      // Convert tensor to array if needed
+      return Array.from(output.data);
+    } catch (error) {
+      console.error(`[HuggingFace] Error generating embedding: ${error}`);
+      throw error;
+    }
+  }
+
+  getName(): string {
+    return `HuggingFace (${this.modelName})`;
+  }
+
+  isInitialized(): boolean {
+    return this.pipeline !== null;
+  }
+
+  async cleanup(): Promise<void> {
+    this.pipeline = null;
+    console.log("[HuggingFace] Adapter cleaned up");
+  }
+}

src/adapters/embeddings/huggingface.ts

@@ -0,0 +1,95 @@
+import { EmbeddingAdapter, EmbeddingConfig } from "../embeddings";
+
+export interface OpenAIConfig extends EmbeddingConfig {
+  apiKey: string;
+  model?: string;
+  baseURL?: string;
+}
+
+export class OpenAIEmbeddingAdapter extends EmbeddingAdapter {
+  private apiKey: string;
+  private model: string;
+  private baseURL: string;
+  private initialized: boolean = false;
+
+  constructor(config: OpenAIConfig) {
+    super(config);
+    if (!config.apiKey) {
+      throw new Error("OpenAI API key is required");
+    }
+    this.apiKey = config.apiKey;
+    this.model = config.model || "text-embedding-3-small";
+    this.baseURL = config.baseURL || "https://api.openai.com/v1";
+  }
+
+  async initialize(): Promise<void> {
+    console.log(`[OpenAI] Initializing with model: ${this.model}...`);
+    this.isInitializing = true;
+    if (!this.apiKey.startsWith("sk-")) {
+      console.warn(
+        "[OpenAI] API key doesn't start with 'sk-', this might cause issues",
+      );
+    }
+    this.initialized = true;
+    this.isInitializing = false;
+    console.log("[OpenAI] Adapter initialized successfully");
+  }
+
+  async generateEmbedding(text: string): Promise<number[]> {
+    if (!this.initialized) {
+      throw new Error(
+        "OpenAI adapter not initialized. Call initialize() first.",
+      );
+    }
+
+    if (this.isInitializing) {
+      throw new Error(
+        "OpenAI adapter is initializing. Please wait a moment and try again.",
+      );
+    }
+
+    try {
+      const response = await fetch(`${this.baseURL}/embeddings`, {
+        method: "POST",
+        headers: {
+          Authorization: `Bearer ${this.apiKey}`,
+          "Content-Type": "application/json",
+        },
+        body: JSON.stringify({
+          input: text,
+          model: this.model,
+        }),
+      });
+
+      if (!response.ok) {
+        throw new Error(
+          `OpenAI API error: ${response.status} ${response.statusText}`,
+        );
+      }
+
+      const data = (await response.json()) as any;
+
+      if (!data.data || !data.data[0] || !data.data[0].embedding) {
+        throw new Error("Invalid response format from OpenAI API");
+      }
+
+      return data.data[0].embedding;
+    } catch (error) {
+      console.error(`[OpenAI] Error generating embedding: ${error}`);
+      throw error;
+    }
+  }
+
+  getName(): string {
+    return `OpenAI (${this.model})`;
+  }
+
+  isInitialized(): boolean {
+    return this.initialized;
+  }
+
+  async cleanup(): Promise<void> {
+    this.initialized = false;
+    console.log("[OpenAI] Adapter cleaned up");
+  }
+}

src/adapters/embeddings/openai.ts

@@ -0,0 +1,6 @@
+export { EmbeddingAdapter } from "./embeddings";
+export type { EmbeddingConfig } from "./embeddings";
+export { HuggingFaceEmbeddingAdapter } from "./embeddings/huggingface";
+export type { HuggingFaceConfig } from "./embeddings/huggingface";
+export { OpenAIEmbeddingAdapter } from "./embeddings/openai";
+export type { OpenAIConfig } from "./embeddings/openai";

src/adapters/index.ts

@@ -1,7 +1,8 @@
+export { Appwrite } from "./services/appwrite";
 export { Code } from "./services/code";
+export { Embeddings } from "./services/embeddings";
 export { Filesystem } from "./services/filesystem";
 export { Git } from "./services/git";
 export { System } from "./services/system";
 export { Terminal } from "./services/terminal";
 export { Synapse } from "./synapse";
-export { Appwrite } from "./services/appwrite";