Atlas/internal/service/invoke/service.go

package invoke

import (
	"context"
	"encoding/json"
	"fmt"
	"time"

	"github.com/rose_cat707/Atlas/internal/domain"
	"github.com/rose_cat707/Atlas/internal/provider"
	"github.com/rose_cat707/Atlas/internal/service/taskevents"
	"github.com/rose_cat707/Atlas/internal/store"
)

type Service struct {
	store      *store.Store
	router     *provider.Router
	models     *provider.ModelsService
	taskEvents *taskevents.Hub
	specs      map[domain.ServiceType]*domain.ServiceTypeSpec
	worker     *asyncWorker
}

func NewService(st *store.Store, router *provider.Router, models *provider.ModelsService, taskEvents *taskevents.Hub) *Service {
	s := &Service{
		store:      st,
		router:     router,
		models:     models,
		taskEvents: taskEvents,
		specs:      defaultSpecs(),
	}
	s.worker = newAsyncWorker(s)
	go s.worker.run()
	return s
}

func (s *Service) Spec(t domain.ServiceType) (*domain.ServiceTypeSpec, bool) {
	spec, ok := s.specs[t]
	return spec, ok
}

func (s *Service) Specs() map[domain.ServiceType]*domain.ServiceTypeSpec {
	return s.specs
}

func (s *Service) Run(ctx context.Context, req domain.InvokeRequest) (*domain.InvokeResult, error) {
	invokeType := req.ServiceType
	if domain.IsImageInvokeService(req.ServiceType) {
		invokeType = domain.ServiceImageGeneration
	}
	spec, ok := s.specs[invokeType]
	if !ok {
		return nil, fmt.Errorf("unknown service_type: %s", req.ServiceType)
	}
	switch req.Mode {
	case domain.ModeStream:
		if spec.Capabilities.Stream == nil {
			return nil, fmt.Errorf("stream not supported for %s", req.ServiceType)
		}
	case domain.ModeSync:
		if spec.Capabilities.Sync == nil {
			return nil, fmt.Errorf("sync not supported for %s", req.ServiceType)
		}
	case domain.ModeAsync:
		if spec.Capabilities.Async == nil {
			return nil, fmt.Errorf("async not supported for %s", req.ServiceType)
		}
	default:
		return nil, fmt.Errorf("invalid mode: %s", req.Mode)
	}

	var params map[string]any
	if len(req.Params) > 0 {
		if err := json.Unmarshal(req.Params, &params); err != nil {
			return nil, fmt.Errorf("invalid params: %w", err)
		}
	}
	if spec.ValidateParams != nil {
		if err := spec.ValidateParams(params); err != nil {
			return nil, err
		}
	}

	cfg, err := s.router.Resolve(ctx, string(domain.ServiceImageGeneration), req.ProviderID)
	if err != nil {
		return nil, err
	}
	if req.Model == "" {
		models, _ := s.models.ListEffective(ctx, cfg)
		req.Model = provider.DefaultModel(models)
	}

	var modelDetail *provider.ModelDetail
	if domain.IsImageInvokeService(req.ServiceType) {
		detail, err := s.models.FindModelDetail(ctx, cfg, req.Model)
		if err != nil {
			return nil, err
		}
		modelDetail = detail
		if err := provider.ValidateModelInput(detail.InputSchema, params); err != nil {
			return nil, err
		}
	}

	switch req.Mode {
	case domain.ModeAsync:
		return s.runAsync(ctx, req, cfg, params, modelDetail)
	case domain.ModeSync:
		return s.runSync(ctx, req, cfg, params, modelDetail)
	default:
		return nil, fmt.Errorf("stream mode handled by chat service")
	}
}

func recordServiceType(req domain.InvokeRequest, detail *provider.ModelDetail) string {
	if detail != nil && detail.Kind != "" {
		return detail.Kind
	}
	return string(req.ServiceType)
}

