feat: integrated open response api w and w/0 streaming

+9 -1

server/api.py

··· 94 94 95 95 global _messages 96 96 97 - return await runtime.backend.generate_response_chat(request) 97 + if request.stream: 98 + # Streaming response 99 + return StreamingResponse( 100 + runtime.backend.generate_response_chat_stream(request), 101 + media_type="text/plain", 102 + headers={"Cache-Control": "no-cache"}, 103 + ) 104 + else: 105 + return await runtime.backend.generate_response_chat(request)

+265 -20

server/backend/mlx.py

··· 24 24 _model_cache: Dict[str, MLXRunner] = {} 25 25 _default_max_tokens: Optional[int] = None # Use dynamic model-aware limits by default 26 26 _current_model_path: Optional[str] = None 27 + # Store generated responses for follow-up support (previous_response_id) 28 + _responses: Dict[str, ResponsesResponse] = {} 27 29 28 30 29 31 def download_model(model_name: str): ··· 200 202 return [{"role": msg.role, "content": msg.content} for msg in messages] 201 203 202 204 205 + def _prepend_previous_response(user_input: str, prev_id: Optional[str]) -> str: 206 + """If prev_id points to a stored response, prepend its output text as context.""" 207 + if not prev_id: 208 + return user_input 209 + prev = _responses.get(prev_id) 210 + if not prev or not getattr(prev, "output", None): 211 + return user_input 212 + prev_text_parts: List[str] = [] 213 + for out in prev.output: 214 + for c in out.get("content", []): 215 + if c.get("type") == "output_text": 216 + prev_text_parts.append(c.get("text", "")) 217 + if prev_text_parts: 218 + return "\n".join(prev_text_parts) + "\n\n" + user_input 219 + return user_input 220 + 221 + 222 + def _calc_usage(runner: MLXRunner, input_text: str, generated_text: str) -> Dict[str, int]: 223 + """Calculate token usage using the runner tokenizer; fall back to zeros on error.""" 224 + try: 225 + input_tokens = len(runner.tokenizer.encode(input_text)) 226 + output_tokens = len(runner.tokenizer.encode(generated_text)) 227 + return {"input_tokens": input_tokens, "output_tokens": output_tokens} 228 + except Exception: 229 + return {"input_tokens": 0, "output_tokens": 0} 230 + 231 + 232 + def _store_response( 233 + response_id: str, 234 + created: int, 235 + completed_at: Optional[int], 236 + model: str, 237 + status: str, 238 + output: List[Dict[str, Any]], 239 + usage: Dict[str, int], 240 + metrics: Optional[Dict[str, Any]] = None, 241 + error: Optional[Dict[str, Any]] = None, 242 + ) -> ResponsesResponse: 243 + """Create a ResponsesResponse, attach metrics to metadata and store it in `_responses`.""" 244 + resp = ResponsesResponse( 245 + id=response_id, 246 + created_at=created, 247 + completed_at=completed_at, 248 + model=model, 249 + status=status, 250 + object="response", 251 + error=error, 252 + output=output, 253 + usage=usage, 254 + ) 255 + if metrics: 256 + try: 257 + resp.metadata["metrics"] = metrics 258 + except Exception: 259 + pass 260 + try: 261 + _responses[response_id] = resp 262 + except Exception: 263 + pass 264 + return resp 265 + 266 + 203 267 def count_tokens(text: str) -> int: 204 268 """Rough token count estimation.""" 205 269 return int(len(text.split()) * 1.3) # Approximation, convert to int 206 270 207 271 208 - async def generate_response_chat(request: ResponsesRequest): 209 - """Generate chat responses""" 272 + async def generate_response_chat_stream( 273 + request: ResponsesRequest 274 + ) -> AsyncGenerator[str, None]: 275 + """Generate streaming chat responses for Responses API.""" 210 276 211 277 model = request.model or "mlx-community/gpt-oss-20b-MXFP4-Q4" 212 - input = request.input or "" 278 + user_input = request.input or "" 213 279 response_id = f"resp-{uuid.uuid4()}" 214 280 msg_id = f"msg_{uuid.uuid4()}" 215 281 created = int(time.time()) 216 282 runner = get_or_load_model(model) 217 - generated_text = runner.generate_batch( 218 - prompt=input, 219 - max_tokens=runner.get_effective_max_tokens(request.max_output_tokens), 220 - temperature=request.temperature or 1, 221 - top_p=request.top_p or 1, 222 - use_chat_template=True, # Already applied in _format_conversation 223 - ) 283 + metrics = None 284 + # If a previous_response_id is provided, prepend its text to the prompt 285 + prev_id = getattr(request, "previous_response_id", None) 286 + user_input = _prepend_previous_response(user_input, prev_id) 287 + 288 + # Calculate input tokens once 289 + input_tokens = len(runner.tokenizer.encode(user_input)) 290 + 291 + # Initial chunk 292 + initial_chunk = { 293 + "id": response_id, 294 + "object": "response.chunk", 295 + "created_at": created, 296 + "model": model, 297 + "status": "in_progress", 298 + "output": [ 299 + { 300 + "type": "message", 301 + "id": msg_id, 302 + "status": "in_progress", 303 + "role": "assistant", 304 + "content": [], 305 + } 306 + ], 307 + "usage": {"input_tokens": input_tokens, "output_tokens": 0}, 308 + } 309 + yield f"data: {json.dumps(initial_chunk)}\n\n" 310 + 311 + # Stream tokens 312 + accumulated_text = "" 313 + output_tokens = 0 314 + try: 315 + for token in runner.generate_streaming( 316 + prompt=user_input, 317 + max_tokens=runner.get_effective_max_tokens(request.max_output_tokens), 318 + temperature=request.temperature or 1, 319 + top_p=request.top_p or 1, 320 + use_chat_template=True, 321 + ): 322 + if isinstance(token, GenerationMetrics): 323 + metrics = token 324 + continue 325 + 326 + accumulated_text += token 327 + output_tokens += 1 # Each yield is one token 328 + 329 + chunk = { 330 + "id": response_id, 331 + "object": "response.chunk", 332 + "created_at": created, 333 + "model": model, 334 + "status": "in_progress", 335 + "output": [ 336 + { 337 + "type": "message", 338 + "id": msg_id, 339 + "status": "in_progress", 340 + "role": "assistant", 341 + "content": [ 342 + { 343 + "type": "output_text", 344 + "text": token, 345 + "annotations": [], 346 + } 347 + ], 348 + } 349 + ], 350 + "usage": {"input_tokens": input_tokens, "output_tokens": output_tokens}, 351 + } 352 + yield f"data: {json.dumps(chunk)}\n\n" 353 + 354 + except Exception as e: 355 + error_chunk = { 356 + "id": response_id, 357 + "object": "response.chunk", 358 + "created_at": created, 359 + "model": model, 360 + "status": "failed", 361 + "error": {"message": str(e), "type": "internal_error"}, 362 + "output": [], 363 + "usage": {"input_tokens": input_tokens, "output_tokens": output_tokens}, 364 + } 365 + yield f"data: {json.dumps(error_chunk)}\n\n" 366 + return 367 + 368 + # Final chunk 224 369 completed_at = int(time.time()) 225 - return ResponsesResponse( 226 - id=response_id, 227 - created_at=created, 228 - completed_at=completed_at, 229 - model=model, 230 - status="completed", 231 - object="response", 232 - output=[ 370 + # Build final chunk with accumulated text and store response for follow-ups 371 + final_chunk = { 372 + "id": response_id, 373 + "object": "response.chunk", 374 + "created_at": created, 375 + "completed_at": completed_at, 376 + "model": model, 377 + "status": "completed", 378 + "output": [ 233 379 { 234 380 "type": "message", 235 381 "id": msg_id, ··· 238 384 "content": [ 239 385 { 240 386 "type": "output_text", 241 - "text": generated_text, 387 + "text": "", 242 388 "annotations": [], 243 389 } 244 390 ], 245 391 } 246 392 ], 247 - usage={"input_tokens": 36}, 393 + "usage": {"input_tokens": input_tokens, "output_tokens": output_tokens}, 394 + } 395 + # Store and return a typed ResponsesResponse for follow-ups 396 + metrics_obj = None 397 + if metrics: 398 + metrics_obj = { 399 + "ttft_ms": metrics.ttft_ms, 400 + "total_tokens": metrics.total_tokens, 401 + "tokens_per_second": metrics.tokens_per_second, 402 + "total_latency_s": metrics.total_latency_s, 403 + } 404 + final_chunk["metrics"] = metrics_obj 405 + 406 + _store_response( 407 + response_id=response_id, 408 + created=created, 409 + completed_at=completed_at, 410 + model=model, 411 + status="completed", 412 + output=final_chunk["output"], 413 + usage={"input_tokens": input_tokens, "output_tokens": output_tokens}, 414 + metrics=metrics_obj, 248 415 ) 416 + yield f"data: {json.dumps(final_chunk)}\n\n" 417 + yield "data: [DONE]\n\n" 418 + 419 + 420 + async def generate_response_chat(request: ResponsesRequest): 421 + """Generate chat responses""" 422 + 423 + model = request.model or "mlx-community/gpt-oss-20b-MXFP4-Q4" 424 + user_input = request.input or "" 425 + response_id = f"resp-{uuid.uuid4()}" 426 + msg_id = f"msg_{uuid.uuid4()}" 427 + created = int(time.time()) 428 + runner = get_or_load_model(model) 429 + 430 + # If a previous_response_id is provided, prepend its text to the prompt 431 + prev_id = getattr(request, "previous_response_id", None) 432 + user_input = _prepend_previous_response(user_input, prev_id) 433 + 434 + metrics_obj = None 435 + try: 436 + start_time = time.time() 437 + generated_text = runner.generate_batch( 438 + prompt=user_input, 439 + max_tokens=runner.get_effective_max_tokens(request.max_output_tokens), 440 + temperature=request.temperature or 1, 441 + top_p=request.top_p or 1, 442 + use_chat_template=True, 443 + ) 444 + 445 + # Metrics for batch generation (approximate) 446 + generation_time = time.time() - start_time 447 + 448 + completed_at = int(time.time()) 449 + status = "completed" 450 + error = None 451 + 452 + # Calculate token usage 453 + usage = _calc_usage(runner, user_input, generated_text) 454 + output_tokens = usage.get("output_tokens", 0) 455 + metrics_obj = { 456 + "ttft_ms": generation_time * 1000.0, 457 + "total_tokens": output_tokens, 458 + "tokens_per_second": (output_tokens / generation_time) if generation_time > 0 else 0.0, 459 + "total_latency_s": generation_time, 460 + } 461 + 462 + except Exception as e: 463 + completed_at = None 464 + status = "failed" 465 + error = {"message": str(e), "type": "internal_error"} 466 + generated_text = "" 467 + usage = {"input_tokens": 0, "output_tokens": 0} 468 + 469 + output_block = [ 470 + { 471 + "type": "message", 472 + "id": msg_id, 473 + "status": "completed" if status == "completed" else "failed", 474 + "role": "assistant", 475 + "content": [ 476 + {"type": "output_text", "text": generated_text, "annotations": []} 477 + ], 478 + } 479 + ] if status == "completed" else [] 480 + 481 + resp = _store_response( 482 + response_id=response_id, 483 + created=created, 484 + completed_at=completed_at, 485 + model=model, 486 + status=status, 487 + output=output_block, 488 + usage=usage, 489 + metrics=(metrics_obj if status == "completed" else None), 490 + error=error, 491 + ) 492 + 493 + return resp

