66 "fmt"
77 "net/http"
88 "net/url"
9+ "regexp"
910
1011 "github.com/appleboy/CodeGPT/core"
1112
@@ -54,25 +55,36 @@ func (c *Client) Completion(ctx context.Context, content string) (*core.Response
5455 return & core.Response {
5556 Content : resp .Content ,
5657 Usage : core.Usage {
57- PromptTokens : resp .Usage .PromptTokens ,
58- CompletionTokens : resp .Usage .CompletionTokens ,
59- TotalTokens : resp .Usage .TotalTokens ,
58+ PromptTokens : resp .Usage .PromptTokens ,
59+ CompletionTokens : resp .Usage .CompletionTokens ,
60+ TotalTokens : resp .Usage .TotalTokens ,
61+ CompletionTokensDetails : resp .Usage .CompletionTokensDetails ,
6062 },
6163 }, nil
6264}
6365
6466// GetSummaryPrefix is an API call to get a summary prefix using function call.
6567func (c * Client ) GetSummaryPrefix (ctx context.Context , content string ) (* core.Response , error ) {
66- resp , err := c .CreateFunctionCall (ctx , content , SummaryPrefixFunc )
67- if err != nil || len (resp .Choices ) != 1 {
68- return nil , err
68+ var resp openai.ChatCompletionResponse
69+ var err error
70+ if checkO1Serial .MatchString (c .model ) {
71+ resp , err = c .CreateChatCompletion (ctx , content )
72+ if err != nil || len (resp .Choices ) != 1 {
73+ return nil , err
74+ }
75+ } else {
76+ resp , err = c .CreateFunctionCall (ctx , content , SummaryPrefixFunc )
77+ if err != nil || len (resp .Choices ) != 1 {
78+ return nil , err
79+ }
6980 }
7081
7182 msg := resp .Choices [0 ].Message
7283 usage := core.Usage {
73- PromptTokens : resp .Usage .PromptTokens ,
74- CompletionTokens : resp .Usage .CompletionTokens ,
75- TotalTokens : resp .Usage .TotalTokens ,
84+ PromptTokens : resp .Usage .PromptTokens ,
85+ CompletionTokens : resp .Usage .CompletionTokens ,
86+ TotalTokens : resp .Usage .TotalTokens ,
87+ CompletionTokensDetails : resp .Usage .CompletionTokensDetails ,
7688 }
7789 if len (msg .ToolCalls ) == 0 {
7890 return & core.Response {
@@ -88,6 +100,8 @@ func (c *Client) GetSummaryPrefix(ctx context.Context, content string) (*core.Re
88100 }, nil
89101}
90102
103+ var checkO1Serial = regexp .MustCompile (`o1-(mini|preview)` )
104+
91105// CreateChatCompletion is an API call to create a function call for a chat message.
92106func (c * Client ) CreateFunctionCall (
93107 ctx context.Context ,
@@ -108,7 +122,7 @@ func (c *Client) CreateFunctionCall(
108122 PresencePenalty : c .presencePenalty ,
109123 Messages : []openai.ChatCompletionMessage {
110124 {
111- Role : openai .ChatMessageRoleSystem ,
125+ Role : openai .ChatMessageRoleAssistant ,
112126 Content : "You are a helpful assistant." ,
113127 },
114128 {
@@ -125,6 +139,11 @@ func (c *Client) CreateFunctionCall(
125139 },
126140 }
127141
142+ if checkO1Serial .MatchString (c .model ) {
143+ req .MaxTokens = 0
144+ req .MaxCompletionsTokens = c .maxTokens
145+ }
146+
128147 return c .client .CreateChatCompletion (ctx , req )
129148}
130149
@@ -142,7 +161,7 @@ func (c *Client) CreateChatCompletion(
142161 PresencePenalty : c .presencePenalty ,
143162 Messages : []openai.ChatCompletionMessage {
144163 {
145- Role : openai .ChatMessageRoleSystem ,
164+ Role : openai .ChatMessageRoleAssistant ,
146165 Content : "You are a helpful assistant." ,
147166 },
148167 {
@@ -152,6 +171,11 @@ func (c *Client) CreateChatCompletion(
152171 },
153172 }
154173
174+ if checkO1Serial .MatchString (c .model ) {
175+ req .MaxTokens = 0
176+ req .MaxCompletionsTokens = c .maxTokens
177+ }
178+
155179 return c .client .CreateChatCompletion (ctx , req )
156180}
157181
0 commit comments