chore: fix data races

.github/workflows/ci.yaml

@@ -45,5 +45,5 @@ jobs:
         env:
           OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
           GENAI_API_KEY: ${{ secrets.GENAI_API_KEY }}
-        run: go test -v ./...
+        run: go test -v ./... -race
 

chains/chains_test.go

@@ -20,6 +20,7 @@ type testLanguageModel struct {
 	simulateWork time.Duration
 	// record the prompt that was passed to the language model
 	recordedPrompt []llms.PromptValue
+	mu             sync.Mutex
 }
 
 type stringPromptValue struct {
@@ -46,9 +47,11 @@ func (l *testLanguageModel) GenerateContent(_ context.Context, mc []llms.Message
 	} else {
 		return nil, fmt.Errorf("passed non-text part")
 	}
+	l.mu.Lock()
 	l.recordedPrompt = []llms.PromptValue{
 		stringPromptValue{s: prompt},
 	}
+	l.mu.Unlock()
 
 	if l.simulateWork > 0 {
 		time.Sleep(l.simulateWork)

llms/ernie/internal/ernieclient/ernieclient.go

@@ -10,6 +10,7 @@ import (
 	"log"
 	"net/http"
 	"strings"
+	"sync"
 	"time"
 )
 
@@ -26,6 +27,7 @@ type Client struct {
 	apiKey      string
 	secretKey   string
 	accessToken string
+	mu          sync.RWMutex
 	httpClient  Doer
 	Model       string
 	ModelPath   ModelPath
@@ -175,7 +177,9 @@ func autoRefresh(c *Client) error {
 				time.Sleep(tryPeriod * time.Minute) // try
 				continue
 			}
+			c.mu.Lock()
 			c.accessToken = authResp.AccessToken
+			c.mu.Unlock()
 			time.Sleep(10 * 24 * time.Hour)
 		}
 	}()
@@ -188,8 +192,11 @@ func (c *Client) CreateCompletion(ctx context.Context, modelPath ModelPath, r *C
 		modelPath = DefaultCompletionModelPath
 	}
 
+	c.mu.RLock()
+	accessToken := c.accessToken
+	c.mu.RUnlock()
 	url := "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/" + string(modelPath) +
-		"?access_token=" + c.accessToken
+		"?access_token=" + accessToken
 	body, e := json.Marshal(r)
 	if e != nil {
 		return nil, e
@@ -219,8 +226,11 @@ func (c *Client) CreateCompletion(ctx context.Context, modelPath ModelPath, r *C
 
 // CreateEmbedding use ernie Embedding-V1.
 func (c *Client) CreateEmbedding(ctx context.Context, texts []string) (*EmbeddingResponse, error) {
+	c.mu.RLock()
+	accessToken := c.accessToken
+	c.mu.RUnlock()
 	url := "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/embeddings/embedding-v1?access_token=" +
-		c.accessToken
+		accessToken
 
 	payload := make(map[string]any)
 	payload["input"] = texts
@@ -346,8 +356,11 @@ func (c *Client) buildURL(modelpath ModelPath) string {
 
 	// ernie example url:
 	// /wenxinworkshop/chat/eb-instant
+	c.mu.RLock()
+	accessToken := c.accessToken
+	c.mu.RUnlock()
 	return fmt.Sprintf("%s/wenxinworkshop/chat/%s?access_token=%s",
-		baseURL, modelpath, c.accessToken,
+		baseURL, modelpath, accessToken,
 	)
 }
 

llms/ernie/internal/ernieclient/ernieclient_test.go

@@ -1,6 +1,30 @@
 package ernieclient
 
-import "testing"
+import (
+	"encoding/json"
+	"io"
+	"net/http"
+	"strings"
+	"testing"
+)
+
+type mockHTTPClient struct{}
+
+// implement ernieclient.Doer interface.
+func (m *mockHTTPClient) Do(_ *http.Request) (*http.Response, error) {
+	authResponse := &authResponse{
+		AccessToken: "test",
+	}
+	b, err := json.Marshal(authResponse)
+	if err != nil {
+		return nil, err
+	}
+	response := &http.Response{
+		StatusCode: http.StatusOK,
+		Body:       io.NopCloser(strings.NewReader(string(b))),
+	}
+	return response, nil
+}
 
 func TestClient_buildURL(t *testing.T) {
 	t.Parallel()
@@ -9,8 +33,6 @@ func TestClient_buildURL(t *testing.T) {
 		secretKey   string
 		accessToken string
 		httpClient  Doer
-		Model       string
-		ModelPath   ModelPath
 	}
 	type args struct {
 		modelpath ModelPath
@@ -22,25 +44,35 @@ func TestClient_buildURL(t *testing.T) {
 		want   string
 	}{
 		{
-			name: "one",
+			name: "with access token",
 			fields: fields{
 				accessToken: "token",
 			},
 			args: args{modelpath: "eb-instant"},
 			want: "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/eb-instant?access_token=token",
 		},
+		{
+			name: "with client, aksk",
+			fields: fields{
+				apiKey:     "test",
+				secretKey:  "test",
+				httpClient: &mockHTTPClient{},
+			},
+			args: args{modelpath: "eb-instant"},
+			want: "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/eb-instant?access_token=test",
+		},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			t.Parallel()
-			c := &Client{
-				apiKey:      tt.fields.apiKey,
-				secretKey:   tt.fields.secretKey,
-				accessToken: tt.fields.accessToken,
-				httpClient:  tt.fields.httpClient,
-				Model:       tt.fields.Model,
-				ModelPath:   tt.fields.ModelPath,
+			c, err := New(
+				WithAKSK(tt.fields.apiKey, tt.fields.secretKey),
+				WithAccessToken(tt.fields.accessToken),
+				WithHTTPClient(tt.fields.httpClient),
+			)
+			if err != nil {
+				t.Errorf("New got error. %v", err)
 			}
 			if got := c.buildURL(tt.args.modelpath); got != tt.want {
 				t.Errorf("buildURL() = %v, want %v", got, tt.want)

llms/openai/internal/openaiclient/chat.go

@@ -381,9 +381,6 @@ func (c *Client) createChat(ctx context.Context, payload *ChatRequest) (*ChatCom
 
 	// Build request
 	body := bytes.NewReader(payloadBytes)
-	if c.baseURL == "" {
-		c.baseURL = defaultBaseURL
-	}
 	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.buildURL("/chat/completions", payload.Model), body)
 	if err != nil {
 		return nil, err

llms/openai/internal/openaiclient/openaiclient.go

@@ -65,6 +65,9 @@ func New(token string, model string, baseURL string, organization string,
 		httpClient:     httpClient,
 		ResponseFormat: responseFormat,
 	}
+	if c.baseURL == "" {
+		c.baseURL = defaultBaseURL
+	}
 
 	for _, opt := range opts {
 		if err := opt(c); err != nil {