diff options
Diffstat (limited to '')
| -rw-r--r-- | internal/httpapi/api.go | 369 |
1 files changed, 369 insertions, 0 deletions
diff --git a/internal/httpapi/api.go b/internal/httpapi/api.go new file mode 100644 index 0000000..ed7939f --- /dev/null +++ b/internal/httpapi/api.go @@ -0,0 +1,369 @@ +package httpapi + +import ( + "context" + "crypto/rand" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "io" + "log/slog" + "net/http" + "strings" + "time" + + "aigw/internal/apierror" + "aigw/internal/auth" + "aigw/internal/catalog" + "aigw/internal/domain" + "aigw/internal/provider" + "aigw/internal/routing" + "aigw/internal/telemetry" + "aigw/internal/usage" +) + +type requestIDKey struct{} + +type API struct { + authenticator auth.Authenticator + catalog *catalog.Catalog + router *routing.Router + forwarder *provider.Forwarder + usageSink telemetry.UsageSink + metrics *telemetry.Metrics + logger *slog.Logger + maxBodyBytes int64 + exposeMetrics bool +} + +type Options struct { + Authenticator auth.Authenticator + Catalog *catalog.Catalog + Router *routing.Router + Forwarder *provider.Forwarder + UsageSink telemetry.UsageSink + Metrics *telemetry.Metrics + Logger *slog.Logger + MaxBodyBytes int64 + ExposeMetrics bool +} + +func New(options Options) *API { + return &API{ + authenticator: options.Authenticator, + catalog: options.Catalog, + router: options.Router, + forwarder: options.Forwarder, + usageSink: options.UsageSink, + metrics: options.Metrics, + logger: options.Logger, + maxBodyBytes: options.MaxBodyBytes, + exposeMetrics: options.ExposeMetrics, + } +} + +func (a *API) Handler() http.Handler { + mux := http.NewServeMux() + mux.HandleFunc("GET /healthz", a.health) + mux.HandleFunc("GET /readyz", a.health) + if a.exposeMetrics { + mux.Handle("GET /metrics", a.metrics) + } + + mux.HandleFunc("GET /v1/models", a.openAIModels) + mux.HandleFunc("GET /api/v1/models", a.openAIModels) + mux.HandleFunc("POST /v1/chat/completions", a.openAIChat) + mux.HandleFunc("POST /api/v1/chat/completions", a.openAIChat) + + mux.HandleFunc("GET /anthropic/v1/models", a.anthropicModels) + mux.HandleFunc("GET /api/anthropic/v1/models", a.anthropicModels) + mux.HandleFunc("POST /anthropic/v1/messages", a.anthropicMessages) + mux.HandleFunc("POST /api/anthropic/v1/messages", a.anthropicMessages) + + return a.withRequestID(a.recoverPanics(mux)) +} + +func (a *API) health(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"status":"ok"}`+"\n") +} + +func (a *API) openAIChat(w http.ResponseWriter, r *http.Request) { + a.serveInference(w, r, domain.ProtocolOpenAI) +} + +func (a *API) anthropicMessages(w http.ResponseWriter, r *http.Request) { + a.serveInference(w, r, domain.ProtocolAnthropic) +} + +func (a *API) serveInference(w http.ResponseWriter, r *http.Request, protocol domain.Protocol) { + requestID := requestIDFrom(r.Context()) + startedAt := time.Now().UTC() + a.metrics.RequestStarted() + success := false + defer func() { a.metrics.RequestFinished(success) }() + + principal, authErr := a.authenticator.Authenticate(r) + if authErr != nil { + apierror.Write(w, apierror.Error{Status: http.StatusForbidden, Type: "access_denied", Message: "Invalid or missing API key"}, requestID) + return + } + if !hasScope(principal, "inference") { + apierror.Write(w, apierror.Error{Status: http.StatusForbidden, Type: "access_denied", Message: "API key does not have inference permission"}, requestID) + return + } + + body, err := readBody(w, r, a.maxBodyBytes) + if err != nil { + status := http.StatusBadRequest + message := "Invalid request body" + if errors.As(err, new(*http.MaxBytesError)) { + status = http.StatusRequestEntityTooLarge + message = "Request body is too large" + } + apierror.Write(w, apierror.Error{Status: status, Type: "invalid_params", Message: message}, requestID) + return + } + + var envelope struct { + Model string `json:"model"` + Stream bool `json:"stream"` + } + if err := json.Unmarshal(body, &envelope); err != nil || strings.TrimSpace(envelope.Model) == "" { + apierror.Write(w, apierror.Error{Status: http.StatusBadRequest, Type: "invalid_params", Message: "Parameter model is required and the body must be valid JSON"}, requestID) + return + } + + routes, err := a.router.Plan(envelope.Model, protocol) + if err != nil { + if errors.Is(err, routing.ErrNoRoute) { + apierror.Write(w, apierror.Error{Status: http.StatusNotFound, Type: "model_not_supported", Message: "Model does not support this API protocol"}, requestID) + } else { + apierror.Write(w, apierror.Error{Status: http.StatusNotFound, Type: "invalid_model", Message: "Model does not exist"}, requestID) + } + return + } + + result, err := a.forwarder.Forward(r.Context(), protocol, requestID, body, r.Header, routes) + if err != nil { + errorType := "no_provider_available" + status := http.StatusBadGateway + if errors.Is(err, context.Canceled) { + errorType = "client_disconnected" + status = 499 + } else { + apierror.Write(w, apierror.Error{Status: status, Type: errorType, Message: "No upstream provider is currently available"}, requestID) + } + a.publishUsage(domain.UsageEvent{ + RequestID: requestID, KeyID: principal.KeyID, TenantID: principal.TenantID, ProjectID: principal.ProjectID, + PublicModel: envelope.Model, Protocol: protocol, Stream: envelope.Stream, StatusCode: status, + Success: false, ErrorType: errorType, Attempts: result.Attempts, StartedAt: startedAt, DurationMS: time.Since(startedAt).Milliseconds(), + }) + return + } + defer result.Response.Body.Close() + + if result.Response.StatusCode >= 400 { + gatewayError := normalizeProviderError(result.Response) + apierror.Write(w, gatewayError, requestID) + a.publishUsage(domain.UsageEvent{ + RequestID: requestID, KeyID: principal.KeyID, TenantID: principal.TenantID, ProjectID: principal.ProjectID, + PublicModel: envelope.Model, ProviderID: result.Route.Provider.ID, UpstreamModel: result.Route.UpstreamModel, + Protocol: protocol, Stream: envelope.Stream, StatusCode: gatewayError.Status, Success: false, + ErrorType: gatewayError.Type, Attempts: result.Attempts, StartedAt: startedAt, DurationMS: time.Since(startedAt).Milliseconds(), + }) + return + } + + stream := envelope.Stream || strings.HasPrefix(strings.ToLower(result.Response.Header.Get("Content-Type")), "text/event-stream") + copyResponseHeaders(w.Header(), result.Response.Header, stream) + w.Header().Set("X-AIGW-Request-ID", requestID) + if stream { + w.Header().Set("X-Accel-Buffering", "no") + } + w.WriteHeader(result.Response.StatusCode) + + observer := usage.NewObserver(protocol, stream) + copyErr := copyResponse(w, result.Response.Body, observer, stream) + usageResult := observer.Usage() + success = copyErr == nil + errorType := "" + if copyErr != nil { + errorType = "stream_interrupted" + } + a.publishUsage(domain.UsageEvent{ + RequestID: requestID, KeyID: principal.KeyID, TenantID: principal.TenantID, ProjectID: principal.ProjectID, + PublicModel: envelope.Model, ProviderID: result.Route.Provider.ID, UpstreamModel: result.Route.UpstreamModel, + Protocol: protocol, Stream: stream, StatusCode: result.Response.StatusCode, Success: success, + ErrorType: errorType, Attempts: result.Attempts, StartedAt: startedAt, DurationMS: time.Since(startedAt).Milliseconds(), Usage: usageResult, + }) + a.logger.Info("inference_request", + "request_id", requestID, + "tenant_id", principal.TenantID, + "project_id", principal.ProjectID, + "model", envelope.Model, + "provider", result.Route.Provider.ID, + "status", result.Response.StatusCode, + "attempts", result.Attempts, + "duration_ms", time.Since(startedAt).Milliseconds(), + ) +} + +func (a *API) openAIModels(w http.ResponseWriter, r *http.Request) { + if !a.authorize(w, r) { + return + } + models := a.catalog.Models(domain.ProtocolOpenAI) + data := make([]map[string]any, 0, len(models)) + for _, model := range models { + data = append(data, map[string]any{"id": model.ID, "object": "model", "created": 0, "owned_by": model.OwnedBy}) + } + writeJSON(w, map[string]any{"object": "list", "data": data}) +} + +func (a *API) anthropicModels(w http.ResponseWriter, r *http.Request) { + if !a.authorize(w, r) { + return + } + models := a.catalog.Models(domain.ProtocolAnthropic) + data := make([]map[string]any, 0, len(models)) + for _, model := range models { + data = append(data, map[string]any{"id": model.ID, "display_name": model.ID, "created_at": "1970-01-01T00:00:00Z", "type": "model"}) + } + response := map[string]any{"data": data, "has_more": false, "first_id": nil, "last_id": nil} + if len(models) > 0 { + response["first_id"] = models[0].ID + response["last_id"] = models[len(models)-1].ID + } + writeJSON(w, response) +} + +func (a *API) authorize(w http.ResponseWriter, r *http.Request) bool { + principal, err := a.authenticator.Authenticate(r) + if err != nil || !hasScope(principal, "inference") { + apierror.Write(w, apierror.Error{Status: http.StatusForbidden, Type: "access_denied", Message: "Invalid API key or insufficient permission"}, requestIDFrom(r.Context())) + return false + } + return true +} + +func (a *API) publishUsage(event domain.UsageEvent) { + if a.usageSink != nil { + a.usageSink.Publish(event) + } +} + +func (a *API) withRequestID(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + id := newRequestID() + w.Header().Set("X-AIGW-Request-ID", id) + next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), requestIDKey{}, id))) + }) +} + +func (a *API) recoverPanics(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + defer func() { + if recovered := recover(); recovered != nil { + a.logger.Error("request_panic", "request_id", requestIDFrom(r.Context()), "error", recovered) + apierror.Write(w, apierror.Error{Status: http.StatusInternalServerError, Type: "internal_server_error", Message: "Internal server error"}, requestIDFrom(r.Context())) + } + }() + next.ServeHTTP(w, r) + }) +} + +func newRequestID() string { + var value [16]byte + if _, err := rand.Read(value[:]); err != nil { + return fmt.Sprintf("req_%d", time.Now().UnixNano()) + } + return "req_" + hex.EncodeToString(value[:]) +} + +func requestIDFrom(ctx context.Context) string { + id, _ := ctx.Value(requestIDKey{}).(string) + return id +} + +func hasScope(principal domain.Principal, wanted string) bool { + if len(principal.Scopes) == 0 { + return true + } + for _, scope := range principal.Scopes { + if scope == wanted || scope == "*" { + return true + } + } + return false +} + +func readBody(w http.ResponseWriter, r *http.Request, limit int64) ([]byte, error) { + r.Body = http.MaxBytesReader(w, r.Body, limit) + defer r.Body.Close() + return io.ReadAll(r.Body) +} + +func copyResponseHeaders(target, source http.Header, stream bool) { + for _, name := range []string{"Content-Type", "Cache-Control", "Retry-After"} { + if value := source.Get(name); value != "" { + target.Set(name, value) + } + } + if !stream { + if value := source.Get("Content-Length"); value != "" { + target.Set("Content-Length", value) + } + } +} + +func copyResponse(w http.ResponseWriter, body io.Reader, observer io.Writer, stream bool) error { + var destination io.Writer = w + if stream { + destination = &flushingWriter{writer: w, controller: http.NewResponseController(w)} + } + _, err := io.CopyBuffer(io.MultiWriter(destination, observer), body, make([]byte, 32<<10)) + return err +} + +type flushingWriter struct { + writer io.Writer + controller *http.ResponseController +} + +func (w *flushingWriter) Write(p []byte) (int, error) { + n, err := w.writer.Write(p) + if err == nil { + _ = w.controller.Flush() + } + return n, err +} + +func normalizeProviderError(response *http.Response) apierror.Error { + payload, _ := io.ReadAll(io.LimitReader(response.Body, 64<<10)) + message := "Upstream provider rejected the request" + var common struct { + Error struct { + Message string `json:"message"` + } `json:"error"` + } + if json.Unmarshal(payload, &common) == nil && common.Error.Message != "" { + message = common.Error.Message + } + switch response.StatusCode { + case http.StatusBadRequest: + return apierror.Error{Status: http.StatusUnprocessableEntity, Type: "provider_unprocessable_entity_error", Message: message} + case http.StatusRequestEntityTooLarge: + return apierror.Error{Status: http.StatusRequestEntityTooLarge, Type: "invalid_params", Message: message} + case http.StatusTooManyRequests: + return apierror.Error{Status: http.StatusTooManyRequests, Type: "rate_limit", Message: message} + default: + return apierror.Error{Status: http.StatusBadGateway, Type: "provider_error", Message: message} + } +} + +func writeJSON(w http.ResponseWriter, value any) { + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(value) +} |