moteva/server/internal/infrastructure/objectrecognition/client.go

package objectrecognition

import (
	"bytes"
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"

	"img_infinite_canvas/internal/domain/design"
)

type ClientOptions struct {
	BaseURL        string
	APIKey         string
	Model          string
	TimeoutSeconds int
}

type Client struct {
	baseURL string
	apiKey  string
	model   string
	client  *http.Client
}

func NewClient(opts ClientOptions) *Client {
	if strings.TrimSpace(opts.BaseURL) == "" || strings.TrimSpace(opts.APIKey) == "" {
		return nil
	}
	timeout := time.Duration(opts.TimeoutSeconds) * time.Second
	if timeout <= 0 {
		timeout = 60 * time.Second
	}
	model := strings.TrimSpace(opts.Model)
	if model == "" {
		model = "gpt-5.4-mini"
	}
	return &Client{
		baseURL: strings.TrimRight(opts.BaseURL, "/"),
		apiKey:  opts.APIKey,
		model:   model,
		client:  &http.Client{Timeout: timeout},
	}
}

func (c *Client) RecognizeObject(ctx context.Context, req design.ObjectRecognitionRequest) (design.ObjectRecognition, error) {
	imageURL := c.imagePayloadURL(ctx, req)
	if strings.TrimSpace(imageURL) == "" {
		return design.ObjectRecognition{}, fmt.Errorf("%w: image_base64 or image_url is required", design.ErrInvalidInput)
	}
	box := clampBox(req.BoundingBox)
	lang := strings.TrimSpace(req.Lang)
	if lang == "" {
		lang = "zh"
	}

	body, err := json.Marshal(map[string]any{
		"model": c.model,
		"messages": []map[string]any{
			{
				"role": "system",
				"content": strings.Join([]string{
					"You are a concise object recognition engine for an image editor.",
					"Identify only the visible object inside the provided normalized bounding box.",
					"Return strict JSON only: {\"id\":\"\",\"label\":\"short object label\",\"confidence\":0.0,\"bbox\":{\"x\":0,\"y\":0,\"width\":0,\"height\":0}}.",
					"Use the requested language for label: Chinese for zh, English for en.",
					"Keep labels short and concrete. Do not describe the whole scene.",
					"Return the input bbox unless you can confidently tighten it to the same object.",
				}, " "),
			},
			{
				"role": "user",
				"content": []map[string]any{
					{
						"type": "text",
						"text": fmt.Sprintf("Recognize the object inside bbox x=%.6f y=%.6f width=%.6f height=%.6f. Language: %s. Extra hint: %s", box.X, box.Y, box.Width, box.Height, lang, strings.TrimSpace(req.Prompt)),
					},
					{
						"type": "image_url",
						"image_url": map[string]any{
							"url": imageURL,
						},
					},
				},
			},
		},
		"response_format": map[string]any{"type": "json_object"},
		"max_tokens":      220,
		"temperature":     0.1,
	})
	if err != nil {
		return design.ObjectRecognition{}, err
	}

	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"/v1/chat/completions", bytes.NewReader(body))
	if err != nil {
		return design.ObjectRecognition{}, err
	}
	httpReq.Header.Set("Authorization", "Bearer "+c.apiKey)
	httpReq.Header.Set("Content-Type", "application/json")

	resp, err := c.client.Do(httpReq)
	if err != nil {
		return design.ObjectRecognition{}, err
	}
	defer resp.Body.Close()

	respBody, err := io.ReadAll(io.LimitReader(resp.Body, 4<<20))
	if err != nil {
		return design.ObjectRecognition{}, err
	}
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return design.ObjectRecognition{}, fmt.Errorf("object recognition failed: %s: %s", resp.Status, trimErrorBody(respBody))
	}

	var payload struct {
		Choices []struct {
			Message struct {
				Content json.RawMessage `json:"content"`
			} `json:"message"`
		} `json:"choices"`
	}
	if err := json.Unmarshal(respBody, &payload); err != nil {
		return design.ObjectRecognition{}, err
	}
	if len(payload.Choices) == 0 {
		return design.ObjectRecognition{}, fmt.Errorf("object recognition returned no choices")
	}
	result, err := parseObjectRecognition(decodeChatContent(payload.Choices[0].Message.Content), box, lang)
	if err != nil {
		return design.ObjectRecognition{}, err
	}
	return result, nil
}

