@@ -30,6 +30,15 @@ import (
3030 "github.com/jpillora/backoff"
3131)
3232
33+ type commandModel interface {
34+ Command (ctx context.Context , content string ) (string , error )
35+ }
36+
37+ type commandClient struct {
38+ client * cohereClient.Client
39+ config * commandProcessorConfig
40+ }
41+
3342//go:generate paramgen -output=paramgen_command.go commandProcessorConfig
3443
3544type commandProcessor struct {
@@ -40,11 +49,11 @@ type commandProcessor struct {
4049 logger log.CtxLogger
4150 config commandProcessorConfig
4251 backoffCfg * backoff.Backoff
43- client * cohereClient. Client
52+ client commandModel
4453}
4554
4655type commandProcessorConfig struct {
47- // Model is one of the Cohere model ( command,embed,rerank) .
56+ // Model is one of the name of a compatible command model version .
4857 Model string `json:"model" default:"command"`
4958 // APIKey is the API key for Cohere api calls.
5059 APIKey string `json:"apiKey" validate:"required"`
@@ -91,8 +100,10 @@ func (p *commandProcessor) Configure(ctx context.Context, cfg config.Config) err
91100 }
92101 p .responseBodyRef = & responseBodyRef
93102
94- // new cohere client
95- p .client = cohereClient .NewClient ()
103+ p .client = & commandClient {
104+ client : cohereClient .NewClient (),
105+ config : & p .config ,
106+ }
96107
97108 p .backoffCfg = & backoff.Backoff {
98109 Factor : p .config .BackoffRetryFactor ,
@@ -132,22 +143,9 @@ func (p *commandProcessor) Process(ctx context.Context, records []opencdc.Record
132143 return append (out , sdk.ErrorRecord {Error : fmt .Errorf ("failed to resolve reference %v: %w" , p .config .RequestBodyRef , err )})
133144 }
134145
146+ content := fmt .Sprintf (p .config .Prompt , p .getInput (requestRef .Get ()))
135147 for {
136- resp , err := p .client .V2 .Chat (
137- ctx ,
138- & cohere.V2ChatRequest {
139- Model : p .config .Model ,
140- Messages : cohere.ChatMessages {
141- {
142- Role : "user" ,
143- User : & cohere.UserMessage {Content : & cohere.UserMessageContent {
144- String : fmt .Sprintf (p .config .Prompt , p .getInput (requestRef .Get ())),
145- }},
146- },
147- },
148- },
149- cohereClient .WithToken (p .config .APIKey ),
150- )
148+ resp , err := p .client .Command (ctx , content )
151149 attempt := p .backoffCfg .Attempt ()
152150 duration := p .backoffCfg .Duration ()
153151
@@ -182,40 +180,48 @@ func (p *commandProcessor) Process(ctx context.Context, records []opencdc.Record
182180
183181 p .backoffCfg .Reset ()
184182
185- chatResponse , err := unmarshalChatResponse ([] byte ( resp . String ()) )
183+ err = p . setField ( & record , p . responseBodyRef , resp )
186184 if err != nil {
187- return append (out , sdk.ErrorRecord {Error : err })
185+ return append (out , sdk.ErrorRecord {Error : fmt . Errorf ( "failed setting response body: %w" , err ) })
188186 }
187+ out = append (out , sdk .SingleRecord (record ))
189188
190- if len (chatResponse .Message .Content ) == 1 {
191- err = p .setField (& record , p .responseBodyRef , chatResponse .Message .Content [0 ].Text )
192- if err != nil {
193- return append (out , sdk.ErrorRecord {Error : fmt .Errorf ("failed setting response body: %w" , err )})
194- }
195- out = append (out , sdk .SingleRecord (record ))
196- }
197189 break
198190 }
199191 }
200192 return out
201193}
202194
203- func (p * commandProcessor ) setField (r * opencdc.Record , refRes * sdk.ReferenceResolver , data any ) error {
204- if refRes == nil {
205- return nil
195+ func (cc * commandClient ) Command (ctx context.Context , content string ) (string , error ) {
196+ resp , err := cc .client .V2 .Chat (
197+ ctx ,
198+ & cohere.V2ChatRequest {
199+ Model : cc .config .Model ,
200+ Messages : cohere.ChatMessages {
201+ {
202+ Role : "user" ,
203+ User : & cohere.UserMessage {Content : & cohere.UserMessageContent {
204+ String : content ,
205+ }},
206+ },
207+ },
208+ },
209+ cohereClient .WithToken (cc .config .APIKey ),
210+ )
211+ if err != nil {
212+ return "" , err
206213 }
207214
208- ref , err := refRes . Resolve ( r )
215+ chatResponse , err := unmarshalChatResponse ([] byte ( resp . String ()) )
209216 if err != nil {
210- return fmt .Errorf ("error reference resolver : %w" , err )
217+ return "" , fmt .Errorf ("error unmarshalling chat response : %w" , err )
211218 }
212219
213- err = ref .Set (data )
214- if err != nil {
215- return fmt .Errorf ("error reference set: %w" , err )
220+ if len (chatResponse .Message .Content ) != 1 {
221+ return "" , fmt .Errorf ("invalid chat content" )
216222 }
217223
218- return nil
224+ return chatResponse . Message . Content [ 0 ]. Text , nil
219225}
220226
221227type ChatResponse struct {
@@ -249,3 +255,21 @@ func (p *commandProcessor) getInput(val any) string {
249255 return fmt .Sprintf ("%v" , v )
250256 }
251257}
258+
259+ func (p * commandProcessor ) setField (r * opencdc.Record , refRes * sdk.ReferenceResolver , data any ) error {
260+ if refRes == nil {
261+ return nil
262+ }
263+
264+ ref , err := refRes .Resolve (r )
265+ if err != nil {
266+ return fmt .Errorf ("error reference resolver: %w" , err )
267+ }
268+
269+ err = ref .Set (data )
270+ if err != nil {
271+ return fmt .Errorf ("error reference set: %w" , err )
272+ }
273+
274+ return nil
275+ }
0 commit comments