1  import type {
  2    ImageBlockParam,
  3    TextBlockParam,
  4    ToolResultBlockParam,
  5  } from '@anthropic-ai/sdk/resources/index.mjs'
  6  import { BASH_TOOL_NAME } from '../tools/BashTool/toolName.js'
  7  import { formatOutput } from '../tools/BashTool/utils.js'
  8  import type {
  9    NotebookCell,
 10    NotebookCellOutput,
 11    NotebookCellSource,
 12    NotebookCellSourceOutput,
 13    NotebookContent,
 14    NotebookOutputImage,
 15  } from '../types/notebook.js'
 16  import { getFsImplementation } from './fsOperations.js'
 17  import { expandPath } from './path.js'
 18  import { jsonParse } from './slowOperations.js'
 19  
 20  const LARGE_OUTPUT_THRESHOLD = 10000
 21  
 22  function isLargeOutputs(
 23    outputs: (NotebookCellSourceOutput | undefined)[],
 24  ): boolean {
 25    let size = 0
 26    for (const o of outputs) {
 27      if (!o) continue
 28      size += (o.text?.length ?? 0) + (o.image?.image_data.length ?? 0)
 29      if (size > LARGE_OUTPUT_THRESHOLD) return true
 30    }
 31    return false
 32  }
 33  
 34  function processOutputText(text: string | string[] | undefined): string {
 35    if (!text) return ''
 36    const rawText = Array.isArray(text) ? text.join('') : text
 37    const { truncatedContent } = formatOutput(rawText)
 38    return truncatedContent
 39  }
 40  
 41  function extractImage(
 42    data: Record<string, unknown>,
 43  ): NotebookOutputImage | undefined {
 44    if (typeof data['image/png'] === 'string') {
 45      return {
 46        image_data: data['image/png'].replace(/\s/g, ''),
 47        media_type: 'image/png',
 48      }
 49    }
 50    if (typeof data['image/jpeg'] === 'string') {
 51      return {
 52        image_data: data['image/jpeg'].replace(/\s/g, ''),
 53        media_type: 'image/jpeg',
 54      }
 55    }
 56    return undefined
 57  }
 58  
 59  function processOutput(output: NotebookCellOutput) {
 60    switch (output.output_type) {
 61      case 'stream':
 62        return {
 63          output_type: output.output_type,
 64          text: processOutputText(output.text),
 65        }
 66      case 'execute_result':
 67      case 'display_data':
 68        return {
 69          output_type: output.output_type,
 70          text: processOutputText(output.data?.['text/plain']),
 71          image: output.data && extractImage(output.data),
 72        }
 73      case 'error':
 74        return {
 75          output_type: output.output_type,
 76          text: processOutputText(
 77            `${output.ename}: ${output.evalue}\n${output.traceback.join('\n')}`,
 78          ),
 79        }
 80    }
 81  }
 82  
 83  function processCell(
 84    cell: NotebookCell,
 85    index: number,
 86    codeLanguage: string,
 87    includeLargeOutputs: boolean,
 88  ): NotebookCellSource {
 89    const cellId = cell.id ?? `cell-${index}`
 90    const cellData: NotebookCellSource = {
 91      cellType: cell.cell_type,
 92      source: Array.isArray(cell.source) ? cell.source.join('') : cell.source,
 93      execution_count:
 94        cell.cell_type === 'code' ? cell.execution_count || undefined : undefined,
 95      cell_id: cellId,
 96    }
 97    // Avoid giving text cells the code language.
 98    if (cell.cell_type === 'code') {
 99      cellData.language = codeLanguage
100    }
101  
102    if (cell.cell_type === 'code' && cell.outputs?.length) {
103      const outputs = cell.outputs.map(processOutput)
104      if (!includeLargeOutputs && isLargeOutputs(outputs)) {
105        cellData.outputs = [
106          {
107            output_type: 'stream',
108            text: `Outputs are too large to include. Use ${BASH_TOOL_NAME} with: cat <notebook_path> | jq '.cells[${index}].outputs'`,
109          },
110        ]
111      } else {
112        cellData.outputs = outputs
113      }
114    }
115  
116    return cellData
117  }
118  
119  function cellContentToToolResult(cell: NotebookCellSource): TextBlockParam {
120    const metadata = []
121    if (cell.cellType !== 'code') {
122      metadata.push(`<cell_type>${cell.cellType}</cell_type>`)
123    }
124    if (cell.language !== 'python' && cell.cellType === 'code') {
125      metadata.push(`<language>${cell.language}</language>`)
126    }
127    const cellContent = `<cell id="${cell.cell_id}">${metadata.join('')}${cell.source}</cell id="${cell.cell_id}">`
128    return {
129      text: cellContent,
130      type: 'text',
131    }
132  }
133  
134  function cellOutputToToolResult(output: NotebookCellSourceOutput) {
135    const outputs: (TextBlockParam | ImageBlockParam)[] = []
136    if (output.text) {
137      outputs.push({
138        text: `\n${output.text}`,
139        type: 'text',
140      })
141    }
142    if (output.image) {
143      outputs.push({
144        type: 'image',
145        source: {
146          data: output.image.image_data,
147          media_type: output.image.media_type,
148          type: 'base64',
149        },
150      })
151    }
152    return outputs
153  }
154  
155  function getToolResultFromCell(cell: NotebookCellSource) {
156    const contentResult = cellContentToToolResult(cell)
157    const outputResults = cell.outputs?.flatMap(cellOutputToToolResult)
158    return [contentResult, ...(outputResults ?? [])]
159  }
160  
161  /**
162   * Reads and parses a Jupyter notebook file into processed cell data
163   */
164  export async function readNotebook(
165    notebookPath: string,
166    cellId?: string,
167  ): Promise<NotebookCellSource[]> {
168    const fullPath = expandPath(notebookPath)
169    const buffer = await getFsImplementation().readFileBytes(fullPath)
170    const content = buffer.toString('utf-8')
171    const notebook = jsonParse(content) as NotebookContent
172    const language = notebook.metadata.language_info?.name ?? 'python'
173    if (cellId) {
174      const cell = notebook.cells.find(c => c.id === cellId)
175      if (!cell) {
176        throw new Error(`Cell with ID "${cellId}" not found in notebook`)
177      }
178      return [processCell(cell, notebook.cells.indexOf(cell), language, true)]
179    }
180    return notebook.cells.map((cell, index) =>
181      processCell(cell, index, language, false),
182    )
183  }
184  
185  /**
186   * Maps notebook cell data to tool result block parameters with sophisticated text block merging
187   */
188  export function mapNotebookCellsToToolResult(
189    data: NotebookCellSource[],
190    toolUseID: string,
191  ): ToolResultBlockParam {
192    const allResults = data.flatMap(getToolResultFromCell)
193  
194    // Merge adjacent text blocks
195    return {
196      tool_use_id: toolUseID,
197      type: 'tool_result' as const,
198      content: allResults.reduce<(TextBlockParam | ImageBlockParam)[]>(
199        (acc, curr) => {
200          if (acc.length === 0) return [curr]
201  
202          const prev = acc[acc.length - 1]
203          if (prev && prev.type === 'text' && curr.type === 'text') {
204            // Merge the text blocks
205            prev.text += '\n' + curr.text
206            return acc
207          }
208  
209          acc.push(curr)
210          return acc
211        },
212        [],
213      ),
214    }
215  }
216  
217  export function parseCellId(cellId: string): number | undefined {
218    const match = cellId.match(/^cell-(\d+)$/)
219    if (match && match[1]) {
220      const index = parseInt(match[1], 10)
221      return isNaN(index) ? undefined : index
222    }
223    return undefined
224  }