@illuma-ai/agents 1.4.0-alpha.0 → 1.4.0-alpha.2

package/src/tools/fileSearch/__tests__/tool.test.ts

@@ -0,0 +1,251 @@
+/**
+ * Unit tests for the file_search library tool.
+ *
+ * The real RAG backend is mocked through the `RagClient` interface so
+ * these tests verify the tool's own logic: target_files filtering,
+ * bounded-concurrency querying, per-file error isolation, formatter
+ * handoff, and empty/no-file edge cases.
+ */
+
+import { createFileSearchTool } from '../tool';
+import { plainTextFormatter, createCitationAnchorFormatter } from '../formatter';
+import type { RagClient, RagQueryParams, RagChunk, FileSearchFile } from '../types';
+
+// Build a mock RagClient that records every query it receives and returns
+// a deterministic chunk set per file.
+function makeRagClient(opts: {
+  chunksByFile?: Record<string, RagChunk[]>;
+  failFileIds?: Set<string>;
+  hang?: boolean;
+} = {}) {
+  const calls: RagQueryParams[] = [];
+  const client: RagClient = {
+    async query(params) {
+      calls.push({ ...params });
+      if (opts.hang) {
+        // Never resolves — used to test timeout handling.
+        await new Promise(() => {});
+      }
+      if (opts.failFileIds?.has(params.file_id)) {
+        throw new Error(`simulated failure for ${params.file_id}`);
+      }
+      return opts.chunksByFile?.[params.file_id] ?? [];
+    },
+  };
+  return { client, calls };
+}
+
+const files: FileSearchFile[] = [
+  { file_id: 'f-alpha', filename: 'alpha.pdf', isCurrentMessage: true },
+  { file_id: 'f-beta', filename: 'beta-report.docx' },
+  { file_id: 'f-gamma', filename: 'gamma_notes.txt' },
+];
+
+describe('createFileSearchTool', () => {
+  it('returns a no-files message when factory was seeded with zero files', async () => {
+    const { client } = makeRagClient();
+    const t = createFileSearchTool({ ragClient: client, files: [] });
+    const result = await t.invoke({ query: 'anything' });
+    // Dual-format tools return [string, artifact?]; LangChain surfaces the
+    // string directly in some invocation paths — handle both shapes.
+    const text = Array.isArray(result) ? result[0] : result;
+    expect(String(text)).toMatch(/no files to search/i);
+  });
+
+  it('queries every file when target_files is omitted', async () => {
+    const { client, calls } = makeRagClient({
+      chunksByFile: {
+        'f-alpha': [chunk('f-alpha', 'alpha page 1', 0.1)],
+        'f-beta': [chunk('f-beta', 'beta text', 0.2)],
+        'f-gamma': [chunk('f-gamma', 'gamma note', 0.3)],
+      },
+    });
+    const t = createFileSearchTool({ ragClient: client, files });
+    await t.invoke({ query: 'test' });
+    expect(calls.map((c) => c.file_id).sort()).toEqual([
+      'f-alpha',
+      'f-beta',
+      'f-gamma',
+    ]);
+  });
+
+  it('filters files by target_files substring (case-insensitive)', async () => {
+    const { client, calls } = makeRagClient({
+      chunksByFile: { 'f-beta': [chunk('f-beta', 'beta', 0.2)] },
+    });
+    const t = createFileSearchTool({ ragClient: client, files });
+    await t.invoke({ query: 'q', target_files: ['BETA-REPORT'] });
+    expect(calls.map((c) => c.file_id)).toEqual(['f-beta']);
+  });
+
+  it('falls back to all files when target_files matches nothing', async () => {
+    const { client, calls } = makeRagClient({
+      chunksByFile: {
+        'f-alpha': [chunk('f-alpha', 'a', 0.1)],
+        'f-beta': [chunk('f-beta', 'b', 0.2)],
+        'f-gamma': [chunk('f-gamma', 'c', 0.3)],
+      },
+    });
+    const warn = jest.fn();
+    const t = createFileSearchTool({
+      ragClient: client,
+      files,
+      logger: { debug: jest.fn(), info: jest.fn(), warn, error: jest.fn() },
+    });
+    await t.invoke({ query: 'q', target_files: ['no-such-file'] });
+    expect(calls.length).toBe(3);
+    expect(warn).toHaveBeenCalled();
+  });
+
+  it('isolates per-file failures so one bad file does not fail the call', async () => {
+    const { client, calls } = makeRagClient({
+      chunksByFile: {
+        'f-alpha': [chunk('f-alpha', 'alpha good', 0.1)],
+        'f-gamma': [chunk('f-gamma', 'gamma good', 0.2)],
+      },
+      failFileIds: new Set(['f-beta']),
+    });
+    const onFileError = jest.fn();
+    const t = createFileSearchTool({
+      ragClient: client,
+      files,
+      callbacks: { onFileError },
+      logger: silentLogger(),
+    });
+    const result = await t.invoke({ query: 'q' });
+    const text = Array.isArray(result) ? result[0] : result;
+    expect(String(text)).toMatch(/alpha good/);
+    expect(String(text)).toMatch(/gamma good/);
+    expect(calls.length).toBe(3);
+    expect(onFileError).toHaveBeenCalledWith(
+      expect.objectContaining({ file_id: 'f-beta' }),
+      expect.any(Error),
+    );
+  });
+
+  it('forwards entity_id, scope, and authHeaders on every query', async () => {
+    const { client, calls } = makeRagClient({
+      chunksByFile: { 'f-alpha': [chunk('f-alpha', 'x', 0.1)] },
+    });
+    const t = createFileSearchTool({
+      ragClient: client,
+      files: [files[0]],
+      entity_id: 'tenant-42',
+      scope: 'user:alice',
+      getAuthHeaders: () => ({ Authorization: 'Bearer TOKEN' }),
+    });
+    await t.invoke({ query: 'q' });
+    expect(calls[0]).toEqual(
+      expect.objectContaining({
+        file_id: 'f-alpha',
+        query: 'q',
+        entity_id: 'tenant-42',
+        scope: 'user:alice',
+        authHeaders: { Authorization: 'Bearer TOKEN' },
+      }),
+    );
+  });
+
+  it('prioritizes current-turn files even when stale files have closer distances', async () => {
+    const { client } = makeRagClient({
+      chunksByFile: {
+        'f-alpha': [chunk('f-alpha', 'current-turn', 0.5)],
+        'f-beta': [chunk('f-beta', 'older-turn', 0.1)],
+      },
+    });
+    const t = createFileSearchTool({
+      ragClient: client,
+      files: [files[0], files[1]], // alpha=current, beta=not
+    });
+    const result = await t.invoke({ query: 'q' });
+    const text = Array.isArray(result) ? String(result[0]) : String(result);
+    const currentIdx = text.indexOf('current-turn');
+    const olderIdx = text.indexOf('older-turn');
+    expect(currentIdx).toBeGreaterThanOrEqual(0);
+    expect(olderIdx).toBeGreaterThan(currentIdx);
+  });
+
+  it('uses plainTextFormatter by default (no citation anchors)', async () => {
+    const { client } = makeRagClient({
+      chunksByFile: { 'f-alpha': [chunk('f-alpha', 'hello', 0.1)] },
+    });
+    const t = createFileSearchTool({
+      ragClient: client,
+      files: [files[0]],
+    });
+    const result = await t.invoke({ query: 'q' });
+    const text = Array.isArray(result) ? String(result[0]) : String(result);
+    expect(text).not.toMatch(/\\ue202turn0file/);
+    expect(text).toMatch(/File: alpha\.pdf/);
+  });
+
+  it('uses citation anchors when createCitationAnchorFormatter is supplied', async () => {
+    const { client } = makeRagClient({
+      chunksByFile: { 'f-alpha': [chunk('f-alpha', 'hello', 0.1)] },
+    });
+    let offset = 0;
+    const formatter = createCitationAnchorFormatter({
+      getSourceOffset: () => offset,
+      advanceSourceOffset: (by) => {
+        offset += by;
+      },
+    });
+    const t = createFileSearchTool({
+      ragClient: client,
+      files: [files[0]],
+      formatter,
+    });
+    const first = await t.invoke({ query: 'q' });
+    const text1 = Array.isArray(first) ? String(first[0]) : String(first);
+    expect(text1).toMatch(/Source 0/);
+    expect(text1).toMatch(/\\ue202turn0file0/);
+
+    // Second call in the same turn: offset should have advanced.
+    const second = await t.invoke({ query: 'q2' });
+    const text2 = Array.isArray(second) ? String(second[0]) : String(second);
+    expect(text2).toMatch(/Source 1/);
+    expect(text2).toMatch(/\\ue202turn0file1/);
+  });
+
+  it('extracts 1-indexed page numbers from metadata.page (rag_api 0-indexed)', async () => {
+    const { client } = makeRagClient({
+      chunksByFile: {
+        'f-alpha': [
+          {
+            file_id: 'f-alpha',
+            page_content: 'page two content',
+            distance: 0.1,
+            metadata: { page: 1 }, // rag_api 0-indexed → display = 2
+          },
+        ],
+      },
+    });
+    const t = createFileSearchTool({
+      ragClient: client,
+      files: [files[0]],
+    });
+    const result = await t.invoke({ query: 'q' });
+    const text = Array.isArray(result) ? String(result[0]) : String(result);
+    expect(text).toMatch(/Page: 2/);
+  });
+});
+
+// ── helpers ──────────────────────────────────────────────────────────────
+
+function chunk(
+  file_id: string,
+  text: string,
+  distance: number,
+  metadata?: Record<string, unknown>,
+): RagChunk {
+  return { file_id, page_content: text, distance, metadata };
+}
+
+function silentLogger() {
+  return {
+    debug: jest.fn(),
+    info: jest.fn(),
+    warn: jest.fn(),
+    error: jest.fn(),
+  };
+}

