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

package/src/tools/fileSearch/tool.ts

@@ -0,0 +1,207 @@
+/**
+ * file_search tool factory — library-native equivalent of the CodeExecutor
+ * pattern. Runtimes supply a `RagClient`, the file list for this turn, and
+ * an optional formatter (ranger uses citation anchors; CLI/A2A use plain
+ * text).
+ *
+ * The tool itself:
+ *   1. Accepts `{ query, target_files? }` from the LLM.
+ *   2. Filters files by `target_files` substring match when provided.
+ *   3. Queries each file in bounded concurrent batches.
+ *   4. Enforces per-file timeouts (failures isolated per file).
+ *   5. Flattens chunks, deprioritizes stale-turn files, caps results.
+ *   6. Hands formatted output to the runtime's formatter for final shape.
+ */
+
+import { tool, DynamicStructuredTool } from '@langchain/core/tools';
+import {
+  fileSearchInputSchema,
+  type FileSearchInput,
+  FileSearchToolName,
+} from './schema';
+import type {
+  FileSearchToolConfig,
+  FileSearchFile,
+  RagChunk,
+  RagQueryParams,
+} from './types';
+import { plainTextFormatter } from './formatter';
+
+const DEFAULT_QUERY_TIMEOUT_MS = 15_000;
+const DEFAULT_CONCURRENCY = 10;
+const DEFAULT_TOP_K = 10;
+
+/**
+ * Build the tool description. Runtimes that use citation anchors supply
+ * `fileCitations: true` (via the formatter); the description includes the
+ * citation ruleset only when that's on.
+ */
+function buildDescription(opts: { fileCitations: boolean }): string {
+  const core = `Performs semantic search across the attached "${FileSearchToolName}" documents using natural language queries. Analyzes the content of loaded files to find relevant information, quotes, and passages matching the query.
+
+**Use target_files to narrow the search:**
+When you know which file(s) contain the relevant information, ALWAYS pass target_files. This is faster and returns more focused results. Pass partial filenames — they match via substring.
+
+**Multiple searches for thorough analysis:**
+For summaries/overviews, call this tool MULTIPLE times with DIFFERENT queries targeting different aspects (intro, methodology, results, conclusions). A single search only returns chunks from one part of the document.`;
+
+  if (!opts.fileCitations) return core;
+
+  return `${core}
+
+**CITING FILE SEARCH RESULTS — MANDATORY:**
+Cite EVERY statement derived from file content. Place the citation anchor IMMEDIATELY after each paragraph using that source. Each search result has a unique source index — use DIFFERENT indices for different claims; do not reuse the same anchor for all paragraphs. Format: \`\\ue202turn0fileN\`. With a page: include \`(p. N)\` inline. Multiple sources: \`\\ue200\\ue202turn0file0\\ue202turn0file1\\ue201\`. NEVER substitute with footnotes, brackets, or symbols.`;
+}
+
+export function createFileSearchTool(
+  config: FileSearchToolConfig,
+): DynamicStructuredTool {
+  const {
+    ragClient,
+    files,
+    entity_id,
+    scope,
+    getAuthHeaders,
+    formatter = plainTextFormatter,
+    queryTimeoutMs = DEFAULT_QUERY_TIMEOUT_MS,
+    concurrencyLimit = DEFAULT_CONCURRENCY,
+    topK = DEFAULT_TOP_K,
+    resultCap,
+    callbacks,
+    logger,
+  } = config;
+
+  // Monotonic call counter used by citation-style formatters to keep source
+  // indices unique across multiple invocations within a single turn.
+  let callIndex = 0;
+
+  // Infer whether the formatter wants citations from the artifact it emits
+  // on an empty-chunk format. This keeps the description/behavior aligned
+  // without forcing the host to declare `fileCitations` twice.
+  const fileCitations = formatter !== plainTextFormatter;
+
+  return tool(
+    async (rawInput: FileSearchInput) => {
+      const { query, target_files } = rawInput;
+
+      if (files.length === 0) {
+        return [
+          'No files to search. Instruct the user to add files for the search.',
+          undefined,
+        ];
+      }
+
+      // target_files: case-insensitive substring match, fallback to all
+      // files with a warning if the filter excludes everything.
+      let filesToQuery: FileSearchFile[] = files;
+      if (target_files && target_files.length > 0) {
+        const lowerTargets = target_files.map((t) => t.toLowerCase());
+        const matched = files.filter((f) =>
+          lowerTargets.some((t) => f.filename.toLowerCase().includes(t)),
+        );
+        if (matched.length === 0) {
+          logger?.warn(
+            `[file_search] No files matched target_files ${target_files.join(', ')}; falling back to all files`,
+          );
+          filesToQuery = files;
+        } else {
+          logger?.info(
+            `[file_search] Filtered to ${matched.length}/${files.length} via target_files`,
+          );
+          filesToQuery = matched;
+        }
+      }
+
+      const authHeaders = getAuthHeaders ? await getAuthHeaders() : undefined;
+
+      const queryOne = async (file: FileSearchFile): Promise<RagChunk[]> => {
+        const params: RagQueryParams = {
+          file_id: file.file_id,
+          query,
+          k: topK,
+          entity_id,
+          scope,
+          authHeaders,
+          timeoutMs: queryTimeoutMs,
+        };
+        try {
+          const chunks = await ragClient.query(params);
+          callbacks?.onFileQueried?.(file, chunks.length);
+          return chunks;
+        } catch (err) {
+          const e = err instanceof Error ? err : new Error(String(err));
+          logger?.error(
+            `[file_search] Query failed for ${file.filename}: ${e.message}`,
+          );
+          callbacks?.onFileError?.(file, e);
+          return [];
+        }
+      };
+
+      // Bounded-concurrency batching. Server-side rerankers handle their
+      // own concurrency; this protects the HTTP connection pool when the
+      // agent has many files.
+      const allChunks: RagChunk[] = [];
+      for (let i = 0; i < filesToQuery.length; i += concurrencyLimit) {
+        const batch = filesToQuery.slice(i, i + concurrencyLimit);
+        const batchResults = await Promise.all(batch.map(queryOne));
+        for (const chunks of batchResults) allChunks.push(...chunks);
+      }
+
+      if (allChunks.length === 0) {
+        return [
+          'No content found in the files. The files may not have been processed correctly or the query may need refinement.',
+          undefined,
+        ];
+      }
+
+      // Build annotated results: attach filename + isCurrentMessage via
+      // a file-id lookup (metadata wins, factory list is fallback).
+      const fileById = new Map(files.map((f) => [f.file_id, f]));
+      const annotated = allChunks.map((c) => {
+        const matched = fileById.get(c.file_id);
+        const filename =
+          (c.metadata?.source
+            ? String(c.metadata.source).split(/[/\\]/).pop()
+            : undefined) ??
+          matched?.filename ??
+          'Unknown';
+        return {
+          ...c,
+          filename,
+          isCurrentMessage: matched?.isCurrentMessage === true,
+        };
+      });
+
+      // Sort: current-turn files first, then by relevance (lower distance).
+      annotated.sort((a, b) => {
+        if (a.isCurrentMessage !== b.isCurrentMessage)
+          return a.isCurrentMessage ? -1 : 1;
+        return a.distance - b.distance;
+      });
+
+      const cap = resultCap ?? Math.max(10, filesToQuery.length * 3);
+      const limited = annotated.slice(0, cap);
+
+      const { message, artifact } = formatter.format(limited, {
+        callIndex,
+        files,
+      });
+      callIndex += 1;
+
+      // Suppress unused-variable warning for fileCitations (currently only
+      // used to gate description; kept in case formatters need it).
+      void fileCitations;
+
+      return [message, artifact];
+    },
+    {
+      name: FileSearchToolName,
+      responseFormat: 'content_and_artifact',
+      description: buildDescription({ fileCitations }),
+      schema: fileSearchInputSchema,
+    },
+  );
+}
+
+export { FileSearchToolName } from './schema';

