| 1 | import { type UIMessage } from "ai"; |
| 2 | |
| 3 | import { isWebSearchToolName } from "@/lib/ai/tool-names"; |
| 4 | import { |
| 5 | getToolInputArgs, |
| 6 | getToolName, |
| 7 | getToolOutput, |
| 8 | getToolState, |
| 9 | isToolPart, |
| 10 | } from "@/lib/ai/uiMessageParts"; |
| 11 | |
| 12 | type NotebookAgentToolCallState = "call" | "partial-call" | "result"; |
| 13 | |
| 14 | export interface NotebookAgentToolCall { |
| 15 | id: string; |
| 16 | toolName: string; |
| 17 | state: NotebookAgentToolCallState; |
| 18 | args?: Record<string, unknown>; |
| 19 | result?: unknown; |
| 20 | } |
| 21 | |
| 22 | export interface NotebookAgentAttachmentContext { |
| 23 | fileAssetId?: string; |
| 24 | fileName: string | null; |
| 25 | fileUrl: string; |
| 26 | processingStatus?: string | null; |
| 27 | ragId?: string | null; |
| 28 | } |
| 29 | |
| 30 | export interface NotebookAgentSelectedChunk { |
| 31 | chunkId: string; |
| 32 | ragId?: string; |
| 33 | slideNumber?: number | null; |
| 34 | content?: string; |
| 35 | } |
| 36 | |
| 37 | export const NOTEBOOK_AGENT_ACTIVITY_TOOL_NAMES = { |
| 38 | USER_INPUT: "__agent.user_input", |
| 39 | ASSISTANT_OUTPUT: "__agent.assistant_output", |
| 40 | STEP_START: "__agent.step_start", |
| 41 | SOURCE: "__agent.source", |
| 42 | REASONING: "__agent.reasoning", |
| 43 | DOCUMENT_CONTEXT: "__agent.document_context", |
| 44 | } as const; |
| 45 | |
| 46 | export function isNotebookAgentActivityEvent(toolName: string): boolean { |
| 47 | return toolName.startsWith("__agent."); |
| 48 | } |
| 49 | |
| 50 | export function parseNotebookAgentToolResult(result: unknown): unknown { |
| 51 | if (typeof result !== "string") { |
| 52 | return result; |
| 53 | } |
| 54 | |
| 55 | try { |
| 56 | return JSON.parse(result) as unknown; |
| 57 | } catch { |
| 58 | return result; |
| 59 | } |
| 60 | } |
| 61 | |
| 62 | export function collectNotebookAgentToolCalls( |
| 63 | messages: UIMessage[], |
| 64 | ): NotebookAgentToolCall[] { |
| 65 | const toolCallMap = new Map<string, NotebookAgentToolCall>(); |
| 66 | |
| 67 | messages.forEach((message, messageIndex) => { |
| 68 | message.parts?.forEach((part, partIndex) => { |
| 69 | if (!isToolPart(part)) { |
| 70 | return; |
| 71 | } |
| 72 | |
| 73 | const state = getToolState(part); |
| 74 | const toolCallId = |
| 75 | part.toolCallId ?? `${getToolName(part)}-${messageIndex}-${partIndex}`; |
| 76 | const existing = toolCallMap.get(toolCallId); |
| 77 | |
| 78 | if (existing?.state === "result" && state !== "result") { |
| 79 | return; |
| 80 | } |
| 81 | |
| 82 | const result = getToolOutput(part); |
| 83 | const parsedResult = |
| 84 | result === undefined |
| 85 | ? existing?.result |
| 86 | : parseNotebookAgentToolResult(result); |
| 87 | |
| 88 | toolCallMap.set(toolCallId, { |
| 89 | id: toolCallId, |
| 90 | toolName: getToolName(part), |
| 91 | state, |
| 92 | args: getToolInputArgs(part), |
| 93 | result: parsedResult, |
| 94 | }); |
| 95 | }); |
| 96 | }); |
| 97 | |
| 98 | return Array.from(toolCallMap.values()); |
| 99 | } |
| 100 | |
| 101 | function isLikelyFileName(value: string): boolean { |
| 102 | return /.+\.[a-z0-9]{2,8}$/i.test(value.trim()); |
| 103 | } |
| 104 | |
| 105 | function isLikelyOpaqueFileName(value: string): boolean { |
| 106 | const trimmed = value.trim(); |
| 107 | if (!trimmed) { |
| 108 | return false; |
| 109 | } |
| 110 | |
| 111 | if (isLikelyFileName(trimmed)) { |
| 112 | return false; |
| 113 | } |
| 114 | |
| 115 | return trimmed.length >= 20; |
| 116 | } |
| 117 | |
| 118 | function isRecord(value: unknown): value is Record<string, unknown> { |
| 119 | return typeof value === "object" && value !== null && !Array.isArray(value); |
| 120 | } |
| 121 | |
| 122 | function getLikelyUploadedFileNames( |
| 123 | toolCalls: NotebookAgentToolCall[], |
| 124 | ): string[] { |
| 125 | const fileNames: string[] = []; |
| 126 | |
| 127 | for (const call of toolCalls) { |
| 128 | if (call.toolName !== NOTEBOOK_AGENT_ACTIVITY_TOOL_NAMES.USER_INPUT) { |
| 129 | continue; |
| 130 | } |
| 131 | |
| 132 | const parsed = parseNotebookAgentToolResult(call.result); |
| 133 | if (!isRecord(parsed)) { |
| 134 | continue; |
| 135 | } |
| 136 | |
| 137 | const text = parsed.text; |
| 138 | if (typeof text !== "string") { |
| 139 | continue; |
| 140 | } |
| 141 | |
| 142 | const normalized = text.trim(); |
| 143 | if (!isLikelyFileName(normalized)) { |
| 144 | continue; |
| 145 | } |
| 146 | |
| 147 | fileNames.push(normalized); |
| 148 | } |
| 149 | |
| 150 | return fileNames; |
| 151 | } |
| 152 | |
| 153 | export function withNotebookAgentDocumentContext( |
| 154 | toolCalls: NotebookAgentToolCall[], |
| 155 | attachments: NotebookAgentAttachmentContext[], |
| 156 | ): NotebookAgentToolCall[] { |
| 157 | if (attachments.length === 0) { |
| 158 | return toolCalls; |
| 159 | } |
| 160 | |
| 161 | const nextCalls = [...toolCalls]; |
| 162 | const existingIds = new Set(nextCalls.map((call) => call.id)); |
| 163 | const fallbackFileNames = getLikelyUploadedFileNames(toolCalls); |
| 164 | let fallbackFileNameIndex = 0; |
| 165 | |
| 166 | for (const [attachmentIndex, attachment] of attachments.entries()) { |
| 167 | const fileIdentity = |
| 168 | attachment.ragId ?? |
| 169 | attachment.fileAssetId ?? |
| 170 | attachment.fileUrl ?? |
| 171 | attachment.fileName ?? |
| 172 | `attachment-${attachmentIndex}`; |
| 173 | const id = `event-document-context-${fileIdentity}`; |
| 174 | |
| 175 | if (existingIds.has(id)) { |
| 176 | continue; |
| 177 | } |
| 178 | |
| 179 | const normalizedStoredFileName = attachment.fileName?.trim() ?? ""; |
| 180 | const shouldUseFallbackFileName = |
| 181 | !normalizedStoredFileName || |
| 182 | normalizedStoredFileName === attachment.ragId || |
| 183 | isLikelyOpaqueFileName(normalizedStoredFileName); |
| 184 | const fallbackFileName = shouldUseFallbackFileName |
| 185 | ? fallbackFileNames[fallbackFileNameIndex] |
| 186 | : undefined; |
| 187 | |
| 188 | if (fallbackFileName) { |
| 189 | fallbackFileNameIndex += 1; |
| 190 | } |
| 191 | |
| 192 | const displayFileName = |
| 193 | fallbackFileName ?? |
| 194 | (shouldUseFallbackFileName ? "Document" : normalizedStoredFileName); |
| 195 | |
| 196 | nextCalls.push({ |
| 197 | id, |
| 198 | toolName: NOTEBOOK_AGENT_ACTIVITY_TOOL_NAMES.DOCUMENT_CONTEXT, |
| 199 | state: "result", |
| 200 | args: { |
| 201 | fileName: displayFileName, |
| 202 | }, |
| 203 | result: { |
| 204 | ragId: attachment.ragId ?? null, |
| 205 | fileAssetId: attachment.fileAssetId ?? null, |
| 206 | fileName: displayFileName, |
| 207 | fileUrl: attachment.fileUrl, |
| 208 | processingStatus: attachment.processingStatus ?? null, |
| 209 | }, |
| 210 | }); |
| 211 | existingIds.add(id); |
| 212 | } |
| 213 | |
| 214 | return nextCalls; |
| 215 | } |
| 216 | |
| 217 | export type NotebookAgentSearchResult = { |
| 218 | query: string; |
| 219 | results: unknown[]; |
| 220 | }; |
| 221 | |
| 222 | export function deriveNotebookAgentSearchResults( |
| 223 | toolCalls: NotebookAgentToolCall[], |
| 224 | ): { |
| 225 | webSearchResults: NotebookAgentSearchResult[]; |
| 226 | documentSearchResults: NotebookAgentSearchResult[]; |
| 227 | } { |
| 228 | const webSearchResults: NotebookAgentSearchResult[] = []; |
| 229 | const documentSearchResults: NotebookAgentSearchResult[] = []; |
| 230 | |
| 231 | for (const call of toolCalls) { |
| 232 | if (call.state !== "result") { |
| 233 | continue; |
| 234 | } |
| 235 | |
| 236 | const parsed = parseNotebookAgentToolResult(call.result) as |
| 237 | | { query?: string; results?: unknown[] } |
| 238 | | undefined; |
| 239 | |
| 240 | const query = |
| 241 | typeof call.args?.query === "string" |
| 242 | ? call.args.query |
| 243 | : typeof parsed?.query === "string" |
| 244 | ? parsed.query |
| 245 | : "Search"; |
| 246 | const results = Array.isArray(parsed?.results) ? parsed.results : []; |
| 247 | |
| 248 | if (call.toolName === "searchDocuments") { |
| 249 | documentSearchResults.push({ query, results }); |
| 250 | continue; |
| 251 | } |
| 252 | |
| 253 | if (isWebSearchToolName(call.toolName)) { |
| 254 | webSearchResults.push({ query, results }); |
| 255 | } |
| 256 | } |
| 257 | |
| 258 | return { |
| 259 | webSearchResults, |
| 260 | documentSearchResults, |
| 261 | }; |
| 262 | } |
| 263 |