Build with Laki.ai headless MCP

import type { AuthProvider } from '@modelcontextprotocol/client'
import {
  Client,
  StreamableHTTPClientTransport,
} from '@modelcontextprotocol/client'

interface SearchData {
  results: Array<{ items: Array<{ id: string }> }>
}

interface ReadError {
  id?: string
  reason: string
  message?: string
}

interface ReadData {
  documents: Array<{
    id: string
    citation: string
    contentBlocks: Array<{
      id: string
      text: string
      citation?: string
    }>
  }>
  errors: ReadError[]
}

interface EvidenceBlock {
  sourceId: string
  text: string
  citation: string
}

interface CitedDraft<T> {
  answer: T
  citedSourceIds: string[]
}

type AnswerBuilder<T> = (
  evidence: EvidenceBlock[],
  readErrors: ReadError[]
) => Promise<CitedDraft<T>>

const token = process.env.LAKI_API_KEY
if (!token) throw new Error('LAKI_API_KEY is required')

const authProvider: AuthProvider = {
  token: async () => token,
}

const endpoint = new URL('https://api.laki.ai/mcp/headless')
const resultMeta = { 'ai.laki/result-format': 'structured' } as const

export const runLegalResearch = async <T>(
  queries: string[],
  buildAnswer: AnswerBuilder<T>
): Promise<T> => {
  const client = new Client({
    name: 'my-backend',
    version: '1.0.0',
  }, {
    versionNegotiation: { mode: { pin: '2026-07-28' } },
  })
  const transport = new StreamableHTTPClientTransport(endpoint, {
    authProvider,
  })

  const callTool = async <TResult>(
    name: string,
    args: Record<string, unknown>
  ): Promise<TResult> => {
    const result = await client.callTool({
      name,
      arguments: args,
      _meta: resultMeta,
    })
    if (result.isError || !result.structuredContent) {
      const message = result.content?.find(
        (item): item is { type: 'text'; text: string } =>
          item.type === 'text'
      )?.text
      throw new Error(message ?? name + ' failed')
    }
    return result.structuredContent as TResult
  }

  try {
    await client.connect(transport)
    await client.listTools({ _meta: resultMeta })

    const search = await callTool<SearchData>('search_legal_sources', {
      queries,
      mode: 'hybrid',
      limit: 5,
    })
    const sourceIds = [
      ...new Set(
        search.results.flatMap(group =>
          group.items.map(item => item.id)
        )
      ),
    ].slice(0, 10)
    if (sourceIds.length === 0) throw new Error('No legal sources found')

    const read = await callTool<ReadData>('read_document', {
      ids: sourceIds,
      includeRelatedSourceStubs: false,
    })
    const evidence = read.documents.flatMap(document =>
      document.contentBlocks.map(block => ({
        sourceId: block.citation ? block.id : document.id,
        text: block.text,
        citation: block.citation ?? document.citation,
      }))
    )
    if (evidence.length === 0) throw new Error('No readable evidence found')

    const draft = await buildAnswer(evidence, read.errors)
    const readSourceIds = new Set(evidence.map(block => block.sourceId))
    const citedSourceIds = [...new Set(draft.citedSourceIds)]
    if (citedSourceIds.some(id => !readSourceIds.has(id))) {
      throw new Error('The answer cited a source that was not read')
    }

    return draft.answer
  } finally {
    await client.close()
  }
}