chat-agent-toolkit 1.2.2 → 1.2.3

package/src/tools/search/doc-utils.ts

@@ -3,8 +3,13 @@
  * @description Document utilities: fallback docs, reranking, and formatting.
  */
 import type { Document } from "./document";
-import path from "node:path";
-import fs from "node:fs";
+
+export interface R2CredentialsInput {
+  accountId: string;
+  accessKeyId: string;
+  secretAccessKey: string;
+  bucket: string;
+}
 
 export function buildFallbackDocs(query: string): Document[] {
   const trimmedQuery = (query || "").trim();
@@ -31,22 +36,45 @@ export function normalizeSourcesOutput(output: unknown, query: string): Document
   return buildFallbackDocs(query);
 }
 
+async function downloadExtractedContent(fileId: string, r2Credentials: R2CredentialsInput): Promise<{ title: string; content: string } | null> {
+  try {
+    const { manageStorage } = await import("manage-storage");
+    const config = {
+      provider: "cloudflare" as const,
+      BUCKET_NAME: r2Credentials.bucket,
+      ACCESS_KEY_ID: r2Credentials.accessKeyId,
+      SECRET_ACCESS_KEY: r2Credentials.secretAccessKey,
+      BUCKET_URL: `https://${r2Credentials.accountId}.r2.cloudflarestorage.com`,
+    };
+    const extractedKey = `${fileId}-extracted.json`;
+    const data = await manageStorage("download", { ...config, key: extractedKey });
+    const parsed = JSON.parse(data);
+    return { title: parsed.title || "Uploaded Document", content: parsed.content || "" };
+  } catch (error) {
+    console.error(`[rerankDocs] Failed to download extracted content for fileId ${fileId}:`, error);
+    return null;
+  }
+}
+
 export async function rerankDocs(
   query: string,
   docs: Document[],
   fileIds: string[],
   optimizationMode: "speed" | "balanced" | "quality",
+  r2Credentials?: R2CredentialsInput,
 ): Promise<Document[]> {
   if (docs.length === 0 && fileIds.length === 0) {
     return docs;
   }
 
-  const filesData = fileIds.map((file) => {
-    const filePath = path.join(process.cwd(), "uploads", file);
-    const contentPath = filePath + "-extracted.json";
-    const content = JSON.parse(fs.readFileSync(contentPath, "utf8"));
-    return { fileName: content.title, content: content.content };
-  });
+  let filesData: { fileName: string; content: string }[] = [];
+
+  if (fileIds.length > 0 && r2Credentials) {
+    const results = await Promise.all(
+      fileIds.map((fileId) => downloadExtractedContent(fileId, r2Credentials))
+    );
+    filesData = results.filter((r) => r !== null) as { fileName: string; content: string }[];
+  }
 
   if (query.toLocaleLowerCase() === "summarize") {
     return docs.slice(0, 15);

package/src/tools/search/metaSearchAgent.ts

@@ -35,6 +35,7 @@ import {
   rerankDocs,
   processDocs,
   normalizeSourcesOutput,
+  type R2CredentialsInput,
 } from "./doc-utils";
 import { groupAndSummarizeDocs } from "./link-summarizer";
 
@@ -350,11 +351,21 @@ class MetaSearchAgent implements MetaSearchAgentType {
         docs = result.docs;
       }
 
+      const r2Credentials: R2CredentialsInput | undefined = process.env.R2_ACCOUNT_ID
+        ? {
+            accountId: process.env.R2_ACCOUNT_ID,
+            accessKeyId: process.env.R2_ACCESS_KEY_ID || "",
+            secretAccessKey: process.env.R2_SECRET_ACCESS_KEY || "",
+            bucket: process.env.R2_UPLOADS_BUCKET || "qwksearch-uploads",
+          }
+        : undefined;
+
       const sortedDocs = await rerankDocs(
         query,
         docs ?? [],
         fileIds,
         optimizationMode,
+        r2Credentials,
       );
 
       const sources = normalizeSourcesOutput(sortedDocs, message);

package/src/tools/search/search-handlers.ts

@@ -7,9 +7,21 @@
  * See extract-webpage/src/search/index.ts for an example.
  */
 import MetaSearchAgent from "./metaSearchAgent";
-import prompts from "write-language/prompts/search-prompts";
+import {
+  webSearchRetrieverPrompt,
+  webSearchResponsePrompt,
+  webSearchRetrieverFewShots,
+  writingAssistantPrompt,
+} from "write-language";
 import type { Config } from "./meta-search-types";
 
+const prompts = {
+  webSearchRetrieverPrompt,
+  webSearchResponsePrompt,
+  webSearchRetrieverFewShots,
+  writingAssistantPrompt,
+};
+
 /**
  * Creates search handler instances with provided search functions.
  * Pass in search functions from extract-webpage to enable web search.

package/src/tools/search/suggestionGeneratorAgent.ts

@@ -8,8 +8,8 @@ import { LineListOutputParser } from "../../utils/outputParser";
 import { formatChatHistoryAsString } from "../../utils";
 import type { ChatTurnMessage } from "./meta-search-types";
 
-const suggestionGeneratorPrompt = `
-You are an AI suggestion generator for an AI powered search engine. You will be given a conversation below. You need to generate 4-5 suggestions based on the conversation. The suggestion should be relevant to the conversation that can be used by the user to ask the chat model for more information.
+const createSuggestionGeneratorPrompt = (maxQuestions: number = 4) => `
+You are an AI suggestion generator for an AI powered search engine. You will be given a conversation below. You need to generate ${maxQuestions} suggestions based on the conversation. The suggestion should be relevant to the conversation that can be used by the user to ask the chat model for more information.
 You need to make sure the suggestions are relevant to the conversation and are helpful to the user. Keep a note that the user might use these suggestions to ask a chat model for more information.
 Make sure the suggestions are medium in length and are informative and relevant to the conversation.
 
@@ -27,6 +27,7 @@ Conversation:
 
 type SuggestionGeneratorInput = {
   chat_history: ChatTurnMessage[];
+  maxQuestions?: number;
 };
 
 const outputParser = new LineListOutputParser({
@@ -37,7 +38,8 @@ const generateSuggestions = async (
   input: SuggestionGeneratorInput,
   llm: LanguageModel,
 ): Promise<string[]> => {
-  const prompt = suggestionGeneratorPrompt.replace(
+  const maxQuestions = input.maxQuestions ?? 4;
+  const prompt = createSuggestionGeneratorPrompt(maxQuestions).replace(
     "{chat_history}",
     formatChatHistoryAsString(input.chat_history),
   );

package/src/tools/composio-mcp.ts

@@ -1,205 +0,0 @@
-/**
- * @fileoverview Composio MCP Integration for AI SDK v5
- *
- * Composio provides dynamic MCP endpoints through sessions/Tool Router.
- * This module creates MCP clients connected to Composio's HTTP endpoints,
- * fetches available tools, and passes them to AI SDK models.
- *
- * Architecture:
- * 1. Create Composio session → Get MCP URL + headers
- * 2. Connect with createMCPClient (HTTP transport)
- * 3. Fetch tools with client.tools()
- * 4. Pass to streamText/generateText
- *
- * @see https://composio.dev/toolkits/ai_ml_api/framework/ai-sdk
- * @see https://ai-sdk.dev/docs/ai-sdk-core/mcp-tools
- */
-
-import { Composio } from '@composio/core';
-import { createMCPClient } from '@ai-sdk/mcp';
-import type { MCPClient } from '@ai-sdk/mcp';
-
-/**
- * Configuration for Composio MCP session
- */
-export interface ComposioMCPConfig {
-  /** Composio API key */
-  apiKey: string;
-  /** User ID to scope the session */
-  userId: string;
-  /** Toolkits to enable (e.g., ['GMAIL', 'NOTION', 'SLACK']) */
-  toolkits: string[];
-  /** Optional: Additional MCP headers (e.g., x-api-key if org requires it) */
-  additionalHeaders?: Record<string, string>;
-}
-
-/**
- * Composio MCP session wrapper
- * Manages the lifecycle of MCP client connections to Composio
- */
-export class ComposioMCPSession {
-  private composio: Composio;
-  private client: MCPClient | null = null;
-  private mcpEndpoint: { url: string; headers: Record<string, string> } | null = null;
-
-  constructor(private config: ComposioMCPConfig) {
-    this.composio = new Composio({
-      apiKey: config.apiKey,
-    });
-  }
-
-  /**
-   * Create Composio Tool Router session and get MCP endpoint
-   * This returns the HTTP URL and auth headers for the MCP server
-   */
-  private async getOrCreateEndpoint() {
-    if (this.mcpEndpoint) {
-      return this.mcpEndpoint;
-    }
-
-    // Create dynamic Composio session with specified toolkits
-    const session = await this.composio.toolkits.authorize({
-      userId: this.config.userId,
-      toolkits: this.config.toolkits,
-    });
-
-    // Extract MCP endpoint info
-    this.mcpEndpoint = {
-      url: session.mcp.url,
-      headers: {
-        ...session.mcp.headers,
-        ...this.config.additionalHeaders,
-      },
-    };
-
-    return this.mcpEndpoint;
-  }
-
-  /**
-   * Get or create the MCP client connection
-   * Uses HTTP transport as recommended by Composio
-   */
-  async getClient(): Promise<MCPClient> {
-    if (this.client) {
-      return this.client;
-    }
-
-    const endpoint = await this.getOrCreateEndpoint();
-
-    // Connect to Composio MCP endpoint via HTTP
-    this.client = await createMCPClient({
-      transport: {
-        type: 'http',
-        url: endpoint.url,
-        headers: endpoint.headers,
-        redirect: 'error',
-      },
-    });
-
-    return this.client;
-  }
-
-  /**
-   * Get all available tools from Composio MCP
-   * These tools can be passed directly to streamText/generateText
-   *
-   * @returns Tools object compatible with AI SDK
-   */
-  async getTools() {
-    const client = await this.getClient();
-    return await client.tools();
-  }
-
-  /**
-   * Close the MCP connection
-   * Should be called after request completion or in finally block
-   */
-  async close() {
-    if (this.client) {
-      await this.client.close();
-      this.client = null;
-    }
-  }
-
-  /**
-   * Get list of available toolkits from Composio
-   * Useful for discovering what integrations are available
-   */
-  static async listAvailableToolkits(apiKey: string): Promise<string[]> {
-    const composio = new Composio({ apiKey });
-    const toolkits = await composio.toolkits.list();
-    return toolkits.map((t) => t.name);
-  }
-
-  /**
-   * Check if a user has authenticated a specific toolkit
-   * Returns true if the user can access the toolkit's tools
-   */
-  async isToolkitAuthenticated(toolkit: string): Promise<boolean> {
-    try {
-      const integrations = await this.composio.integrations.list({
-        appName: toolkit,
-      });
-      return integrations.length > 0;
-    } catch {
-      return false;
-    }
-  }
-}
-
-/**
- * Helper function to create a one-shot Composio MCP session
- * For quick tool usage without managing session lifecycle
- *
- * @example
- * ```ts
- * const tools = await getComposioTools({
- *   apiKey: process.env.COMPOSIO_API_KEY!,
- *   userId: 'user-123',
- *   toolkits: ['GMAIL', 'SLACK'],
- * });
- *
- * const result = await generateText({
- *   model: openai('gpt-4o'),
- *   prompt: 'Check my unread emails',
- *   tools,
- * });
- * ```
- */
-export async function getComposioTools(config: ComposioMCPConfig) {
-  const session = new ComposioMCPSession(config);
-  try {
-    return await session.getTools();
-  } finally {
-    await session.close();
-  }
-}
-
-/**
- * Create a reusable Composio MCP session
- * Better for multiple requests with the same toolkits
- *
- * @example
- * ```ts
- * const session = await createComposioSession({
- *   apiKey: process.env.COMPOSIO_API_KEY!,
- *   userId: 'user-123',
- *   toolkits: ['GMAIL', 'NOTION'],
- * });
- *
- * try {
- *   const tools = await session.getTools();
- *   const result = await streamText({ model, messages, tools });
- * } finally {
- *   await session.close();
- * }
- * ```
- */
-export async function createComposioSession(
-  config: ComposioMCPConfig
-): Promise<ComposioMCPSession> {
-  const session = new ComposioMCPSession(config);
-  // Pre-initialize the connection
-  await session.getClient();
-  return session;
-}