@@ -5,15 +5,15 @@ import (
55 "encoding/base64"
66 "errors"
77 "fmt"
8- "io"
98 "os"
109 "strings"
1110
12- "github.com/revrost/go-openrouter"
11+ openrouter "github.com/OpenRouterTeam/go-sdk"
12+ "github.com/OpenRouterTeam/go-sdk/models/components"
1313)
1414
1515type OpenRouterLLM struct {
16- Client * openrouter.Client
16+ Client * openrouter.OpenRouter
1717 Model string
1818 PromptPrefix string
1919}
@@ -24,81 +24,138 @@ func NewOpenRouterLLM(model string, promptPrefix string) *OpenRouterLLM {
2424 apiKey = os .Getenv ("OPENROUTER_KEY" ) // fallback
2525 }
2626
27- client := openrouter .NewClient (apiKey )
27+ client := openrouter .New (
28+ openrouter .WithSecurity (apiKey ),
29+ )
2830 return & OpenRouterLLM {
2931 Client : client ,
3032 Model : model ,
3133 PromptPrefix : promptPrefix ,
3234 }
3335}
3436
35- func (o * OpenRouterLLM ) Generate (ctx context.Context , prompt string ) (any , error ) {
36- fullPrompt := prompt
37+ func (o * OpenRouterLLM ) buildPrompt (prompt string ) string {
3738 if o .PromptPrefix != "" {
38- fullPrompt = o .PromptPrefix + "\n " + prompt
39+ return o .PromptPrefix + "\n " + prompt
40+ }
41+ return prompt
42+ }
43+
44+ // firstChoiceText extracts the assistant's text out of a non-streaming
45+ // ChatResult. Content is normally a plain string (ChatUserMessageContentTypeStr
46+ // equivalent on the response side), but defensively also handles the
47+ // content-parts array shape by concatenating any text parts.
48+ func firstChoiceText (result * components.ChatResult ) (string , error ) {
49+ if result == nil || len (result .Choices ) == 0 {
50+ return "" , errors .New ("no response from OpenRouter" )
51+ }
52+
53+ content , ok := result .Choices [0 ].Message .Content .GetOrZero ()
54+ if ! ok {
55+ return "" , errors .New ("empty response content from OpenRouter" )
56+ }
57+
58+ if content .Str != nil {
59+ return * content .Str , nil
60+ }
61+
62+ if len (content .ArrayOfChatContentItems ) > 0 {
63+ var sb strings.Builder
64+ for _ , item := range content .ArrayOfChatContentItems {
65+ if item .ChatContentText != nil {
66+ sb .WriteString (item .ChatContentText .Text )
67+ }
68+ }
69+ return sb .String (), nil
3970 }
4071
41- resp , err := o .Client .CreateChatCompletion (ctx , openrouter.ChatCompletionRequest {
42- Model : o .Model ,
43- Messages : []openrouter.ChatCompletionMessage {
44- openrouter .UserMessage (fullPrompt ),
72+ return "" , errors .New ("unsupported response content shape from OpenRouter" )
73+ }
74+
75+ func (o * OpenRouterLLM ) Generate (ctx context.Context , prompt string ) (any , error ) {
76+ fullPrompt := o .buildPrompt (prompt )
77+
78+ res , err := o .Client .Chat .Send (ctx , components.ChatRequest {
79+ Model : openrouter .String (o .Model ),
80+ Messages : []components.ChatMessages {
81+ components .CreateChatMessagesUser (components.ChatUserMessage {
82+ Content : components .CreateChatUserMessageContentStr (fullPrompt ),
83+ }),
4584 },
46- })
85+ }, nil )
4786 if err != nil {
4887 return nil , err
4988 }
50- if len ( resp . Choices ) == 0 {
89+ if res == nil || res . ChatResult == nil {
5190 return nil , errors .New ("no response from OpenRouter" )
5291 }
53- return resp .Choices [0 ].Message .Content .Text , nil
92+
93+ return firstChoiceText (res .ChatResult )
5494}
5595
56- // GenerateStream uses OpenRouter's streaming chat completion API
96+ // GenerateStream uses OpenRouter's streaming chat completion API.
5797func (o * OpenRouterLLM ) GenerateStream (ctx context.Context , prompt string ) (<- chan StreamChunk , error ) {
58- fullPrompt := prompt
59- if o .PromptPrefix != "" {
60- fullPrompt = o .PromptPrefix + "\n " + prompt
61- }
98+ fullPrompt := o .buildPrompt (prompt )
6299
63- stream , err := o .Client .CreateChatCompletionStream (ctx , openrouter.ChatCompletionRequest {
64- Model : o .Model ,
65- Messages : []openrouter.ChatCompletionMessage {
66- openrouter .UserMessage (fullPrompt ),
100+ res , err := o .Client .Chat .Send (ctx , components.ChatRequest {
101+ Model : openrouter .String (o .Model ),
102+ Stream : openrouter .Bool (true ),
103+ Messages : []components.ChatMessages {
104+ components .CreateChatMessagesUser (components.ChatUserMessage {
105+ Content : components .CreateChatUserMessageContentStr (fullPrompt ),
106+ }),
67107 },
68- })
108+ }, nil )
69109 if err != nil {
70110 return nil , err
71111 }
112+ if res == nil || res .EventStream == nil {
113+ return nil , errors .New ("no streaming response from OpenRouter" )
114+ }
72115
73116 ch := make (chan StreamChunk , 16 )
74117 go func () {
75118 defer close (ch )
76- defer stream .Close ()
119+ defer res .EventStream .Close ()
120+
77121 var sb strings.Builder
78- for {
79- resp , err := stream .Recv ()
80- if err != nil {
81- if errors .Is (err , io .EOF ) {
82- ch <- StreamChunk {Done : true , FullText : sb .String ()}
83- return
122+ for res .EventStream .Next () {
123+ event := res .EventStream .Value ()
124+ if event == nil {
125+ continue
126+ }
127+ chunk := event .Data
128+
129+ if chunk .Error != nil {
130+ ch <- StreamChunk {
131+ Done : true ,
132+ FullText : sb .String (),
133+ Err : fmt .Errorf ("openrouter stream error: %s" , chunk .Error .Message ),
84134 }
85- ch <- StreamChunk {Done : true , FullText : sb .String (), Err : err }
86135 return
87136 }
88- if len (resp .Choices ) > 0 {
89- delta := resp .Choices [0 ].Delta .Content
90- if delta != "" {
91- sb .WriteString (delta )
92- ch <- StreamChunk {Delta : delta }
93- }
137+ if len (chunk .Choices ) == 0 {
138+ continue
139+ }
140+
141+ delta , ok := chunk .Choices [0 ].Delta .Content .GetOrZero ()
142+ if ok && delta != "" {
143+ sb .WriteString (delta )
144+ ch <- StreamChunk {Delta : delta }
94145 }
95146 }
147+
148+ if err := res .EventStream .Err (); err != nil {
149+ ch <- StreamChunk {Done : true , FullText : sb .String (), Err : err }
150+ return
151+ }
152+ ch <- StreamChunk {Done : true , FullText : sb .String ()}
96153 }()
97154
98155 return ch , nil
99156}
100157
101- // getOpenRouterMimeType converts normalized MIME types to OpenRouter's expected format
158+ // getOpenRouterMimeType converts normalized MIME types to OpenRouter's expected format.
102159func getOpenRouterMimeType (mt string ) string {
103160 mt = strings .ToLower (strings .TrimSpace (mt ))
104161 switch {
@@ -125,10 +182,7 @@ func getOpenRouterMimeType(mt string) string {
125182}
126183
127184func (o * OpenRouterLLM ) GenerateWithFiles (ctx context.Context , prompt string , files []File ) (any , error ) {
128- fullPrompt := prompt
129- if o .PromptPrefix != "" {
130- fullPrompt = o .PromptPrefix + "\n " + prompt
131- }
185+ fullPrompt := o .buildPrompt (prompt )
132186
133187 // Separate files by type
134188 var textFiles []File
@@ -159,36 +213,58 @@ func (o *OpenRouterLLM) GenerateWithFiles(ctx context.Context, prompt string, fi
159213 textPrompt = combinePromptWithFiles (fullPrompt , textFiles )
160214 }
161215
162- var msg openrouter. ChatCompletionMessage
216+ var content components. ChatUserMessageContent
163217
164- // Handle image files using UserMessageWithImage
165- if len (imageFiles ) > 0 {
166- // Use first image with UserMessageWithImage
218+ switch {
219+ case len (imageFiles ) > 0 :
220+ // Only the first image is attached, matching the original function's behavior.
167221 firstImage := imageFiles [0 ]
168222 encoded := base64 .StdEncoding .EncodeToString (firstImage .Data )
169- dataURL := fmt .Sprintf ("data:%s;base64,%s" , getOpenRouterMimeType (normalizeMIME (firstImage .Name , firstImage .MIME )), encoded )
170- msg = openrouter .UserMessageWithImage (textPrompt , dataURL )
223+ dataURL := fmt .Sprintf (
224+ "data:%s;base64,%s" ,
225+ getOpenRouterMimeType (normalizeMIME (firstImage .Name , firstImage .MIME )),
226+ encoded ,
227+ )
228+ content = components .CreateChatUserMessageContentArrayOfChatContentItems ([]components.ChatContentItems {
229+ components .CreateChatContentItemsText (components.ChatContentText {Text : textPrompt }),
230+ components .CreateChatContentItemsImageURL (components.ChatContentImage {
231+ ImageURL : components.ChatContentImageImageURL {URL : dataURL },
232+ }),
233+ })
171234
172- // For additional images, we would need to append to the message
173- // This depends on go-openrouter's content structure
174- } else if len (pdfFiles ) > 0 {
175- // Handle PDF files using UserMessageWithPDF
235+ case len (pdfFiles ) > 0 :
176236 firstPDF := pdfFiles [0 ]
177- msg = openrouter .UserMessageWithPDF (textPrompt , firstPDF .Name , string (firstPDF .Data ))
178- } else {
179- // Fallback to text-only message
180- msg = openrouter .UserMessage (textPrompt )
237+ encoded := base64 .StdEncoding .EncodeToString (firstPDF .Data )
238+ dataURL := fmt .Sprintf ("data:application/pdf;base64,%s" , encoded )
239+ filename := firstPDF .Name
240+ content = components .CreateChatUserMessageContentArrayOfChatContentItems ([]components.ChatContentItems {
241+ components .CreateChatContentItemsText (components.ChatContentText {Text : textPrompt }),
242+ components .CreateChatContentItemsFile (components.ChatContentFile {
243+ File : components.File {
244+ FileData : & dataURL ,
245+ Filename : & filename ,
246+ },
247+ }),
248+ })
249+
250+ default :
251+ content = components .CreateChatUserMessageContentStr (textPrompt )
181252 }
182253
183- resp , err := o .Client .CreateChatCompletion (ctx , openrouter.ChatCompletionRequest {
184- Model : o .Model ,
185- Messages : []openrouter.ChatCompletionMessage {msg },
186- })
254+ res , err := o .Client .Chat .Send (ctx , components.ChatRequest {
255+ Model : openrouter .String (o .Model ),
256+ Messages : []components.ChatMessages {
257+ components .CreateChatMessagesUser (components.ChatUserMessage {
258+ Content : content ,
259+ }),
260+ },
261+ }, nil )
187262 if err != nil {
188263 return nil , err
189264 }
190- if len ( resp . Choices ) == 0 {
265+ if res == nil || res . ChatResult == nil {
191266 return nil , errors .New ("no response from OpenRouter" )
192267 }
193- return resp .Choices [0 ].Message .Content .Text , nil
268+
269+ return firstChoiceText (res .ChatResult )
194270}
0 commit comments