Luigit
repositories / will

will

owned by admin

src/agent/zai/tools.ts

Raw
import { randomUUID } from 'node:crypto'
import { defineTool } from '@earendil-works/pi-coding-agent'
import { Type } from 'typebox'
import type { ServiceClient } from '../service-client.ts'
import type { ZaiRestClient } from './client.ts'

/**
 * Curated ZAI capability tools with stable will-owned names and schemas
 * (spec AC19/AC20: no dynamic MCP passthrough; results cached by
 * digest + capability + model + schema version).
 *
 * Provider limits enforced explicitly (docs 2026-09):
 * - transcription: wav/mp3, <=25 MB, <=30 s audio
 * - layout parsing: images <=10 MB, PDF <=50 MB / 30 pages
 * - vision MCP video: <=8 MB (enforced in mcp.ts)
 */

const ANALYSIS_SCHEMA_VERSION = 1

export interface ZaiToolDeps {
  client: ZaiRestClient
  service: ServiceClient
  analysisModel: string
}

export class CapabilityError extends Error {}

function dataUrl(_digest: string, type: string, bytes: Uint8Array): string {
  const b64 = Buffer.from(bytes).toString('base64')
  return `data:${type};base64,${b64}`
}

async function cachedAnalysis(
  deps: ZaiToolDeps,
  key: { digest: string; capability: string; model: string },
  run: () => Promise<{ result: unknown; usage?: unknown }>,
): Promise<{ fromCache: boolean; result: unknown }> {
  const found = await deps.service.getAnalysis({
    digest: key.digest,
    capability: key.capability,
    model: key.model,
    schemaVersion: ANALYSIS_SCHEMA_VERSION,
  })
  if (found) return { fromCache: true, result: found.result }
  const { result, usage } = await run()
  await deps.service.putAnalysis({
    digest: key.digest,
    capability: key.capability,
    model: key.model,
    schemaVersion: ANALYSIS_SCHEMA_VERSION,
    result,
    ...(usage !== undefined ? { usage } : {}),
  })
  return { fromCache: false, result }
}

function textResult(payload: unknown, isError = false) {
  return {
    content: [{ type: 'text' as const, text: JSON.stringify(payload) }],
    ...(isError ? { isError: true } : {}),
    details: {},
  }
}