package/src/tools/fileSearch/formatter.ts

@@ -0,0 +1,131 @@
+/**
+ * Default result formatters.
+ *
+ * - `plainTextFormatter`: CLI / A2A / generic output. No citation anchors.
+ * - `citationAnchorFormatter`: ranger-style `\ue202turn0fileN` anchors with
+ *   a monotonic `sourceOffset` so multi-call turns stay globally unique.
+ *
+ * Runtimes can supply their own `FileSearchResultFormatter` to override.
+ */
+
+import type {
+  FileSearchResultFormatter,
+  FileSearchFile,
+  RagChunk,
+} from './types';
+
+type AnnotatedChunk = RagChunk & {
+  filename: string;
+  isCurrentMessage: boolean;
+};
+
+export const plainTextFormatter: FileSearchResultFormatter = {
+  format(chunks, { files: _files }) {
+    if (chunks.length === 0) {
+      return { message: 'No relevant results found in the available files.' };
+    }
+    const body = chunks
+      .map((c) => {
+        const page = getPage(c);
+        const rel = (1 - c.distance).toFixed(4);
+        return (
+          `File: ${c.filename}` +
+          (page != null ? `\nPage: ${page}` : '') +
+          `\nRelevance: ${rel}\nContent: ${c.page_content}\n`
+        );
+      })
+      .join('\n---\n');
+
+    const sources = chunks.map((c) => ({
+      type: 'file' as const,
+      fileId: c.file_id,
+      content: c.page_content,
+      fileName: c.filename,
+      relevance: 1 - c.distance,
+      pages: getPage(c) != null ? [getPage(c) as number] : [],
+    }));
+
+    return { message: body, artifact: { file_search: { sources } } };
+  },
+};
+
+export interface CitationAnchorFormatterOptions {
+  /** Tool name used in the `file_search` artifact wrapper. Defaults to `'file_search'`. */
+  toolName?: string;
+  /**
+   * Monotonic counter for source indices within a turn. Pass the SAME
+   * function to the formatter across multiple calls in the same turn so
+   * anchors stay globally unique.
+   */
+  getSourceOffset?: () => number;
+  /** Called after formatting to advance the offset. */
+  advanceSourceOffset?: (by: number) => void;
+}
+
+export function createCitationAnchorFormatter(
+  opts: CitationAnchorFormatterOptions = {},
+): FileSearchResultFormatter {
+  const toolName = opts.toolName ?? 'file_search';
+  const getOffset = opts.getSourceOffset ?? (() => 0);
+  const advance = opts.advanceSourceOffset ?? (() => {});
+
+  return {
+    format(chunks) {
+      if (chunks.length === 0) {
+        return {
+          message:
+            'No results found or errors occurred while searching the files.',
+        };
+      }
+      const base = getOffset();
+      const body = chunks
+        .map((c, i) => {
+          const globalIndex = base + i;
+          const page = getPage(c);
+          const rel = (1 - c.distance).toFixed(4);
+          return (
+            `[Source ${globalIndex}] File: ${c.filename} | Anchor: \\ue202turn0file${globalIndex}` +
+            (page != null ? ` | Page: ${page}` : '') +
+            ` | Relevance: ${rel}\nContent: ${c.page_content}\n` +
+            `↑ Cite this source using: \\ue202turn0file${globalIndex}`
+          );
+        })
+        .join('\n---\n');
+
+      const sources = chunks.map((c) => ({
+        type: 'file' as const,
+        fileId: c.file_id,
+        content: c.page_content,
+        fileName: c.filename,
+        relevance: 1 - c.distance,
+        pages: getPage(c) != null ? [getPage(c) as number] : [],
+        pageRelevance:
+          getPage(c) != null
+            ? { [getPage(c) as number]: 1 - c.distance }
+            : {},
+      }));
+
+      advance(chunks.length);
+      return {
+        message: body,
+        artifact: { [toolName]: { sources, fileCitations: true } },
+      };
+    },
+  };
+}
+
+/** Extract a 1-indexed page number from the chunk metadata, or null. */
+function getPage(chunk: AnnotatedChunk | RagChunk): number | null {
+  const raw =
+    (chunk.metadata?.page as unknown) ??
+    (chunk.metadata?.page_number as unknown) ??
+    null;
+  if (raw == null) return null;
+  const parsed = typeof raw === 'number' ? raw : parseInt(String(raw), 10);
+  if (Number.isNaN(parsed) || parsed < 0) return null;
+  // rag_api stores 0-indexed; display is 1-indexed
+  return parsed + 1;
+}
+
+// Re-export so consumers only import from the formatter module.
+export type { FileSearchResultFormatter, FileSearchFile, RagChunk };

