105 lines
2.7 KiB
Go
105 lines
2.7 KiB
Go
package mistral
|
|
|
|
import (
|
|
schemas "github.com/maximhq/bifrost/core/schemas"
|
|
)
|
|
|
|
// ToMistralOCRRequest converts a Bifrost OCR request to a Mistral OCR request.
|
|
func ToMistralOCRRequest(req *schemas.BifrostOCRRequest) *MistralOCRRequest {
|
|
if req == nil {
|
|
return nil
|
|
}
|
|
|
|
mistralReq := &MistralOCRRequest{
|
|
Model: req.Model,
|
|
Document: MistralOCRDocument{
|
|
Type: string(req.Document.Type),
|
|
},
|
|
}
|
|
|
|
if req.ID != nil {
|
|
mistralReq.ID = *req.ID
|
|
}
|
|
|
|
switch req.Document.Type {
|
|
case schemas.OCRDocumentTypeDocumentURL:
|
|
if req.Document.DocumentURL != nil {
|
|
mistralReq.Document.DocumentURL = *req.Document.DocumentURL
|
|
}
|
|
case schemas.OCRDocumentTypeImageURL:
|
|
if req.Document.ImageURL != nil {
|
|
mistralReq.Document.ImageURL = *req.Document.ImageURL
|
|
}
|
|
}
|
|
|
|
if req.Params != nil {
|
|
mistralReq.IncludeImageBase64 = req.Params.IncludeImageBase64
|
|
mistralReq.Pages = req.Params.Pages
|
|
mistralReq.ImageLimit = req.Params.ImageLimit
|
|
mistralReq.ImageMinSize = req.Params.ImageMinSize
|
|
mistralReq.TableFormat = req.Params.TableFormat
|
|
mistralReq.ExtractHeader = req.Params.ExtractHeader
|
|
mistralReq.ExtractFooter = req.Params.ExtractFooter
|
|
mistralReq.BBoxAnnotationFormat = req.Params.BBoxAnnotationFormat
|
|
mistralReq.DocumentAnnotationFormat = req.Params.DocumentAnnotationFormat
|
|
mistralReq.DocumentAnnotationPrompt = req.Params.DocumentAnnotationPrompt
|
|
mistralReq.ExtraParams = req.Params.ExtraParams
|
|
}
|
|
|
|
return mistralReq
|
|
}
|
|
|
|
// ToBifrostOCRResponse converts a Mistral OCR response to a Bifrost OCR response.
|
|
func (r *MistralOCRResponse) ToBifrostOCRResponse() *schemas.BifrostOCRResponse {
|
|
if r == nil {
|
|
return nil
|
|
}
|
|
|
|
resp := &schemas.BifrostOCRResponse{
|
|
Model: r.Model,
|
|
DocumentAnnotation: r.DocumentAnnotation,
|
|
}
|
|
|
|
// Convert pages
|
|
if len(r.Pages) > 0 {
|
|
resp.Pages = make([]schemas.OCRPage, len(r.Pages))
|
|
for i, p := range r.Pages {
|
|
page := schemas.OCRPage{
|
|
Index: p.Index,
|
|
Markdown: p.Markdown,
|
|
}
|
|
if len(p.Images) > 0 {
|
|
page.Images = make([]schemas.OCRPageImage, len(p.Images))
|
|
for j, img := range p.Images {
|
|
page.Images[j] = schemas.OCRPageImage{
|
|
ID: img.ID,
|
|
TopLeftX: img.TopLeftX,
|
|
TopLeftY: img.TopLeftY,
|
|
BottomRightX: img.BottomRightX,
|
|
BottomRightY: img.BottomRightY,
|
|
ImageBase64: img.ImageBase64,
|
|
}
|
|
}
|
|
}
|
|
if p.Dimensions != nil {
|
|
page.Dimensions = &schemas.OCRPageDimensions{
|
|
DPI: p.Dimensions.DPI,
|
|
Height: p.Dimensions.Height,
|
|
Width: p.Dimensions.Width,
|
|
}
|
|
}
|
|
resp.Pages[i] = page
|
|
}
|
|
}
|
|
|
|
// Convert usage info
|
|
if r.UsageInfo != nil {
|
|
resp.UsageInfo = &schemas.OCRUsageInfo{
|
|
PagesProcessed: r.UsageInfo.PagesProcessed,
|
|
DocSizeBytes: r.UsageInfo.DocSizeBytes,
|
|
}
|
|
}
|
|
|
|
return resp
|
|
}
|