export function createZaiTools(deps: ZaiToolDeps) {
  const zaiAnalyzeMedia = defineTool({
    name: 'zai_analyze_media',
    label: 'Analyze media',
    description:
      'Analyze an image (or image-bearing document) with the multimodal coding model. ' +
      'Input is a media digest from chat or the media store; ask a concrete question. ' +
      'Cached by digest and question-independent schema: repeated calls are cheap.',
    parameters: Type.Object({
      digest: Type.String({ description: 'sha256 media digest' }),
      question: Type.String({ description: 'What to look at or describe' }),
    }),
    async execute(_id, params) {
      try {
        const meta = await deps.service.recall({
          kind: 'analysis',
          digest: params.digest,
          limit: 1,
        })
        void meta
        const bytes = await deps.service.getMedia(params.digest)
        if (bytes.byteLength > 10 << 20)
          throw new CapabilityError('image exceeds 10 MB analysis limit')
        const cached = await cachedAnalysis(
          deps,
          {
            digest: params.digest,
            capability: 'zai_analyze_media',
            model: deps.analysisModel,
          },
          async () => {
            const res = await deps.client.chatCompletion({
              model: deps.analysisModel,
              messages: [
                {
                  role: 'user',
                  content: [
                    { type: 'text', text: params.question },
                    {
                      type: 'image_url',
                      image_url: { url: dataUrl(params.digest, 'image/png', bytes) },
                    },
                  ],
                },
              ],
            })
            return {
              result: { answer: res.content },
              ...(res.usage !== undefined ? { usage: res.usage } : {}),
            }
          },
        )
        return textResult({ ok: true, fromCache: cached.fromCache, ...(cached.result as object) })
      } catch (err) {
        return textResult(
          { ok: false, error: err instanceof Error ? err.message : String(err) },
          true,
        )
      }
    },
  })

  const zaiTranscribe = defineTool({
    name: 'zai_transcribe',
    label: 'Transcribe audio',
    description:
      'Transcribe a voice or audio message to text (glm-asr). Supported: wav/mp3, ' +
      'up to 25 MB and 30 seconds; longer audio fails explicitly.',
    parameters: Type.Object({
      digest: Type.String({ description: 'sha256 media digest of the audio' }),
    }),
    async execute(_id, params) {
      try {
        const bytes = await deps.service.getMedia(params.digest)
        if (bytes.byteLength > 25 << 20) throw new CapabilityError('audio exceeds 25 MB limit')
        const cached = await cachedAnalysis(
          deps,
          { digest: params.digest, capability: 'zai_transcribe', model: 'glm-asr-2512' },
          async () => {
            const res = await deps.client.transcribe(bytes, {
              fileName: 'audio.wav',
              mimeType: 'audio/wav',
            })
            return { result: { text: res.text } }
          },
        )
        return textResult({ ok: true, fromCache: cached.fromCache, ...(cached.result as object) })
      } catch (err) {
        return textResult(
          { ok: false, error: err instanceof Error ? err.message : String(err) },
          true,
        )
      }
    },
  })

  const zaiParseDocument = defineTool({
    name: 'zai_parse_document',
    label: 'Parse document',
    description:
      'Extract text, tables, and layout from a PDF or document image as markdown ' +
      '(glm-ocr). Images <=10 MB, PDFs <=50 MB and 30 pages.',
    parameters: Type.Object({
      digest: Type.String({ description: 'sha256 media digest of the document' }),
    }),
    async execute(_id, params) {
      try {
        const bytes = await deps.service.getMedia(params.digest)
        if (bytes.byteLength > 50 << 20) throw new CapabilityError('document exceeds 50 MB limit')
        const cached = await cachedAnalysis(
          deps,
          { digest: params.digest, capability: 'zai_parse_document', model: 'glm-ocr' },
          async () => {
            const res = await deps.client.layoutParsing(bytes, { fileName: 'document.pdf' })
            return {
              result: {
                markdown: res.markdown,
                ...(res.pages !== undefined ? { pages: res.pages } : {}),
              },
            }
          },
        )
        return textResult({ ok: true, fromCache: cached.fromCache, ...(cached.result as object) })
      } catch (err) {
        return textResult(
          { ok: false, error: err instanceof Error ? err.message : String(err) },
          true,
        )
      }
    },
  })

  const zaiGenerateImage = defineTool({
    name: 'zai_generate_image',
    label: 'Generate image',
    description:
      'Generate an image from a prompt and store it in the media store. ' +
      'Returns the new digest; send it to the group with chat_send media.',
    parameters: Type.Object({
      prompt: Type.String({ description: 'Image description' }),
    }),
    async execute(_id, params) {
      try {
        const gen = await deps.client.generateImage(params.prompt)
        let bytes: Uint8Array | undefined
        if (gen.url) bytes = await deps.client.fetchBinary(gen.url)
        else if (gen.b64) bytes = new Uint8Array(Buffer.from(gen.b64, 'base64'))
        if (!bytes) throw new CapabilityError('image generation returned no data')
        const meta = await deps.service.putMedia(bytes, { type: 'image/png' })
        return textResult({ ok: true, digest: meta.digest, size: meta.size })
      } catch (err) {
        return textResult(
          { ok: false, error: err instanceof Error ? err.message : String(err) },
          true,
        )
      }
    },
  })

  const zaiGenerateVideo = defineTool({
    name: 'zai_generate_video',
    label: 'Generate video',
    description:
      'Generate a short video from a prompt (async, polled). ' +
      'Returns the stored digest on success; explicit failure otherwise.',
    parameters: Type.Object({
      prompt: Type.String({ description: 'Video description' }),
    }),
    async execute(_id, params) {
      try {
        const task = await deps.client.generateVideo(params.prompt)
        if (!task.id) throw new CapabilityError('video generation returned no task id')
        const done = await deps.client.pollTask('video', task.id)
        if (!done.url) return textResult({ ok: false, status: done.status }, true)
        const bytes = await deps.client.fetchBinary(done.url)
        const meta = await deps.service.putMedia(bytes, { type: 'video/mp4' })
        return textResult({ ok: true, digest: meta.digest })
      } catch (err) {
        return textResult(
          { ok: false, error: err instanceof Error ? err.message : String(err) },
          true,
        )
      }
    },
  })

  return [zaiAnalyzeMedia, zaiTranscribe, zaiParseDocument, zaiGenerateImage, zaiGenerateVideo]
}

export function analysisCacheKey(digest: string, capability: string, model: string) {
  return { digest, capability, model, schemaVersion: ANALYSIS_SCHEMA_VERSION }
}

export function newCapabilityEventId(): string {
  return randomUUID()
}