func (s *Service) runSync(ctx context.Context, req domain.InvokeRequest, cfg *provider.Config, params map[string]any, detail *provider.ModelDetail) (*domain.InvokeResult, error) {
	body, err := buildPredictionBody(predictionVersion(req.Model, detail), params)
	if err != nil {
		return nil, err
	}
	meta := provider.CallMeta{
		SessionID: req.SessionID,
		Method:    "POST",
		Path:      "/v1/predictions",
		Model:     req.Model,
	}
	res, err := s.router.ProxyWithMeta(ctx, cfg, "POST", "/v1/predictions", body, map[string]string{"Prefer": "wait"}, meta)
	if err != nil {
		return nil, err
	}
	if res.StatusCode >= 400 {
		return nil, fmt.Errorf("upstream error %d: %s", res.StatusCode, string(res.Body))
	}
	output := parsePredictionOutput(res.Body)
	pid := cfg.ID
	gen := &store.Generation{
		SessionID:   req.SessionID,
		ServiceType: recordServiceType(req, detail),
		ProviderID:  &pid,
		Model:       req.Model,
		InputJSON:   string(req.Params),
		OutputJSON:  string(output),
	}
	id, err := s.store.InsertGeneration(ctx, gen)
	if err != nil {
		return nil, err
	}
	return &domain.InvokeResult{GenerationID: id, Output: output}, nil
}

func (s *Service) runAsync(ctx context.Context, req domain.InvokeRequest, cfg *provider.Config, params map[string]any, detail *provider.ModelDetail) (*domain.InvokeResult, error) {
	body, err := buildPredictionBody(predictionVersion(req.Model, detail), params)
	if err != nil {
		return nil, err
	}
	meta := provider.CallMeta{
		SessionID: req.SessionID,
		Method:    "POST",
		Path:      "/v1/predictions",
		Model:     req.Model,
	}
	res, err := s.router.ProxyWithMeta(ctx, cfg, "POST", "/v1/predictions", body, nil, meta)
	if err != nil {
		return nil, err
	}
	if res.StatusCode >= 400 {
		return nil, fmt.Errorf("upstream error %d: %s", res.StatusCode, string(res.Body))
	}
	var pred struct {
		ID string `json:"id"`
	}
	_ = json.Unmarshal(res.Body, &pred)
	task := &store.Task{
		SessionID:   req.SessionID,
		ServiceType: recordServiceType(req, detail),
		Protocol:    cfg.Protocol,
		ProviderID:  cfg.ID,
		Status:      "running",
		Progress:    0,
		InputJSON:   string(req.Params),
		ExternalID:  pred.ID,
	}
	taskID, err := s.store.InsertTask(ctx, task)
	if err != nil {
		return nil, err
	}
	task.ID = taskID
	s.notifyTask(*task)
	return &domain.InvokeResult{TaskID: &taskID}, nil
}

func (s *Service) notifyTask(task store.Task) {
	s.store.Cache.Set(fmt.Sprintf("task:%d", task.ID), &task)
	if s.taskEvents != nil {
		s.taskEvents.Publish(task)
	}
}

func predictionVersion(model string, detail *provider.ModelDetail) string {
	if detail != nil {
		return detail.PredictionVersion()
	}
	return model
}

func buildPredictionBody(version string, params map[string]any) ([]byte, error) {
	input := map[string]any{}
	for k, v := range params {
		input[k] = v
	}
	payload := map[string]any{
		"version": version,
		"input":   input,
	}
	return json.Marshal(payload)
}

func parsePredictionOutput(body []byte) json.RawMessage {
	var pred map[string]any
	if err := json.Unmarshal(body, &pred); err != nil {
		return json.RawMessage(`{}`)
	}
	out := map[string]any{}
	if output, ok := pred["output"]; ok {
		out["output"] = output
	}
	if urls, ok := pred["urls"]; ok {
		out["urls"] = urls
	}
	data, _ := json.Marshal(out)
	return data
}

func defaultSpecs() map[domain.ServiceType]*domain.ServiceTypeSpec {
	return map[domain.ServiceType]*domain.ServiceTypeSpec{
		domain.ServiceChat: {
			Type: domain.ServiceChat,
			Capabilities: domain.InvokeCapabilities{
				Stream: &domain.StreamCapability{},
			},
			DefaultParams: func() map[string]any { return map[string]any{"temperature": 0.7} },
			ValidateParams: func(params map[string]any) error { return nil },
		},
		domain.ServiceImageGeneration: {
			Type: domain.ServiceImageGeneration,
			Capabilities: domain.InvokeCapabilities{
				Sync:  &domain.SyncCapability{},
				Async: &domain.AsyncCapability{Driver: "replicate", PollInterval: 2 * time.Second},
			},
			DefaultParams: func() map[string]any { return map[string]any{} },
		},
	}
}