iop/apps/edge/internal/openai/server.go

package openai

import (
	"context"
	"errors"
	"fmt"
	"net"
	"net/http"
	"time"

	"go.uber.org/zap"

	edgeservice "iop/apps/edge/internal/service"
	"iop/packages/go/config"
)

type runService interface {
	SubmitRun(context.Context, edgeservice.SubmitRunRequest) (edgeservice.RunResult, error)
	OllamaAPI(context.Context, edgeservice.OllamaAPIRequest) (edgeservice.OllamaAPIView, error)
}

type Server struct {
	cfg     config.EdgeOpenAIConf
	service runService
	logger  *zap.Logger
	server  *http.Server
}

func NewServer(cfg config.EdgeOpenAIConf, svc runService, logger *zap.Logger) *Server {
	if logger == nil {
		logger = zap.NewNop()
	}
	return &Server{cfg: cfg, service: svc, logger: logger}
}

func (s *Server) Enabled() bool {
	return s != nil && s.cfg.Enabled
}

func (s *Server) Start(ctx context.Context) error {
	if !s.Enabled() {
		return nil
	}
	if s.cfg.Listen == "" {
		s.cfg.Listen = "0.0.0.0:18081"
	}

	mux := s.routes()

	s.server = &http.Server{
		Addr:              s.cfg.Listen,
		Handler:           mux,
		ReadHeaderTimeout: 5 * time.Second,
	}
	ln, err := net.Listen("tcp", s.cfg.Listen)
	if err != nil {
		return fmt.Errorf("openai server listen %s: %w", s.cfg.Listen, err)
	}
	go func() {
		<-ctx.Done()
		shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		_ = s.server.Shutdown(shutdownCtx)
	}()
	go func() {
		s.logger.Info("openai-compatible server listening", zap.String("addr", s.cfg.Listen))
		if err := s.server.Serve(ln); err != nil && !errors.Is(err, http.ErrServerClosed) {
			s.logger.Warn("openai-compatible server exited", zap.Error(err))
		}
	}()
	return nil
}

func (s *Server) Stop(ctx context.Context) error {
	if s == nil || s.server == nil {
		return nil
	}
	return s.server.Shutdown(ctx)
}