fix: model selection and throwing exceptions (#302)

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
 	}
 }

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)