Cyrene_v2/internal/runtime/multimodal/vision/adapter.go

package vision

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"strings"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
)

// visionMaxTokens caps the vision/OCR response length.
const visionMaxTokens = 512

// VisionAdapter understands an image and returns a description of it. It is the
// backend behind the vision tool; real providers implement it without exposing
// the image to the main-model context.
type VisionAdapter interface {
	// Describe returns a description of img. prompt is an optional instruction
	// overriding the default "describe this image" ask.
	Describe(ctx context.Context, img ImageRef, prompt string) (string, error)
}

// OCRAdapter extracts text from an image. It is the backend behind the OCR tool.
type OCRAdapter interface {
	// Extract returns the text found in img.
	Extract(ctx context.Context, img ImageRef) (string, error)
}

// OpenAICompatibleAdapter implements both VisionAdapter and OCRAdapter by
// calling an OpenAI-compatible /chat/completions endpoint with the image
// embedded as a content part. It is the ready default implementation; real
// cloud vision/OCR providers expose the same protocol.
type OpenAICompatibleAdapter struct {
	cfg    config.LLMConfig
	model  string
	client *http.Client
	log    *slog.Logger
}

// OpenAIOption customizes an OpenAICompatibleAdapter.
type OpenAIOption func(*OpenAICompatibleAdapter)

// WithModel overrides the adapter model (defaults to cfg.Model).
func WithModel(m string) OpenAIOption {
	return func(a *OpenAICompatibleAdapter) {
		if m != "" {
			a.model = m
		}
	}
}

// WithHTTPClient sets the *http.Client used for the request; nil keeps default.
func WithHTTPClient(c *http.Client) OpenAIOption {
	return func(a *OpenAICompatibleAdapter) {
		if c != nil {
			a.client = c
		}
	}
}

// WithLogger sets the logger used on failures; nil keeps a discard logger.
func WithLogger(l *slog.Logger) OpenAIOption {
	return func(a *OpenAICompatibleAdapter) {
		if l != nil {
			a.log = l
		}
	}
}

// NewOpenAICompatible builds an adapter from cfg, using cfg.Model unless
// overridden.
func NewOpenAICompatible(cfg config.LLMConfig, opts ...OpenAIOption) *OpenAICompatibleAdapter {
	a := &OpenAICompatibleAdapter{
		cfg:    cfg,
		model:  cfg.Model,
		client: defaultHTTPClient(),
		log:    slog.New(slog.NewTextHandler(io.Discard, nil)),
	}
	for _, o := range opts {
		o(a)
	}
	return a
}

// Describe returns an image description from the underlying chat endpoint.
func (a *OpenAICompatibleAdapter) Describe(ctx context.Context, img ImageRef, prompt string) (string, error) {
	if prompt == "" {
		prompt = "Describe this image in detail."
	}
	return a.ask(ctx, img, prompt)
}

// Extract returns text found in the image via the same endpoint.
func (a *OpenAICompatibleAdapter) Extract(ctx context.Context, img ImageRef) (string, error) {
	return a.ask(ctx, img, "Extract all text visible in this image. Return only the text.")
}

// ask resolves the image, builds a chat request and returns the assistant text.
func (a *OpenAICompatibleAdapter) ask(ctx context.Context, img ImageRef, prompt string) (string, error) {
	if strings.TrimSpace(a.cfg.BaseURL) == "" {
		return "", ErrNotConfigured
	}
	if strings.TrimSpace(a.model) == "" {
		return "", ErrNotConfigured
	}
	if err := ctx.Err(); err != nil {
		return "", err
	}
	imageURL, err := a.imageURL(ctx, img)
	if err != nil {
		return "", err
	}
	payload, err := json.Marshal(a.request(imageURL, prompt))
	if err != nil {
		return "", fmt.Errorf("vision: %w: marshal request: %v", ErrUpstream, err)
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, strings.TrimRight(a.cfg.BaseURL, "/")+"/chat/completions", bytes.NewReader(payload))
	if err != nil {
		return "", fmt.Errorf("vision: %w: build request: %v", ErrUpstream, err)
	}
	req.Header.Set("Content-Type", "application/json")
	if a.cfg.APIKey != "" {
		req.Header.Set("Authorization", "Bearer "+a.cfg.APIKey)
	}
	resp, err := a.client.Do(req)
	if err != nil {
		return "", fmt.Errorf("vision: %w: %v", ErrUpstream, err)
	}
	defer func() { _ = resp.Body.Close() }()
	if resp.StatusCode != http.StatusOK {
		snippet, _ := io.ReadAll(io.LimitReader(resp.Body, maxFetchedBody))
		return "", fmt.Errorf("vision: %w: status %d: %s", ErrUpstream, resp.StatusCode, strings.TrimSpace(string(snippet)))
	}
	return a.decode(resp.Body)
}

// imageURL converts an ImageRef to the image_url value the wire expects: an
// http(s) URL passes through, anything else becomes a base64 data: URI.
func (a *OpenAICompatibleAdapter) imageURL(ctx context.Context, img ImageRef) (string, error) {
	if img.Source == SourceURL {
		return img.Data, nil
	}
	mime, data, err := Resolve(ctx, img, a.client)
	if err != nil {
		return "", err
	}
	return DataURI(mime, data), nil
}

// openAIVisionRequest is the JSON body sent to /chat/completions.
type openAIVisionRequest struct {
	Model     string          `json:"model"`
	Messages  []openAIMessage `json:"messages"`
	MaxTokens int             `json:"max_tokens"`
}

type openAIMessage struct {
	Role    string `json:"role"`
	Content any    `json:"content"`
}

type openAIPart struct {
	Type     string          `json:"type"`
	Text     string          `json:"text,omitempty"`
	ImageURL *openAIImageURL `json:"image_url,omitempty"`
}

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

type openAIVisionResponse struct {
	Choices []openAIChoice `json:"choices"`
}

type openAIChoice struct {
	Message openAIResponseMessage `json:"message"`
}

type openAIResponseMessage struct {
	Content string `json:"content"`
}

// request builds the wire body for an image+prompt.
func (a *OpenAICompatibleAdapter) request(imageURL, prompt string) openAIVisionRequest {
	return openAIVisionRequest{
		Model: a.model,
		Messages: []openAIMessage{
			{Role: "user", Content: []openAIPart{
				{Type: "text", Text: prompt},
				{Type: "image_url", ImageURL: &openAIImageURL{URL: imageURL}},
			}},
		},
		MaxTokens: visionMaxTokens,
	}
}

// decode extracts the first choice's text from a response body.
func (a *OpenAICompatibleAdapter) decode(r io.Reader) (string, error) {
	var wire openAIVisionResponse
	if err := json.NewDecoder(r).Decode(&wire); err != nil {
		return "", fmt.Errorf("vision: %w: decode response: %v", ErrUpstream, err)
	}
	if len(wire.Choices) == 0 {
		return "", fmt.Errorf("vision: %w: empty choices", ErrUpstream)
	}
	return wire.Choices[0].Message.Content, nil
}

// interface assertion: one concrete adapter serves both directions.
var _ VisionAdapter = (*OpenAICompatibleAdapter)(nil)
var _ OCRAdapter = (*OpenAICompatibleAdapter)(nil)