package/src/tools/fileSearch/types.ts

@@ -0,0 +1,147 @@
+/**
+ * File-search tool types. Mirrors the web_search / code_executor split —
+ * the library owns tool logic; the host supplies a `RagClient` + config
+ * shaped to its own deployment (auth strategy, scope identity, file set).
+ */
+
+export interface FileSearchFile {
+  /** Stable identifier within the RAG backend. */
+  file_id: string;
+  /** Human-readable name surfaced in tool results and prompts. */
+  filename: string;
+  /**
+   * Hint that this file arrived on the *current* conversation turn (as
+   * opposed to an earlier turn). Hosts that don't distinguish leave this
+   * undefined; the formatter deprioritizes older files when set.
+   */
+  isCurrentMessage?: boolean;
+}
+
+/**
+ * A single chunk returned by the RAG backend for a query.
+ * Shape is normalized here so the library stays independent of any
+ * specific RAG service's response format — the `RagClient` is responsible
+ * for translating the backend response into this shape.
+ */
+export interface RagChunk {
+  file_id: string;
+  page_content: string;
+  distance: number; // 0 = perfect match, 1 = orthogonal
+  metadata?: Record<string, unknown>;
+}
+
+export interface RagQueryParams {
+  file_id: string;
+  query: string;
+  /** Top-K chunks to return per file. Default 10. */
+  k?: number;
+  /** Optional tenant/entity ID — forwarded to the backend verbatim. */
+  entity_id?: string;
+  /**
+   * Scope identifier. Hosts use this for per-tenant isolation in the RAG
+   * backend (ranger → userId, cli → agentId, a2a → task-id).
+   */
+  scope?: string;
+  /**
+   * Per-request auth headers — the host builds these (e.g., ranger sends
+   * `Authorization: Bearer <short-lived-JWT-for-userId>`). Allows
+   * different runtimes to use different auth strategies without the
+   * library knowing.
+   */
+  authHeaders?: Record<string, string>;
+  /** Optional per-call timeout override. */
+  timeoutMs?: number;
+}
+
+/**
+ * Pluggable RAG backend. Runtimes provide an implementation that speaks
+ * to whatever vector DB / search service they've deployed (rag_api is the
+ * default; azure search, pinecone, etc. are valid alternates).
+ */
+export interface RagClient {
+  query(params: RagQueryParams): Promise<RagChunk[]>;
+}
+
+/**
+ * Formatter callback that shapes raw chunks into the runtime's preferred
+ * presentation. Ranger uses citation anchors (`\ue202turn0fileN`), CLI
+ * uses plain text, A2A server can return structured parts.
+ */
+export interface FileSearchResultFormatter {
+  format(
+    chunks: Array<RagChunk & { filename: string; isCurrentMessage: boolean }>,
+    context: {
+      /**
+       * Monotonic call index within a single turn. Ranger uses this to
+       * keep citation source indices globally unique across multiple
+       * file_search invocations.
+       */
+      callIndex: number;
+      /** Same files list the factory was seeded with (for lookups). */
+      files: FileSearchFile[];
+    },
+  ): {
+    /** Message returned to the LLM (content). */
+    message: string;
+    /** Optional artifact payload (sources, metadata) returned alongside. */
+    artifact?: unknown;
+  };
+}
+
+export interface FileSearchToolCallbacks {
+  /**
+   * Fires when a RAG query completes successfully for one file. Hosts can
+   * use this for telemetry, streaming UI updates, etc.
+   */
+  onFileQueried?: (file: FileSearchFile, chunkCount: number) => void;
+  /** Fires when a RAG query fails. Hosts can log or surface via UI. */
+  onFileError?: (file: FileSearchFile, error: Error) => void;
+}
+
+export interface FileSearchToolLogger {
+  debug: (msg: string, ...args: unknown[]) => void;
+  info: (msg: string, ...args: unknown[]) => void;
+  warn: (msg: string, ...args: unknown[]) => void;
+  error: (msg: string, ...args: unknown[]) => void;
+}
+
+export interface FileSearchToolConfig {
+  /** The RAG backend client. Required. */
+  ragClient: RagClient;
+  /** Files the agent has access to this turn. Empty = tool self-reports so. */
+  files: FileSearchFile[];
+  /** Tenant/entity ID forwarded to the backend per-query. */
+  entity_id?: string;
+  /**
+   * Per-turn scope identity. Most runtimes pin this once (userId/agentId);
+   * if omitted, no scope is sent to the backend.
+   */
+  scope?: string;
+  /**
+   * Per-call auth header builder. Called on every tool invocation so
+   * the host can mint fresh short-lived tokens (ranger's JWT pattern).
+   * When omitted, no auth headers are sent.
+   */
+  getAuthHeaders?: () => Record<string, string> | Promise<Record<string, string>>;
+  /**
+   * Result formatter. When omitted, the default plain-text formatter is
+   * used (suitable for CLI/A2A runtimes that don't need citation anchors).
+   */
+  formatter?: FileSearchResultFormatter;
+  /** Per-file query timeout in ms. Default 15_000. */
+  queryTimeoutMs?: number;
+  /**
+   * Max concurrent in-flight RAG queries. Protects HTTP connection pools
+   * under heavy file counts. Default 10.
+   */
+  concurrencyLimit?: number;
+  /** Top-K chunks per file. Default 10. */
+  topK?: number;
+  /**
+   * Cap on total chunks returned to the LLM after sorting. Default is
+   * `max(10, filesCount * 3)`.
+   */
+  resultCap?: number;
+  callbacks?: FileSearchToolCallbacks;
+  logger?: FileSearchToolLogger;
+}