mirror of
https://github.com/bendtherules/ask262.git
synced 2026-08-18 13:21:55 +00:00
feat: Support Fireworks embedding for ingest and tool
This commit is contained in:
@@ -0,0 +1,36 @@
|
||||
# Ask262 Environment Configuration
|
||||
# Copy this file to .env and fill in your actual values
|
||||
# Bun automatically loads .env files - no dotenv package needed!
|
||||
|
||||
# =============================================================================
|
||||
# EMBEDDING PROVIDER CONFIGURATION
|
||||
# =============================================================================
|
||||
|
||||
# Choose your embedding provider: "ollama" (local) or "fireworks" (cloud)
|
||||
# Default: "ollama"
|
||||
ASK262_EMBEDDING_PROVIDER=ollama
|
||||
|
||||
# =============================================================================
|
||||
# OLLAMA (LOCAL) CONFIGURATION
|
||||
# =============================================================================
|
||||
|
||||
# Ollama server URL (optional, defaults to http://localhost:11434)
|
||||
# OLLAMA_HOST=http://localhost:11434
|
||||
|
||||
# =============================================================================
|
||||
# FIREWORKS.AI (CLOUD) CONFIGURATION
|
||||
# =============================================================================
|
||||
|
||||
# Required if using Fireworks embeddings
|
||||
# Get your API key from: https://app.fireworks.ai/models/fireworks/qwen3-embedding-8b
|
||||
# FIREWORKS_API_KEY=fw_xxxxxxxxxxxxxxxxxxxxxxxx
|
||||
|
||||
# Fireworks base URL (optional, defaults to https://api.fireworks.ai/inference/v1)
|
||||
# FIREWORKS_BASE_URL=https://api.fireworks.ai/inference/v1
|
||||
|
||||
# =============================================================================
|
||||
# HTTP SERVER CONFIGURATION
|
||||
# =============================================================================
|
||||
|
||||
# Port for the HTTP MCP server (optional, defaults to 3000)
|
||||
# ASK262_PORT=3000
|
||||
@@ -16,7 +16,33 @@ MCP server for exploring the ECMAScript specification and its implementation in
|
||||
|
||||
### Environment Variables
|
||||
|
||||
- `OLLAMA_HOST` (optional): Ollama server URL. Defaults to `http://localhost:11434`
|
||||
Ask262 uses Bun's built-in `.env` support (no dotenv package needed).
|
||||
|
||||
**Setup:**
|
||||
```bash
|
||||
# Copy the example file
|
||||
cp .env.example .env
|
||||
|
||||
# Edit with your values
|
||||
nano .env # or use your preferred editor
|
||||
```
|
||||
|
||||
**Key variables:**
|
||||
- `ASK262_EMBEDDING_PROVIDER`: Choose `ollama` (local) or `fireworks` (cloud). Default: `ollama`
|
||||
- `OLLAMA_HOST`: Ollama server URL. Default: `http://localhost:11434`
|
||||
- `FIREWORKS_API_KEY`: Required if using Fireworks. Get from https://app.fireworks.ai
|
||||
- `ASK262_PORT`: HTTP server port. Default: `3000`
|
||||
|
||||
**Example `.env`:**
|
||||
```bash
|
||||
# Use Fireworks for embeddings (faster, cloud-based)
|
||||
ASK262_EMBEDDING_PROVIDER=fireworks
|
||||
FIREWORKS_API_KEY=fw_your_key_here
|
||||
|
||||
# Or use local Ollama (default)
|
||||
# ASK262_EMBEDDING_PROVIDER=ollama
|
||||
# OLLAMA_HOST=http://localhost:11434
|
||||
```
|
||||
|
||||
## Installation
|
||||
|
||||
|
||||
@@ -13,6 +13,7 @@
|
||||
"@modelcontextprotocol/sdk": "^1.0.4",
|
||||
"acorn": "^8.16.0",
|
||||
"cheerio": "^1.2.0",
|
||||
"commander": "^14.0.3",
|
||||
"glob": "^13.0.6",
|
||||
"graphology": "^0.26.0",
|
||||
"hono": "^4.12.14",
|
||||
@@ -153,7 +154,7 @@
|
||||
|
||||
"command-line-usage": ["command-line-usage@7.0.4", "", { "dependencies": { "array-back": "^6.2.2", "chalk-template": "^0.4.0", "table-layout": "^4.1.1", "typical": "^7.3.0" } }, "sha512-85UdvzTNx/+s5CkSgBm/0hzP80RFHAa7PsfeADE5ezZF3uHz3/Tqj9gIKGT9PTtpycc3Ua64T0oVulGfKxzfqg=="],
|
||||
|
||||
"commander": ["commander@10.0.1", "", {}, "sha512-y4Mg2tXshplEbSGzx7amzPwKKOCGuoSRP/CjEdwwk0FOGlUbq6lKuoyDZTNZkmxHdJtp54hdfY/JUrdL7Xfdug=="],
|
||||
"commander": ["commander@14.0.3", "", {}, "sha512-H+y0Jo/T1RZ9qPP4Eh1pkcQcLRglraJaSLoyOtHxu6AapkjWVCy2Sit1QQ4x3Dng8qDlSsZEet7g5Pq06MvTgw=="],
|
||||
|
||||
"content-disposition": ["content-disposition@1.1.0", "", {}, "sha512-5jRCH9Z/+DRP7rkvY83B+yGIGX96OYdJmzngqnw2SBSxqCFPd0w2km3s5iawpGX8krnwSGmF0FW5Nhr0Hfai3g=="],
|
||||
|
||||
@@ -515,6 +516,8 @@
|
||||
|
||||
"langchain/uuid": ["uuid@10.0.0", "", { "bin": { "uuid": "dist/bin/uuid" } }, "sha512-8XkAphELsDnEGrDxUOHB3RGvXz6TeuYSGEZBOjtTtPm2lwhGBjLgOzLHB63IUWfBpNucQjND6d3AOudO+H3RWQ=="],
|
||||
|
||||
"langsmith/commander": ["commander@10.0.1", "", {}, "sha512-y4Mg2tXshplEbSGzx7amzPwKKOCGuoSRP/CjEdwwk0FOGlUbq6lKuoyDZTNZkmxHdJtp54hdfY/JUrdL7Xfdug=="],
|
||||
|
||||
"langsmith/uuid": ["uuid@10.0.0", "", { "bin": { "uuid": "dist/bin/uuid" } }, "sha512-8XkAphELsDnEGrDxUOHB3RGvXz6TeuYSGEZBOjtTtPm2lwhGBjLgOzLHB63IUWfBpNucQjND6d3AOudO+H3RWQ=="],
|
||||
|
||||
"openai/@types/node": ["@types/node@18.19.130", "", { "dependencies": { "undici-types": "~5.26.4" } }, "sha512-GRaXQx6jGfL8sKfaIDD6OupbIHBr9jv7Jnaml9tB7l4v068PAOXqfcujMMo5PhbIs6ggR1XODELqahT2R8v0fg=="],
|
||||
|
||||
@@ -57,6 +57,7 @@
|
||||
"@modelcontextprotocol/sdk": "^1.0.4",
|
||||
"acorn": "^8.16.0",
|
||||
"cheerio": "^1.2.0",
|
||||
"commander": "^14.0.3",
|
||||
"glob": "^13.0.6",
|
||||
"graphology": "^0.26.0",
|
||||
"hono": "^4.12.14",
|
||||
|
||||
@@ -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
|
||||
|
||||
+7
-1
@@ -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";
|
||||
|
||||
@@ -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";
|
||||
}
|
||||
@@ -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;
|
||||
};
|
||||
}
|
||||
+4
-10
@@ -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;
|
||||
|
||||
+4
-11
@@ -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
|
||||
|
||||
+46
-10
@@ -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
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user