{"version":3,"sources":["../src/ocr.ts"],"sourcesContent":["/**\n * Copyright 2024 Bloom Labs Inc\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\nimport { Mistral } from '@mistralai/mistralai';\nimport {\n  OCRRequest,\n  Document as OCRDDocument,\n  OCRResponse,\n} from '@mistralai/mistralai/models/components';\nimport type { GenerateRequest, GenerateResponseData, Genkit } from 'genkit';\nimport { GenerationCommonConfigSchema, z } from 'genkit';\nimport type { ModelAction } from 'genkit/model';\nimport { modelRef } from 'genkit/model';\n\nconst ImageURLSchema = z.object({\n  url: z.string(),\n  detail: z.string().nullable().optional(),\n});\n\nconst OCRSchema = GenerationCommonConfigSchema.extend({\n  document: z.object({\n    type: z.enum(['document_url', 'image_url']).optional(),\n    documentUrl: z.string(),\n    documentName: z.string().nullable().optional(),\n    imageUrl: z.union([z.string(), ImageURLSchema]).optional(),\n  }),\n});\n\nexport const ocr = modelRef({\n  name: 'mistral/ocr',\n  info: {\n    versions: ['mistral-ocr-latest'],\n    label: 'Mistral - OCR',\n    supports: {\n      media: false,\n      output: ['text', 'json'],\n      multiturn: false,\n      systemRole: false,\n      tools: false,\n    },\n  },\n  configSchema: OCRSchema,\n});\n\nfunction toOCRRequest(request: GenerateRequest<typeof OCRSchema>): OCRRequest {\n  const docType = request.config?.document?.type || 'document_url';\n\n  const options: OCRRequest = {\n    model: 'mistral-ocr-latest',\n    document:\n      docType === 'document_url'\n        ? {\n            type: 'document_url',\n            documentUrl: request.config?.document?.documentUrl || '',\n            documentName: request.config?.document?.documentName || null,\n          }\n        : {\n            type: 'image_url',\n            imageUrl:\n              typeof request.config?.document?.imageUrl === 'string'\n                ? request.config.document.imageUrl\n                : request.config?.document?.imageUrl?.url || '',\n          },\n  };\n\n  (Object.keys(options) as (keyof OCRRequest)[]).forEach((k) => {\n    if (options[k] === undefined) {\n      delete options[k];\n    }\n  });\n  return options;\n}\n\nfunction toGenerateResponse(result: OCRResponse): GenerateResponseData {\n  const candidates: GenerateResponseData['candidates'] = result.pages.map(\n    (page, index) => ({\n      index: index,\n      finishReason: 'stop',\n      message: {\n        role: 'model',\n        content: [\n          {\n            text: page.markdown,\n          },\n          ...page.images.map((image) => ({\n            media: {\n              contentType: 'image/png',\n              url: `data:image/png;base64,${image.imageBase64}`,\n            },\n          })),\n        ],\n      },\n    })\n  );\n  return { candidates };\n}\n\nexport function OCRModel(\n  ai: Genkit,\n  client: Mistral\n): ModelAction<typeof OCRSchema> {\n  return ai.defineModel<typeof OCRSchema>(\n    {\n      name: ocr.name,\n      ...ocr.info,\n      configSchema: ocr.configSchema,\n    },\n    async (request) => {\n      const result = await client.ocr.process(toOCRRequest(request));\n      return toGenerateResponse(result);\n    }\n  );\n}\n"],"mappings":";;;;;AAsBA,SAAS,8BAA8B,SAAS;AAEhD,SAAS,gBAAgB;AAEzB,MAAM,iBAAiB,EAAE,OAAO;AAAA,EAC9B,KAAK,EAAE,OAAO;AAAA,EACd,QAAQ,EAAE,OAAO,EAAE,SAAS,EAAE,SAAS;AACzC,CAAC;AAED,MAAM,YAAY,6BAA6B,OAAO;AAAA,EACpD,UAAU,EAAE,OAAO;AAAA,IACjB,MAAM,EAAE,KAAK,CAAC,gBAAgB,WAAW,CAAC,EAAE,SAAS;AAAA,IACrD,aAAa,EAAE,OAAO;AAAA,IACtB,cAAc,EAAE,OAAO,EAAE,SAAS,EAAE,SAAS;AAAA,IAC7C,UAAU,EAAE,MAAM,CAAC,EAAE,OAAO,GAAG,cAAc,CAAC,EAAE,SAAS;AAAA,EAC3D,CAAC;AACH,CAAC;AAEM,MAAM,MAAM,SAAS;AAAA,EAC1B,MAAM;AAAA,EACN,MAAM;AAAA,IACJ,UAAU,CAAC,oBAAoB;AAAA,IAC/B,OAAO;AAAA,IACP,UAAU;AAAA,MACR,OAAO;AAAA,MACP,QAAQ,CAAC,QAAQ,MAAM;AAAA,MACvB,WAAW;AAAA,MACX,YAAY;AAAA,MACZ,OAAO;AAAA,IACT;AAAA,EACF;AAAA,EACA,cAAc;AAChB,CAAC;AAED,SAAS,aAAa,SAAwD;AAxD9E;AAyDE,QAAM,YAAU,mBAAQ,WAAR,mBAAgB,aAAhB,mBAA0B,SAAQ;AAElD,QAAM,UAAsB;AAAA,IAC1B,OAAO;AAAA,IACP,UACE,YAAY,iBACR;AAAA,MACE,MAAM;AAAA,MACN,eAAa,mBAAQ,WAAR,mBAAgB,aAAhB,mBAA0B,gBAAe;AAAA,MACtD,gBAAc,mBAAQ,WAAR,mBAAgB,aAAhB,mBAA0B,iBAAgB;AAAA,IAC1D,IACA;AAAA,MACE,MAAM;AAAA,MACN,UACE,SAAO,mBAAQ,WAAR,mBAAgB,aAAhB,mBAA0B,cAAa,WAC1C,QAAQ,OAAO,SAAS,aACxB,yBAAQ,WAAR,mBAAgB,aAAhB,mBAA0B,aAA1B,mBAAoC,QAAO;AAAA,IACnD;AAAA,EACR;AAEA,EAAC,OAAO,KAAK,OAAO,EAA2B,QAAQ,CAAC,MAAM;AAC5D,QAAI,QAAQ,CAAC,MAAM,QAAW;AAC5B,aAAO,QAAQ,CAAC;AAAA,IAClB;AAAA,EACF,CAAC;AACD,SAAO;AACT;AAEA,SAAS,mBAAmB,QAA2C;AACrE,QAAM,aAAiD,OAAO,MAAM;AAAA,IAClE,CAAC,MAAM,WAAW;AAAA,MAChB;AAAA,MACA,cAAc;AAAA,MACd,SAAS;AAAA,QACP,MAAM;AAAA,QACN,SAAS;AAAA,UACP;AAAA,YACE,MAAM,KAAK;AAAA,UACb;AAAA,UACA,GAAG,KAAK,OAAO,IAAI,CAAC,WAAW;AAAA,YAC7B,OAAO;AAAA,cACL,aAAa;AAAA,cACb,KAAK,yBAAyB,MAAM,WAAW;AAAA,YACjD;AAAA,UACF,EAAE;AAAA,QACJ;AAAA,MACF;AAAA,IACF;AAAA,EACF;AACA,SAAO,EAAE,WAAW;AACtB;AAEO,SAAS,SACd,IACA,QAC+B;AAC/B,SAAO,GAAG;AAAA,IACR;AAAA,MACE,MAAM,IAAI;AAAA,OACP,IAAI,OAFT;AAAA,MAGE,cAAc,IAAI;AAAA,IACpB;AAAA,IACA,CAAO,YAAY;AACjB,YAAM,SAAS,MAAM,OAAO,IAAI,QAAQ,aAAa,OAAO,CAAC;AAC7D,aAAO,mBAAmB,MAAM;AAAA,IAClC;AAAA,EACF;AACF;","names":[]}