@corgelabs/corge 0.0.1 → 1.0.0

package/src/ai/service.ts

@@ -0,0 +1,317 @@
+import { createGoogleVertex } from '@ai-sdk/google-vertex'
+import { createVertexAnthropic } from '@ai-sdk/google-vertex/anthropic'
+import { HTTPClient, Mistral } from '@mistralai/mistralai'
+import { TypeSafeClient } from '@typesafe-ai/sdk'
+import { embed, generateText, Output, type LanguageModel, type ModelMessage } from 'ai'
+import type { Core } from '../index'
+import {
+  DEFAULT_EXTRACTION_PROMPT,
+  envelopeSchema,
+  type AiGenerateOptions,
+  type AiGenerateResult,
+  type AiUsageSummary,
+  type DecideResult,
+  type JevQuestion,
+  type VisionOptions,
+  type VisionResult
+} from './schema'
+import { formatOcrPages, normalizeUsage } from './utils'
+
+let cachedVertexKey = ''
+let cachedVertexAnthropic: ReturnType<typeof createVertexAnthropic> | undefined
+let cachedVertexEmbedding: ReturnType<typeof createGoogleVertex> | undefined
+
+let cachedJevKey = ''
+let cachedJev: TypeSafeClient | undefined
+
+let cachedMistralKey = ''
+let cachedMistral: Mistral | undefined
+
+export class AiService {
+  readonly #core: Core
+
+  constructor(coreClient: Core) {
+    this.#core = coreClient
+  }
+
+  #vertexProviders = () => {
+    const project = this.#core.config.gcpProject
+    const location = this.#core.config.models.vertexLocation
+    const key = `${project}:${location}`
+    if (!cachedVertexAnthropic || !cachedVertexEmbedding || cachedVertexKey !== key) {
+      cachedVertexKey = key
+      cachedVertexAnthropic = createVertexAnthropic({ project, location })
+      cachedVertexEmbedding = createGoogleVertex({
+        project,
+        location: location === 'global' ? 'us-central1' : location
+      })
+    }
+    return { vertexAnthropic: cachedVertexAnthropic, vertexEmbeddingProvider: cachedVertexEmbedding }
+  }
+
+  languageModel = (modelId: string = this.#core.config.models.vertex): LanguageModel =>
+    this.#vertexProviders().vertexAnthropic(modelId)
+
+  #jevClient = (apiKey: string): TypeSafeClient =>
+    cachedJev && cachedJevKey === apiKey
+      ? cachedJev
+      : ((cachedJevKey = apiKey),
+        (cachedJev = new TypeSafeClient({ apiKey, defaultModel: this.#core.config.models.jev })))
+
+  #mistralClient = (): Mistral => {
+    const apiKey = this.#core.env.MISTRAL_API_KEY
+    if (cachedMistral && cachedMistralKey === apiKey) return cachedMistral
+    const httpClient = new HTTPClient()
+    httpClient.addHook('response', async (res) => {
+      const limit = Number(res.headers.get('x-ratelimit-limit-ocr-pages-minute') ?? '0')
+      const remHeader = res.headers.get('x-ratelimit-remaining-ocr-pages-minute')
+      const remaining = res.status === 429 ? 0 : remHeader !== null ? Number(remHeader) : undefined
+      if (limit > 0 || remaining !== undefined) {
+        await this.#core.internals.resilience.syncLimit('mistral', remaining, limit)
+      }
+    })
+    cachedMistralKey = apiKey
+    return (cachedMistral = new Mistral({
+      apiKey,
+      httpClient,
+      retryConfig: {
+        strategy: 'backoff',
+        backoff: { initialInterval: 5000, maxInterval: 20000, exponent: 1.5, maxElapsedTime: 90000 },
+        retryConnectionErrors: true
+      }
+    }))
+  }
+
+  describeImage = async (documentUrl: string, signal?: AbortSignal): Promise<string> =>
+    this.#core.internals.resilience.circuitBreaker('vertex', async () => {
+      await this.#core.internals.resilience.acquireLimit('ai')
+      const match = /^data:([^;,]+);base64,(.+)$/.exec(documentUrl)
+      const mediaType = match?.[1] ?? 'image/jpeg'
+      const base64 = match?.[2] ?? documentUrl
+      const modelId = this.#core.config.models.vertex
+      const { text } = await generateText({
+        model: this.languageModel(modelId),
+        abortSignal: signal,
+        messages: [
+          {
+            role: 'user',
+            content: [
+              { type: 'file', data: base64, mediaType },
+              { type: 'text', text: 'Describe what is illustrated on this image concisely in markdown.' }
+            ]
+          }
+        ]
+      })
+      const trimmed = text.trim()
+      if (!trimmed) throw new Error('vision model returned empty description')
+      return trimmed
+    })
+
+  #vertexVisionFallback = async (documentUrl: string, opts?: VisionOptions): Promise<VisionResult> => {
+    const markdown = await this.describeImage(documentUrl, opts?.signal)
+    let annotation: Record<string, unknown> = {}
+    const modelId = this.#core.config.models.vertex
+    let usage: AiUsageSummary | undefined
+    if (opts?.schema) {
+      await this.#core.internals.resilience.acquireLimit('ai')
+      const match = /^data:([^;,]+);base64,(.+)$/.exec(documentUrl)
+      const mediaType = match?.[1] ?? 'application/pdf'
+      const base64 = match?.[2] ?? documentUrl
+      const res = await this.#core.internals.resilience.circuitBreaker('vertex', () =>
+        generateText({
+          model: this.languageModel(modelId),
+          abortSignal: opts.signal,
+          messages: [
+            {
+              role: 'user',
+              content: [
+                { type: 'file', data: base64, mediaType },
+                {
+                  type: 'text',
+                  text: `${opts.prompt ?? DEFAULT_EXTRACTION_PROMPT}\n\nJSON Schema:\n${JSON.stringify(opts.schema)}\n\nRespond with ONLY a valid JSON object.`
+                }
+              ]
+            }
+          ]
+        })
+      )
+      usage = normalizeUsage(modelId, res.usage)
+      const jsonMatch = /\{[\s\S]*\}/.exec(res.text)
+      if (jsonMatch?.[0]) {
+        try {
+          annotation = JSON.parse(jsonMatch[0]) as Record<string, unknown>
+        } catch {}
+      }
+    }
+    return { markdown, annotation, ...(usage ? { usage } : {}) }
+  }
+
+  vision = async (documentUrl: string, opts?: VisionOptions): Promise<VisionResult> => {
+    const isImage = documentUrl.startsWith('data:image')
+    const hasSchema = Boolean(opts?.schema)
+    const allowFallback = opts?.fallbackVision !== false
+    await this.#core.internals.payments.enforceAiQuota()
+
+    return this.#core.internals.resilience.circuitBreaker(
+      'mistral',
+      async () => {
+        await this.#core.internals.resilience.acquireLimit('mistral')
+        const modelId = this.#core.config.models.mistralOcr
+        const result = await this.#mistralClient().ocr.process(
+          {
+            model: modelId,
+            confidenceScoresGranularity: 'block',
+            pages: !hasSchema && !isImage ? Array.from({ length: 10 }, (_, i) => i) : undefined,
+            document: isImage
+              ? { type: 'image_url' as const, imageUrl: documentUrl }
+              : { type: 'document_url' as const, documentUrl },
+            ...(opts?.schema && {
+              documentAnnotationFormat: {
+                type: 'json_schema' as const,
+                jsonSchema: { name: 'fields', schemaDefinition: opts.schema, strict: true }
+              },
+              documentAnnotationPrompt: opts.prompt ?? DEFAULT_EXTRACTION_PROMPT
+            })
+          },
+          opts?.signal ? { fetchOptions: { signal: opts.signal } } : undefined
+        )
+        const pagesProcessed = result.usageInfo?.pagesProcessed ?? result.pages.length
+        if (pagesProcessed > 1) {
+          await this.#core.internals.resilience.chargeLimit('mistral', pagesProcessed - 1)
+        }
+        let markdown = formatOcrPages(result.pages)
+        const annotation = JSON.parse(result.documentAnnotation ?? '{}') as Record<string, unknown>
+
+        if (!markdown && (opts?.fallbackVision ?? !hasSchema)) {
+          markdown = await this.describeImage(documentUrl, opts?.signal)
+        }
+
+        const usage = normalizeUsage(modelId, null, pagesProcessed)
+        await this.#core.internals.payments.trackAiUsage({
+          model: usage.model,
+          inputTokens: usage.inputTokens,
+          outputTokens: usage.outputTokens,
+          totalTokens: usage.totalTokens,
+          source: 'vision'
+        })
+        return { markdown, annotation, usage }
+      },
+      allowFallback ? () => this.#vertexVisionFallback(documentUrl, opts) : undefined,
+      { isDev: this.#core.env.isDev }
+    )
+  }
+
+  decide = async (state: unknown, questions: Record<string, JevQuestion>): Promise<DecideResult> =>
+    this.#core.internals.resilience.circuitBreaker(
+      'jev',
+      async () => {
+        await this.#core.internals.payments.enforceAiQuota()
+        await this.#core.internals.resilience.acquireLimit('jev')
+        const raw = await this.#jevClient(this.#core.env.TYPESAFE_API_KEY).systemOne({
+          state: state as never,
+          questions: questions as never
+        })
+        const parsed = envelopeSchema.safeParse(raw)
+        if (parsed.success && Object.keys(questions).every((q) => parsed.data.answers[q]?.type === questions[q]?.type)) {
+          const usage = normalizeUsage(parsed.data.model || this.#core.config.models.jev, null)
+          await this.#core.internals.payments.trackAiUsage({
+            model: usage.model,
+            inputTokens: usage.inputTokens,
+            outputTokens: usage.outputTokens,
+            totalTokens: usage.totalTokens,
+            source: 'decide'
+          })
+          return {
+            ...parsed.data,
+            usage
+          }
+        }
+        throw new Error('Jev decision returned no valid answer envelope')
+      },
+      undefined,
+      { isDev: this.#core.env.isDev }
+    )
+
+  generate = async <O = string>(
+    promptOrMessages: string | ModelMessage[],
+    opts?: AiGenerateOptions<O>
+  ): Promise<AiGenerateResult<O>> =>
+    this.#core.internals.resilience.circuitBreaker('vertex', async () => {
+      await this.#core.internals.payments.enforceAiQuota()
+      await this.#core.internals.resilience.acquireLimit('ai')
+      const modelId = opts?.model ?? this.#core.config.models.vertex
+      const model = this.languageModel(modelId)
+      const messages: ModelMessage[] =
+        typeof promptOrMessages === 'string'
+          ? [{ role: 'user', content: [{ type: 'text', text: promptOrMessages }] }]
+          : promptOrMessages
+
+      if (opts?.schema) {
+        const res = await generateText({
+          model,
+          system: opts.system,
+          messages,
+          output: Output.object({ schema: opts.schema }),
+          abortSignal: opts.signal
+        })
+        const norm = normalizeUsage(modelId, res.usage)
+        await this.#core.internals.payments.trackAiUsage({
+          model: norm.model,
+          inputTokens: norm.inputTokens,
+          outputTokens: norm.outputTokens,
+          totalTokens: norm.totalTokens,
+          source: 'generate'
+        })
+        return {
+          text: res.text,
+          output: res.output as O,
+          usage: res.usage
+        }
+      }
+
+      const res = await generateText({
+        model,
+        system: opts?.system,
+        messages,
+        abortSignal: opts?.signal
+      })
+      const norm = normalizeUsage(modelId, res.usage)
+      await this.#core.internals.payments.trackAiUsage({
+        model: norm.model,
+        inputTokens: norm.inputTokens,
+        outputTokens: norm.outputTokens,
+        totalTokens: norm.totalTokens,
+        source: 'generate'
+      })
+      return {
+        text: res.text,
+        output: res.text as unknown as O,
+        usage: res.usage
+      }
+    })
+
+  embed = async (text: string): Promise<number[] | null> => {
+    const normalized = text.trim()
+    if (normalized.length < 15) return null
+    try {
+      return await this.#core.internals.resilience.circuitBreaker('vertex', async () => {
+        await this.#core.internals.resilience.acquireLimit('ai')
+        const embeddingModel = this.#vertexProviders().vertexEmbeddingProvider.textEmbeddingModel(
+          this.#core.config.models.vertexEmbedding
+        )
+        const { embedding } = await embed({
+          model: embeddingModel,
+          value: normalized.slice(0, 8000),
+          dimensions: 768,
+          abortSignal: AbortSignal.timeout(15_000),
+          providerOptions: {
+            vertex: { outputDimensionality: 768 }
+          }
+        })
+        return Array.isArray(embedding) && embedding.length === 768 ? embedding : null
+      })
+    } catch {
+      return null
+    }
+  }
+}

