Skip to content

Commit

Permalink
feat(chat_completions): support partial mode
Browse files Browse the repository at this point in the history
  • Loading branch information
northes committed May 30, 2024
1 parent 12b8194 commit 62e528b
Show file tree
Hide file tree
Showing 3 changed files with 58 additions and 3 deletions.
15 changes: 15 additions & 0 deletions api_chat_completions.go
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,8 @@ func (c *Client) Chat() IChat {
type ChatCompletionsMessage struct {
Role ChatCompletionsMessageRole `json:"role"`
Content string `json:"content"`
Partial bool `json:"partial,omitempty"`
Name string `json:"name,omitempty"`
}

type ChatCompletionsRequest struct {
Expand Down Expand Up @@ -90,6 +92,19 @@ func (c *chat) Completions(ctx context.Context, req *ChatCompletionsRequest) (*C
return chatCompletionsResp, nil
}

func (c *ChatCompletionsResponse) GetMessage() (*ChatCompletionsMessage, error) {
if len(c.Choices) == 0 {
return nil, fmt.Errorf("empty choices")

Check warning on line 97 in api_chat_completions.go

View check run for this annotation

Codecov / codecov/patch

api_chat_completions.go#L97

Added line #L97 was not covered by tests
}
for _, choice := range c.Choices {
if choice.Message != nil {
return choice.Message, nil
}
}

return nil, fmt.Errorf("no such choice")

Check warning on line 105 in api_chat_completions.go

View check run for this annotation

Codecov / codecov/patch

api_chat_completions.go#L105

Added line #L105 was not covered by tests
}

type ChatCompletionsStreamResponse struct {
resp *httpx.Response
}
Expand Down
34 changes: 34 additions & 0 deletions api_chat_completions_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,11 +3,14 @@ package moonshot_test
import (
"context"
"errors"
"fmt"
"io"
"strings"
"testing"

"github.com/northes/go-moonshot"
"github.com/northes/go-moonshot/test"
"github.com/stretchr/testify/require"
)

func TestChat(t *testing.T) {
Expand Down Expand Up @@ -99,3 +102,34 @@ func TestChatWithContext(t *testing.T) {
}
t.Log(test.MarshalJsonToStringX(resp))
}

func TestPartialMode(t *testing.T) {
cli, err := NewTestClient()
if err != nil {
t.Fatal(err)
}
ctx := context.Background()

const leadingText = "{"

builder := moonshot.NewChatCompletionsBuilder()
builder.AddSystemContent("请从产品描述中提取名称、尺寸、价格和颜色,并在一个 JSON 对象中输出。")
builder.AddUserContent("大米 SmartHome Mini 是一款小巧的智能家居助手,有黑色和银色两种颜色,售价为 998 元,尺寸为 256 x 128 x 128mm。可让您通过语音或应用程序控制灯光、恒温器和其他联网设备,无论您将它放在家中的任何位置。")
builder.AddAssistantContent(leadingText, true)
builder.SetTemperature(0.3)

resp, err := cli.Chat().Completions(ctx, builder.ToRequest())
if err != nil {
t.Fatal(err)
}
message, err := resp.GetMessage()
if err != nil {
t.Fatal(err)
}

require.Equal(t, strings.HasPrefix(message.Content, leadingText), false, "message content should not start with '{'")

jsonStr := fmt.Sprintf("%s%s", leadingText, message.Content)

t.Log(jsonStr)
}
12 changes: 9 additions & 3 deletions chat_completions_builder.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ package moonshot
type IChatCompletionsBuilder interface {
AddUserContent(content string) IChatCompletionsBuilder
AddSystemContent(content string) IChatCompletionsBuilder
AddAssistantContent(content string) IChatCompletionsBuilder
AddAssistantContent(content string, partialMode ...bool) IChatCompletionsBuilder
AddPrompt(prompt string) IChatCompletionsBuilder
AddMessage(message *ChatCompletionsMessage) IChatCompletionsBuilder

Expand Down Expand Up @@ -66,11 +66,17 @@ func (c *chatCompletionsBuilder) AddSystemContent(content string) IChatCompletio
return c
}

// AddAssistantContent adds a message with the role of assistant
func (c *chatCompletionsBuilder) AddAssistantContent(content string) IChatCompletionsBuilder {
// AddAssistantContent add a message with the role of assistant, and partial mode
func (c *chatCompletionsBuilder) AddAssistantContent(content string, partialMode ...bool) IChatCompletionsBuilder {
var partial bool
if len(partialMode) == 1 {
partial = partialMode[0]
}

c.req.Messages = append(c.req.Messages, &ChatCompletionsMessage{
Role: RoleAssistant,
Content: content,
Partial: partial,
})
return c
}
Expand Down

0 comments on commit 62e528b

Please sign in to comment.