@juspay/neurolink
Version:
The pipe layer of an AI nervous system: one interface connecting provider neurons to your application, across three inference types — generate, stream and decide. `decide` returns typed, calibrated judgements from a non-generative model (~400ms, ~$0.00002
511 lines (510 loc) • 23.8 kB
JavaScript
import { existsSync, readFileSync } from "fs";
import { extname, resolve } from "path";
import { z } from "zod";
import { SpanSerializer, SpanType, SpanStatus, getMetricsAggregator, } from "../observability/index.js";
import { logger } from "../utils/logger.js";
import { redactUrlForError } from "../utils/logSanitize.js";
import { ImageProcessor } from "../utils/imageProcessor.js";
import { createChunker } from "./ChunkerFactory.js";
import { createVectorQueryTool, InMemoryVectorStore, } from "./retrieval/vectorQueryTool.js";
import { ImageLoader } from "./document/imageLoader.js";
import { IMAGE_EXTENSIONS } from "../processors/config/index.js";
import { withSpan } from "../telemetry/withSpan.js";
import { tracers } from "../telemetry/tracers.js";
/**
* Maps file extensions to recommended chunking strategies
*/
const EXTENSION_TO_STRATEGY = {
".md": "markdown",
".mdx": "markdown",
".html": "html",
".htm": "html",
".json": "json",
".tex": "latex",
".latex": "latex",
".txt": "recursive",
".csv": "recursive",
".xml": "recursive",
".yaml": "recursive",
".yml": "recursive",
".ts": "recursive",
".js": "recursive",
".py": "recursive",
".java": "recursive",
".go": "recursive",
".rs": "recursive",
".c": "recursive",
".cpp": "recursive",
".rb": "recursive",
".php": "recursive",
".swift": "recursive",
".kt": "recursive",
};
/**
* Detect the best chunking strategy from file extension
*/
function detectStrategy(filePath) {
const ext = extname(filePath).toLowerCase();
return EXTENSION_TO_STRATEGY[ext] || "recursive";
}
/** Embedding dimension used by {@link generateSimpleEmbedding}. */
const EMBEDDING_DIMENSION = 128;
/**
* Simple hash function for strings (FNV-1a variant).
* Maps a word to a bucket index deterministically.
*/
function hashWord(word, buckets) {
let hash = 2166136261;
for (let i = 0; i < word.length; i++) {
hash ^= word.charCodeAt(i);
hash = (hash * 16777619) >>> 0;
}
return hash % buckets;
}
/**
* Generate deterministic embeddings for chunks.
* Combines character-frequency (40%) with word-level hash features (60%)
* for better semantic discrimination than pure character frequency.
* When a real embedding provider is configured, it will be used instead.
*/
function generateSimpleEmbedding(text, dimension) {
const charEmbedding = new Array(dimension).fill(0);
const wordEmbedding = new Array(dimension).fill(0);
// Character-frequency features
for (let i = 0; i < text.length; i++) {
const charCode = text.charCodeAt(i);
const idx = charCode % dimension;
charEmbedding[idx] += 1;
}
// Word-level hash features (TF-IDF-like)
const words = text
.toLowerCase()
.replace(/[^a-z0-9\s]/g, "")
.split(/\s+/)
.filter((w) => w.length > 1);
for (const word of words) {
const idx = hashWord(word, dimension);
wordEmbedding[idx] += 1;
}
// Combine: 40% character, 60% word
const combined = new Array(dimension);
for (let i = 0; i < dimension; i++) {
combined[i] = 0.4 * charEmbedding[i] + 0.6 * wordEmbedding[i];
}
// Normalize to unit vector
const magnitude = Math.sqrt(combined.reduce((sum, v) => sum + v * v, 0));
if (magnitude > 0) {
for (let i = 0; i < dimension; i++) {
combined[i] /= magnitude;
}
}
return combined;
}
/**
* Diversify retrieval results via round-robin across source files.
* Ensures at least one chunk per source file appears in the top-K results,
* preventing any single file from dominating retrieval.
*/
function diversifyResults(results, topK) {
// Group by source file
const byFile = new Map();
for (const r of results) {
const source = r.metadata?.source || "unknown";
if (!byFile.has(source)) {
byFile.set(source, []);
}
const sourceGroup = byFile.get(source);
if (sourceGroup) {
sourceGroup.push(r);
}
}
// If only one source file, no diversification needed
if (byFile.size <= 1) {
return results.slice(0, topK);
}
// Round-robin selection from each source file group
const diversified = [];
const iterators = [...byFile.values()].map((arr) => ({ arr, idx: 0 }));
while (diversified.length < topK &&
iterators.some((it) => it.idx < it.arr.length)) {
for (const it of iterators) {
if (it.idx < it.arr.length && diversified.length < topK) {
diversified.push(it.arr[it.idx++]);
}
}
}
return diversified;
}
/**
* Prepare RAG tools from the provided configuration.
*
* This function:
* 1. Loads and reads all specified files
* 2. Chunks them using the configured (or auto-detected) strategy
* 3. Generates embeddings for each chunk
* 4. Stores them in an in-memory vector store
* 5. Creates a tool the AI model can use to search the documents
*
* @param ragConfig - RAG configuration from generate/stream options
* @param fallbackProvider - Provider to use for embeddings if not specified in ragConfig
* @returns Prepared RAG tool to inject into the tools record
*/
export async function prepareRAGTool(ragConfig, fallbackProvider) {
const span = SpanSerializer.createSpan(SpanType.RAG, "rag.prepare", {
"rag.operation": "prepare",
"rag.files_count": ragConfig.files?.length ?? 0,
"rag.strategy": ragConfig.strategy ?? "auto",
"rag.chunk_size": ragConfig.chunkSize ?? 1000,
});
const startTime = Date.now();
try {
const result = await _prepareRAGToolInner(ragConfig, fallbackProvider);
span.durationMs = Date.now() - startTime;
const endedSpan = SpanSerializer.endSpan(span, SpanStatus.OK);
endedSpan.attributes = {
...endedSpan.attributes,
"rag.chunks_indexed": result.chunksIndexed,
"rag.files_loaded": result.filesLoaded,
};
getMetricsAggregator().recordSpan(endedSpan);
return result;
}
catch (error) {
span.durationMs = Date.now() - startTime;
const endedSpan = SpanSerializer.endSpan(span, SpanStatus.ERROR);
endedSpan.statusMessage =
error instanceof Error ? error.message : String(error);
getMetricsAggregator().recordSpan(endedSpan);
throw error;
}
}
async function _prepareRAGToolInner(ragConfig, fallbackProvider) {
const { files, strategy: userStrategy, chunkSize = 1000, chunkOverlap = 200, topK: userTopK = 5, toolName = "search_knowledge_base", toolDescription = "REQUIRED: Search through pre-loaded local documents to find relevant information. Use this tool FIRST before any web search or other tools. This searches an indexed knowledge base of documents the user has provided.", embeddingProvider, embeddingModel, } = ragConfig;
if (!files || files.length === 0) {
throw new Error("RAG config requires at least one file path in 'files'");
}
// 1. Load files — separate text and image files
const fileContents = [];
const imageFiles = [];
for (const filePath of files) {
// HTTP(S) sources are handled by ImageLoader directly and must not go
// through resolve()/existsSync(), which would mangle the URL into a
// bogus local path and always report it as missing.
if (/^https?:\/\//i.test(filePath)) {
// `new URL()` throws on a string the guarding regex still accepts —
// "https://" is the minimal case. Uncaught, it leaves this loop, is
// rethrown by prepareRAGTool, and disables RAG for the WHOLE call, so
// one bad entry discards every file already loaded. Skip the source
// instead, matching how an unreadable local file is handled below.
let ext;
try {
ext = extname(new URL(filePath).pathname).toLowerCase();
}
catch {
logger.warn(`[RAG] Malformed URL source, skipping: ${redactUrlForError(filePath)}`);
continue;
}
if (ext === ".svg") {
// IMAGE_EXTENSIONS includes .svg, but every consumer downstream of
// here treats an image as raster bytes: ImageLoader returns the file
// verbatim, captioning base64s it into a vision call, and the Bedrock
// embed path ships it as-is — Nova rejects the format outright and
// Titan accepts markup it cannot embed meaningfully. Skipping is
// honest; supporting it needs a rasterize-or-sanitize step in
// ImageLoader, not a classification change here.
logger.warn(`[RAG] SVG is not supported as a RAG image source, skipping: ${redactUrlForError(filePath)}`);
continue;
}
if (IMAGE_EXTENSIONS.includes(ext)) {
imageFiles.push({ path: filePath, ext });
}
else {
logger.warn(`[RAG] Unsupported URL source, skipping: ${redactUrlForError(filePath)}`);
}
continue;
}
const resolvedPath = resolve(filePath);
if (!existsSync(resolvedPath)) {
logger.warn(`[RAG] File not found, skipping: ${resolvedPath}`);
continue;
}
// Detect image files by extension
const ext = extname(resolvedPath).toLowerCase();
if (ext === ".svg") {
// See the URL branch above for why SVG is excluded from the raster path.
logger.warn(`[RAG] SVG is not supported as a RAG image source, skipping: ${resolvedPath}`);
continue;
}
if (IMAGE_EXTENSIONS.includes(ext)) {
imageFiles.push({ path: resolvedPath, ext });
continue;
}
try {
const content = readFileSync(resolvedPath, "utf-8");
const strategy = userStrategy || detectStrategy(resolvedPath);
fileContents.push({ path: resolvedPath, content, strategy });
}
catch (error) {
logger.warn(`[RAG] Failed to read file: ${resolvedPath}: ${error instanceof Error ? error.message : String(error)}`);
}
}
// Auto-increase topK for multi-file scenarios to ensure coverage
// (computed after loading so it reflects only files that actually exist).
// Image files count toward the file total so image-only RAG gets a topK
// that covers every source.
const totalSources = fileContents.length + imageFiles.length;
const topK = totalSources > 1 ? Math.max(userTopK, totalSources * 3) : userTopK;
if (fileContents.length === 0 && imageFiles.length === 0) {
throw new Error("RAG: No files could be loaded. Check that file paths exist and are readable.");
}
logger.info(`[RAG] Loaded ${fileContents.length} files for indexing`);
// 2. Chunk all files
const allChunks = [];
for (const { path, content, strategy } of fileContents) {
try {
const chunker = await createChunker(strategy, {
maxSize: chunkSize,
overlap: Math.min(chunkOverlap, Math.floor(chunkSize * 0.5)),
});
const chunks = await chunker.chunk(content, {
metadata: { source: path },
});
for (const chunk of chunks) {
allChunks.push({
text: chunk.text,
metadata: { ...chunk.metadata, source: path },
});
}
}
catch (error) {
logger.warn(`[RAG] Chunking failed for ${path}, using fallback: ${error instanceof Error ? error.message : String(error)}`);
// Fallback: treat entire file as one chunk
allChunks.push({
text: content.slice(0, chunkSize),
metadata: { source: path, fallback: true },
});
}
}
logger.info(`[RAG] Created ${allChunks.length} chunks from ${fileContents.length} text files`);
// 2b. Load image files (multi-modal RAG).
//
// Deliberately NOT embedded here. Images are vectorised below through the
// same `embedFn` as the text chunks, because index-wide a single embedding
// space is the whole precondition for similarity to mean anything. Embedding
// them here — as this path originally did, with a hash of the caption —
// produced a 128-dimension image vector sitting in an index whose text
// vectors came from a provider at 768 or 1536, which is not a worse ranking
// but an incomparable one.
const imageChunks = [];
if (imageFiles.length > 0) {
const imageLoader = new ImageLoader();
for (const { path: imgPath } of imageFiles) {
try {
const imageDoc = await imageLoader.load(imgPath);
// The two `ext === ".svg"` checks above are a cheap early-out, not the
// boundary: both match the NAME exactly, and the name is not what
// decides the type. ImageLoader resolves it two ways that disagree
// with a bare `.svg` comparison — `.svgz` maps to image/svg+xml
// through EXTENSION_MIME_MAP, and loadFromURL ignores the extension
// entirely and sniffs the bytes, so any image-extension URL actually
// serving SVG resolves to image/svg+xml here. A local raster suffix
// has the opposite problem: loadFromPath trusts the extension before
// the bytes, so inspect both. Otherwise `hasImage: true` below sends
// SVG markup to the raster embed call. Skip rather than throw,
// matching the early-outs and the identical guard
// in RAGPipeline.ingestImages — the two entry points are independent
// (prepareRAGTool builds its chunks here, not via ingestImages), so
// each needs its own.
if (imageDoc.mimeType === "image/svg+xml" ||
ImageProcessor.detectImageType(imageDoc.image) === "image/svg+xml") {
logger.warn(`[RAG] SVG is not supported as a RAG image source, skipping: ${redactUrlForError(imgPath)}`);
continue;
}
// Use filename-based text representation for simple embedding
const imageText = imageDoc.text;
imageChunks.push({
id: `img-chunk-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`,
text: imageText,
metadata: {
text: imageText,
source: imgPath,
hasImage: true,
mimeType: imageDoc.mimeType,
imageWidth: imageDoc.metadata.width,
imageHeight: imageDoc.metadata.height,
},
});
}
catch (error) {
logger.warn(`[RAG] Failed to load image: ${imgPath}: ${error instanceof Error ? error.message : String(error)}`);
}
}
logger.info(`[RAG] Loaded ${imageChunks.length} image files for multi-modal RAG`);
}
// 3. Generate embeddings and store in vector store
const vectorStore = new InMemoryVectorStore();
const indexName = "rag-index";
// When the caller configured an embedding provider/model, embed BOTH the
// index chunks and (below) the queries through that provider — previously
// those config fields had no runtime effect and retrieval always used the
// deterministic hash embedding, which is a lexical fingerprint rather than
// a semantic space. Index and query must share one embedding space, so the
// provider path replaces the hash path wholesale; a provider failure while
// creating the provider or building the index falls back to the hash for
// both sides, but one at query time fails that search.
const wantProviderEmbeddings = Boolean(embeddingProvider || embeddingModel);
const embedProviderName = embeddingProvider || fallbackProvider || "vertex";
const embedModelName = embeddingModel || "gemini-2.5-flash";
let embedFn = (text) => Promise.resolve(generateSimpleEmbedding(text, EMBEDDING_DIMENSION));
if (wantProviderEmbeddings) {
try {
const { AIProviderFactory } = await import("../core/factory.js");
const embedderProvider = (await AIProviderFactory.createProvider(embedProviderName, embedModelName));
if (typeof embedderProvider.embed === "function") {
const providerEmbed = embedderProvider.embed.bind(embedderProvider);
embedFn = (text) => providerEmbed(text, embedModelName);
}
else {
logger.warn(`[RAG] Embedding provider '${embedProviderName}' has no embed(); falling back to hash embeddings`);
}
}
catch (error) {
logger.warn("[RAG] Failed to create embedding provider; falling back to hash embeddings", { error: error instanceof Error ? error.message : String(error) });
}
}
// Text and image captions embed together, in one call and under one
// fallback. Splitting them into two try/catch blocks would let the text side
// succeed on the provider while the image side fell back to the hash, which
// is the mixed-space index this is trying to prevent.
let chunkVectors;
let imageVectors;
try {
const allVectors = await Promise.all([
...allChunks.map((chunk) => embedFn(chunk.text)),
...imageChunks.map((chunk) => embedFn(chunk.text)),
]);
chunkVectors = allVectors.slice(0, allChunks.length);
imageVectors = allVectors.slice(allChunks.length);
}
catch (error) {
// One failed chunk must not leave a mixed-space index — flip the whole
// index AND all queries back to the hash space together.
logger.warn("[RAG] Provider embedding failed mid-index; falling back to hash embeddings for index and queries", { error: error instanceof Error ? error.message : String(error) });
embedFn = (text) => Promise.resolve(generateSimpleEmbedding(text, EMBEDDING_DIMENSION));
chunkVectors = allChunks.map((chunk) => generateSimpleEmbedding(chunk.text, EMBEDDING_DIMENSION));
imageVectors = imageChunks.map((chunk) => generateSimpleEmbedding(chunk.text, EMBEDDING_DIMENSION));
}
const items = [
...allChunks.map((chunk, i) => ({
id: `rag-chunk-${i}`,
vector: chunkVectors[i],
metadata: {
text: chunk.text,
...chunk.metadata,
},
})),
...imageChunks.map((chunk, i) => ({
id: chunk.id,
vector: imageVectors[i],
metadata: chunk.metadata,
})),
];
// The guard near the top of this function fires when no source was found
// at all. This one covers the other way to end up with an empty index:
// sources WERE found and every one of them failed to load — an image over
// the size ceiling, a corrupt file, an unreachable URL. Each of those is
// caught and warned individually, so without this the call returns
// `chunksIndexed: 0` alongside a positive `filesLoaded` and reports
// success, and the caller gets a silently empty knowledge base.
if (items.length === 0) {
throw new Error("RAG: every source failed to load, so nothing was indexed. " +
"Check the warnings above for the per-source reason.");
}
await vectorStore.upsert(indexName, items);
logger.info(`[RAG] Indexed ${items.length} chunks (${allChunks.length} text + ${imageChunks.length} images) in vector store`);
// 4. Create the search tool
// Determine embedding provider/model for the query tool
const provider = embeddingProvider || fallbackProvider || "vertex";
const model = embeddingModel || "gemini-2.5-flash";
const queryTool = createVectorQueryTool({
id: toolName,
description: toolDescription,
indexName,
embeddingModel: { provider, modelName: model },
topK,
includeSources: true,
}, vectorStore);
// Convert to Vercel AI SDK Tool format
const aiTool = {
description: queryTool.description,
inputSchema: z.object({
query: z
.string()
.describe("The search query to find relevant information"),
}),
execute: async ({ query }) => {
return withSpan({
name: "neurolink.rag.search",
tracer: tracers.rag,
attributes: {
"rag.query_length": query ? String(query).length : 0,
"rag.top_k": topK ?? 5,
},
}, async (span) => {
// Query through the same embedding space the index was built in.
// A failure here must fail the search: swapping in a hash vector
// would be scored against provider-space vectors of another
// dimension, which cosine reads as 0 for every chunk, so the tool
// would return arbitrary chunks as if they were relevant. A hash-built
// index never reaches this — its embedFn cannot reject.
const queryEmbedding = await embedFn(query);
// Fetch more candidates than needed so diversity can select across files and images
const loadedSources = fileContents.length + imageChunks.length;
const fetchK = loadedSources > 1 ? topK * 3 : topK;
const rawResults = await vectorStore.query({
indexName,
queryVector: queryEmbedding,
topK: fetchK,
});
// Apply source-file diversity for multi-file/multi-image RAG
const results = loadedSources > 1
? diversifyResults(rawResults, topK)
: rawResults.slice(0, topK);
if (results.length === 0) {
span.setAttribute("rag.results_count", 0);
return {
relevantContext: "No relevant documents found for the query.",
sources: [],
totalResults: 0,
};
}
const relevantContext = results
.map((r, i) => `[${i + 1}] ${r.metadata?.text || r.text || ""}` +
(r.metadata?.hasImage ? " [image]" : ""))
.join("\n\n");
span.setAttribute("rag.results_count", results.length);
return {
relevantContext,
sources: results.map((r) => ({
id: r.id,
score: r.score,
source: r.metadata?.source,
text: (r.metadata?.text || r.text || "").slice(0, 200),
hasImage: r.metadata?.hasImage ?? false,
})),
totalResults: results.length,
};
});
},
};
return {
tool: aiTool,
toolName,
chunksIndexed: allChunks.length + imageChunks.length,
filesLoaded: fileContents.length + imageChunks.length,
};
}
/** @internal Exported for testing only */
export { generateSimpleEmbedding as _generateSimpleEmbedding };
/** @internal Exported for testing only */
export { diversifyResults as _diversifyResults };