package/src/tools/fileSearch/index.ts

@@ -0,0 +1,23 @@
+export { createFileSearchTool, FileSearchToolName } from './tool';
+export {
+  HttpRagClient,
+  getRagBaseUrl,
+  RAG_API_URL_ENV,
+  type HttpRagClientOptions,
+} from './ragClient';
+export {
+  plainTextFormatter,
+  createCitationAnchorFormatter,
+  type CitationAnchorFormatterOptions,
+} from './formatter';
+export { fileSearchInputSchema, type FileSearchInput } from './schema';
+export type {
+  FileSearchFile,
+  RagChunk,
+  RagClient,
+  RagQueryParams,
+  FileSearchResultFormatter,
+  FileSearchToolCallbacks,
+  FileSearchToolConfig,
+  FileSearchToolLogger,
+} from './types';

package/src/tools/fileSearch/ragClient.ts

@@ -0,0 +1,141 @@
+/**
+ * Default HTTP RAG client. Posts to `${baseUrl}/query` with the shape
+ * rag_api expects (`{ file_id, query, k, entity_id? }`). Runtimes that
+ * use a different vector backend implement their own `RagClient`.
+ *
+ * Auth is runtime-provided per call (via `authHeaders` on the params) so
+ * short-lived tokens can be minted per request without the client
+ * caching stale credentials.
+ */
+
+import fetch from 'node-fetch';
+import { getEnvironmentVariable } from '@langchain/core/utils/env';
+import type {
+  RagClient,
+  RagQueryParams,
+  RagChunk,
+  FileSearchToolLogger,
+} from './types';
+
+export const RAG_API_URL_ENV = 'RAG_API_URL';
+
+/** Resolve base URL at call time so env-var changes propagate. */
+export function getRagBaseUrl(override?: string): string {
+  const url =
+    override ??
+    getEnvironmentVariable(RAG_API_URL_ENV) ??
+    '';
+  if (!url) {
+    throw new Error(
+      `file_search: ${RAG_API_URL_ENV} is not configured. ` +
+        `Set the env var or pass baseUrl to HttpRagClient.`,
+    );
+  }
+  return url.replace(/\/$/, '');
+}
+
+export interface HttpRagClientOptions {
+  /** Base URL of the RAG service (no trailing slash). Falls back to env. */
+  baseUrl?: string;
+  /** Default headers sent on every request (e.g., a static API key). */
+  defaultHeaders?: Record<string, string>;
+  /** Default timeout if params don't override. Default 15_000. */
+  defaultTimeoutMs?: number;
+  logger?: FileSearchToolLogger;
+}
+
+/**
+ * Expected rag_api response shape: `[[{ page_content, metadata }, distance], ...]`
+ * — an array of [doc, score] tuples. Normalized here into `RagChunk[]`.
+ */
+type RagApiResponse = Array<
+  [
+    {
+      page_content: string;
+      metadata?: Record<string, unknown>;
+    },
+    number,
+  ]
+>;
+
+export class HttpRagClient implements RagClient {
+  private readonly baseUrlOverride?: string;
+  private readonly defaultHeaders: Record<string, string>;
+  private readonly defaultTimeoutMs: number;
+  private readonly logger?: FileSearchToolLogger;
+
+  constructor(opts: HttpRagClientOptions = {}) {
+    this.baseUrlOverride = opts.baseUrl;
+    this.defaultHeaders = opts.defaultHeaders ?? {};
+    this.defaultTimeoutMs = opts.defaultTimeoutMs ?? 15_000;
+    this.logger = opts.logger;
+  }
+
+  async query(params: RagQueryParams): Promise<RagChunk[]> {
+    const baseUrl = getRagBaseUrl(this.baseUrlOverride);
+    const url = `${baseUrl}/query`;
+
+    const body: Record<string, unknown> = {
+      file_id: params.file_id,
+      query: params.query,
+      k: params.k ?? 10,
+    };
+    if (params.entity_id) body.entity_id = params.entity_id;
+    if (params.scope) body.scope = params.scope;
+
+    const headers: Record<string, string> = {
+      'Content-Type': 'application/json',
+      ...this.defaultHeaders,
+      ...(params.authHeaders ?? {}),
+    };
+
+    const timeoutMs = params.timeoutMs ?? this.defaultTimeoutMs;
+    const controller =
+      typeof AbortController !== 'undefined' ? new AbortController() : null;
+    const timer = controller
+      ? setTimeout(() => controller.abort(), timeoutMs)
+      : null;
+
+    this.logger?.debug('[file_search] RAG query', {
+      url,
+      file_id: params.file_id,
+      k: body.k,
+    });
+
+    try {
+      const res = await fetch(url, {
+        method: 'POST',
+        headers,
+        body: JSON.stringify(body),
+        signal: controller?.signal as unknown as undefined,
+      });
+      if (!res.ok) {
+        const text = await res.text().catch(() => '');
+        throw new Error(
+          `RAG query failed: ${res.status} ${res.statusText} — ${text.slice(0, 200)}`,
+        );
+      }
+      const json = (await res.json()) as RagApiResponse;
+      return this.normalize(params.file_id, json);
+    } finally {
+      if (timer) clearTimeout(timer);
+    }
+  }
+
+  /** Convert rag_api's tuple format into the library's normalized shape. */
+  private normalize(file_id: string, resp: RagApiResponse): RagChunk[] {
+    if (!Array.isArray(resp)) {
+      this.logger?.warn('[file_search] RAG response not an array', { resp });
+      return [];
+    }
+    return resp
+      .filter((row) => Array.isArray(row) && row.length === 2)
+      .map(([doc, distance]) => ({
+        file_id:
+          (doc?.metadata?.file_id as string | undefined) ?? file_id,
+        page_content: doc?.page_content ?? '',
+        distance: typeof distance === 'number' ? distance : 1,
+        metadata: doc?.metadata,
+      }));
+  }
+}

package/src/tools/fileSearch/schema.ts

@@ -0,0 +1,19 @@
+import { z } from 'zod';
+
+export const fileSearchInputSchema = z.object({
+  query: z
+    .string()
+    .describe(
+      "A natural language query to search for relevant information in the files. Be SPECIFIC and TARGETED — use keywords for the specific section or topic you need. For comprehensive tasks (summaries, overviews), call this tool multiple times with different targeted queries (e.g., 'introduction', 'methodology', 'results', 'conclusions') rather than one broad query.",
+    ),
+  target_files: z
+    .array(z.string())
+    .optional()
+    .describe(
+      'Optional list of filenames (or partial names) to limit the search to. When provided, only files whose name contains one of these strings will be searched. Use this to avoid searching irrelevant files. Omit to search all available files.',
+    ),
+});
+
+export type FileSearchInput = z.infer<typeof fileSearchInputSchema>;
+
+export const FileSearchToolName = 'file_search';