366 lines
12 KiB
Go
366 lines
12 KiB
Go
package openai
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
edgeservice "iop/apps/edge/internal/service"
|
|
)
|
|
|
|
type singleRequestChatPolicy struct {
|
|
status int
|
|
finishReason string
|
|
errorType string
|
|
message string
|
|
silent bool
|
|
error bool
|
|
}
|
|
|
|
func singleRequestChatTerminalPolicy(disposition edgeservice.SingleRequestTerminalDisposition) singleRequestChatPolicy {
|
|
if disposition.Kind == "" {
|
|
disposition.Kind = edgeservice.SingleRequestTerminalEndTurn
|
|
}
|
|
if disposition.Validate() != nil {
|
|
disposition = edgeservice.SingleRequestTerminalDisposition{
|
|
Kind: edgeservice.SingleRequestTerminalError, ErrorClass: edgeservice.SingleRequestTerminalErrorProvider,
|
|
}
|
|
}
|
|
switch disposition.Kind {
|
|
case edgeservice.SingleRequestTerminalEndTurn:
|
|
return singleRequestChatPolicy{status: http.StatusOK, finishReason: "stop"}
|
|
case edgeservice.SingleRequestTerminalLength:
|
|
return singleRequestChatPolicy{status: http.StatusOK, finishReason: "length"}
|
|
case edgeservice.SingleRequestTerminalCancelled:
|
|
return singleRequestChatPolicy{silent: true}
|
|
case edgeservice.SingleRequestTerminalError:
|
|
switch disposition.ErrorClass {
|
|
case edgeservice.SingleRequestTerminalErrorValidation, edgeservice.SingleRequestTerminalErrorContext:
|
|
return singleRequestChatPolicy{status: http.StatusBadRequest, errorType: "invalid_request_error", message: "single-request execution was rejected", error: true}
|
|
case edgeservice.SingleRequestTerminalErrorTimeout:
|
|
return singleRequestChatPolicy{status: http.StatusBadGateway, errorType: "run_error", message: "single-request execution timed out", error: true}
|
|
default:
|
|
return singleRequestChatPolicy{status: http.StatusBadGateway, errorType: "run_error", message: "single-request execution failed", error: true}
|
|
}
|
|
default:
|
|
return singleRequestChatPolicy{status: http.StatusBadGateway, errorType: "run_error", message: "single-request execution failed", error: true}
|
|
}
|
|
}
|
|
|
|
func (s *Server) handleChatSingleRequestStream(
|
|
w http.ResponseWriter,
|
|
r *http.Request,
|
|
capability singleRequestService,
|
|
dispatch routeDispatch,
|
|
body []byte,
|
|
) {
|
|
requestID, err := newLogicalRequestRandomID()
|
|
if err != nil {
|
|
writeError(w, http.StatusServiceUnavailable, "run_error", "single-request execution is unavailable")
|
|
return
|
|
}
|
|
requestID = "req_" + requestID
|
|
stream, err := newSingleRequestChatStream(w, requestID, dispatch.SingleRequest.PublicModel)
|
|
if err != nil {
|
|
writeError(w, http.StatusInternalServerError, "run_error", "single-request streaming is unavailable")
|
|
return
|
|
}
|
|
execution, err := capability.StartSingleRequest(r.Context(), edgeservice.SingleRequestRequest{
|
|
RequestID: requestID,
|
|
Binding: dispatch.SingleRequest.Clone(),
|
|
Prompt: string(append([]byte(nil), body...)),
|
|
})
|
|
if err != nil || execution == nil {
|
|
if errors.Is(err, edgeservice.ErrSingleRequestExecutorUnavailable) {
|
|
writeError(w, http.StatusServiceUnavailable, "run_error", "single-request execution is unavailable")
|
|
return
|
|
}
|
|
writeError(w, http.StatusBadGateway, "run_error", "single-request execution could not be started")
|
|
return
|
|
}
|
|
defer execution.Cancel()
|
|
_ = pumpSingleRequestChatStream(r.Context(), execution, stream, newWallClockSingleRequestAnthropicTicker)
|
|
}
|
|
|
|
func (s *Server) handleChatSingleRequest(
|
|
w http.ResponseWriter,
|
|
r *http.Request,
|
|
capability singleRequestService,
|
|
dispatch routeDispatch,
|
|
body []byte,
|
|
) {
|
|
requestID, err := newLogicalRequestRandomID()
|
|
if err != nil {
|
|
writeError(w, http.StatusServiceUnavailable, "run_error", "single-request execution is unavailable")
|
|
return
|
|
}
|
|
requestID = "req_" + requestID
|
|
execution, err := capability.StartSingleRequest(r.Context(), edgeservice.SingleRequestRequest{
|
|
RequestID: requestID,
|
|
Binding: dispatch.SingleRequest.Clone(),
|
|
Prompt: string(append([]byte(nil), body...)),
|
|
})
|
|
if err != nil || execution == nil {
|
|
if errors.Is(err, edgeservice.ErrSingleRequestExecutorUnavailable) {
|
|
writeError(w, http.StatusServiceUnavailable, "run_error", "single-request execution is unavailable")
|
|
return
|
|
}
|
|
writeError(w, http.StatusBadGateway, "run_error", "single-request execution could not be started")
|
|
return
|
|
}
|
|
defer execution.Cancel()
|
|
for {
|
|
select {
|
|
case <-r.Context().Done():
|
|
execution.Cancel()
|
|
return
|
|
case progress, ok := <-execution.Progress():
|
|
if !ok {
|
|
if r.Context().Err() == nil && execution.State() != edgeservice.SingleRequestStateCancelled {
|
|
writeError(w, http.StatusBadGateway, "run_error", "single-request execution failed")
|
|
}
|
|
return
|
|
}
|
|
switch progress.Stage {
|
|
case edgeservice.SingleRequestStateFinalizing:
|
|
if progress.Result == nil {
|
|
_ = execution.AcknowledgeTerminal(false)
|
|
writeError(w, http.StatusBadGateway, "run_error", "single-request execution failed")
|
|
return
|
|
}
|
|
writeErr := writeChatSingleRequestTerminal(w, requestID, dispatch.SingleRequest.PublicModel, *progress.Result)
|
|
_ = execution.AcknowledgeTerminal(writeErr == nil)
|
|
return
|
|
case edgeservice.SingleRequestStateFailed:
|
|
policy := singleRequestChatTerminalPolicy(singleRequestProgressTerminal(progress, edgeservice.SingleRequestTerminalDisposition{Kind: edgeservice.SingleRequestTerminalError, ErrorClass: edgeservice.SingleRequestTerminalErrorProvider}))
|
|
writeError(w, policy.status, policy.errorType, policy.message)
|
|
return
|
|
case edgeservice.SingleRequestStateCancelled:
|
|
return
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func writeChatSingleRequestTerminal(w http.ResponseWriter, requestID, model string, result edgeservice.SingleRequestResult) error {
|
|
policy := singleRequestChatTerminalPolicy(result.Terminal)
|
|
if policy.silent || policy.error || policy.finishReason == "" {
|
|
return errors.New("single-request result has no Chat terminal")
|
|
}
|
|
content := result.Output
|
|
if policy.finishReason == "length" {
|
|
content = ""
|
|
}
|
|
response := chatCompletionResponse{
|
|
ID: "chatcmpl_iop_" + strings.TrimPrefix(requestID, "req_"), Object: "chat.completion",
|
|
Created: time.Now().Unix(), Model: model,
|
|
Choices: []chatCompletionChoice{{Index: 0, Message: chatMessage{Role: "assistant", Content: content}, FinishReason: policy.finishReason}},
|
|
}
|
|
encoded, err := json.Marshal(response)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
encoded = append(encoded, '\n')
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(policy.status)
|
|
n, err := w.Write(encoded)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if n != len(encoded) {
|
|
return io.ErrShortWrite
|
|
}
|
|
return nil
|
|
}
|
|
|
|
type singleRequestChatStream struct {
|
|
mu sync.Mutex
|
|
w http.ResponseWriter
|
|
flusher http.Flusher
|
|
id string
|
|
model string
|
|
created int64
|
|
started bool
|
|
terminal bool
|
|
terminalErr error
|
|
}
|
|
|
|
func newSingleRequestChatStream(w http.ResponseWriter, requestID, model string) (*singleRequestChatStream, error) {
|
|
flusher, ok := w.(http.Flusher)
|
|
if !ok || strings.TrimSpace(requestID) == "" || strings.TrimSpace(model) == "" {
|
|
return nil, errors.New("single-request Chat stream is unavailable")
|
|
}
|
|
return &singleRequestChatStream{w: w, flusher: flusher, id: "chatcmpl_iop_" + strings.TrimPrefix(requestID, "req_"), model: model, created: time.Now().Unix()}, nil
|
|
}
|
|
|
|
func (s *singleRequestChatStream) startLocked() error {
|
|
if s.started {
|
|
return nil
|
|
}
|
|
s.w.Header().Set("Content-Type", "text/event-stream")
|
|
s.w.Header().Set("Cache-Control", "no-cache")
|
|
s.w.WriteHeader(http.StatusOK)
|
|
s.flusher.Flush()
|
|
s.started = true
|
|
return nil
|
|
}
|
|
|
|
func (s *singleRequestChatStream) Start() error {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
if s.terminal {
|
|
return s.terminalErr
|
|
}
|
|
return s.startLocked()
|
|
}
|
|
|
|
func (s *singleRequestChatStream) Ping() error {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
if s.terminal {
|
|
return s.terminalErr
|
|
}
|
|
if err := s.startLocked(); err != nil {
|
|
return err
|
|
}
|
|
_, err := io.WriteString(s.w, ": ping\n\n")
|
|
if err == nil {
|
|
s.flusher.Flush()
|
|
}
|
|
return err
|
|
}
|
|
|
|
func (s *singleRequestChatStream) writeSSELocked(value any) error {
|
|
payload, err := json.Marshal(value)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if _, err = fmt.Fprintf(s.w, "data: %s\n\n", payload); err != nil {
|
|
return err
|
|
}
|
|
s.flusher.Flush()
|
|
return nil
|
|
}
|
|
|
|
func (s *singleRequestChatStream) Final(result edgeservice.SingleRequestResult) error {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
if s.terminal {
|
|
return s.terminalErr
|
|
}
|
|
if err := s.startLocked(); err != nil {
|
|
return err
|
|
}
|
|
policy := singleRequestChatTerminalPolicy(result.Terminal)
|
|
if policy.silent || policy.error || policy.finishReason == "" {
|
|
return errors.New("single-request result has no Chat stream terminal")
|
|
}
|
|
s.terminal = true
|
|
if policy.finishReason != "length" && result.Output != "" {
|
|
if err := s.writeSSELocked(chatCompletionChunk{ID: s.id, Object: "chat.completion.chunk", Created: s.created, Model: s.model, Choices: []chatCompletionChunkChoice{{Index: 0, Delta: chatDelta{Content: result.Output}}}}); err != nil {
|
|
s.terminalErr = err
|
|
return err
|
|
}
|
|
}
|
|
if err := s.writeSSELocked(chatCompletionChunk{ID: s.id, Object: "chat.completion.chunk", Created: s.created, Model: s.model, Choices: []chatCompletionChunkChoice{{Index: 0, Delta: chatDelta{}, FinishReason: policy.finishReason}}}); err != nil {
|
|
s.terminalErr = err
|
|
return err
|
|
}
|
|
if _, err := io.WriteString(s.w, "data: [DONE]\n\n"); err != nil {
|
|
s.terminalErr = err
|
|
return err
|
|
}
|
|
s.flusher.Flush()
|
|
return nil
|
|
}
|
|
|
|
func (s *singleRequestChatStream) TerminalError(disposition edgeservice.SingleRequestTerminalDisposition) error {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
if s.terminal {
|
|
return s.terminalErr
|
|
}
|
|
policy := singleRequestChatTerminalPolicy(disposition)
|
|
if policy.silent {
|
|
s.terminal = true
|
|
return nil
|
|
}
|
|
if !policy.error {
|
|
policy = singleRequestChatTerminalPolicy(edgeservice.SingleRequestTerminalDisposition{Kind: edgeservice.SingleRequestTerminalError, ErrorClass: edgeservice.SingleRequestTerminalErrorProvider})
|
|
}
|
|
if err := s.startLocked(); err != nil {
|
|
return err
|
|
}
|
|
s.terminal = true
|
|
if err := s.writeSSELocked(errorResponse{Error: errorBody{Type: policy.errorType, Message: policy.message}}); err != nil {
|
|
s.terminalErr = err
|
|
return err
|
|
}
|
|
if _, err := io.WriteString(s.w, "data: [DONE]\n\n"); err != nil {
|
|
s.terminalErr = err
|
|
return err
|
|
}
|
|
s.flusher.Flush()
|
|
return nil
|
|
}
|
|
|
|
func pumpSingleRequestChatStream(ctx context.Context, execution edgeservice.SingleRequestExecution, stream *singleRequestChatStream, tickerFactory singleRequestAnthropicTickerFactory) error {
|
|
if execution == nil || stream == nil || tickerFactory == nil {
|
|
return errors.New("single-request Chat stream is unavailable")
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
execution.Cancel()
|
|
return err
|
|
}
|
|
if err := stream.Start(); err != nil {
|
|
execution.Cancel()
|
|
return err
|
|
}
|
|
ticker := tickerFactory()
|
|
if ticker == nil {
|
|
execution.Cancel()
|
|
return errors.New("single-request Chat stream is unavailable")
|
|
}
|
|
defer ticker.Stop()
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
execution.Cancel()
|
|
return ctx.Err()
|
|
case <-ticker.Ticks():
|
|
if err := stream.Ping(); err != nil {
|
|
execution.Cancel()
|
|
return err
|
|
}
|
|
case progress, ok := <-execution.Progress():
|
|
if !ok {
|
|
if ctx.Err() != nil {
|
|
return ctx.Err()
|
|
}
|
|
if execution.State() == edgeservice.SingleRequestStateCompleted {
|
|
return nil
|
|
}
|
|
return stream.TerminalError(edgeservice.SingleRequestTerminalDisposition{Kind: edgeservice.SingleRequestTerminalError, ErrorClass: edgeservice.SingleRequestTerminalErrorProvider})
|
|
}
|
|
switch progress.Stage {
|
|
case edgeservice.SingleRequestStateFinalizing:
|
|
if progress.Result == nil {
|
|
err := stream.TerminalError(edgeservice.SingleRequestTerminalDisposition{Kind: edgeservice.SingleRequestTerminalError, ErrorClass: edgeservice.SingleRequestTerminalErrorProvider})
|
|
return errors.Join(err, execution.AcknowledgeTerminal(false))
|
|
}
|
|
err := stream.Final(*progress.Result)
|
|
return errors.Join(err, execution.AcknowledgeTerminal(err == nil))
|
|
case edgeservice.SingleRequestStateFailed:
|
|
return stream.TerminalError(singleRequestProgressTerminal(progress, edgeservice.SingleRequestTerminalDisposition{Kind: edgeservice.SingleRequestTerminalError, ErrorClass: edgeservice.SingleRequestTerminalErrorProvider}))
|
|
case edgeservice.SingleRequestStateCancelled:
|
|
return stream.TerminalError(edgeservice.SingleRequestTerminalDisposition{Kind: edgeservice.SingleRequestTerminalCancelled})
|
|
}
|
|
}
|
|
}
|
|
}
|