@@ -11,6 +11,7 @@ import (
1111 "os"
1212 "sort"
1313 "strings"
14+ "sync/atomic"
1415
1516 "github.com/acorn-io/gptscript/pkg/cache"
1617 "github.com/acorn-io/gptscript/pkg/hash"
@@ -27,11 +28,14 @@ const (
2728)
2829
2930var (
30- key = os .Getenv ("OPENAI_API_KEY" )
31- url = os .Getenv ("OPENAI_URL" )
31+ key = os .Getenv ("OPENAI_API_KEY" )
32+ url = os .Getenv ("OPENAI_URL" )
33+ transactionID int64
3234)
3335
3436type Client struct {
37+ url string
38+ key string
3539 c * openai.Client
3640 cache * cache.Client
3741}
@@ -107,7 +111,11 @@ func (c *Client) ListModules(ctx context.Context) (result []string, _ error) {
107111}
108112
109113func (c * Client ) cacheKey (request openai.ChatCompletionRequest ) string {
110- return hash .Encode (request )
114+ return hash .Encode (map [string ]any {
115+ "url" : c .url ,
116+ "key" : c .key ,
117+ "request" : request ,
118+ })
111119}
112120
113121func (c * Client ) seed (request openai.ChatCompletionRequest ) int {
@@ -231,11 +239,12 @@ func toMessages(ctx context.Context, cache *cache.Client, request types.Completi
231239}
232240
233241type Status struct {
234- Request any
235- Response any
236- Cached bool
237- Chunks any
238- PartialResponse * types.CompletionMessage
242+ OpenAITransactionID string
243+ Request any
244+ Response any
245+ Cached bool
246+ Chunks any
247+ PartialResponse * types.CompletionMessage
239248}
240249
241250func (c * Client ) Call (ctx context.Context , messageRequest types.CompletionRequest , status chan <- Status ) (* types.CompletionMessage , error ) {
@@ -281,8 +290,10 @@ func (c *Client) Call(ctx context.Context, messageRequest types.CompletionReques
281290 }
282291 }
283292
293+ id := fmt .Sprint (atomic .AddInt64 (& transactionID , 1 ))
284294 status <- Status {
285- Request : request ,
295+ OpenAITransactionID : id ,
296+ Request : request ,
286297 }
287298
288299 var cacheResponse bool
@@ -291,7 +302,7 @@ func (c *Client) Call(ctx context.Context, messageRequest types.CompletionReques
291302 if err != nil {
292303 return nil , err
293304 } else if ! ok {
294- response , err = c .call (ctx , request , status )
305+ response , err = c .call (ctx , request , id , status )
295306 if err != nil {
296307 return nil , err
297308 }
@@ -305,9 +316,10 @@ func (c *Client) Call(ctx context.Context, messageRequest types.CompletionReques
305316 }
306317
307318 status <- Status {
308- Chunks : response ,
309- Response : result ,
310- Cached : cacheResponse ,
319+ OpenAITransactionID : id ,
320+ Chunks : response ,
321+ Response : result ,
322+ Cached : cacheResponse ,
311323 }
312324
313325 return & result , nil
@@ -391,7 +403,7 @@ func (c *Client) store(ctx context.Context, key string, responses []openai.ChatC
391403 return c .cache .Store (ctx , key , buf .Bytes ())
392404}
393405
394- func (c * Client ) call (ctx context.Context , request openai.ChatCompletionRequest , partial chan <- Status ) (responses []openai.ChatCompletionStreamResponse , _ error ) {
406+ func (c * Client ) call (ctx context.Context , request openai.ChatCompletionRequest , transactionID string , partial chan <- Status ) (responses []openai.ChatCompletionStreamResponse , _ error ) {
395407 cacheKey := c .cacheKey (request )
396408 request .Stream = true
397409
@@ -404,6 +416,7 @@ func (c *Client) call(ctx context.Context, request openai.ChatCompletionRequest,
404416 }
405417
406418 partial <- Status {
419+ OpenAITransactionID : transactionID ,
407420 PartialResponse : & types.CompletionMessage {
408421 Role : types .CompletionMessageRoleTypeAssistant ,
409422 Content : types .Text (msg + "Waiting for model response...\n " ),
@@ -429,7 +442,8 @@ func (c *Client) call(ctx context.Context, request openai.ChatCompletionRequest,
429442 if partial != nil {
430443 partialMessage = appendMessage (partialMessage , response )
431444 partial <- Status {
432- PartialResponse : & partialMessage ,
445+ OpenAITransactionID : transactionID ,
446+ PartialResponse : & partialMessage ,
433447 }
434448 }
435449 responses = append (responses , response )
0 commit comments