Repository: tmc/langchaingo
Pull request #302 (2023-09-27 05:52:22)

fix: model selection and throwing exceptions
### PR Checklist

- [*] Read the [Contributing documentation](https://github.com/tmc/langchaingo/blob/main/CONTRIBUTING.md).
- [*] Read the [Code of conduct documentation](https://github.com/tmc/langchaingo/blob/main/CODE_OF_CONDUCT.md).
- [*] Name your Pull Request title clearly, concisely, and prefixed with the name of the primarily affected package you changed according to [Good commit messages](https://go.dev/doc/contribute#commit_messages) (such as `memory: add interfaces for X, Y` or `util: add whizzbang helpers`).
- [*] Check that there isn't already a PR that solves the problem the same way to avoid creating a duplicate.
- [*] Provide a description in this PR that addresses **what** the PR is solving, or reference the issue that it solves (e.g. `Fixes #123`).
- [*] Describes the source of new concepts.
- [*] References existing implementations as appropriate.
- [*] Contains test coverage for new functions.
- [*] Passes all [`golangci-lint`](https://golangci-lint.run/) checks.



Patch:
diff --git a/llms/ernie/erniellm.go b/llms/ernie/erniellm.go
index 406c1006..0b7abc15 100644
--- a/llms/ernie/erniellm.go
+++ b/llms/ernie/erniellm.go
@@ -28,14 +28,6 @@ var (
 
 // New returns a new Anthropic LLM.
 func New(opts ...Option) (*LLM, error) {
-	c, err := newClient(opts...)
-
-	return &LLM{
-		client: c,
-	}, err
-}
-
-func newClient(opts ...Option) (*ernieclient.Client, error) {
 	options := &options{
 		apiKey:    os.Getenv(ernieAPIKey),
 		secretKey: os.Getenv(ernieSecretKey),
@@ -45,7 +37,16 @@ func newClient(opts ...Option) (*ernieclient.Client, error) {
 		opt(options)
 	}
 
-	if options.accessToken == "" && (options.apiKey == "" || options.secretKey == "") {
+	c, err := newClient(options)
+
+	return &LLM{
+		client: c,
+		model:  options.modelName,
+	}, err
+}
+
+func newClient(opts *options) (*ernieclient.Client, error) {
+	if opts.accessToken == "" && (opts.apiKey == "" || opts.secretKey == "") {
 		return nil, fmt.Errorf(`%w
 You can pass auth info by use ernie.New(ernie.WithAKSK("{api Key}","{serect Key}")) ,
 or
@@ -55,8 +56,8 @@ doc: https://cloud.baidu.com/doc/WENXINWORKSHOP/s/flfmc9do2`, ernieclient.ErrNot
 	}
 
 	return ernieclient.New(
-		ernieclient.WithAccessToken(options.accessToken),
-		ernieclient.WithAKSK(options.apiKey, options.secretKey))
+		ernieclient.WithAccessToken(opts.accessToken),
+		ernieclient.WithAKSK(opts.apiKey, opts.secretKey))
 }
 
 // GeneratePrompt implements llms.LanguageModel.
@@ -107,6 +108,11 @@ func (l *LLM) Generate(ctx context.Context, prompts []string, options ...llms.Ca
 		if err != nil {
 			return nil, err
 		}
+		if result.ErrorCode > 0 {
+			return nil, fmt.Errorf("%w, error_code:%v, erro_msg:%v, id:%v",
+				ErrCodeResponse, result.ErrorCode, result.ErrorMsg, result.ID)
+		}
+
 		generations = append(generations, &llms.Generation{
 			Text: result.Result,
 		})
@@ -126,8 +132,8 @@ func (l *LLM) CreateEmbedding(ctx context.Context, texts []string) ([][]float32,
 	}
 
 	if resp.ErrorCode > 0 {
-		return nil, fmt.Errorf("%w, error_code:%v, erro_msg:%v",
-			ErrCodeResponse, resp.ErrorCode, resp.ErrorMsg)
+		return nil, fmt.Errorf("%w, error_code:%v, erro_msg:%v, id:%v",
+			ErrCodeResponse, resp.ErrorCode, resp.ErrorMsg, resp.ID)
 	}
 
 	emb := make([][]float32, 0, len(texts))
@@ -159,6 +165,7 @@ func (l *LLM) getModelPath(opts llms.CallOptions) ernieclient.ModelPath {
 	case ModelNameLlama2_70BChat:
 		return "llama_2_70b"
 	default:
+
 		return ernieclient.DefaultCompletionModelPath
 	}
 }
diff --git a/llms/ernie/internal/ernieclient/ernieclient.go b/llms/ernie/internal/ernieclient/ernieclient.go
index 0ebf5a16..a8d771d5 100644
--- a/llms/ernie/internal/ernieclient/ernieclient.go
+++ b/llms/ernie/internal/ernieclient/ernieclient.go
@@ -281,11 +281,14 @@ func parseStreamingCompletionResponse(ctx context.Context, resp *http.Response,
 			if line == "" {
 				continue
 			}
-			if !strings.HasPrefix(line, dataPrefix) {
+
+			if !strings.HasPrefix(line, dataPrefix) && !strings.HasPrefix(line, "{") {
 				continue
 			}
+
 			data := strings.TrimPrefix(line, dataPrefix)
 			streamPayload := &Completion{}
+
 			err := json.NewDecoder(bytes.NewReader([]byte(data))).Decode(&streamPayload)
 			if err != nil {
 				log.Fatalf("failed to decode stream payload: %v", err)
diff --git a/llms/openai/internal/openaiclient/chat.go b/llms/openai/internal/openaiclient/chat.go
index 50001c9b..ea660856 100644
--- a/llms/openai/internal/openaiclient/chat.go
+++ b/llms/openai/internal/openaiclient/chat.go
@@ -221,6 +221,7 @@ func parseStreamingChatResponse(ctx context.Context, r *http.Response, payload *
 		}
 		chunk := []byte(streamResponse.Choices[0].Delta.Content)
 		response.Choices[0].Message.Content += streamResponse.Choices[0].Delta.Content
+		response.Choices[0].FinishReason = streamResponse.Choices[0].FinishReason
 		if streamResponse.Choices[0].Delta.FunctionCall != nil {
 			if response.Choices[0].Message.FunctionCall == nil {
 				response.Choices[0].Message.FunctionCall = streamResponse.Choices[0].Delta.FunctionCall


Test patch:
diff --git a/llms/openai/internal/openaiclient/chat_test.go b/llms/openai/internal/openaiclient/chat_test.go
new file mode 100644
index 00000000..fc6bc977
--- /dev/null
+++ b/llms/openai/internal/openaiclient/chat_test.go
@@ -0,0 +1,32 @@
+package openaiclient
+
+import (
+	"bytes"
+	"context"
+	"io"
+	"net/http"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestParseStreamingChatResponse_FinishReason(t *testing.T) {
+	t.Parallel()
+	mockBody := `data: {"choices":[{"index":0,"delta":{"role":"assistant","content":"hello"},"finish_reason":"stop"}]}`
+	r := &http.Response{
+		StatusCode: http.StatusOK,
+		Body:       io.NopCloser(bytes.NewBufferString(mockBody)),
+	}
+
+	req := &ChatRequest{
+		StreamingFunc: func(ctx context.Context, chunk []byte) error {
+			return nil
+		},
+	}
+
+	resp, err := parseStreamingChatResponse(context.Background(), r, req)
+
+	assert.NoError(t, err)
+	assert.NotNil(t, resp)
+	assert.Equal(t, "stop", resp.Choices[0].FinishReason)
+}