MangoPig
136 lines
3.1 KiB
Go
136 lines
3.1 KiB
Go
// Path: Code/Tracker/clients/docling.go
|
|
|
|
package clients
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"mime/multipart"
|
|
"net/http"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
type DoclingClient struct {
|
|
baseURL string
|
|
httpClient *http.Client
|
|
}
|
|
|
|
type DoclingResponse struct {
|
|
Document struct {
|
|
MdContent string `json:"md_content"`
|
|
} `json:"document"`
|
|
Status string `json:"status"`
|
|
ProcessingTime float64 `json:"processing_time"`
|
|
}
|
|
|
|
func NewDoclingClient(baseURL string) *DoclingClient {
|
|
return &DoclingClient{
|
|
baseURL: strings.TrimSuffix(baseURL, "/"),
|
|
httpClient: &http.Client{
|
|
Timeout: 5 * time.Minute, // OCR can take a while
|
|
},
|
|
}
|
|
}
|
|
|
|
func (c *DoclingClient) Health() error {
|
|
resp, err := c.httpClient.Get(c.baseURL + "/health")
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
return fmt.Errorf("docling unhealthy: status %d", resp.StatusCode)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (c *DoclingClient) Convert(filePath string) (string, error) {
|
|
// Open file
|
|
f, err := os.Open(filePath)
|
|
if err != nil {
|
|
return "", fmt.Errorf("failed to open file: %w", err)
|
|
}
|
|
defer f.Close()
|
|
|
|
// Create multipart form
|
|
var buf bytes.Buffer
|
|
writer := multipart.NewWriter(&buf)
|
|
|
|
// Add file
|
|
part, err := writer.CreateFormFile("files", filepath.Base(filePath))
|
|
if err != nil {
|
|
return "", fmt.Errorf("failed to create form file: %w", err)
|
|
}
|
|
if _, err := io.Copy(part, f); err != nil {
|
|
return "", fmt.Errorf("failed to copy file: %w", err)
|
|
}
|
|
|
|
// Add options
|
|
writer.WriteField("to_formats", "md")
|
|
writer.WriteField("do_ocr", "true")
|
|
|
|
// For images, force OCR
|
|
ext := strings.ToLower(filepath.Ext(filePath))
|
|
if ext == ".png" || ext == ".jpg" || ext == ".jpeg" {
|
|
writer.WriteField("force_ocr", "true")
|
|
writer.WriteField("ocr_lang", "en")
|
|
writer.WriteField("ocr_lang", "ch_tra")
|
|
}
|
|
|
|
writer.Close()
|
|
|
|
// Make request
|
|
req, err := http.NewRequest("POST", c.baseURL+"/v1/convert/file", &buf)
|
|
if err != nil {
|
|
return "", fmt.Errorf("failed to create request: %w", err)
|
|
}
|
|
req.Header.Set("Content-Type", writer.FormDataContentType())
|
|
|
|
resp, err := c.httpClient.Do(req)
|
|
if err != nil {
|
|
return "", fmt.Errorf("request failed: %w", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
body, _ := io.ReadAll(resp.Body)
|
|
return "", fmt.Errorf("docling error: status %d, body: %s", resp.StatusCode, string(body))
|
|
}
|
|
|
|
// Parse response
|
|
var result DoclingResponse
|
|
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
|
return "", fmt.Errorf("failed to decode response: %w", err)
|
|
}
|
|
|
|
return result.Document.MdContent, nil
|
|
}
|
|
|
|
func getMimeType(path string) string {
|
|
ext := strings.ToLower(filepath.Ext(path))
|
|
switch ext {
|
|
case ".pdf":
|
|
return "application/pdf"
|
|
case ".docx":
|
|
return "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
|
|
case ".doc":
|
|
return "application/msword"
|
|
case ".xlsx":
|
|
return "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
|
|
case ".xls":
|
|
return "application/vnd.ms-excel"
|
|
case ".png":
|
|
return "image/png"
|
|
case ".jpg", ".jpeg":
|
|
return "image/jpeg"
|
|
default:
|
|
return "application/octet-stream"
|
|
}
|
|
}
|