11package proxy
22
33import (
4+ "bufio"
5+ "bytes"
46 "context"
57 "encoding/json"
8+ "errors"
69 "fmt"
710 "io"
811 "net/http"
@@ -11,7 +14,6 @@ import (
1114 "github.com/bricks-cloud/bricksllm/internal/telemetry"
1215 "github.com/bricks-cloud/bricksllm/internal/util"
1316 "github.com/gin-gonic/gin"
14-
1517 responsesOpenai "github.com/openai/openai-go/responses"
1618 goopenai "github.com/sashabaranov/go-openai"
1719)
@@ -103,8 +105,21 @@ func getResponsesHandler(prod, private bool, client http.Client, e estimator) gi
103105 }
104106
105107 c .Set ("costInUsd" , cost )
106- c .Set ("promptTokenCount" , resp .Usage .InputTokens )
107- c .Set ("completionTokenCount" , resp .Usage .OutputTokens )
108+
109+ promptTokenCount , err := int64ToInt (resp .Usage .InputTokens )
110+ if err != nil {
111+ telemetry .Incr ("bricksllm.proxy.get_responses_handler.int64_to_int_error" , nil , 1 )
112+ logError (log , "error when converting int64 to int for prompt token count" , prod , err )
113+ }
114+
115+ completionTokenCount , err := int64ToInt (resp .Usage .OutputTokens )
116+ if err != nil {
117+ telemetry .Incr ("bricksllm.proxy.get_responses_handler.int64_to_int_error" , nil , 1 )
118+ logError (log , "error when converting int64 to int for completion token count" , prod , err )
119+ }
120+
121+ c .Set ("promptTokenCount" , promptTokenCount )
122+ c .Set ("completionTokenCount" , completionTokenCount )
108123
109124 c .Data (res .StatusCode , "application/json" , bytes )
110125 return
@@ -134,9 +149,114 @@ func getResponsesHandler(prod, private bool, client http.Client, e estimator) gi
134149 return
135150 }
136151
137- // handle streaming response
152+ buffer := bufio .NewReader (res .Body )
153+ content := ""
154+ streamingResponse := [][]byte {}
155+
156+ streamCost := 0.0
157+ streamPromptTokenCount := 0
158+ streamCompletionTokenCount := 0
159+
160+ defer func () {
161+ c .Set ("content" , content )
162+ c .Set ("streaming_response" , bytes .Join (streamingResponse , []byte {'\n' }))
163+
164+ c .Set ("costInUsd" , streamCost )
165+ c .Set ("promptTokenCount" , streamPromptTokenCount )
166+ c .Set ("completionTokenCount" , streamCompletionTokenCount )
167+ }()
168+
138169 telemetry .Incr ("bricksllm.proxy.get_responses_handler.streaming_requests" , nil , 1 )
139- // TODO: implement the actual streaming logic here
170+ c .Stream (func (w io.Writer ) bool {
171+ raw , err := buffer .ReadBytes ('\n' )
172+ if err != nil {
173+ if err == io .EOF {
174+ return false
175+ }
176+
177+ if errors .Is (err , context .DeadlineExceeded ) {
178+ telemetry .Incr ("bricksllm.proxy.get_responses_handler.context_deadline_exceeded_error" , nil , 1 )
179+ logError (log , "context deadline exceeded when reading bytes from openai responses api response" , prod , err )
180+
181+ return false
182+ }
183+
184+ telemetry .Incr ("bricksllm.proxy.get_responses_handler.read_bytes_error" , nil , 1 )
185+ logError (log , "error when reading bytes from openai responses api response" , prod , err )
186+
187+ apiErr := & goopenai.ErrorResponse {
188+ Error : & goopenai.APIError {
189+ Type : "bricksllm_error" ,
190+ Message : err .Error (),
191+ },
192+ }
193+
194+ errBytes , err := json .Marshal (apiErr )
195+ if err != nil {
196+ telemetry .Incr ("bricksllm.proxy.get_responses_handler.json_marshal_error" , nil , 1 )
197+ logError (log , "error when marshalling bytes for openai streaming responses api error response" , prod , err )
198+ return false
199+ }
200+
201+ c .SSEvent ("" , string (errBytes ))
202+ c .SSEvent ("" , " [DONE]" )
203+ return false
204+ }
205+
206+ streamingResponse = append (streamingResponse , raw )
207+
208+ noSpaceLine := bytes .TrimSpace (raw )
209+ if ! bytes .HasPrefix (noSpaceLine , headerData ) {
210+ return true
211+ }
212+
213+ noPrefixLine := bytes .TrimPrefix (noSpaceLine , headerData )
214+ c .SSEvent ("" , " " + string (noPrefixLine ))
215+
216+ if string (noPrefixLine ) == "[DONE]" {
217+ return false
218+ }
219+
220+ responsesStreamResp := & responsesOpenai.ResponseStreamEventUnion {}
221+ err = json .Unmarshal (noPrefixLine , responsesStreamResp )
222+ if err != nil {
223+ telemetry .Incr ("bricksllm.proxy.get_responses_handler.completion_response_unmarshall_error" , nil , 1 )
224+ logError (log , "error when unmarshalling openai responses api stream response" , prod , err )
225+ }
226+ if err == nil {
227+ textDelta := responsesStreamResp .AsResponseOutputTextDelta ().Delta
228+ if len (textDelta ) > 0 {
229+ content += textDelta
230+ }
231+
232+ if responsesStreamResp .Response .Status == "completed" {
233+ streamCost , err = e .EstimateResponseApiTotalCost (model , responsesStreamResp .Response .Usage )
234+ if err != nil {
235+ telemetry .Incr ("bricksllm.proxy.get_chat_completion_handler.estimate_total_cost_error" , nil , 1 )
236+ logError (log , "error when estimating openai cost" , prod , err )
237+ }
238+ streamPromptTokenCount , err = int64ToInt (responsesStreamResp .Response .Usage .InputTokens )
239+ if err != nil {
240+ telemetry .Incr ("bricksllm.proxy.get_responses_handler.int64_to_int_error" , nil , 1 )
241+ logError (log , "error when converting int64 to int for prompt token count" , prod , err )
242+ }
243+
244+ streamCompletionTokenCount , err = int64ToInt (responsesStreamResp .Response .Usage .OutputTokens )
245+ if err != nil {
246+ telemetry .Incr ("bricksllm.proxy.get_responses_handler.int64_to_int_error" , nil , 1 )
247+ logError (log , "error when converting int64 to int for completion token count" , prod , err )
248+ }
249+ }
250+ }
251+ return true
252+ })
140253 telemetry .Timing ("bricksllm.proxy.get_chat_completion_handler.streaming_latency" , time .Since (start ), nil , 1 )
141254 }
142255}
256+
257+ func int64ToInt (src int64 ) (int , error ) {
258+ if src > int64 (int (^ uint (0 )>> 1 )) {
259+ return 0 , fmt .Errorf ("int64 value %d overflows int" , src )
260+ }
261+ return int (src ), nil
262+ }
0 commit comments