first commit

core/providers/mistral/ocr.go

@@ -0,0 +1,104 @@
+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
+}