ocr-mistral.go

ocr-mistral.go

package invoicing

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"mime/multipart"
	"net/http"
	"net/url"
	"os"
)

// Maximum file size allowed by Mistral API (52.4 MB)
const MistralMaxFileSize = 52 * 1024 * 1024

type mistralClient struct {
	apiKey string
}

func newMistralClient(apiKey string) *mistralClient {
	return &mistralClient{
		apiKey: apiKey,
	}
}

func (c *mistralClient) GetContent(ctx context.Context, filepath string) (*mistralOCRResponse, error) {
	fileInfo, err := os.Stat(filepath)
	if err != nil {
		return nil, fmt.Errorf("error checking file size: %v", err)
	}

	if fileInfo.Size() > MistralMaxFileSize {
		return nil, fmt.Errorf("file is too large (%.2f MB). Maximum allowed size is %.2f MB",
			float64(fileInfo.Size())/1024/1024, float64(MistralMaxFileSize)/1024/1024)
	}

	id, err := c.Upload(ctx, filepath)
	if err != nil {
		return nil, fmt.Errorf("error uploading file: %v", err)
	}

	url, err := c.GetSignedURL(ctx, id)
	if err != nil {
		return nil, fmt.Errorf("error getting signed URL: %v", err)
	}

	ocr, err := c.GetOCR(ctx, url)
	if err != nil {
		return nil, fmt.Errorf("error getting OCR: %v", err)
	}

	return ocr, nil
}

func (c *mistralClient) Upload(ctx context.Context, filepath string) (string, error) {
	file, err := os.Open(filepath)
	if err != nil {
		return "", fmt.Errorf("error opening file: %v", err)
	}
	defer file.Close()

	body := &bytes.Buffer{}
	writer := multipart.NewWriter(body)

	if err := writer.WriteField("purpose", "ocr"); err != nil {
		return "", fmt.Errorf("error writing purpose field: %v", err)
	}

	part, err := writer.CreateFormFile("file", filepath)
	if err != nil {
		return "", fmt.Errorf("error creating form file: %v", err)
	}

	if _, err := io.Copy(part, file); err != nil {
		return "", fmt.Errorf("error copying file: %v", err)
	}

	if err := writer.Close(); err != nil {
		return "", fmt.Errorf("error closing writer: %v", err)
	}

	req, err := http.NewRequestWithContext(ctx, "POST", "https://api.mistral.ai/v1/files", body)
	if err != nil {
		return "", fmt.Errorf("error creating request: %v", err)
	}

	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", c.apiKey))
	req.Header.Set("Content-Type", writer.FormDataContentType())

	type uploadResponse struct {
		ID string `json:"id"`
	}

	var ur uploadResponse
	if err := sendRequest(ctx, req, &ur); err != nil {
		return "", ErrRateLimited
	}

	return ur.ID, nil
}

func (c *mistralClient) GetSignedURL(ctx context.Context, fileID string) (string, error) {
	req, err := http.NewRequestWithContext(ctx, "GET", fmt.Sprintf("https://api.mistral.ai/v1/files/%s/url", fileID), nil)
	if err != nil {
		return "", fmt.Errorf("error creating request: %v", err)
	}

	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", c.apiKey))
	params := url.Values{}
	params.Add("expiry", "24")
	req.URL.RawQuery = params.Encode()

	type urlResponse struct {
		URL string `json:"url"`
	}

	var ur urlResponse
	if err := sendRequest(ctx, req, &ur); err != nil {
		return "", fmt.Errorf("error getting signed URL: %v", err)
	}

	return ur.URL, nil
}

type mistralOCRRequest struct {
	Model              string          `json:"model"`
	Document           mistralDocument `json:"document"`
	IncludeImageBase64 bool            `json:"include_image_base64"`
}

type mistralDocument struct {
	Type        string `json:"type"`
	DocumentURL string `json:"document_url"`
}

type mistralOCRPage struct {
	Index      int                  `json:"index"`
	Markdown   string               `json:"markdown"`
	Images     []interface{}        `json:"images"`
	Dimensions mistralOCRDimensions `json:"dimensions"`
}

type mistralOCRDimensions struct {
	Dpi    int `json:"dpi"`
	Height int `json:"height"`
	Width  int `json:"width"`
}

type mistralOCRUsageInfo struct {
	PagesProcessed int `json:"pages_processed"`
	DocSizeBytes   int `json:"doc_size_bytes"`
}

type mistralOCRResponse struct {
	Pages     []mistralOCRPage    `json:"pages"`
	Model     string              `json:"model"`
	UsageInfo mistralOCRUsageInfo `json:"usage_info"`
}

func (c *mistralClient) GetOCR(ctx context.Context, fileURL string) (*mistralOCRResponse, error) {
	reqBody := mistralOCRRequest{
		Model: "mistral-ocr-latest",
		Document: mistralDocument{
			Type:        "document_url",
			DocumentURL: fileURL,
		},
		IncludeImageBase64: false,
	}

	jsonReqBody, err := json.Marshal(reqBody)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal request body: %v", err)
	}

	req, err := http.NewRequestWithContext(ctx, "POST", "https://api.mistral.ai/v1/ocr", bytes.NewBuffer(jsonReqBody))
	if err != nil {
		return nil, fmt.Errorf("error creating request: %v", err)
	}
	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", c.apiKey))
	req.Header.Set("Content-Type", "application/json")

	var ocrResponse mistralOCRResponse
	if err := sendRequest(ctx, req, &ocrResponse); err != nil {
		return nil, fmt.Errorf("error sending request: %v", err)
	}

	return &ocrResponse, nil
}

func sendRequest(ctx context.Context, req *http.Request, v any) error {
	resp, err := http.DefaultClient.Do(req.WithContext(ctx))
	if err != nil {
		return fmt.Errorf("error sending request: %v", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("unexpected status code: %d", resp.StatusCode)
	}

	b, err := io.ReadAll(resp.Body)
	if err != nil {
		return fmt.Errorf("error reading response body: %v", err)
	}

	if err := json.Unmarshal(b, v); err != nil {
		return fmt.Errorf("error unmarshalling response: %v", err)
	}

	return nil
}