func (c *Client) imagePayloadURL(ctx context.Context, req design.ObjectRecognitionRequest) string {
	imageBase64 := strings.TrimSpace(req.ImageBase64)
	if imageBase64 != "" {
		if strings.HasPrefix(imageBase64, "data:image/") {
			return imageBase64
		}
		return "data:image/jpeg;base64," + imageBase64
	}

	imageURL := strings.TrimSpace(req.ImageURL)
	if imageURL == "" || strings.HasPrefix(imageURL, "data:image/") {
		return imageURL
	}
	if !strings.HasPrefix(imageURL, "http://") && !strings.HasPrefix(imageURL, "https://") {
		return imageURL
	}

	httpReq, err := http.NewRequestWithContext(ctx, http.MethodGet, imageURL, nil)
	if err != nil {
		return imageURL
	}
	resp, err := c.client.Do(httpReq)
	if err != nil {
		return imageURL
	}
	defer resp.Body.Close()
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return imageURL
	}

	data, err := io.ReadAll(io.LimitReader(resp.Body, 20<<20))
	if err != nil || len(data) == 0 {
		return imageURL
	}
	contentType := strings.TrimSpace(strings.Split(resp.Header.Get("Content-Type"), ";")[0])
	if contentType == "" || !strings.HasPrefix(contentType, "image/") {
		contentType = http.DetectContentType(data)
	}
	if !strings.HasPrefix(contentType, "image/") {
		return imageURL
	}
	return "data:" + contentType + ";base64," + base64.StdEncoding.EncodeToString(data)
}

func parseObjectRecognition(content string, fallbackBox design.NormalizedBox, lang string) (design.ObjectRecognition, error) {
	content = strings.TrimSpace(content)
	if content == "" {
		return design.ObjectRecognition{}, fmt.Errorf("object recognition returned empty content")
	}
	if strings.HasPrefix(content, "```") {
		content = strings.Trim(content, "`")
		content = strings.TrimSpace(strings.TrimPrefix(content, "json"))
	}
	if start := strings.Index(content, "{"); start >= 0 {
		if end := strings.LastIndex(content, "}"); end >= start {
			content = content[start : end+1]
		}
	}

	var parsed struct {
		ID         string               `json:"id"`
		Label      string               `json:"label"`
		Confidence float64              `json:"confidence"`
		BBox       design.NormalizedBox `json:"bbox"`
		Box        design.NormalizedBox `json:"box"`
	}
	if err := json.Unmarshal([]byte(content), &parsed); err != nil {
		return design.ObjectRecognition{}, err
	}
	label := strings.TrimSpace(parsed.Label)
	if label == "" {
		if strings.HasPrefix(strings.ToLower(strings.TrimSpace(lang)), "en") {
			label = "selected object"
		} else {
			label = "选中对象"
		}
	}
	box := parsed.BBox
	if box.Width <= 0 || box.Height <= 0 {
		box = parsed.Box
	}
	if box.Width <= 0 || box.Height <= 0 {
		box = fallbackBox
	}
	return design.ObjectRecognition{
		ID:         strings.TrimSpace(parsed.ID),
		Label:      label,
		Confidence: clampFloat(parsed.Confidence, 0, 1),
		BBox:       clampBox(box),
	}, nil
}

func decodeChatContent(raw json.RawMessage) string {
	var text string
	if err := json.Unmarshal(raw, &text); err == nil {
		return text
	}
	var parts []struct {
		Type string `json:"type"`
		Text string `json:"text"`
	}
	if err := json.Unmarshal(raw, &parts); err == nil {
		var builder strings.Builder
		for _, part := range parts {
			if strings.TrimSpace(part.Text) == "" {
				continue
			}
			if builder.Len() > 0 {
				builder.WriteString("\n")
			}
			builder.WriteString(part.Text)
		}
		return builder.String()
	}
	return string(raw)
}

func trimErrorBody(body []byte) string {
	value := strings.TrimSpace(string(body))
	if value == "" {
		return "empty error body"
	}
	if len(value) > 1200 {
		return value[:1200]
	}
	return value
}

func clampBox(box design.NormalizedBox) design.NormalizedBox {
	width := clampFloat(box.Width, 0, 1)
	height := clampFloat(box.Height, 0, 1)
	return design.NormalizedBox{
		X:      clampFloat(box.X, 0, 1-width),
		Y:      clampFloat(box.Y, 0, 1-height),
		Width:  width,
		Height: height,
	}
}

func clampFloat(value float64, minValue float64, maxValue float64) float64 {
	if value < minValue {
		return minValue
	}
	if value > maxValue {
		return maxValue
	}
	return value
}