UNPKG

genkitx-mistral

Version:

Firebase Genkit AI framework plugin for Mistral AI APIs.

1 lines 5.67 kB
{"version":3,"sources":["../src/ocr.ts"],"sourcesContent":["/**\n * Copyright 2024 The Fire Company\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":[]}