+6 -4

tiles/src/runtime/mlx.rs

··· 448 448 println!( 449 449 "{}", 450 450 format!( 451 - "\n{} {:.1} tok/s | {} tokens | {:.0}ms TTFT", 451 + "\n{} {:.1} tok/s | {} tokens | {:.0}s TTFT", 452 452 "💡".yellow(), 453 453 bench_metrics.total_tokens as f64 454 454 / bench_metrics.total_latency_s, 455 455 bench_metrics.total_tokens, 456 - bench_metrics.ttft_ms 456 + bench_metrics.ttft_ms / 1000.0 457 457 ) 458 458 .dimmed() 459 459 ); ··· 536 536 537 537 let body = json!({ 538 538 "model": model_name, 539 + "input": input, 539 540 "chat_start": chat_start, 540 541 "stream": true, 541 542 "python_code": python_code, 542 543 "messages": [{"role": "assistant", "content": g_reply}, {"role": "user", "content": input}] 543 544 }); 544 545 let res = client 545 - .post("http://127.0.0.1:6969/v1/chat/completions") 546 + .post("http://127.0.0.1:6969/v1/responses") 546 547 .json(&body) 547 548 .send() 548 549 .await ··· 573 574 574 575 // Parse JSON 575 576 let v: Value = serde_json::from_str(data).unwrap(); 577 + // println!("{:?}", v); 576 578 // Check for metrics in the response 577 579 if let Some(metrics_obj) = v.get("metrics") { 578 580 metrics = serde_json::from_value(metrics_obj.clone()).ok(); 579 581 } 580 - if let Some(delta) = v["choices"][0]["delta"]["content"].as_str() { 582 + if let Some(delta) = v["output"][0]["content"][0]["text"].as_str() { 581 583 accumulated.push_str(delta); 582 584 if !run_args.memory && delta.contains("**[Answer]**") { 583 585 is_answer_start = true;

Configure Feed

Configure Feed