feat: Support Fireworks embedding for ingest and tool

src/agent-tools/searchSpecSections.ts

@@ -4,7 +4,7 @@
  */
 
 import type { Table } from "@lancedb/lancedb";
-import type { OllamaEmbeddings } from "@langchain/ollama";
+import type { Embeddings } from "@langchain/core/embeddings";
 import { z } from "zod";
 
 // #region Zod schemas (not exported)
@@ -62,12 +62,12 @@ export type SearchSpecInput = z.infer<typeof inputSchema>;
  * Creates the search spec sections tool function.
  * Performs semantic vector search to find relevant spec sections.
  * @param table - LanceDB table containing spec vectors
- * @param embeddings - Ollama embeddings instance
+ * @param embeddings - Embeddings instance (Ollama or Fireworks)
  * @returns Function that performs the search and returns structured output
  */
 export function createSearchSpecSectionsTool(
   table: Table,
-  embeddings: OllamaEmbeddings,
+  embeddings: Embeddings,
 ) {
   return async ({ query }: SearchSpecInput): Promise<SearchSpecOutput> => {
     // Generate embedding for the query

src/constants.ts

@@ -4,5 +4,11 @@ export const CODE_DIR = "./engine262/src";
 export const GRAPH_FILE = "./graphology/graph.json";
 
 // Model configurations
-export const EMBEDDING_MODEL = "qwen3-embedding:0.6b";
+export const OLLAMA_EMBEDDING_MODEL = "qwen3-embedding:0.6b";
 export const RERANKER_MODEL = "dengcao/Qwen3-Reranker-0.6B:Q8_0";
+
+// Embedding provider configuration
+export const EMBEDDING_PROVIDER =
+  process.env.ASK262_EMBEDDING_PROVIDER ?? "ollama";
+export const FIREWORKS_EMBEDDING_MODEL = "fireworks/qwen3-embedding-8b";
+export const FIREWORKS_BASE_URL = "https://api.fireworks.ai/inference/v1";

src/lib/embeddings-factory.ts

@@ -0,0 +1,76 @@
+import type { Embeddings } from "@langchain/core/embeddings";
+import { OllamaEmbeddings } from "@langchain/ollama";
+import {
+  EMBEDDING_PROVIDER,
+  FIREWORKS_BASE_URL,
+  FIREWORKS_EMBEDDING_MODEL,
+  OLLAMA_EMBEDDING_MODEL,
+} from "../constants.js";
+import { FireworksEmbeddings } from "./fireworks-embeddings.js";
+
+/**
+ * Type for supported embedding providers.
+ */
+export type EmbeddingProvider = "ollama" | "fireworks";
+
+/**
+ * Create an embeddings instance based on the configured provider.
+ *
+ * @param provider - The embedding provider to use. Defaults to EMBEDDING_PROVIDER env var or "ollama"
+ * @returns Embeddings instance (OllamaEmbeddings or FireworksEmbeddings)
+ * @throws Error if provider is invalid or required credentials are missing
+ *
+ * @example
+ * ```typescript
+ * // Use default provider from env
+ * const embeddings = createEmbeddings();
+ *
+ * // Explicitly use Fireworks
+ * const embeddings = createEmbeddings("fireworks");
+ *
+ * // Explicitly use Ollama
+ * const embeddings = createEmbeddings("ollama");
+ * ```
+ */
+export function createEmbeddings(provider?: EmbeddingProvider): Embeddings {
+  const selectedProvider =
+    provider ?? (EMBEDDING_PROVIDER as EmbeddingProvider);
+
+  switch (selectedProvider) {
+    case "ollama": {
+      console.log("[Embeddings] Using Ollama provider");
+      return new OllamaEmbeddings({
+        model: OLLAMA_EMBEDDING_MODEL,
+        baseUrl: process.env.OLLAMA_HOST,
+      });
+    }
+
+    case "fireworks": {
+      const apiKey = process.env.FIREWORKS_API_KEY;
+      if (!apiKey) {
+        throw new Error("FIREWORKS_API_KEY environment variable is required");
+      }
+      console.log("[Embeddings] Using Fireworks provider");
+      return new FireworksEmbeddings({
+        apiKey,
+        modelName: FIREWORKS_EMBEDDING_MODEL,
+        baseUrl: FIREWORKS_BASE_URL,
+      });
+    }
+
+    default: {
+      throw new Error(
+        `Unknown embedding provider: ${selectedProvider}. Use 'ollama' or 'fireworks'.`,
+      );
+    }
+  }
+}
+
+/**
+ * Get the currently configured embedding provider.
+ *
+ * @returns The active provider name
+ */
+export function getEmbeddingProvider(): EmbeddingProvider {
+  return (EMBEDDING_PROVIDER as EmbeddingProvider) ?? "ollama";
+}

src/lib/fireworks-embeddings.ts

@@ -0,0 +1,199 @@
+import { Embeddings, type EmbeddingsParams } from "@langchain/core/embeddings";
+
+/**
+ * Interface for FireworksEmbeddings parameters.
+ */
+export interface FireworksEmbeddingsParams extends EmbeddingsParams {
+  /**
+   * API key for Fireworks.ai
+   * Can also be set via FIREWORKS_API_KEY env var
+   */
+  apiKey?: string;
+
+  /**
+   * Model name to use
+   * @default "fireworks/qwen3-embedding-8b"
+   */
+  modelName?: string;
+
+  /**
+   * Base URL for Fireworks API
+   * @default "https://api.fireworks.ai/inference/v1"
+   */
+  baseUrl?: string;
+
+  /**
+   * Maximum number of documents to embed in a single request
+   * @default 100
+   */
+  batchSize?: number;
+
+  /**
+   * Maximum retries for rate limit errors
+   * @default 3
+   */
+  maxRetries?: number;
+
+  /**
+   * Initial wait time in ms for rate limit retries (doubles each retry)
+   * @default 1000
+   */
+  initialRetryDelayMs?: number;
+}
+
+/**
+ * Fireworks.ai embeddings implementation for LangChain.
+ * Uses the qwen3-embedding-8b model via Fireworks inference API.
+ *
+ * @example
+ * ```typescript
+ * const embeddings = new FireworksEmbeddings({
+ *   apiKey: process.env.FIREWORKS_API_KEY,
+ *   modelName: "fireworks/qwen3-embedding-8b",
+ * });
+ *
+ * const vectors = await embeddings.embedDocuments(["hello", "world"]);
+ * ```
+ */
+export class FireworksEmbeddings extends Embeddings {
+  private apiKey: string;
+  private modelName: string;
+  private baseUrl: string;
+  private batchSize: number;
+  private maxRetries: number;
+  private initialRetryDelayMs: number;
+
+  constructor(params?: FireworksEmbeddingsParams) {
+    super(params ?? {});
+
+    this.apiKey = params?.apiKey ?? process.env.FIREWORKS_API_KEY ?? "";
+    if (!this.apiKey) {
+      throw new Error(
+        "Fireworks API key is required. Set FIREWORKS_API_KEY env var or pass apiKey parameter.",
+      );
+    }
+
+    this.modelName = params?.modelName ?? "fireworks/qwen3-embedding-8b";
+    this.baseUrl = params?.baseUrl ?? "https://api.fireworks.ai/inference/v1";
+    this.batchSize = params?.batchSize ?? 100;
+    this.maxRetries = params?.maxRetries ?? 3;
+    this.initialRetryDelayMs = params?.initialRetryDelayMs ?? 1000;
+  }
+
+  /**
+   * Embed a single document (query).
+   * Uses the embeddings endpoint optimized for search queries.
+   */
+  async embedQuery(document: string): Promise<number[]> {
+    const vectors = await this.embedDocuments([document]);
+    return vectors[0];
+  }
+
+  /**
+   * Embed multiple documents in batches with rate limit handling.
+   */
+  async embedDocuments(documents: string[]): Promise<number[][]> {
+    if (documents.length === 0) {
+      return [];
+    }
+
+    const allEmbeddings: number[][] = [];
+
+    // Process in batches
+    for (let i = 0; i < documents.length; i += this.batchSize) {
+      const batch = documents.slice(i, i + this.batchSize);
+      const batchEmbeddings = await this.embedBatchWithRetry(batch);
+      allEmbeddings.push(...batchEmbeddings);
+    }
+
+    return allEmbeddings;
+  }
+
+  /**
+   * Embed a single batch with retry logic for rate limits.
+   */
+  private async embedBatchWithRetry(
+    documents: string[],
+    attempt = 1,
+  ): Promise<number[][]> {
+    try {
+      return await this.embedBatch(documents);
+    } catch (error) {
+      // Check if it's a rate limit error (429)
+      const isRateLimit =
+        error instanceof Error &&
+        (error.message.includes("429") || error.message.includes("rate limit"));
+
+      if (isRateLimit && attempt < this.maxRetries) {
+        const delay = this.initialRetryDelayMs * 2 ** (attempt - 1);
+        console.error(
+          `[Fireworks] Rate limit hit. Waiting ${delay}ms before retry ${attempt}/${this.maxRetries}...`,
+        );
+        await sleep(delay);
+        return this.embedBatchWithRetry(documents, attempt + 1);
+      }
+
+      // Fail fast for other errors or if retries exhausted
+      throw error;
+    }
+  }
+
+  /**
+   * Make the actual API call to Fireworks for embeddings.
+   */
+  private async embedBatch(documents: string[]): Promise<number[][]> {
+    const url = `${this.baseUrl}/embeddings`;
+
+    const response = await fetch(url, {
+      method: "POST",
+      headers: {
+        Authorization: `Bearer ${this.apiKey}`,
+        "Content-Type": "application/json",
+        Accept: "application/json",
+      },
+      body: JSON.stringify({
+        model: this.modelName,
+        input: documents,
+      }),
+    });
+
+    if (!response.ok) {
+      const errorText = await response.text();
+      throw new Error(
+        `Fireworks API error: ${response.status} ${response.statusText} - ${errorText}`,
+      );
+    }
+
+    const data = (await response.json()) as FireworksEmbeddingResponse;
+
+    // Extract embeddings from response
+    // Fireworks returns embeddings in the same order as input
+    const embeddings = data.data.map((item) => item.embedding);
+
+    return embeddings;
+  }
+}
+
+/**
+ * Sleep utility for rate limit retries.
+ */
+function sleep(ms: number): Promise<void> {
+  return new Promise((resolve) => setTimeout(resolve, ms));
+}
+
+/**
+ * Fireworks API response structure for embeddings.
+ */
+interface FireworksEmbeddingResponse {
+  object: "list";
+  data: Array<{
+    object: "embedding";
+    embedding: number[];
+    index: number;
+  }>;
+  model: string;
+  usage: {
+    prompt_tokens: number;
+    total_tokens: number;
+  };
+}

src/mcp-server-http.ts

@@ -10,7 +10,6 @@ import path from "node:path";
 import { fileURLToPath } from "node:url";
 import { serve } from "@hono/node-server";
 import * as lancedbSdk from "@lancedb/lancedb";
-import { OllamaEmbeddings } from "@langchain/ollama";
 import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
 import { WebStandardStreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/webStandardStreamableHttp.js";
 import { Hono } from "hono";
@@ -33,21 +32,16 @@ import {
   sectionContentToolMetadata,
   sectionContentToolName,
 } from "./agent-tools/index.js";
-import {
-  EMBEDDING_MODEL,
-  STORAGE_DIR as STORAGE_DIR_REL,
-} from "./constants.js";
+import { STORAGE_DIR as STORAGE_DIR_REL } from "./constants.js";
+import { createEmbeddings } from "./lib/embeddings-factory.js";
 
 // Resolve storage path relative to this script's directory
 const __filename = fileURLToPath(import.meta.url);
 const __dirname = path.dirname(__filename);
 const STORAGE_DIR = path.resolve(__dirname, "..", STORAGE_DIR_REL);
 
-// Initialize embeddings
-const embeddings = new OllamaEmbeddings({
-  model: EMBEDDING_MODEL,
-  baseUrl: process.env.OLLAMA_HOST,
-});
+// Initialize embeddings based on ASK262_EMBEDDING_PROVIDER env var
+const embeddings = createEmbeddings();
 
 // Server port (default: 3000)
 const PORT = Number(process.env.ASK262_PORT) || 3000;

src/mcp-server.ts

@@ -8,7 +8,6 @@
 import path from "node:path";
 import { fileURLToPath } from "node:url";
 import * as lancedbSdk from "@lancedb/lancedb";
-import { OllamaEmbeddings } from "@langchain/ollama";
 import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
 import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
 import { z } from "zod";
@@ -35,10 +34,8 @@ import {
   sectionContentToolMetadata,
   sectionContentToolName,
 } from "./agent-tools/index.js";
-import {
-  EMBEDDING_MODEL,
-  STORAGE_DIR as STORAGE_DIR_REL,
-} from "./constants.js";
+import { STORAGE_DIR as STORAGE_DIR_REL } from "./constants.js";
+import { createEmbeddings } from "./lib/embeddings-factory.js";
 
 // Resolve storage path relative to this script's directory
 const __filename = fileURLToPath(import.meta.url);
@@ -87,12 +84,8 @@ export interface SearchSpecMCPOutput extends McpToolOutputBase {
 
 // #endregion
 
-// Initialize embeddings
-// OLLAMA_HOST env var is optional - @langchain/ollama defaults to http://localhost:11434
-const embeddings = new OllamaEmbeddings({
-  model: EMBEDDING_MODEL,
-  baseUrl: process.env.OLLAMA_HOST,
-});
+// Initialize embeddings based on ASK262_EMBEDDING_PROVIDER env var
+const embeddings = createEmbeddings();
 
 export async function main() {
   // Connect to LanceDB

src/setup/ingest.ts

@@ -4,18 +4,19 @@ import readline from "node:readline";
 import * as lancedbSdk from "@lancedb/lancedb";
 import { Index } from "@lancedb/lancedb";
 import { Document } from "@langchain/core/documents";
-import { OllamaEmbeddings } from "@langchain/ollama";
+import type { Embeddings } from "@langchain/core/embeddings";
 import * as cheerio from "cheerio";
+import { Command } from "commander";
 import { glob } from "glob";
 import ora from "ora";
-import { EMBEDDING_MODEL, SPEC_DIR, STORAGE_DIR } from "../constants.js";
+import { SPEC_DIR, STORAGE_DIR } from "../constants.js";
+import {
+  createEmbeddings,
+  type EmbeddingProvider,
+} from "../lib/embeddings-factory.js";
 import { HTMLTextSplitter } from "./text-splitters/index.js";
 import { formatForIngestion } from "./utils/formatHTMLForIngestion.js";
 
-const embeddings = new OllamaEmbeddings({
-  model: EMBEDDING_MODEL,
-});
-
 const htmlSplitter = new HTMLTextSplitter({
   chunkSize: 8192,
   maxChunkSize: 12288,
@@ -48,6 +49,7 @@ interface ChunkInfo {
 
 async function generateEmbeddingsWithProgress(
   documents: Document[],
+  embeddings: Embeddings,
 ): Promise<number[][]> {
   const total = documents.length;
   const vectors: number[][] = [];
@@ -275,6 +277,37 @@ async function buildSpecDocuments(): Promise<Document[]> {
 }
 
 async function main() {
+  // Parse command line arguments using Commander
+  const program = new Command()
+    .name("ingest")
+    .description("Ingest ECMAScript specification into vector database")
+    .version("1.0.0")
+    .option(
+      "-p, --embedding-provider <provider>",
+      "Embedding provider to use (ollama or fireworks)",
+      process.env.ASK262_EMBEDDING_PROVIDER ?? "ollama",
+    )
+    .option(
+      "-y, --yes",
+      "Automatically overwrite existing vector store without prompting",
+      false,
+    )
+    .parse();
+
+  const options = program.opts();
+  const provider = options.embeddingProvider as EmbeddingProvider;
+
+  // Validate provider
+  if (provider !== "ollama" && provider !== "fireworks") {
+    console.error(`Error: Unknown embedding provider "${provider}"`);
+    console.error('Use "ollama" or "fireworks"');
+    process.exit(1);
+  }
+
+  // Create embeddings instance based on provider
+  console.log(`Initializing embeddings provider: ${provider}`);
+  const embeddings = createEmbeddings(provider);
+
   console.log("Building specification documents...");
   const specDocs = await buildSpecDocuments();
   console.log(`Built ${specDocs.length} specification documents.`);
@@ -306,9 +339,12 @@ async function main() {
   }
 
   if (tableExists) {
-    const shouldOverwrite = await askUser(
-      "Do you want to overwrite the existing vector store?",
-    );
+    let shouldOverwrite = options.yes;
+    if (!shouldOverwrite) {
+      shouldOverwrite = await askUser(
+        "Do you want to overwrite the existing vector store?",
+      );
+    }
     if (!shouldOverwrite) {
       console.log("Ingest cancelled by user.");
       process.exit(0);
@@ -318,7 +354,7 @@ async function main() {
   }
 
   console.log("Generating embeddings...");
-  const vectors = await generateEmbeddingsWithProgress(specDocs);
+  const vectors = await generateEmbeddingsWithProgress(specDocs, embeddings);
 
   console.log("Creating table with documents...");
   // Prepare data records with vector, text, and metadata

src/test/manual/test-search-spec-sections.ts

@@ -14,7 +14,7 @@
 import * as lancedbSdk from "@lancedb/lancedb";
 import { OllamaEmbeddings } from "@langchain/ollama";
 import { createSearchSpecSectionsTool } from "../../agent-tools/index.js";
-import { EMBEDDING_MODEL, STORAGE_DIR } from "../../constants.js";
+import { OLLAMA_EMBEDDING_MODEL, STORAGE_DIR } from "../../constants.js";
 
 async function main() {
   // Get query from command line or use default
@@ -28,7 +28,7 @@ async function main() {
 
   try {
     const embeddings = new OllamaEmbeddings({
-      model: EMBEDDING_MODEL,
+      model: OLLAMA_EMBEDDING_MODEL,
     });
 
     const db = await lancedbSdk.connect(STORAGE_DIR);

src/test/manual/verify-db.ts

@@ -21,10 +21,10 @@ import type { Table } from "@lancedb/lancedb";
 import * as lancedbSdk from "@lancedb/lancedb";
 import { OllamaEmbeddings } from "@langchain/ollama";
 import { Command } from "commander";
-import { EMBEDDING_MODEL, STORAGE_DIR } from "../../constants.js";
+import { OLLAMA_EMBEDDING_MODEL, STORAGE_DIR } from "../../constants.js";
 
 const embeddings = new OllamaEmbeddings({
-  model: EMBEDDING_MODEL,
+  model: OLLAMA_EMBEDDING_MODEL,
 });
 
 interface DocumentRecord {