@@ -181,14 +181,12 @@ impl BackendModel for LLamaEdgeApiServer {
181
181
let ( wasm_module, listen_addr) = if let Some ( old_model) = & old_model {
182
182
let listen_addr = load_model_options. override_server_address . clone ( ) . map_or (
183
183
old_model. listen_addr ,
184
- |addr| {
185
- match std:: net:: TcpListener :: bind ( & addr) {
186
- Ok ( listener) => listener. local_addr ( ) . unwrap ( ) ,
187
- Err ( _) => {
188
- eprintln ! ( "Failed to start the model on address {}" , addr) ;
189
- eprintln ! ( "Using the previous one {}" , old_model. listen_addr) ;
190
- old_model. listen_addr
191
- }
184
+ |addr| match std:: net:: TcpListener :: bind ( & addr) {
185
+ Ok ( listener) => listener. local_addr ( ) . unwrap ( ) ,
186
+ Err ( _) => {
187
+ eprintln ! ( "Failed to start the model on address {}" , addr) ;
188
+ eprintln ! ( "Using the previous one {}" , old_model. listen_addr) ;
189
+ old_model. listen_addr
192
190
}
193
191
} ,
194
192
) ;
@@ -205,10 +203,23 @@ impl BackendModel for LLamaEdgeApiServer {
205
203
( old_model. wasm_module . clone ( ) , listen_addr)
206
204
} else {
207
205
let addr = std:: env:: var ( "MOLY_API_SERVER_ADDR" ) . unwrap_or ( "localhost:0" . to_string ( ) ) ;
208
- let new_addr = std:: net:: TcpListener :: bind ( & addr)
209
- . unwrap ( )
210
- . local_addr ( )
211
- . unwrap ( ) ;
206
+
207
+ let listen_addr = load_model_options
208
+ . override_server_address
209
+ . clone ( )
210
+ . map ( |addr| match std:: net:: TcpListener :: bind ( & addr) {
211
+ Ok ( listener) => Some ( listener. local_addr ( ) . unwrap ( ) ) ,
212
+ Err ( _) => None ,
213
+ } )
214
+ . flatten ( ) ;
215
+
216
+ let new_addr = match listen_addr {
217
+ Some ( addr) => addr,
218
+ None => {
219
+ let listener = std:: net:: TcpListener :: bind ( & addr) . unwrap ( ) ;
220
+ listener. local_addr ( ) . unwrap ( )
221
+ }
222
+ } ;
212
223
213
224
( Module :: from_bytes ( None , WASM ) . unwrap ( ) , new_addr)
214
225
} ;
@@ -312,15 +323,24 @@ impl BackendModel for LLamaEdgeApiServer {
312
323
313
324
async_rt. spawn ( async move {
314
325
let request_body = serde_json:: to_string ( & data) . unwrap ( ) ;
315
- let resp = reqwest:: ClientBuilder :: new ( )
326
+ let request = reqwest:: ClientBuilder :: new ( )
316
327
. no_proxy ( )
317
328
. build ( )
318
329
. unwrap ( )
319
330
. post ( url)
320
- . body ( request_body)
321
- . send ( )
322
- . await
323
- . map_err ( |e| anyhow ! ( e) ) ;
331
+ . body ( request_body) ;
332
+
333
+ let resp = tokio:: select! {
334
+ res = request. send( ) => Some ( res. map_err( |e| anyhow!( e) ) ) ,
335
+ _ = cancel. recv( ) => None ,
336
+ } ;
337
+
338
+ let Some ( resp) = resp else {
339
+ let _ = tx. send ( Ok ( ChatResponse :: ChatResponseChunk ( stop_chunk (
340
+ StopReason :: Stop ,
341
+ ) ) ) ) ;
342
+ return ;
343
+ } ;
324
344
325
345
match resp {
326
346
Ok ( resp) => {
@@ -351,8 +371,18 @@ impl BackendModel for LLamaEdgeApiServer {
351
371
StopReason :: Stop ,
352
372
) ) ) ) ;
353
373
} else {
354
- let resp: Result < ChatResponseData , anyhow:: Error > =
355
- resp. json ( ) . await . map_err ( |e| anyhow ! ( e) ) ;
374
+ let resp = tokio:: select! {
375
+ res = resp. json:: <ChatResponseData >( ) => Some ( res. map_err( |e| anyhow!( e) ) ) ,
376
+ _ = cancel. recv( ) => None ,
377
+ } ;
378
+
379
+ let Some ( resp) = resp else {
380
+ let _ = tx. send ( Ok ( ChatResponse :: ChatResponseChunk ( stop_chunk (
381
+ StopReason :: Stop ,
382
+ ) ) ) ) ;
383
+ return ;
384
+ } ;
385
+
356
386
let _ = tx. send ( resp. map ( ChatResponse :: ChatFinalResponseData ) ) ;
357
387
}
358
388
}
0 commit comments