Update GenerateText result struct

Author: h0rv

README.md

@@ -36,7 +36,7 @@ import (
 		wx.WithDecodingMethod(wx.Greedy),
 	)
 
-  println(result)
+  println(result.Text)
 ```
 
 ## Development Setup

models/generate.go

@@ -15,30 +15,45 @@ const (
 	GenerateTextEndpoint string = GenerationEndpoint + "/text"
 )
 
-type GenerateResult struct {
-	GeneratedText string `json:"generated_text"`
-	StopReason    string `json:"stop_reason"`
+type StopReason = string
+
+const (
+	NotFinished        StopReason = "NOT_FINISHED"  // Possibly more tokens to be streamed
+	MaxTokens          StopReason = "MAX_TOKENS"    // Maximum requested tokens reached
+	EndOfSequenceToken StopReason = "EOS_TOKEN"     // End of sequence token encountered
+	Cancelled          StopReason = "CANCELLED"     // Request canceled by the client
+	TimeLimit          StopReason = "TIME_LIMIT"    // Time limit reached
+	StopSequence       StopReason = "STOP_SEQUENCE" // Stop sequence encountered
+	TokenLimit         StopReason = "TOKEN_LIMIT"   // Token limit reached
+	Error              StopReason = "ERROR"         // Error encountered
+)
+
+type GenerateTextResult struct {
+	Text                string     `json:"generated_text"`
+	GeneratedTokenCount int        `json:"generated_token_count"`
+	InputTokenCount     int        `json:"input_token_count"`
+	StopReason          StopReason `json:"stop_reason"`
 }
 
-type GeneratePayload struct {
+type GenerateTextPayload struct {
 	ProjectID  string           `json:"project_id"`
 	Model      string           `json:"model_id"`
 	Prompt     string           `json:"input"`
 	Parameters *GenerateOptions `json:"parameters,omitempty"`
 }
 
-type generateResponse struct {
-	Status     string           `json:"status"`
-	StatusCode int              `json:"status_code"`
-	Results    []GenerateResult `json:"results"`
+type generateTextResponse struct {
+	Status     string               `json:"status"`
+	StatusCode int                  `json:"status_code"`
+	Results    []GenerateTextResult `json:"results"`
 }
 
 // GenerateText generates completion text based on a given prompt and parameters
-func (m *Model) GenerateText(prompt string, options ...GenerateOption) (string, error) {
+func (m *Model) GenerateText(prompt string, options ...GenerateOption) (GenerateTextResult, error) {
 	m.CheckAndRefreshToken()
 
 	if prompt == "" {
-		return "", errors.New("prompt cannot be empty")
+		return GenerateTextResult{}, errors.New("prompt cannot be empty")
 	}
 
 	opts := &GenerateOptions{}
@@ -48,7 +63,7 @@ func (m *Model) GenerateText(prompt string, options ...GenerateOption) (string,
 		}
 	}
 
-	payload := GeneratePayload{
+	payload := GenerateTextPayload{
 		ProjectID:  m.projectID,
 		Model:      m.modelType,
 		Prompt:     prompt,
@@ -57,17 +72,21 @@ func (m *Model) GenerateText(prompt string, options ...GenerateOption) (string,
 
 	response, err := m.generateTextRequest(payload)
 	if err != nil {
-		return "", err
+		return GenerateTextResult{}, err
+	}
+
+	if len(response.Results) == 0 {
+		return GenerateTextResult{}, errors.New("no result recieved")
 	}
 
-	result := response.Results[0].GeneratedText
+	result := response.Results[0]
 
 	return result, nil
 }
 
 // generateTextRequest sends the generate request and handles the response using the http package.
 // Returns error on non-2XX response
-func (m *Model) generateTextRequest(payload GeneratePayload) (generateResponse, error) {
+func (m *Model) generateTextRequest(payload GenerateTextPayload) (generateTextResponse, error) {
 	params := url.Values{
 		"version": {m.apiVersion},
 	}
@@ -81,37 +100,37 @@ func (m *Model) generateTextRequest(payload GeneratePayload) (generateResponse,
 
 	payloadJSON, err := json.Marshal(payload)
 	if err != nil {
-		return generateResponse{}, err
+		return generateTextResponse{}, err
 	}
 
 	req, err := http.NewRequest(http.MethodPost, generateTextURL.String(), bytes.NewBuffer(payloadJSON))
 	if err != nil {
-		return generateResponse{}, err
+		return generateTextResponse{}, err
 	}
 
 	req.Header.Set("Content-Type", "application/json")
 	req.Header.Set("Authorization", "Bearer "+m.token.value)
 
 	res, err := m.httpClient.Do(req)
 	if err != nil {
-		return generateResponse{}, err
+		return generateTextResponse{}, err
 	}
 
 	statusCode := res.StatusCode
 
 	if statusCode < 200 || statusCode >= 300 {
 		body, err := io.ReadAll(res.Body)
 		if err != nil {
-			return generateResponse{}, fmt.Errorf("request failed with status code %d", statusCode)
+			return generateTextResponse{}, fmt.Errorf("request failed with status code %d", statusCode)
 		}
-		return generateResponse{}, fmt.Errorf("request failed with status code %d and error %s", statusCode, body)
+		return generateTextResponse{}, fmt.Errorf("request failed with status code %d and error %s", statusCode, body)
 	}
 	defer res.Body.Close()
 
-	var generateRes generateResponse
+	var generateRes generateTextResponse
 
 	if err := json.NewDecoder(res.Body).Decode(&generateRes); err != nil {
-		return generateResponse{}, err
+		return generateTextResponse{}, err
 	}
 
 	return generateRes, nil

models/test/generate_test.go

@@ -72,7 +72,7 @@ func TestGenerateText(t *testing.T) {
 	if err != nil {
 		t.Fatalf("Expected no error, but got an error: %v", err)
 	}
-	if result == "" {
+	if result.Text == "" {
 		t.Fatal("Expected a result, but got an empty string")
 	}
 }
@@ -88,7 +88,7 @@ func TestGenerateTextWithNilOptions(t *testing.T) {
 	if err != nil {
 		t.Fatalf("Expected no error, but got an error: %v", err)
 	}
-	if result == "" {
+	if result.Text == "" {
 		t.Fatal("Expected a result, but got an empty string")
 	}
 }