From 24c50ebd0ddd895207bc17e8c9275d47d4e0f991 Mon Sep 17 00:00:00 2001 From: Benjamin Shafii Date: Wed, 1 Jan 2025 15:11:55 +0100 Subject: [PATCH] feat(plugin): add leightweight context mode --- .../plugin/views/assistant/ai-chat/chat.tsx | 40 +- .../ai-chat/components/example-prompts.tsx | 42 +- .../ai-chat/context-limit-indicator.tsx | 41 +- .../ai-chat/tool-handlers/onboard-handler.tsx | 204 ++++----- .../assistant/ai-chat/use-context-items.ts | 64 ++- packages/web/app/api/(newai)/chat/route.ts | 411 +++++++++--------- 6 files changed, 443 insertions(+), 359 deletions(-) diff --git a/packages/plugin/views/assistant/ai-chat/chat.tsx b/packages/plugin/views/assistant/ai-chat/chat.tsx index 30a1decd..57271d0b 100644 --- a/packages/plugin/views/assistant/ai-chat/chat.tsx +++ b/packages/plugin/views/assistant/ai-chat/chat.tsx @@ -50,6 +50,7 @@ export const ChatComponent: React.FC = ({ currentFile, screenpipe, textSelections, + isLightweightMode, } = useContextItems(); const uniqueReferences = getUniqueReferences(); @@ -72,9 +73,42 @@ export const ChatComponent: React.FC = ({ }); const contextString = React.useMemo(() => { - const contextString = JSON.stringify(contextItems); - return contextString; - }, [contextItems]); + if (isLightweightMode) { + // In lightweight mode, only include metadata + const lightweightContext = { + files: Object.fromEntries( + Object.entries(files).map(([id, file]) => [ + id, + { ...file, content: '' } + ]) + ), + folders: Object.fromEntries( + Object.entries(folders).map(([id, folder]) => [ + id, + { ...folder, files: folder.files.map(f => ({ ...f, content: '' })) } + ]) + ), + tags: Object.fromEntries( + Object.entries(tags).map(([id, tag]) => [ + id, + { ...tag, files: tag.files.map(f => ({ ...f, content: '' })) } + ]) + ), + searchResults: Object.fromEntries( + Object.entries(searchResults).map(([id, search]) => [ + id, + { ...search, results: search.results.map(r => ({ ...r, content: '' })) } + ]) + ), + // Keep these as is + currentFile: currentFile ? { ...currentFile, content: '' } : null, + screenpipe, + textSelections, + }; + return JSON.stringify(lightweightContext); + } + return JSON.stringify(contextItems); + }, [contextItems, isLightweightMode]); logger.debug("contextString", contextString); const chatBody = { diff --git a/packages/plugin/views/assistant/ai-chat/components/example-prompts.tsx b/packages/plugin/views/assistant/ai-chat/components/example-prompts.tsx index e20d5d18..ce900192 100644 --- a/packages/plugin/views/assistant/ai-chat/components/example-prompts.tsx +++ b/packages/plugin/views/assistant/ai-chat/components/example-prompts.tsx @@ -7,18 +7,18 @@ interface Example { } const examples: Example[] = [ +// { +// prompt: "Move all my untitled notes to an 'Inbox' folder", +// description: "File organization", +// icon: "📁" +// }, +// { +// prompt: "Rename my daily notes to include topics discussed", +// description: "Smart renaming", +// icon: "✏️" +// }, { - prompt: "Move all my untitled notes to an 'Inbox' folder", - description: "File organization", - icon: "📁" - }, - { - prompt: "Rename my daily notes to include topics discussed", - description: "Smart renaming", - icon: "✏️" - }, - { - prompt: "Add a summary section to my meeting notes", + prompt: "Add a summary section at the bottom of my current note", description: "Content editing", icon: "➕" }, @@ -27,16 +27,16 @@ const examples: Example[] = [ description: "Smart search", icon: "🔍" }, - { - prompt: "Tag my book notes with relevant categories", - description: "Auto-tagging", - icon: "🏷️" - }, - { - prompt: "Analyze my vault structure and suggest improvements", - description: "Vault analysis", - icon: "📊" - }, +// { +// prompt: "Tag my book notes with relevant categories", +// description: "Auto-tagging", +// icon: "🏷️" +// }, +// { +// prompt: "Analyze my vault structure and suggest improvements", +// description: "Vault analysis", +// icon: "📊" +// }, { prompt: "Get a summary of my day from Screenpipe", description: "Daily summary", diff --git a/packages/plugin/views/assistant/ai-chat/context-limit-indicator.tsx b/packages/plugin/views/assistant/ai-chat/context-limit-indicator.tsx index 076d98f2..11e71deb 100644 --- a/packages/plugin/views/assistant/ai-chat/context-limit-indicator.tsx +++ b/packages/plugin/views/assistant/ai-chat/context-limit-indicator.tsx @@ -3,6 +3,7 @@ import { init, get_encoding } from "tiktoken/init"; import wasmBinary from "tiktoken/tiktoken_bg.wasm"; import { useDebouncedCallback } from 'use-debounce'; import { logger } from "../../../services/logger"; +import { useContextItems } from "./use-context-items"; interface TokenStats { contextSize: number; @@ -19,10 +20,10 @@ export function ContextLimitIndicator({ const [stats, setStats] = React.useState({ contextSize: 0, percentUsed: 0 }); const [error, setError] = React.useState(); const [tiktokenInitialized, setTiktokenInitialized] = React.useState(false); + const { isLightweightMode, toggleLightweightMode } = useContextItems(); // Initialize encoder once on mount React.useEffect(() => { - async function setup() { try { if (!tiktokenInitialized) { @@ -39,7 +40,6 @@ export function ContextLimitIndicator({ // Debounced token calculation const calculateTokens = useDebouncedCallback((text: string) => { - console.log("calculateTokens", { text, tiktokenInitialized }); if (!text || !tiktokenInitialized) return; const encoder = get_encoding("cl100k_base"); @@ -67,14 +67,39 @@ export function ContextLimitIndicator({ } const isOverLimit = stats.contextSize > maxContextSize; + const shouldWarn = stats.percentUsed > 80; return ( -
- - {isOverLimit ? "Context size exceeds maximum" : "Context used"} - - {stats.percentUsed.toFixed(0)}% +
+
+ + {isOverLimit + ? "Context size exceeds maximum" + : shouldWarn + ? "Context size nearing limit" + : "Context used"} + + {stats.percentUsed.toFixed(0)}% +
+ +
); } \ No newline at end of file diff --git a/packages/plugin/views/assistant/ai-chat/tool-handlers/onboard-handler.tsx b/packages/plugin/views/assistant/ai-chat/tool-handlers/onboard-handler.tsx index 00cfe0cc..74615645 100644 --- a/packages/plugin/views/assistant/ai-chat/tool-handlers/onboard-handler.tsx +++ b/packages/plugin/views/assistant/ai-chat/tool-handlers/onboard-handler.tsx @@ -1,75 +1,75 @@ import React, { useState } from "react"; import { TFile, TFolder } from "obsidian"; import { ToolHandlerProps } from "./types"; - - - +import { useContextItems } from "../use-context-items"; +import { usePlugin } from "../../provider"; export function OnboardHandler({ toolInvocation, handleAddResult, - app, }: ToolHandlerProps) { const [isAnalyzing, setIsAnalyzing] = useState(false); const [isValidated, setIsValidated] = useState(false); + const { toggleLightweightMode } = useContextItems(); + const plugin = usePlugin(); const getFilesFromPath = (path: string): TFile[] => { + const allUserFiles = plugin.getAllUserMarkdownFiles(); + if (path === "/") { - return app.vault.getMarkdownFiles(); - } - - const folder = app.vault.getAbstractFileByPath(path); - if (!folder || !(folder instanceof TFolder)) { - return []; + return allUserFiles; } - const files: TFile[] = []; - folder.children.forEach(child => { - if (child instanceof TFile) { - files.push(child); - } + // Filter files that belong to the specified path + return allUserFiles.filter(file => { + const filePath = file.path; + return filePath.startsWith(path + "/") || filePath === path; }); - return files; }; const analyzeFolderStructure = async ( path: string, depth = 0, - maxDepth = 3, + maxDepth = 3 ) => { + toggleLightweightMode(); + const files = getFilesFromPath(path); const structure = { path, - files: await Promise.all(files.map(async file => { - const fileData = { - name: file.name, - path: file.path, - content: await app.vault.read(file), - type: "file" as const, - depth: depth + 1, - }; - - - - return fileData; - })), + files: await Promise.all( + files.map(async file => { + const fileData = { + name: file.name, + path: file.path, + type: "file" as const, + depth: depth + 1, + }; + return fileData; + }) + ), subfolders: [], depth, }; if (depth < maxDepth && path !== "/") { - const folder = app.vault.getAbstractFileByPath(path) as TFolder; - if (folder && folder instanceof TFolder) { - for (const child of folder.children) { - if (child instanceof TFolder) { - const subStructure = await analyzeFolderStructure( - child.path, - depth + 1, - maxDepth - ); - structure.subfolders.push(subStructure); - } - } + // Get all user folders at current path + const userFolders = plugin.getAllUserFolders().filter(folderPath => { + // Only include direct subfolders of current path + const isSubfolder = folderPath.startsWith(path + "/"); + const folderDepth = folderPath.split("/").length; + const currentDepth = path.split("/").length; + return isSubfolder && folderDepth === currentDepth + 1; + }); + + // Analyze each subfolder + for (const folderPath of userFolders) { + const subStructure = await analyzeFolderStructure( + folderPath, + depth + 1, + maxDepth + ); + structure.subfolders.push(subStructure); } } @@ -78,88 +78,74 @@ export function OnboardHandler({ const handleAnalyze = async () => { setIsAnalyzing(true); + toggleLightweightMode(); try { const { path = "/", maxDepth = 3 } = toolInvocation.args; const structure = await analyzeFolderStructure(path, 0, maxDepth); - + setIsValidated(true); - handleAddResult( - JSON.stringify({ - success: true, - structure, - analyzedPath: path, - message: `Vault structure analyzed successfully for path: ${path}`, - }) - ); + + // Prepare analysis data in the format expected by generateSettings + const analysisData = { + structure, + stats: { + totalFiles: structure.files.length, + fileTypes: structure.files.reduce((acc, file) => { + const ext = file.name.split(".").pop() || "no-extension"; + acc[ext] = (acc[ext] || 0) + 1; + return acc; + }, {} as Record), + folderCount: structure.subfolders.length, + maxDepth: maxDepth, + }, + }; + + console.log("analysisData", analysisData); + handleAddResult("goodbye"); } catch (error) { + console.error("Analysis error:", error); handleAddResult( JSON.stringify({ success: false, - error: error.message, + error: error.message || "An error occurred during analysis", }) ); + } finally { + setIsAnalyzing(false); } - setIsAnalyzing(false); }; - const renderStructurePreview = () => ( -
- This will analyze your vault structure up to 3 levels deep to suggest optimal organization. - The analysis will: -
    -
  • Scan your folder hierarchy
    • -
  • Analyze naming patterns
    • -
  • Identify common groupings
    • -
  • Suggest organizational improvements
    • -
-
- ); - return ( -
-
- Would you like to analyze your vault structure for better organization? +
+
+ This will analyze your vault structure to suggest optimal organization + and settings. The analysis will: +
    +
  • Scan your folder hierarchy
    • +
  • Analyze file naming patterns
    • +
  • Identify common groupings
    • +
  • Generate recommended settings
    • +
- {renderStructurePreview()} - - {!isValidated && !("result" in toolInvocation) && ( -
- -
- )} +
+ +
); -} \ No newline at end of file +} diff --git a/packages/plugin/views/assistant/ai-chat/use-context-items.ts b/packages/plugin/views/assistant/ai-chat/use-context-items.ts index 5acbeeb0..24fee689 100644 --- a/packages/plugin/views/assistant/ai-chat/use-context-items.ts +++ b/packages/plugin/views/assistant/ai-chat/use-context-items.ts @@ -78,6 +78,7 @@ type ContextCollections = { interface ContextItemsState extends ContextCollections { currentFile: FileContextItem | null; includeCurrentFile: boolean; + isLightweightMode: boolean; // Actions for each type addFile: (file: FileContextItem) => void; @@ -93,6 +94,7 @@ interface ContextItemsState extends ContextCollections { setCurrentFile: (file: FileContextItem | null) => void; toggleCurrentFile: () => void; clearAll: () => void; + toggleLightweightMode: () => void; // Processing methods processFolderFiles: ( @@ -116,73 +118,105 @@ export const useContextItems = create((set, get) => ({ textSelections: {}, currentFile: null, includeCurrentFile: true, + isLightweightMode: false, - // Updated add actions with reference checking + // Add toggle function + toggleLightweightMode: () => set(state => ({ isLightweightMode: !state.isLightweightMode })), + + // Update addFile to handle lightweight mode addFile: file => set(state => { const existingItemIndex = Object.values(state.files).findIndex( item => item.reference === file.reference ); + const lightweightFile = state.isLightweightMode ? { + ...file, + content: '', // Remove content in lightweight mode + } : file; + if (existingItemIndex !== -1) { - // Replace existing item return { files: { ...state.files, - [file.id]: { ...file, createdAt: Date.now() }, + [file.id]: { ...lightweightFile, createdAt: Date.now() }, }, }; } return { - files: { ...state.files, [file.id]: file }, + files: { ...state.files, [file.id]: lightweightFile }, }; }), + // Update addFolder to handle lightweight mode addFolder: folder => set(state => { const existingItemIndex = Object.values(state.folders).findIndex( item => item.reference === folder.reference ); + const lightweightFolder = state.isLightweightMode ? { + ...folder, + files: folder.files.map(f => ({ ...f, content: '' })), // Remove content in lightweight mode + } : folder; + if (existingItemIndex !== -1) { return { folders: { ...state.folders, - [folder.id]: { ...folder, createdAt: Date.now() }, + [folder.id]: { ...lightweightFolder, createdAt: Date.now() }, }, }; } return { - folders: { ...state.folders, [folder.id]: folder }, + folders: { ...state.folders, [folder.id]: lightweightFolder }, }; }), + // Add YouTube video without lightweight mode addYouTubeVideo: video => set(state => ({ youtubeVideos: { ...state.youtubeVideos, [video.id]: video }, })), - addTag: tag => - set(state => ({ - tags: { ...state.tags, [tag.id]: tag }, - })), - + // Add Screenpipe data without lightweight mode addScreenpipe: data => set(state => ({ screenpipe: { ...state.screenpipe, [data.id]: data }, })), + // Update addTag to handle lightweight mode + addTag: tag => + set(state => { + const lightweightTag = state.isLightweightMode ? { + ...tag, + files: tag.files.map(f => ({ ...f, content: '' })), // Remove content in lightweight mode + } : tag; + + return { + tags: { ...state.tags, [tag.id]: lightweightTag }, + }; + }), + + // Update addSearchResults to handle lightweight mode addSearchResults: search => - set(state => ({ - searchResults: { ...state.searchResults, [search.id]: search }, - })), + set(state => { + const lightweightSearch = state.isLightweightMode ? { + ...search, + results: search.results.map(r => ({ ...r, content: '' })), // Remove content in lightweight mode + } : search; + + return { + searchResults: { ...state.searchResults, [search.id]: lightweightSearch }, + }; + }), + // Add text selection without lightweight mode addTextSelection: selection => set(state => { const reference = selection.reference; - // Remove any existing items with same reference first get().removeByReference(reference); return { diff --git a/packages/web/app/api/(newai)/chat/route.ts b/packages/web/app/api/(newai)/chat/route.ts index f77ea9b1..e641dfdd 100644 --- a/packages/web/app/api/(newai)/chat/route.ts +++ b/packages/web/app/api/(newai)/chat/route.ts @@ -1,4 +1,4 @@ -import { convertToCoreMessages, streamText, StreamData, tool } from "ai"; +import { convertToCoreMessages, streamText, createDataStreamResponse, generateId } from "ai"; import { NextResponse, NextRequest } from "next/server"; import { incrementAndLogTokenUsage } from "@/lib/incrementAndLogTokenUsage"; import { handleAuthorization } from "@/lib/handleAuthorization"; @@ -22,222 +22,227 @@ const settingsSchema = z.object({ }); export async function POST(req: NextRequest) { - try { - const { userId } = await handleAuthorization(req); - const { - messages, - newUnifiedContext, - enableScreenpipe, - currentDatetime, - unifiedContext: oldUnifiedContext, - model: bodyModel, - } = await req.json(); + return createDataStreamResponse({ + execute: async (dataStream) => { + try { + const { userId } = await handleAuthorization(req); + const { + messages, + newUnifiedContext, + enableScreenpipe, + currentDatetime, + unifiedContext: oldUnifiedContext, + model: bodyModel, + } = await req.json(); - const chosenModelName = bodyModel; - console.log("Chat using model:", chosenModelName); - const model = getModel(chosenModelName); + const chosenModelName = bodyModel; + console.log("Chat using model:", chosenModelName); + const model = getModel(chosenModelName); - const contextString = - newUnifiedContext || - oldUnifiedContext - ?.map((file) => { - return `File: ${file.title}\n\nContent:\n${file.content}\nPath: ${file.path} Reference: ${file.reference}`; - }) - .join("\n\n"); + const contextString = + newUnifiedContext || + oldUnifiedContext + ?.map((file) => { + return `File: ${file.title}\n\nContent:\n${file.content}\nPath: ${file.path} Reference: ${file.reference}`; + }) + .join("\n\n"); - const data = new StreamData(); + dataStream.writeData('initialized call'); - const result = await streamText({ - model, - system: getChatSystemPrompt( - contextString, - enableScreenpipe, - currentDatetime - ), - maxSteps: 3, - messages: convertToCoreMessages(messages), - tools: { - getSearchQuery: { - description: "Extract queries to search for notes", - parameters: z.object({ - query: z - .string() - .describe("The search query to find relevant notes"), - }), - }, - searchByName: { - description: "Search for files by name pattern", - parameters: z.object({ - query: z - .string() - .describe( - "The name pattern to search for (e.g., 'Untitled*' or exact name)" - ), - }), - }, - getYoutubeVideoId: { - description: "Get the YouTube video ID from a URL", - parameters: z.object({ - videoId: z.string().describe("The YouTube video ID"), - }), - }, - getLastModifiedFiles: { - description: "Get the last modified files in the vault", - parameters: z.object({ - count: z - .number() - .describe("The number of last modified files to retrieve"), - }), - }, - onboardUser: { - description: "Onboard the user to the vault", - parameters: z.object({}), - }, - appendContentToFile: { - description: "Append content to a file with user confirmation", - parameters: z.object({ - content: z.string().describe("The content to append to the file"), - message: z - .string() - .describe("Message to show to the user for confirmation"), - fileName: z - .string() - .optional() - .describe("Optional specific file to append to"), - }), - }, - addTextToDocument: { - description: "Adds the text to the current document when the user requests to do so", - parameters: z.object({ - content: z.string().describe("The text content to add to the document"), - path: z.string().optional().describe("Optional path to the document. If not provided, uses current document"), - }), - }, - modifyDocumentText: { - description: "Modifies the text in the current document or selected text according to user's request", - parameters: z.object({ - content: z.string().describe("The new text content to replace the current content or selection"), - path: z.string().optional().describe("Optional path to the document. If not provided, uses current document"), - }), - }, - generateSettings: { - description: - "Generate vault organization settings based on user preferences", - parameters: settingsSchema, - }, - analyzeVaultStructure: { - description: - "Analyze vault structure to suggest organization improvements", - parameters: z.object({ - path: z - .string() - .describe( - "Path to analyze. Use '/' for all files or specific folder path" - ), - maxDepth: z - .number() - .optional() - .describe("Maximum depth to analyze"), - addToContext: z - .boolean() - .optional() - .describe("Whether to add analyzed files to context"), - }), - }, - getScreenpipeDailySummary: { - description: "Get a summary of the user's day using Screenpipe data", - parameters: z.object({ - startTime: z - .string() - .optional() - .describe("Start time in ISO format"), - endTime: z.string().optional().describe("End time in ISO format"), - }), - }, - moveFiles: { - description: "Move files to their designated folders", - parameters: z.object({ - moves: z.array( - z.object({ - sourcePath: z + const result = await streamText({ + model, + system: getChatSystemPrompt( + contextString, + enableScreenpipe, + currentDatetime + ), + maxSteps: 3, + messages: convertToCoreMessages(messages), + tools: { + getSearchQuery: { + description: "Extract queries to search for notes", + parameters: z.object({ + query: z + .string() + .describe("The search query to find relevant notes"), + }), + }, + searchByName: { + description: "Search for files by name pattern", + parameters: z.object({ + query: z + .string() + .describe( + "The name pattern to search for (e.g., 'Untitled*' or exact name)" + ), + }), + }, + getYoutubeVideoId: { + description: "Get the YouTube video ID from a URL", + parameters: z.object({ + videoId: z.string().describe("The YouTube video ID"), + }), + }, + getLastModifiedFiles: { + description: "Get the last modified files in the vault", + parameters: z.object({ + count: z + .number() + .describe("The number of last modified files to retrieve"), + }), + }, + onboardUser: { + description: "Onboard the user to the vault", + parameters: z.object({}), + }, + appendContentToFile: { + description: "Append content to a file with user confirmation", + parameters: z.object({ + content: z.string().describe("The content to append to the file"), + message: z + .string() + .describe("Message to show to the user for confirmation"), + fileName: z + .string() + .optional() + .describe("Optional specific file to append to"), + }), + }, + addTextToDocument: { + description: "Adds the text to the current document when the user requests to do so", + parameters: z.object({ + content: z.string().describe("The text content to add to the document"), + path: z.string().optional().describe("Optional path to the document. If not provided, uses current document"), + }), + }, + modifyDocumentText: { + description: "Modifies the text in the current document or selected text according to user's request", + parameters: z.object({ + content: z.string().describe("The new text content to replace the current content or selection"), + path: z.string().optional().describe("Optional path to the document. If not provided, uses current document"), + }), + }, + generateSettings: { + description: + "Generate vault organization settings based on user preferences", + parameters: settingsSchema, + }, + analyzeVaultStructure: { + description: + "Analyze vault structure to suggest organization improvements", + parameters: z.object({ + path: z .string() .describe( - "Source path (e.g., '/' for root, or specific folder path)" + "Path to analyze. Use '/' for all files or specific folder path" ), - destinationPath: z.string().describe("Destination folder path"), - pattern: z - .object({ - namePattern: z + maxDepth: z + .number() + .optional() + .describe("Maximum depth to analyze"), + addToContext: z + .boolean() + .optional() + .describe("Whether to add analyzed files to context"), + }), + }, + getScreenpipeDailySummary: { + description: "Get a summary of the user's day using Screenpipe data", + parameters: z.object({ + startTime: z + .string() + .optional() + .describe("Start time in ISO format"), + endTime: z.string().optional().describe("End time in ISO format"), + }), + }, + moveFiles: { + description: "Move files to their designated folders", + parameters: z.object({ + moves: z.array( + z.object({ + sourcePath: z .string() - .optional() .describe( - "File name pattern to match (e.g., 'untitled-*')" + "Source path (e.g., '/' for root, or specific folder path)" ), - extension: z + destinationPath: z.string().describe("Destination folder path"), + pattern: z + .object({ + namePattern: z + .string() + .optional() + .describe( + "File name pattern to match (e.g., 'untitled-*')" + ), + extension: z + .string() + .optional() + .describe("File extension to match"), + }) + .optional(), + }) + ), + message: z.string().describe("Confirmation message to show user"), + }), + }, + renameFiles: { + description: "Rename files based on pattern or criteria", + parameters: z.object({ + files: z.array( + z.object({ + oldPath: z + .string() + .describe("The current full path of the file"), + newName: z .string() - .optional() - .describe("File extension to match"), + .describe("Proposed new file name (no directories)"), }) - .optional(), - }) - ), - message: z.string().describe("Confirmation message to show user"), - }), - }, - renameFiles: { - description: "Rename files based on pattern or criteria", - parameters: z.object({ - files: z.array( - z.object({ - oldPath: z + ), + message: z.string().describe("Confirmation message to show user"), + }), + }, + executeActionsOnFileBasedOnPrompt: { + description: + "Analyze file content and apply one of (recommendTags & appendTag), (recommendFolders & moveFile), or (recommendName & moveFile)", + parameters: z.object({ + filePaths: z + .array(z.string()) + .describe("List of file paths to analyze"), + userPrompt: z .string() - .describe("The current full path of the file"), - newName: z - .string() - .describe("Proposed new file name (no directories)"), - }) - ), - message: z.string().describe("Confirmation message to show user"), - }), - }, - executeActionsOnFileBasedOnPrompt: { - description: - "Analyze file content and apply one of (recommendTags & appendTag), (recommendFolders & moveFile), or (recommendName & moveFile)", - parameters: z.object({ - filePaths: z - .array(z.string()) - .describe("List of file paths to analyze"), - userPrompt: z - .string() - .describe( - "User instructions to decide how to rename or re-tag or re-folder the files" - ), - }), - }, - }, - onFinish: async ({ usage, experimental_providerMetadata }) => { - console.log("Token usage:", usage); - const googleMetadata = experimental_providerMetadata?.google as unknown as GoogleGenerativeAIProviderMetadata | undefined; - console.log("Google metadata:", JSON.stringify(googleMetadata, null, 2)); + .describe( + "User instructions to decide how to rename or re-tag or re-folder the files" + ), + }), + }, + }, + onFinish: async ({ usage, experimental_providerMetadata }) => { + console.log("Token usage:", usage); + const googleMetadata = experimental_providerMetadata?.google as unknown as GoogleGenerativeAIProviderMetadata | undefined; + console.log("Google metadata:", JSON.stringify(googleMetadata, null, 2)); - if (googleMetadata?.groundingMetadata) { - data.appendMessageAnnotation({ - type: "search-results", - groundingMetadata: googleMetadata.groundingMetadata - }); - } + if (googleMetadata?.groundingMetadata) { + dataStream.writeMessageAnnotation({ + type: "search-results", + groundingMetadata: googleMetadata.groundingMetadata + }); + } - await incrementAndLogTokenUsage(userId, usage.totalTokens); - data.close(); - }, - }); + await incrementAndLogTokenUsage(userId, usage.totalTokens); + dataStream.writeData('call completed'); + }, + }); - return result.toDataStreamResponse({ data }); - } catch (error) { - console.error("Error in POST request:", error); - return NextResponse.json( - { error: error.message || "An unexpected error occurred" }, - { status: error.status || 500 } - ); - } + result.mergeIntoDataStream(dataStream); + } catch (error) { + console.error("Error in POST request:", error); + throw error; + } + }, + onError: (error) => { + console.error("Error in stream:", error); + return error instanceof Error ? error.message : String(error); + }, + }); }