BCAI/easyai-ai-gateway
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: enrich task record details

Browse Source
This commit is contained in:
wangbo
2026-05-10 22:33:58 +08:00
parent 205a4b625e
commit 53f8edfb67
19 changed files with 1781 additions and 165 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/api/internal/clients/clients_test.go
+58
View File
@@ -17,6 +17,7 @@ func TestOpenAIClientChatContract(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
gotPath = r.URL.Path
gotAuth = r.Header.Get("Authorization")
w.Header().Set("X-Request-Id", "req-chat-test")
var body map[string]any
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
t.Fatalf("decode request: %v", err)
@@ -53,6 +54,55 @@ func TestOpenAIClientChatContract(t *testing.T) {
if response.Usage.TotalTokens != 5 || response.Result["id"] != "chatcmpl-test" {
t.Fatalf("unexpected response: %+v", response)
}
if response.RequestID != "req-chat-test" || response.ResponseStartedAt.IsZero() || response.ResponseFinishedAt.IsZero() {
t.Fatalf("response metadata was not captured: %+v", response)
}
}
func TestOpenAIClientChatStreamContract(t *testing.T) {
var gotStream bool
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
var body map[string]any
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
t.Fatalf("decode request: %v", err)
}
gotStream, _ = body["stream"].(bool)
w.Header().Set("Content-Type", "text/event-stream")
_, _ = w.Write([]byte("data: {\"id\":\"chatcmpl-stream\",\"object\":\"chat.completion.chunk\",\"model\":\"deepseek-v4-flash\",\"choices\":[{\"delta\":{\"content\":\"hello\"}}]}\n\n"))
_, _ = w.Write([]byte("data: {\"id\":\"chatcmpl-stream\",\"object\":\"chat.completion.chunk\",\"model\":\"deepseek-v4-flash\",\"choices\":[{\"delta\":{\"content\":\" world\"},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":1,\"completion_tokens\":2,\"total_tokens\":3}}\n\n"))
_, _ = w.Write([]byte("data: [DONE]\n\n"))
}))
defer server.Close()
response, err := (OpenAIClient{HTTPClient: server.Client()}).Run(context.Background(), Request{
Kind: "chat.completions",
Model: "DeepSeek-V4-Flash",
Body: map[string]any{
"model": "DeepSeek-V4-Flash",
"messages": []any{map[string]any{"role": "user", "content": "ping"}},
"stream": true,
},
Candidate: store.RuntimeModelCandidate{
BaseURL: server.URL,
ModelName: "deepseek-v4-flash",
Credentials: map[string]any{"apiKey": "test-key"},
},
})
if err != nil {
t.Fatalf("run openai stream client: %v", err)
}
if !gotStream {
t.Fatalf("expected upstream stream request")
}
if response.Usage.TotalTokens != 3 {
t.Fatalf("unexpected usage: %+v", response.Usage)
}
choices, _ := response.Result["choices"].([]any)
choice, _ := choices[0].(map[string]any)
message, _ := choice["message"].(map[string]any)
if message["content"] != "hello world" {
t.Fatalf("unexpected stream response: %+v", response.Result)
}
}
func TestGeminiClientChatContract(t *testing.T) {
@@ -111,6 +161,14 @@ func TestGeminiClientChatContract(t *testing.T) {
}
}
func TestGeminiURLAcceptsVersionedBaseURL(t *testing.T) {
got := geminiURL("https://generativelanguage.googleapis.com/v1beta", "gemini-2.5-flash", "test-key")
want := "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent?key=test-key"
if got != want {
t.Fatalf("unexpected gemini url: %s", got)
}
}
func extractText(result map[string]any) string {
choices, _ := result["choices"].([]any)
choice, _ := choices[0].(map[string]any)
apps/api/internal/clients/gemini.go
+21 -2
View File
@@ -8,6 +8,7 @@ import (
"net/http"
"net/url"
"strings"
"time"
)
type GeminiClient struct {
@@ -30,11 +31,26 @@ func (c GeminiClient) Run(ctx context.Context, request Request) (Response, error
if err != nil {
return Response{}, &ClientError{Code: "network", Message: err.Error(), Retryable: true}
}
responseStartedAt := time.Now()
requestID := requestIDFromHTTPResponse(resp)
result, err := decodeHTTPResponse(resp)
responseFinishedAt := time.Now()
if err != nil {
return Response{}, err
return Response{}, annotateResponseError(err, requestID, responseStartedAt, responseFinishedAt)
}
return Response{Result: geminiResult(request, result), Usage: geminiUsage(result), Progress: providerProgress(request)}, nil
output := geminiResult(request, result)
if requestID == "" {
requestID = requestIDFromResult(output)
}
return Response{
Result: output,
RequestID: requestID,
Usage: geminiUsage(result),
Progress: providerProgress(request),
ResponseStartedAt: responseStartedAt,
ResponseFinishedAt: responseFinishedAt,
ResponseDurationMS: responseDurationMS(responseStartedAt, responseFinishedAt),
}, nil
}
func geminiURL(baseURL string, model string, apiKey string) string {
@@ -42,6 +58,9 @@ func geminiURL(baseURL string, model string, apiKey string) string {
if base == "" {
base = "https://generativelanguage.googleapis.com"
}
if strings.HasSuffix(base, "/v1beta") {
base = strings.TrimSuffix(base, "/v1beta")
}
escapedModel := url.PathEscape(model)
return fmt.Sprintf("%s/v1beta/models/%s:generateContent?key=%s", base, escapedModel, url.QueryEscape(apiKey))
}
apps/api/internal/clients/helpers.go
+169
View File
@@ -1,6 +1,8 @@
package clients
import (
"bufio"
"bytes"
"encoding/json"
"fmt"
"io"
@@ -48,6 +50,7 @@ func decodeHTTPResponse(resp *http.Response) (map[string]any, error) {
Code: statusCodeName(resp.StatusCode),
Message: errorMessage(raw, resp.Status),
StatusCode: resp.StatusCode,
RequestID: requestIDFromHTTPResponse(resp),
Retryable: HTTPRetryable(resp.StatusCode),
}
}
@@ -61,6 +64,137 @@ func decodeHTTPResponse(resp *http.Response) (map[string]any, error) {
return out, nil
}
func decodeOpenAIStreamResponse(resp *http.Response, onDelta StreamDelta) (map[string]any, error) {
defer resp.Body.Close()
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
raw, _ := io.ReadAll(io.LimitReader(resp.Body, 16*1024*1024))
return nil, &ClientError{
Code: statusCodeName(resp.StatusCode),
Message: errorMessage(raw, resp.Status),
StatusCode: resp.StatusCode,
RequestID: requestIDFromHTTPResponse(resp),
Retryable: HTTPRetryable(resp.StatusCode),
}
}
if result, ok, err := decodeOpenAIStreamReader(resp.Body, onDelta); ok || err != nil {
return result, err
}
return map[string]any{}, nil
}
func decodeOpenAIStreamReader(reader io.Reader, onDelta StreamDelta) (map[string]any, bool, error) {
scanner := bufio.NewScanner(reader)
scanner.Buffer(make([]byte, 0, 64*1024), 16*1024*1024)
rawLines := make([]string, 0)
parts := make([]string, 0)
var last map[string]any
var usage Usage
for scanner.Scan() {
rawLine := scanner.Text()
rawLines = append(rawLines, rawLine)
line := strings.TrimSpace(rawLine)
if !strings.HasPrefix(line, "data:") {
continue
}
payload := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
if payload == "" || payload == "[DONE]" {
continue
}
var event map[string]any
if err := json.Unmarshal([]byte(payload), &event); err != nil {
continue
}
last = event
if text := streamEventText(event); text != "" {
parts = append(parts, text)
if onDelta != nil {
if err := onDelta(text); err != nil {
return nil, true, err
}
}
}
if eventUsage := usageFromOpenAI(event); eventUsage.TotalTokens > 0 {
usage = eventUsage
}
}
if err := scanner.Err(); err != nil {
return nil, true, &ClientError{Code: "stream_read_error", Message: err.Error(), Retryable: true}
}
if last == nil {
raw := []byte(strings.Join(rawLines, "\n"))
if len(raw) == 0 {
return map[string]any{}, true, nil
}
var out map[string]any
if err := json.Unmarshal(raw, &out); err != nil {
return nil, false, nil
}
return out, true, nil
}
return buildOpenAIStreamResult(last, parts, usage), true, nil
}
func decodeOpenAIStream(raw []byte) (map[string]any, bool) {
if !bytes.Contains(raw, []byte("data:")) {
return nil, false
}
result, ok, err := decodeOpenAIStreamReader(bytes.NewReader(raw), nil)
return result, ok && err == nil
}
func buildOpenAIStreamResult(last map[string]any, parts []string, usage Usage) map[string]any {
if len(parts) == 0 {
return last
}
var out map[string]any
out = map[string]any{
"id": stringFromAny(firstPresent(last["id"], "chatcmpl-stream")),
"object": "chat.completion",
"model": stringFromAny(last["model"]),
"choices": []any{map[string]any{
"index": 0,
"message": map[string]any{
"role": "assistant",
"content": strings.Join(parts, ""),
},
"finish_reason": "stop",
}},
}
if usage.TotalTokens > 0 {
out["usage"] = map[string]any{
"prompt_tokens": usage.InputTokens,
"completion_tokens": usage.OutputTokens,
"total_tokens": usage.TotalTokens,
}
}
return out
}
func streamEventText(event map[string]any) string {
if choices, ok := event["choices"].([]any); ok {
for _, rawChoice := range choices {
choice, _ := rawChoice.(map[string]any)
if delta, ok := choice["delta"].(map[string]any); ok {
if content, ok := delta["content"].(string); ok {
return content
}
}
if message, ok := choice["message"].(map[string]any); ok {
if content, ok := message["content"].(string); ok {
return content
}
}
}
}
if delta, ok := event["delta"].(string); ok {
return delta
}
if text, ok := event["output_text"].(string); ok {
return text
}
return ""
}
func usageFromOpenAI(result map[string]any) Usage {
usage, _ := result["usage"].(map[string]any)
input := intFromAny(firstPresent(usage["prompt_tokens"], usage["input_tokens"]))
@@ -72,6 +206,34 @@ func usageFromOpenAI(result map[string]any) Usage {
return Usage{InputTokens: input, OutputTokens: output, TotalTokens: total}
}
func requestIDFromHTTPResponse(resp *http.Response) string {
if resp == nil {
return ""
}
for _, key := range []string{
"x-request-id",
"x-requestid",
"request-id",
"x-amzn-requestid",
"x-amz-request-id",
"cf-ray",
} {
if value := strings.TrimSpace(resp.Header.Get(key)); value != "" {
return value
}
}
return ""
}
func requestIDFromResult(result map[string]any) string {
for _, key := range []string{"request_id", "requestId", "id", "response_id", "responseId"} {
if value := strings.TrimSpace(stringFromAny(result[key])); value != "" {
return value
}
}
return ""
}
func intFromAny(value any) int {
switch typed := value.(type) {
case float64:
@@ -85,6 +247,13 @@ func intFromAny(value any) int {
}
}
func stringFromAny(value any) string {
if text, ok := value.(string); ok {
return text
}
return ""
}
func firstPresent(values ...any) any {
for _, value := range values {
if value != nil {
apps/api/internal/clients/openai.go
+30 -3
View File
@@ -6,6 +6,7 @@ import (
"encoding/json"
"net/http"
"strings"
"time"
)
type OpenAIClient struct {
@@ -23,6 +24,7 @@ func (c OpenAIClient) Run(ctx context.Context, request Request) (Response, error
}
body := cloneBody(request.Body)
body["model"] = request.Candidate.ModelName
stream := request.Stream || boolValue(body, "stream")
raw, _ := json.Marshal(body)
req, err := http.NewRequestWithContext(ctx, http.MethodPost, joinURL(request.Candidate.BaseURL, endpoint), bytes.NewReader(raw))
if err != nil {
@@ -34,11 +36,36 @@ func (c OpenAIClient) Run(ctx context.Context, request Request) (Response, error
if err != nil {
return Response{}, &ClientError{Code: "network", Message: err.Error(), Retryable: true}
}
result, err := decodeHTTPResponse(resp)
responseStartedAt := time.Now()
requestID := requestIDFromHTTPResponse(resp)
result, err := decodeOpenAIResponse(resp, stream, request.StreamDelta)
responseFinishedAt := time.Now()
if err != nil {
return Response{}, err
return Response{}, annotateResponseError(err, requestID, responseStartedAt, responseFinishedAt)
}
return Response{Result: result, Usage: usageFromOpenAI(result), Progress: providerProgress(request)}, nil
if requestID == "" {
requestID = requestIDFromResult(result)
}
return Response{
Result: result,
RequestID: requestID,
Usage: usageFromOpenAI(result),
Progress: providerProgress(request),
ResponseStartedAt: responseStartedAt,
ResponseFinishedAt: responseFinishedAt,
ResponseDurationMS: responseDurationMS(responseStartedAt, responseFinishedAt),
}, nil
}
func decodeOpenAIResponse(resp *http.Response, stream bool, onDelta StreamDelta) (map[string]any, error) {
if stream {
result, err := decodeOpenAIStreamResponse(resp, onDelta)
if err == nil {
return result, nil
}
return nil, err
}
return decodeHTTPResponse(resp)
}
func openAIEndpoint(kind string) string {
apps/api/internal/clients/simulation.go
+11 -3
View File
@@ -4,6 +4,7 @@ import (
"context"
"fmt"
"strings"
"time"
)
type SimulationClient struct{}
@@ -17,10 +18,17 @@ func (c SimulationClient) Run(ctx context.Context, request Request) (Response, e
if profile == "fatal_failure" || profile == "non_retryable_failure" {
return Response{}, &ClientError{Code: "bad_request", Message: "simulated non-retryable failure", Retryable: false}
}
responseStartedAt := time.Now()
result := simulatedResult(request)
responseFinishedAt := time.Now()
return Response{
Result: simulatedResult(request),
Usage: simulatedUsage(request),
Progress: simulatedProgress(request),
Result: result,
RequestID: requestIDFromResult(result),
Usage: simulatedUsage(request),
Progress: simulatedProgress(request),
ResponseStartedAt: responseStartedAt,
ResponseFinishedAt: responseFinishedAt,
ResponseDurationMS: responseDurationMS(responseStartedAt, responseFinishedAt),
}, nil
}
apps/api/internal/clients/types.go
+79 -13
View File
@@ -4,23 +4,30 @@ import (
"context"
"errors"
"net/http"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
type Request struct {
Kind string
ModelType string
Model string
Body map[string]any
Candidate store.RuntimeModelCandidate
Stream bool
Kind string
ModelType string
Model string
Body map[string]any
Candidate store.RuntimeModelCandidate
Stream bool
StreamDelta StreamDelta
}
type Response struct {
Result map[string]any
Usage Usage
Progress []Progress
Result map[string]any
RequestID string
Usage Usage
Progress []Progress
Metrics map[string]any
ResponseStartedAt time.Time
ResponseFinishedAt time.Time
ResponseDurationMS int64
}
type Usage struct {
@@ -36,15 +43,21 @@ type Progress struct {
Payload map[string]any
}
type StreamDelta func(text string) error
type Client interface {
Run(ctx context.Context, request Request) (Response, error)
}
type ClientError struct {
Code string
Message string
StatusCode int
Retryable bool
Code string
Message string
StatusCode int
RequestID string
ResponseStartedAt time.Time
ResponseFinishedAt time.Time
ResponseDurationMS int64
Retryable bool
}
func (e *ClientError) Error() string {
@@ -67,6 +80,59 @@ func ErrorCode(err error) string {
return "client_error"
}
type ResponseMetadata struct {
RequestID string
ResponseStartedAt time.Time
ResponseFinishedAt time.Time
ResponseDurationMS int64
StatusCode int
}
func ErrorResponseMetadata(err error) ResponseMetadata {
var clientErr *ClientError
if errors.As(err, &clientErr) {
return ResponseMetadata{
RequestID: clientErr.RequestID,
ResponseStartedAt: clientErr.ResponseStartedAt,
ResponseFinishedAt: clientErr.ResponseFinishedAt,
ResponseDurationMS: clientErr.ResponseDurationMS,
StatusCode: clientErr.StatusCode,
}
}
return ResponseMetadata{}
}
func annotateResponseError(err error, requestID string, startedAt time.Time, finishedAt time.Time) error {
var clientErr *ClientError
if !errors.As(err, &clientErr) {
return err
}
if clientErr.RequestID == "" {
clientErr.RequestID = requestID
}
if clientErr.ResponseStartedAt.IsZero() {
clientErr.ResponseStartedAt = startedAt
}
if clientErr.ResponseFinishedAt.IsZero() {
clientErr.ResponseFinishedAt = finishedAt
}
if clientErr.ResponseDurationMS == 0 {
clientErr.ResponseDurationMS = responseDurationMS(clientErr.ResponseStartedAt, clientErr.ResponseFinishedAt)
}
return err
}
func HTTPRetryable(status int) bool {
return status == http.StatusTooManyRequests || status == http.StatusRequestTimeout || status >= 500
}
func responseDurationMS(startedAt time.Time, finishedAt time.Time) int64 {
if startedAt.IsZero() || finishedAt.IsZero() {
return 0
}
duration := finishedAt.Sub(startedAt).Milliseconds()
if duration < 0 {
return 0
}
return duration
}