freedakgmail
140 lines
3.3 KiB
Go
140 lines
3.3 KiB
Go
package handler
|
|
|
|
import (
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
|
|
"github.com/edgeai/gateway/internal/adapter"
|
|
"github.com/edgeai/gateway/pkg/api"
|
|
)
|
|
|
|
// SSEWriter writes Server-Sent Events to an HTTP response.
|
|
type SSEWriter struct {
|
|
w http.ResponseWriter
|
|
flusher http.Flusher
|
|
}
|
|
|
|
// NewSSEWriter creates a new SSEWriter. Returns nil if streaming is not supported.
|
|
func NewSSEWriter(w http.ResponseWriter) *SSEWriter {
|
|
flusher, ok := w.(http.Flusher)
|
|
if !ok {
|
|
return nil
|
|
}
|
|
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
w.Header().Set("Cache-Control", "no-cache")
|
|
w.Header().Set("Connection", "keep-alive")
|
|
w.Header().Set("X-Accel-Buffering", "no")
|
|
|
|
return &SSEWriter{w: w, flusher: flusher}
|
|
}
|
|
|
|
// WriteChunk writes a single SSE data event.
|
|
func (s *SSEWriter) WriteChunk(data any) error {
|
|
jsonData, err := json.Marshal(data)
|
|
if err != nil {
|
|
return fmt.Errorf("marshal sse data: %w", err)
|
|
}
|
|
fmt.Fprintf(s.w, "data: %s\n\n", jsonData)
|
|
s.flusher.Flush()
|
|
return nil
|
|
}
|
|
|
|
// WriteDone writes the [DONE] marker.
|
|
func (s *SSEWriter) WriteDone() {
|
|
fmt.Fprintf(s.w, "data: [DONE]\n\n")
|
|
s.flusher.Flush()
|
|
}
|
|
|
|
// StreamChatCompletion streams chunks from an adapter to the client in OpenAI SSE format.
|
|
func StreamChatCompletion(sse *SSEWriter, ch <-chan adapter.StreamChunk, requestID, taskID, model string) (int, int, error) {
|
|
inputTokens := 0
|
|
outputTokens := 0
|
|
|
|
for chunk := range ch {
|
|
if chunk.Error != nil {
|
|
return inputTokens, outputTokens, chunk.Error
|
|
}
|
|
|
|
if chunk.Done {
|
|
if chunk.InputTokens > 0 {
|
|
inputTokens = chunk.InputTokens
|
|
}
|
|
if chunk.OutputTokens > 0 {
|
|
outputTokens = chunk.OutputTokens
|
|
}
|
|
|
|
// Write final chunk with finish_reason
|
|
sseChunk := map[string]any{
|
|
"id": requestID,
|
|
"object": "chat.completion.chunk",
|
|
"model": model,
|
|
"choices": []map[string]any{
|
|
{
|
|
"index": 0,
|
|
"delta": map[string]any{},
|
|
"finish_reason": chunk.FinishReason,
|
|
},
|
|
},
|
|
}
|
|
if inputTokens > 0 || outputTokens > 0 {
|
|
sseChunk["usage"] = map[string]int{
|
|
"input_tokens": inputTokens,
|
|
"output_tokens": outputTokens,
|
|
"total_tokens": inputTokens + outputTokens,
|
|
}
|
|
}
|
|
sse.WriteChunk(sseChunk)
|
|
sse.WriteDone()
|
|
return inputTokens, outputTokens, nil
|
|
}
|
|
|
|
// Write content delta
|
|
sseChunk := map[string]any{
|
|
"id": requestID,
|
|
"object": "chat.completion.chunk",
|
|
"model": model,
|
|
"choices": []map[string]any{
|
|
{
|
|
"index": 0,
|
|
"delta": map[string]any{
|
|
"content": chunk.Delta,
|
|
},
|
|
"finish_reason": nil,
|
|
},
|
|
},
|
|
}
|
|
sse.WriteChunk(sseChunk)
|
|
}
|
|
|
|
return inputTokens, outputTokens, nil
|
|
}
|
|
|
|
// BuildChatResponse creates a non-streaming ChatResponse from adapter result.
|
|
func BuildChatResponse(requestID, taskID, logicalModel string, resp *adapter.ChatResponse) api.ChatResponse {
|
|
return api.ChatResponse{
|
|
RequestID: requestID,
|
|
TaskID: taskID,
|
|
Status: "completed",
|
|
Model: logicalModel,
|
|
Choices: []api.Choice{
|
|
{
|
|
Index: 0,
|
|
Message: &api.Message{
|
|
Role: "assistant",
|
|
Content: resp.Content,
|
|
},
|
|
FinishReason: resp.FinishReason,
|
|
},
|
|
},
|
|
LogicalModel: logicalModel,
|
|
ActualModel: resp.ActualModel,
|
|
Usage: &api.Usage{
|
|
InputTokens: resp.InputTokens,
|
|
OutputTokens: resp.OutputTokens,
|
|
TotalTokens: resp.InputTokens + resp.OutputTokens,
|
|
},
|
|
}
|
|
}
|