package/src/ai/utils.ts

@@ -0,0 +1,38 @@
+import type { OCRPageObject } from '@mistralai/mistralai/models/components'
+import type { AiUsageSummary } from './schema'
+
+export const formatOcrPages = (pages: OCRPageObject[]): string =>
+  pages
+    .map((p, i) => {
+      const text = p.blocks
+        ? p.blocks
+            .filter(
+              (b) =>
+                b.type !== 'image' &&
+                ('confidenceScores' in b ? (b.confidenceScores?.averageContentConfidenceScore ?? 1) : 1) >= 0.5
+            )
+            .map((b) => ('content' in b ? b.content.trim() : ''))
+            .filter(Boolean)
+            .join('\n\n')
+        : p.markdown.replace(/!\[[^\]]*\]\([^)]*\)/g, '').trim()
+      return text ? `### Page ${i + 1}\n${text}` : ''
+    })
+    .filter(Boolean)
+    .join('\n\n')
+
+export const normalizeUsage = (
+  model: string,
+  raw?: { inputTokens?: number; outputTokens?: number; totalTokens?: number; tokens?: number } | null,
+  pages?: number
+): AiUsageSummary => {
+  const inputTokens = raw?.inputTokens ?? raw?.tokens ?? 0
+  const outputTokens = raw?.outputTokens ?? 0
+  const totalTokens = raw?.totalTokens ?? (inputTokens + outputTokens || (pages ? pages * 1000 : 0))
+  return {
+    model,
+    inputTokens,
+    outputTokens,
+    totalTokens,
+    ...(pages !== undefined ? { pages } : {})
+  }
+}

