From 18c7e7224ba7925e695e7537b283a17c2e50317d Mon Sep 17 00:00:00 2001 From: kardolus Date: Sat, 17 Aug 2024 10:41:18 -0400 Subject: [PATCH] Add debug mode --- client/client.go | 35 +++++++++++++++++++++++++++-- configmanager/configmanager_test.go | 11 +++++++++ http/http.go | 22 +++++++++++++----- http/http_test.go | 7 ++++-- integration/integration_test.go | 16 +++++++++++++ types/config.go | 1 + 6 files changed, 82 insertions(+), 10 deletions(-) diff --git a/client/client.go b/client/client.go index 080bf59..e824193 100644 --- a/client/client.go +++ b/client/client.go @@ -122,7 +122,17 @@ func (c *Client) Query(input string) (string, int, error) { return "", 0, err } - raw, err := c.caller.Post(c.getEndpoint(c.Config.CompletionsPath), body, false) + endpoint := c.getEndpoint(c.Config.CompletionsPath) + + if c.Config.Debug { + c.printRequestDebugInfo(endpoint, body) + } + + raw, err := c.caller.Post(endpoint, body, false) + if c.Config.Debug { + c.printResponseDebugInfo(raw) + } + if err != nil { return "", 0, err } @@ -154,7 +164,13 @@ func (c *Client) Stream(input string) error { return err } - result, err := c.caller.Post(c.getEndpoint(c.Config.CompletionsPath), body, true) + endpoint := c.getEndpoint(c.Config.CompletionsPath) + + if c.Config.Debug { + c.printRequestDebugInfo(endpoint, body) + } + + result, err := c.caller.Post(endpoint, body, true) if err != nil { return err } @@ -317,3 +333,18 @@ func generateUniqueSlug() string { guid := uuid.New() return InteractiveThreadPrefix + guid.String()[:4] } + +func (c *Client) printRequestDebugInfo(endpoint string, body []byte) { + bodyString := strings.ReplaceAll(string(body), "'", "'\"'\"'") // Escape single quotes + + fmt.Printf("\nGenerated cURL command:\n\n") + fmt.Printf("curl --location --insecure --request POST '%s' \\\n", endpoint) + fmt.Printf(" --header \"Authorization: Bearer ${OPENAI_API_KEY}\" \\\n") + fmt.Printf(" --header 'Content-Type: application/json' \\\n") + fmt.Printf(" --data-raw '%s'\n\n", bodyString) +} + +func (c *Client) printResponseDebugInfo(raw []byte) { + fmt.Printf("\nResponse\n\n") + fmt.Printf("%s\n\n", raw) +} diff --git a/configmanager/configmanager_test.go b/configmanager/configmanager_test.go index 0d68dde..d35fcaf 100644 --- a/configmanager/configmanager_test.go +++ b/configmanager/configmanager_test.go @@ -38,6 +38,7 @@ func testConfig(t *testing.T, when spec.G, it spec.S) { defaultOmitHistory = false defaultAutoCreateNewThread = false defaultTrackTokenUsage = false + defaultDebug = false defaultTemperature = 1.1 defaultTopP = 2.2 defaultFrequencyPenalty = 3.3 @@ -78,6 +79,7 @@ func testConfig(t *testing.T, when spec.G, it spec.S) { CommandPrompt: defaultCommandPrompt, AutoCreateNewThread: defaultAutoCreateNewThread, TrackTokenUsage: defaultTrackTokenUsage, + Debug: defaultDebug, } envPrefix = strings.ToUpper(defaultConfig.Name) + "_" @@ -114,6 +116,7 @@ func testConfig(t *testing.T, when spec.G, it spec.S) { Expect(subject.Config.CommandPrompt).To(Equal(defaultCommandPrompt)) Expect(subject.Config.AutoCreateNewThread).To(Equal(defaultAutoCreateNewThread)) Expect(subject.Config.TrackTokenUsage).To(Equal(defaultTrackTokenUsage)) + Expect(subject.Config.Debug).To(Equal(defaultDebug)) }) it("should prioritize user-provided config over defaults", func() { @@ -136,6 +139,7 @@ func testConfig(t *testing.T, when spec.G, it spec.S) { CommandPrompt: "user-command-prompt", AutoCreateNewThread: true, TrackTokenUsage: true, + Debug: true, } mockConfigStore.EXPECT().ReadDefaults().Return(defaultConfig).Times(1) @@ -154,6 +158,7 @@ func testConfig(t *testing.T, when spec.G, it spec.S) { Expect(subject.Config.OmitHistory).To(BeTrue()) Expect(subject.Config.AutoCreateNewThread).To(BeTrue()) Expect(subject.Config.TrackTokenUsage).To(BeTrue()) + Expect(subject.Config.Debug).To(BeTrue()) Expect(subject.Config.Role).To(Equal("user-role")) Expect(subject.Config.Thread).To(Equal("user-thread")) Expect(subject.Config.Temperature).To(Equal(2.5)) @@ -176,6 +181,7 @@ func testConfig(t *testing.T, when spec.G, it spec.S) { os.Setenv(envPrefix+"OMIT_HISTORY", "true") os.Setenv(envPrefix+"AUTO_CREATE_NEW_THREAD", "true") os.Setenv(envPrefix+"TRACK_TOKEN_USAGE", "true") + os.Setenv(envPrefix+"DEBUG", "true") os.Setenv(envPrefix+"ROLE", "env-role") os.Setenv(envPrefix+"THREAD", "env-thread") os.Setenv(envPrefix+"TEMPERATURE", "2.2") @@ -201,6 +207,7 @@ func testConfig(t *testing.T, when spec.G, it spec.S) { Expect(subject.Config.OmitHistory).To(BeTrue()) Expect(subject.Config.AutoCreateNewThread).To(BeTrue()) Expect(subject.Config.TrackTokenUsage).To(BeTrue()) + Expect(subject.Config.Debug).To(BeTrue()) Expect(subject.Config.Role).To(Equal("env-role")) Expect(subject.Config.Thread).To(Equal("env-thread")) Expect(subject.Config.Temperature).To(Equal(2.2)) @@ -223,6 +230,7 @@ func testConfig(t *testing.T, when spec.G, it spec.S) { os.Setenv(envPrefix+"OMIT_HISTORY", "true") os.Setenv(envPrefix+"AUTO_CREATE_NEW_THREAD", "true") os.Setenv(envPrefix+"TRACK_TOKEN_USAGE", "true") + os.Setenv(envPrefix+"DEBUG", "false") os.Setenv(envPrefix+"ROLE", "env-role") os.Setenv(envPrefix+"THREAD", "env-thread") os.Setenv(envPrefix+"TEMPERATURE", "2.2") @@ -244,6 +252,7 @@ func testConfig(t *testing.T, when spec.G, it spec.S) { OmitHistory: false, AutoCreateNewThread: false, TrackTokenUsage: false, + Debug: true, Role: "user-role", Thread: "user-thread", Temperature: 1.5, @@ -270,6 +279,7 @@ func testConfig(t *testing.T, when spec.G, it spec.S) { Expect(subject.Config.OmitHistory).To(BeTrue()) Expect(subject.Config.AutoCreateNewThread).To(BeTrue()) Expect(subject.Config.TrackTokenUsage).To(BeTrue()) + Expect(subject.Config.Debug).To(BeFalse()) Expect(subject.Config.Role).To(Equal("env-role")) Expect(subject.Config.Thread).To(Equal("env-thread")) Expect(subject.Config.Temperature).To(Equal(2.2)) @@ -428,6 +438,7 @@ func unsetEnvironmentVariables(envPrefix string) { "COMMAND_PROMPT", "AUTO_CREATE_NEW_THREAD", "TRACK_TOKEN_USAGE", + "DEBUG", } for _, variable := range variables { diff --git a/http/http.go b/http/http.go index 2087f79..f0d9619 100644 --- a/http/http.go +++ b/http/http.go @@ -56,12 +56,22 @@ func (r *RestCaller) Post(url string, body []byte, stream bool) ([]byte, error) return r.doRequest(http.MethodPost, url, body, stream) } -func ProcessResponse(r io.Reader, w io.Writer) []byte { +func (r *RestCaller) ProcessResponse(reader io.Reader, writer io.Writer) []byte { var result []byte - scanner := bufio.NewScanner(r) + if r.config.Debug { + fmt.Printf("\nResponse\n\n") + } + + scanner := bufio.NewScanner(reader) for scanner.Scan() { line := scanner.Text() + + if r.config.Debug { + fmt.Println(line) + continue + } + if strings.HasPrefix(line, "data:") { line = line[6:] // Skip the "data: " prefix if len(line) < 6 { @@ -69,7 +79,7 @@ func ProcessResponse(r io.Reader, w io.Writer) []byte { } if line == "[DONE]" { - _, _ = w.Write([]byte("\n")) + _, _ = writer.Write([]byte("\n")) result = append(result, []byte("\n")...) break } @@ -77,13 +87,13 @@ func ProcessResponse(r io.Reader, w io.Writer) []byte { var data types.Data err := json.Unmarshal([]byte(line), &data) if err != nil { - _, _ = fmt.Fprintf(w, "Error: %s\n", err.Error()) + _, _ = fmt.Fprintf(writer, "Error: %s\n", err.Error()) continue } for _, choice := range data.Choices { if content, ok := choice.Delta["content"]; ok { - _, _ = w.Write([]byte(content)) + _, _ = writer.Write([]byte(content)) result = append(result, []byte(content)...) } } @@ -119,7 +129,7 @@ func (r *RestCaller) doRequest(method, url string, body []byte, stream bool) ([] } if stream { - return ProcessResponse(response.Body, os.Stdout), nil + return r.ProcessResponse(response.Body, os.Stdout), nil } result, err := io.ReadAll(response.Body) diff --git a/http/http_test.go b/http/http_test.go index bd3dc65..dd91a7f 100644 --- a/http/http_test.go +++ b/http/http_test.go @@ -16,14 +16,17 @@ func TestUnitHTTP(t *testing.T) { } func testHTTP(t *testing.T, when spec.G, it spec.S) { + var subject http.RestCaller + it.Before(func() { RegisterTestingT(t) + subject = http.RestCaller{} }) when("ProcessResponse()", func() { it("parses a stream as expected", func() { buf := &bytes.Buffer{} - http.ProcessResponse(strings.NewReader(stream), buf) + subject.ProcessResponse(strings.NewReader(stream), buf) output := buf.String() Expect(output).To(Equal("a b c\n")) }) @@ -32,7 +35,7 @@ func testHTTP(t *testing.T, when spec.G, it spec.S) { expectedOutput := "Error: unexpected end of JSON input\n" var buf bytes.Buffer - http.ProcessResponse(strings.NewReader(input), &buf) + subject.ProcessResponse(strings.NewReader(input), &buf) output := buf.String() Expect(output).To(Equal(expectedOutput)) diff --git a/integration/integration_test.go b/integration/integration_test.go index 78b808d..f0b642a 100644 --- a/integration/integration_test.go +++ b/integration/integration_test.go @@ -487,6 +487,22 @@ max_tokens: 100 Expect(output).To(ContainSubstring("Token Usage:")) }) + it("prints debug information with the --debug flag", func() { + Expect(os.Setenv("OPENAI_DEBUG", "true")).To(Succeed()) + + output := runCommand("--query", "tell me a joke") + + Expect(output).To(ContainSubstring("Generated cURL command")) + Expect(output).To(ContainSubstring("/v1/chat/completions")) + Expect(output).To(ContainSubstring("--header \"Authorization: Bearer ${OPENAI_API_KEY}\"")) + Expect(output).To(ContainSubstring("--header 'Content-Type: application/json'")) + Expect(output).To(ContainSubstring("\"model\":\"gpt-3.5-turbo\"")) + Expect(output).To(ContainSubstring("\"messages\":")) + Expect(output).To(ContainSubstring("Response")) + + Expect(os.Unsetenv("OPENAI_DEBUG")).To(Succeed()) + }) + it("should assemble http errors as expected", func() { Expect(os.Setenv(apiKeyEnvVar, "wrong-token")).To(Succeed()) diff --git a/types/config.go b/types/config.go index 430bd15..25c4570 100644 --- a/types/config.go +++ b/types/config.go @@ -21,4 +21,5 @@ type Config struct { CommandPrompt string `yaml:"command_prompt"` AutoCreateNewThread bool `yaml:"auto_create_new_thread"` TrackTokenUsage bool `yaml:"track_token_usage"` + Debug bool `yaml:"debug"` }