初始提交:边缘AI算力机统一AI通讯层
This commit is contained in:
@@ -0,0 +1,192 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/edgeai/gateway/internal/adapter"
|
||||
"github.com/edgeai/gateway/internal/auth"
|
||||
"github.com/edgeai/gateway/internal/config"
|
||||
"github.com/edgeai/gateway/internal/handler"
|
||||
"github.com/edgeai/gateway/internal/middleware"
|
||||
"github.com/edgeai/gateway/internal/observability"
|
||||
"github.com/edgeai/gateway/internal/router"
|
||||
"github.com/edgeai/gateway/internal/scheduler"
|
||||
"github.com/edgeai/gateway/internal/session"
|
||||
)
|
||||
|
||||
// Server is the main HTTP server for the AI gateway.
|
||||
type Server struct {
|
||||
cfg *config.Config
|
||||
logger *observability.Logger
|
||||
metrics *observability.Metrics
|
||||
HTTPSrv *http.Server
|
||||
auth *auth.Authenticator
|
||||
registry *adapter.Registry
|
||||
modelMap *router.LogicalModelMapping
|
||||
scheduler *scheduler.Scheduler
|
||||
sessions *session.Store
|
||||
}
|
||||
|
||||
// New creates a new Server instance with all components wired.
|
||||
func New(cfg *config.Config, logger *observability.Logger) (*Server, error) {
|
||||
// Ensure data directory exists
|
||||
dbPath := extractDBPath(cfg.Storage.SessionDB)
|
||||
if dbPath != "" {
|
||||
os.MkdirAll(filepath.Dir(dbPath), 0755)
|
||||
}
|
||||
|
||||
// Initialize auth
|
||||
authPath := filepath.Join(filepath.Dir(dbPath), "auth.db")
|
||||
authenticator, err := auth.NewAuthenticator(authPath, logger)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("init auth: %w", err)
|
||||
}
|
||||
|
||||
// Initialize session store
|
||||
sessionStore, err := session.NewStore(dbPath)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("init session store: %w", err)
|
||||
}
|
||||
|
||||
// Initialize adapter registry
|
||||
registry := adapter.NewRegistry()
|
||||
|
||||
// Initialize logical model mapping
|
||||
modelMap := router.NewLogicalModelMapping(cfg)
|
||||
|
||||
// Register adapters for each unique endpoint
|
||||
registered := make(map[string]bool)
|
||||
for _, mc := range cfg.Models {
|
||||
key := mc.Provider + "|" + mc.Endpoint
|
||||
if !registered[key] {
|
||||
switch mc.Provider {
|
||||
case "ollama":
|
||||
registry.Register(mc.Provider, adapter.NewOllamaAdapter(mc.Endpoint))
|
||||
}
|
||||
registered[key] = true
|
||||
}
|
||||
}
|
||||
|
||||
// Initialize scheduler
|
||||
sched := scheduler.NewScheduler(&cfg.Scheduler, logger)
|
||||
|
||||
// Initialize metrics
|
||||
metrics := observability.NewMetrics()
|
||||
|
||||
s := &Server{
|
||||
cfg: cfg,
|
||||
logger: logger,
|
||||
metrics: metrics,
|
||||
auth: authenticator,
|
||||
registry: registry,
|
||||
modelMap: modelMap,
|
||||
scheduler: sched,
|
||||
sessions: sessionStore,
|
||||
}
|
||||
|
||||
mux := http.NewServeMux()
|
||||
s.registerRoutes(mux)
|
||||
|
||||
// Apply middleware chain (order: Recovery → Logging → RequestID → BodyLimit → Auth → handler)
|
||||
h := middleware.RequestID(mux)
|
||||
h = middleware.BodyLimit(cfg.Server.MaxRequestBodyMB)(h)
|
||||
h = s.auth.Middleware(h)
|
||||
h = middleware.Logging(logger)(h)
|
||||
h = middleware.Recovery(logger)(h)
|
||||
|
||||
s.HTTPSrv = &http.Server{
|
||||
Addr: fmt.Sprintf("%s:%d", cfg.Server.Host, cfg.Server.Port),
|
||||
Handler: h,
|
||||
ReadTimeout: 30 * time.Second,
|
||||
WriteTimeout: 0, // no write timeout for SSE
|
||||
IdleTimeout: 120 * time.Second,
|
||||
}
|
||||
|
||||
return s, nil
|
||||
}
|
||||
|
||||
func (s *Server) registerRoutes(mux *http.ServeMux) {
|
||||
// Health and readiness
|
||||
mux.HandleFunc("/health", s.handleHealth)
|
||||
mux.HandleFunc("/ready", s.handleReady)
|
||||
|
||||
// Metrics
|
||||
mux.HandleFunc(s.cfg.Observability.MetricsPath, s.metrics.Handler())
|
||||
|
||||
// OpenAI-compatible API
|
||||
mux.HandleFunc("/v1/chat/completions", s.handleChatCompletions)
|
||||
mux.HandleFunc("/v1/models", s.handleModels)
|
||||
|
||||
// Session management
|
||||
mux.HandleFunc("/v1/sessions", s.handleSessions)
|
||||
mux.HandleFunc("/v1/sessions/", s.handleSessionByID)
|
||||
}
|
||||
|
||||
// Authenticator returns the authenticator instance (for testing/management).
|
||||
func (s *Server) Authenticator() *auth.Authenticator {
|
||||
return s.auth
|
||||
}
|
||||
|
||||
// Start begins listening for HTTP requests.
|
||||
func (s *Server) Start() error {
|
||||
s.logger.Info("http server starting", observability.F().
|
||||
Event("server_start").
|
||||
Set("addr", s.HTTPSrv.Addr))
|
||||
return s.HTTPSrv.ListenAndServe()
|
||||
}
|
||||
|
||||
// Shutdown gracefully shuts down the server.
|
||||
func (s *Server) Shutdown() error {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
|
||||
s.scheduler.Stop()
|
||||
if s.sessions != nil {
|
||||
s.sessions.Close()
|
||||
}
|
||||
if s.auth != nil {
|
||||
s.auth.Close()
|
||||
}
|
||||
|
||||
return s.HTTPSrv.Shutdown(ctx)
|
||||
}
|
||||
|
||||
func (s *Server) handleHealth(w http.ResponseWriter, r *http.Request) {
|
||||
handler.WriteJSON(w, http.StatusOK, map[string]string{"status": "ok"})
|
||||
}
|
||||
|
||||
func (s *Server) handleReady(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
|
||||
ready := true
|
||||
reasons := []string{}
|
||||
|
||||
for _, name := range s.registry.Names() {
|
||||
a, _ := s.registry.Get(name)
|
||||
if err := a.HealthCheck(r.Context()); err != nil {
|
||||
ready = false
|
||||
reasons = append(reasons, fmt.Sprintf("%s: %v", name, err))
|
||||
}
|
||||
}
|
||||
|
||||
if ready {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
w.Write([]byte(`{"status":"ready"}`))
|
||||
} else {
|
||||
w.WriteHeader(http.StatusServiceUnavailable)
|
||||
fmt.Fprintf(w, `{"status":"not_ready","reasons":["%s"]}`, strings.Join(reasons, `","`))
|
||||
}
|
||||
}
|
||||
|
||||
func extractDBPath(connStr string) string {
|
||||
if strings.HasPrefix(connStr, "sqlite://") {
|
||||
return strings.TrimPrefix(connStr, "sqlite://")
|
||||
}
|
||||
return connStr
|
||||
}
|
||||
Reference in New Issue
Block a user