package/src/attachments/index.ts

@@ -0,0 +1,3 @@
+export * from './schema'
+export { AttachmentsService } from './service'
+export { coreAttachmentsRoutes } from './routes'

package/src/attachments/routes.ts

@@ -0,0 +1,28 @@
+import { core } from '../index'
+import type { StepRunner } from '../tasks'
+import {
+  processAttachmentsInput,
+  processAttachmentsOutput,
+  scanAttachmentInput,
+  scanAttachmentOutput,
+  type ProcessAttachmentsInput,
+  type ScanAttachmentInput
+} from './schema'
+
+export const coreAttachmentsRoutes = {
+  'attachments.scan': {
+    mode: 'workflow' as const,
+    input: scanAttachmentInput,
+    output: scanAttachmentOutput,
+    shortCircuit: ({ input }: { input: ScanAttachmentInput }) => core.internals.context.scanShortCircuit({ input }),
+    handler: ({ input, step }: { input: ScanAttachmentInput; step: StepRunner }) =>
+      core.attachments.scan(input, { step })
+  },
+  'attachments.process': {
+    mode: 'workflow' as const,
+    input: processAttachmentsInput,
+    output: processAttachmentsOutput,
+    handler: ({ input, step }: { input: ProcessAttachmentsInput; step: StepRunner }) =>
+      core.attachments.process(input, { step })
+  }
+}

