diff --git a/cli/cmd/geniex/BUILD.bazel b/cli/cmd/geniex/BUILD.bazel index 1faf8cd17..78b74dbae 100644 --- a/cli/cmd/geniex/BUILD.bazel +++ b/cli/cmd/geniex/BUILD.bazel @@ -10,6 +10,7 @@ go_library( "main.go", "model.go", "run.go", + "run_history.go", "serve.go", "update.go", "version.go", @@ -46,6 +47,8 @@ go_cgo_test( srcs = [ "infer_test.go", "model_test.go", + "run_history_test.go", + "run_test.go", ], embed = [":geniex_lib"], gotags = [ @@ -54,6 +57,9 @@ go_cgo_test( ], deps = [ "//cli/internal/testutil", + "@com_github_openai_openai_go_v3//:openai-go", + "@com_github_openai_openai_go_v3//option", + "@com_github_openai_openai_go_v3//packages/ssestream", ], ) diff --git a/cli/cmd/geniex/run.go b/cli/cmd/geniex/run.go index 8f295243f..39fcfc42a 100644 --- a/cli/cmd/geniex/run.go +++ b/cli/cmd/geniex/run.go @@ -41,9 +41,7 @@ func tagStreamError(err error) error { var body struct { Code *int32 `json:"code"` } - // code < 0 means the server error was not an SDKError (SDKErrorCode - // returns -1); keep the original stream error in that case. - if sonic.Unmarshal(se.Event.Data, &body) == nil && body.Code != nil && *body.Code >= 0 { + if sonic.Unmarshal(se.Event.Data, &body) == nil && body.Code != nil && *body.Code != -1 { return geniex_sdk.SDKError(*body.Code) } } @@ -127,8 +125,10 @@ func runCompletions(ctx context.Context, name string, modelType geniex_sdk.Model // repl var history []openai.ChatCompletionMessageParamUnion + protected := 0 if systemPrompt != "" { history = append(history, openai.SystemMessage(systemPrompt)) + protected = 1 } processor := &common.Processor{ @@ -136,14 +136,20 @@ func runCompletions(ctx context.Context, name string, modelType geniex_sdk.Model Verbose: verbose, TestMode: testMode, Reset: func() error { - history = nil _, err := client.Chat.Completions.New(ctx, openai.ChatCompletionNewParams{ Messages: nil, Model: name, }) + if err == nil { + history = history[:0] + if protected == 1 { + history = append(history, openai.SystemMessage(systemPrompt)) + } + } return err }, Run: func(prompt string, images, audios []string, onToken func(string) bool) (string, geniex_sdk.ProfileData, error) { + var userMessage openai.ChatCompletionMessageParamUnion if len(images) > 0 || len(audios) > 0 { contents := make([]openai.ChatCompletionContentPartUnionParam, 0) contents = append(contents, openai.ChatCompletionContentPartUnionParam{ @@ -169,87 +175,114 @@ func runCompletions(ctx context.Context, name string, modelType geniex_sdk.Model }, }) } - history = append(history, openai.UserMessage(contents)) + userMessage = openai.UserMessage(contents) } else { - history = append(history, openai.UserMessage(prompt)) + userMessage = openai.UserMessage(prompt) + } + pending := append(append([]openai.ChatCompletionMessageParamUnion(nil), history...), userMessage) + type attemptResult struct { + text string + profile geniex_sdk.ProfileData + complete bool } + var emitted bool + result, kept, dropped, err := retryChatHistory(pending, protected, + func(messages []openai.ChatCompletionMessageParamUnion) (attemptResult, error) { + emitted = false - start := time.Now() - acc := openai.ChatCompletionAccumulator{} - stream := client.Chat.Completions.NewStreaming(ctx, openai.ChatCompletionNewParams{ - Messages: history, - Model: name, - StreamOptions: openai.ChatCompletionStreamOptionsParam{IncludeUsage: openai.Opt(true)}, - Temperature: openai.Float(float64(temperature)), - TopP: openai.Float(float64(topP)), - PresencePenalty: openai.Float(float64(presencePenalty)), - FrequencyPenalty: openai.Float(float64(frequencyPenalty)), - Seed: openai.Int(int64(seed)), - MaxCompletionTokens: openai.Int(int64(maxTokens)), - }, + start := time.Now() + acc := openai.ChatCompletionAccumulator{} + stream := client.Chat.Completions.NewStreaming(ctx, openai.ChatCompletionNewParams{ + Messages: messages, + Model: name, + StreamOptions: openai.ChatCompletionStreamOptionsParam{IncludeUsage: openai.Opt(true)}, + Temperature: openai.Float(float64(temperature)), + TopP: openai.Float(float64(topP)), + PresencePenalty: openai.Float(float64(presencePenalty)), + FrequencyPenalty: openai.Float(float64(frequencyPenalty)), + Seed: openai.Int(int64(seed)), + MaxCompletionTokens: openai.Int(int64(maxTokens)), + }, - option.WithJSONSet("enable_think", enableThink), - option.WithJSONSet("top_k", topK), - option.WithJSONSet("min_p", minP), - option.WithJSONSet("repetition_penalty", repetitionPenalty), - option.WithJSONSet("grammar_path", grammarPath), - option.WithJSONSet("grammar_string", grammarString), - option.WithJSONSet("ngl", ngl), - option.WithJSONSet("nctx", nctx), - option.WithJSONSet("compute", computeUnit), - option.WithJSONSet("vit_compute", vitComputeUnit), - option.WithJSONSet("spec_type", specType), - option.WithJSONSet("spec_draft_model", draftModel), - option.WithJSONSet("spec_n_max", draftTokens), - option.WithJSONSet("spec_n_min", draftMin), - option.WithJSONSet("spec_p_min", draftPMin), - option.WithJSONSet("power_mode", powerMode)) + option.WithJSONSet("enable_think", enableThink), + option.WithJSONSet("top_k", topK), + option.WithJSONSet("min_p", minP), + option.WithJSONSet("repetition_penalty", repetitionPenalty), + option.WithJSONSet("grammar_path", grammarPath), + option.WithJSONSet("grammar_string", grammarString), + option.WithJSONSet("ngl", ngl), + option.WithJSONSet("nctx", nctx), + option.WithJSONSet("compute", computeUnit), + option.WithJSONSet("vit_compute", vitComputeUnit), + option.WithJSONSet("spec_type", specType), + option.WithJSONSet("spec_draft_model", draftModel), + option.WithJSONSet("spec_n_max", draftTokens), + option.WithJSONSet("spec_n_min", draftMin), + option.WithJSONSet("spec_p_min", draftPMin), + option.WithJSONSet("power_mode", powerMode)) + defer stream.Close() - var firstToken time.Time - var profileData geniex_sdk.ProfileData - for stream.Next() { - if firstToken.IsZero() { - firstToken = time.Now() - } + var firstToken time.Time + var profileData geniex_sdk.ProfileData + for stream.Next() { + if firstToken.IsZero() { + firstToken = time.Now() + } - chunk := stream.Current() - acc.AddChunk(chunk) - if len(chunk.Choices) > 0 { - if !onToken(chunk.Choices[0].Delta.Content) { - stream.Close() - break + chunk := stream.Current() + acc.AddChunk(chunk) + if len(chunk.Choices) > 0 { + content := chunk.Choices[0].Delta.Content + if content != "" { + emitted = true + } + if !onToken(content) { + stream.Close() + break + } + } + if chunk.Usage.PromptTokens > 0 { + profileData.PromptTokens = chunk.Usage.PromptTokens + profileData.GeneratedTokens = chunk.Usage.CompletionTokens + det := chunk.Usage.CompletionTokensDetails + profileData.DraftNAccepted = det.AcceptedPredictionTokens + profileData.DraftNTotal = det.AcceptedPredictionTokens + det.RejectedPredictionTokens + } } - } - if chunk.Usage.PromptTokens > 0 { - profileData.PromptTokens = chunk.Usage.PromptTokens - profileData.GeneratedTokens = chunk.Usage.CompletionTokens - det := chunk.Usage.CompletionTokensDetails - profileData.DraftNAccepted = det.AcceptedPredictionTokens - profileData.DraftNTotal = det.AcceptedPredictionTokens + det.RejectedPredictionTokens - } - } - // zero token generated - if firstToken.IsZero() { - firstToken = time.Now() - } + // zero token generated + if firstToken.IsZero() { + firstToken = time.Now() + } - end := time.Now() - profileData.TTFT = firstToken.Sub(start).Microseconds() - profileData.PromptTime = profileData.TTFT - profileData.DecodeTime = end.Sub(firstToken).Microseconds() - profileData.DecodingSpeed = float64(profileData.GeneratedTokens) / float64(end.Sub(firstToken).Seconds()) + end := time.Now() + profileData.TTFT = firstToken.Sub(start).Microseconds() + profileData.PromptTime = profileData.TTFT + profileData.DecodeTime = end.Sub(firstToken).Microseconds() + profileData.DecodingSpeed = float64(profileData.GeneratedTokens) / float64(end.Sub(firstToken).Seconds()) - if stream.Err() != nil { - return "", profileData, tagStreamError(stream.Err()) - } + if stream.Err() != nil { + return attemptResult{profile: profileData}, tagStreamError(stream.Err()) + } - if len(acc.Choices) > 0 { - history = append(history, openai.AssistantMessage(acc.Choices[0].Message.Content)) - return acc.Choices[0].Message.Content, profileData, nil - } + if len(acc.Choices) > 0 { + return attemptResult{text: acc.Choices[0].Message.Content, profile: profileData, complete: true}, nil + } - return "", profileData, nil + return attemptResult{profile: profileData}, nil + }, func(err error) bool { + return !emitted && errors.Is(err, geniex_sdk.ErrLlmGenerationPromptTooLong) + }) + if err != nil { + return result.text, result.profile, err + } + if dropped > 0 { + fmt.Println(render.GetTheme().Info.Sprint("Older messages were removed to fit the context window.")) + } + if result.complete { + history = append(kept, openai.AssistantMessage(result.text)) + } + return result.text, result.profile, nil }, } if len(prompt) > 0 || input != "" { diff --git a/cli/cmd/geniex/run_history.go b/cli/cmd/geniex/run_history.go new file mode 100644 index 000000000..fec1b92d0 --- /dev/null +++ b/cli/cmd/geniex/run_history.go @@ -0,0 +1,21 @@ +// Copyright 2024-2026 Qualcomm Technologies, Inc. and/or its subsidiaries. +// SPDX-License-Identifier: BSD-3-Clause + +package main + +import "github.com/openai/openai-go/v3" + +func retryChatHistory[T any](messages []openai.ChatCompletionMessageParamUnion, protected int, + request func([]openai.ChatCompletionMessageParamUnion) (T, error), retryable func(error) bool, +) (T, []openai.ChatCompletionMessageParamUnion, int, error) { + dropped := 0 + for { + result, err := request(messages) + if err == nil || !retryable(err) || len(messages)-protected < 3 || + messages[protected].OfUser == nil || messages[protected+1].OfAssistant == nil { + return result, messages, dropped, err + } + messages = append(messages[:protected:protected], messages[protected+2:]...) + dropped++ + } +} diff --git a/cli/cmd/geniex/run_history_test.go b/cli/cmd/geniex/run_history_test.go new file mode 100644 index 000000000..8e1aba341 --- /dev/null +++ b/cli/cmd/geniex/run_history_test.go @@ -0,0 +1,94 @@ +// Copyright 2024-2026 Qualcomm Technologies, Inc. and/or its subsidiaries. +// SPDX-License-Identifier: BSD-3-Clause + +package main + +import ( + "errors" + "testing" + + "github.com/openai/openai-go/v3" +) + +func TestRetryChatHistory(t *testing.T) { + tooLong := errors.New("prompt too long") + messages := []openai.ChatCompletionMessageParamUnion{ + openai.SystemMessage("rules"), + openai.UserMessage("old"), openai.AssistantMessage("answer"), + openai.UserMessage("new"), + } + calls := 0 + result, kept, dropped, err := retryChatHistory(messages, 1, + func(got []openai.ChatCompletionMessageParamUnion) (string, error) { + calls++ + if calls == 1 { + return "", tooLong + } + if len(got) != 2 || got[0].OfSystem == nil || got[1].OfUser == nil { + t.Fatalf("retry lost system or current user: %+v", got) + } + return "ok", nil + }, func(err error) bool { return errors.Is(err, tooLong) }) + if err != nil || result != "ok" || dropped != 1 || len(kept) != 2 || len(messages) != 4 { + t.Fatalf("result=%q kept=%d dropped=%d calls=%d err=%v", result, len(kept), dropped, calls, err) + } +} + +func TestRetryChatHistoryStopsAtCurrentTurn(t *testing.T) { + tooLong := errors.New("prompt too long") + for _, protected := range []int{0, 1} { + messages := []openai.ChatCompletionMessageParamUnion{openai.UserMessage("current")} + if protected == 1 { + messages = append([]openai.ChatCompletionMessageParamUnion{openai.SystemMessage("rules")}, messages...) + } + calls := 0 + _, kept, dropped, err := retryChatHistory(messages, protected, + func([]openai.ChatCompletionMessageParamUnion) (struct{}, error) { + calls++ + return struct{}{}, tooLong + }, func(err error) bool { return errors.Is(err, tooLong) }) + if !errors.Is(err, tooLong) || calls != 1 || dropped != 0 || len(kept) != len(messages) { + t.Fatalf("protected=%d calls=%d dropped=%d err=%v", protected, calls, dropped, err) + } + } +} + +func TestRetryChatHistoryDropsMultipleTurns(t *testing.T) { + tooLong := errors.New("prompt too long") + messages := []openai.ChatCompletionMessageParamUnion{ + openai.UserMessage("first"), openai.AssistantMessage("one"), + openai.UserMessage("second"), openai.AssistantMessage("two"), + openai.UserMessage("current"), + } + calls := 0 + _, kept, dropped, err := retryChatHistory(messages, 0, + func(got []openai.ChatCompletionMessageParamUnion) (struct{}, error) { + calls++ + if calls < 3 { + return struct{}{}, tooLong + } + if len(got) != 1 || got[0].OfUser == nil { + t.Fatalf("current user turn was lost: %+v", got) + } + return struct{}{}, nil + }, func(err error) bool { return errors.Is(err, tooLong) }) + if err != nil || calls != 3 || dropped != 2 || len(kept) != 1 || len(messages) != 5 { + t.Fatalf("calls=%d dropped=%d kept=%d err=%v", calls, dropped, len(kept), err) + } +} + +func TestRetryChatHistoryKeepsHistoryOnOtherErrors(t *testing.T) { + failure := errors.New("network failure") + messages := []openai.ChatCompletionMessageParamUnion{ + openai.UserMessage("old"), openai.AssistantMessage("answer"), openai.UserMessage("new"), + } + calls := 0 + _, kept, dropped, err := retryChatHistory(messages, 0, + func([]openai.ChatCompletionMessageParamUnion) (struct{}, error) { + calls++ + return struct{}{}, failure + }, func(error) bool { return false }) + if !errors.Is(err, failure) || calls != 1 || dropped != 0 || len(kept) != len(messages) { + t.Fatalf("calls=%d dropped=%d err=%v", calls, dropped, err) + } +} diff --git a/cli/cmd/geniex/run_test.go b/cli/cmd/geniex/run_test.go new file mode 100644 index 000000000..b9aa1bd17 --- /dev/null +++ b/cli/cmd/geniex/run_test.go @@ -0,0 +1,148 @@ +// Copyright 2024-2026 Qualcomm Technologies, Inc. and/or its subsidiaries. +// SPDX-License-Identifier: BSD-3-Clause + +package main + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "net/http/httptest" + "reflect" + "testing" + + "github.com/openai/openai-go/v3" + "github.com/openai/openai-go/v3/option" + "github.com/openai/openai-go/v3/packages/ssestream" + + geniex_sdk "github.com/qualcomm/GenieX/bindings/go" +) + +func TestTagStreamErrorRetainsNegativeSDKCode(t *testing.T) { + streamErr := &ssestream.StreamError{ + Message: "Input prompt too long", + Event: ssestream.Event{Data: []byte(`{"code":-200103,"error":"Input prompt too long"}`)}, + } + if err := tagStreamError(streamErr); !errors.Is(err, geniex_sdk.ErrLlmGenerationPromptTooLong) { + t.Fatalf("got %v, want prompt-too-long SDK error", err) + } +} + +func TestTagStreamErrorKeepsNonSDKError(t *testing.T) { + streamErr := &ssestream.StreamError{ + Message: "other error", + Event: ssestream.Event{Data: []byte(`{"code":-1,"error":"other error"}`)}, + } + if err := tagStreamError(streamErr); err != streamErr { + t.Fatalf("got %v, want original stream error", err) + } +} + +func TestRunHistoryRetriesStreamPromptError(t *testing.T) { + requests := 0 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requests++ + var body struct { + Messages []json.RawMessage `json:"messages"` + } + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + t.Errorf("decode request: %v", err) + return + } + want := 3 + if requests == 2 { + want = 1 + } + if len(body.Messages) != want { + t.Errorf("request %d: got %d messages, want %d", requests, len(body.Messages), want) + } + w.Header().Set("Content-Type", "text/event-stream") + if requests == 1 { + fmt.Fprint(w, "data: {\"error\":\"Input prompt too long\",\"code\":-200103}\n\n") + return + } + fmt.Fprint(w, "data: {\"id\":\"test\",\"object\":\"chat.completion.chunk\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"ok\"}}]}\n\ndata: [DONE]\n\n") + })) + defer server.Close() + api := openai.NewClient(option.WithBaseURL(server.URL + "/v1")) + messages := []openai.ChatCompletionMessageParamUnion{ + openai.UserMessage("old"), openai.AssistantMessage("answer"), openai.UserMessage("current"), + } + text, kept, dropped, err := retryChatHistory(messages, 0, + func(got []openai.ChatCompletionMessageParamUnion) (string, error) { + stream := api.Chat.Completions.NewStreaming(context.Background(), openai.ChatCompletionNewParams{ + Model: "test", Messages: got, + }) + defer stream.Close() + var content string + for stream.Next() { + if chunk := stream.Current(); len(chunk.Choices) > 0 { + content += chunk.Choices[0].Delta.Content + } + } + if stream.Err() != nil { + return "", tagStreamError(stream.Err()) + } + return content, nil + }, func(err error) bool { return errors.Is(err, geniex_sdk.ErrLlmGenerationPromptTooLong) }) + if err != nil || text != "ok" || dropped != 1 || len(kept) != 1 || requests != 2 { + t.Fatalf("text=%q dropped=%d kept=%d requests=%d err=%v", text, dropped, len(kept), requests, err) + } +} + +func TestRunCompletionsContinuesAfterContextOverflow(t *testing.T) { + oldClient, oldPrompt, oldSystemPrompt := client, prompt, systemPrompt + defer func() { + client, prompt, systemPrompt = oldClient, oldPrompt, oldSystemPrompt + }() + prompt = []string{"old", "current", "next"} + systemPrompt = "rules" + + var requests [][]string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var body struct { + Messages []struct { + Role string `json:"role"` + Content string `json:"content"` + } `json:"messages"` + } + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + t.Errorf("decode request: %v", err) + return + } + roles := make([]string, len(body.Messages)) + for i, message := range body.Messages { + roles[i] = message.Role + ":" + message.Content + } + requests = append(requests, roles) + if len(requests) == 1 { + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, "null") + return + } + w.Header().Set("Content-Type", "text/event-stream") + if len(requests) == 3 { + fmt.Fprint(w, "data: {\"error\":\"Input prompt too long\",\"code\":-200103}\n\n") + return + } + fmt.Fprint(w, "data: {\"id\":\"test\",\"object\":\"chat.completion.chunk\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"ok\"}}]}\n\ndata: [DONE]\n\n") + })) + defer server.Close() + client = openai.NewClient(option.WithBaseURL(server.URL + "/v1")) + + if err := runCompletions(context.Background(), "test", geniex_sdk.ModelTypeLLM); err != nil { + t.Fatalf("runCompletions: %v", err) + } + want := [][]string{ + {"system:rules"}, + {"system:rules", "user:old"}, + {"system:rules", "user:old", "assistant:ok", "user:current"}, + {"system:rules", "user:current"}, + {"system:rules", "user:current", "assistant:ok", "user:next"}, + } + if !reflect.DeepEqual(requests, want) { + t.Fatalf("requests = %v, want %v", requests, want) + } +}