package/src/attachments/schema.ts

@@ -0,0 +1,45 @@
+import { z } from 'zod'
+import { CORE_CONFIG } from '../config'
+import { fileRef, fileScan } from '../files'
+
+export const processedAttachment = z.union([
+  z.object({ ref: fileRef, scan: fileScan }),
+  z.object({ ref: fileRef, error: z.string() })
+])
+export type ProcessedAttachment = z.infer<typeof processedAttachment>
+
+export interface AttachmentTypeSpec<S extends z.ZodType = z.ZodType> {
+  description: string
+  schema: S
+}
+export type AttachmentTypesMap = Record<string, AttachmentTypeSpec>
+
+export const defineAttachmentTypes = <T extends AttachmentTypesMap>(types: T): T => types
+
+const registeredAttachmentTypes: AttachmentTypesMap = {}
+
+export const registerAttachmentTypes = (types?: AttachmentTypesMap, reset = false): void => {
+  if (reset) {
+    for (const k of Object.keys(registeredAttachmentTypes)) delete registeredAttachmentTypes[k]
+  }
+  if (!types) return
+  Object.assign(registeredAttachmentTypes, types)
+}
+
+export const getRegisteredAttachmentTypes = (): AttachmentTypesMap => registeredAttachmentTypes
+
+export const scanAttachmentInput = z.object({ ref: fileRef, type: z.string().optional() })
+export const scanAttachmentOutput = fileScan
+export type ScanAttachmentInput = z.input<typeof scanAttachmentInput>
+export type ScanAttachmentOutput = z.infer<typeof scanAttachmentOutput>
+
+export const attachmentInput = z.union([fileRef, scanAttachmentInput])
+export const processAttachmentsInput = z.object({
+  refs: z
+    .array(attachmentInput)
+    .min(1)
+    .refine((arr) => arr.length <= CORE_CONFIG.limits.maxAttachments, 'Too many attachments')
+})
+export const processAttachmentsOutput = z.object({ attachments: z.array(processedAttachment) })
+export type ProcessAttachmentsInput = z.input<typeof processAttachmentsInput>
+export type ProcessAttachmentsOutput = z.infer<typeof processAttachmentsOutput>