()
| 498 | """ |
| 499 | |
| 500 | async def wrapped_generator(): |
| 501 | span = trace.get_current_span() |
| 502 | if span is not None and span.is_recording(): |
| 503 | last_time = None |
| 504 | count = 0 |
| 505 | try: |
| 506 | async for chunk in original_generator: |
| 507 | last_time = time.time() |
| 508 | # 首包捕获 |
| 509 | if count == 0 and span is not None and span.is_recording(): |
| 510 | last_time = time.time() |
| 511 | span.add_event("first_chunk", {"time": last_time}) |
| 512 | count += 1 |
| 513 | yield chunk |
| 514 | except Exception as e: |
| 515 | # 错误捕获 |
| 516 | if span is not None and span.is_recording(): |
| 517 | span.add_event("stream_error", {"time": time.time(), "error": str(e), "total_chunk": count}) |
| 518 | span.record_exception(e) |
| 519 | span.set_status({"code": "ERROR", "description": str(e)}) |
| 520 | raise |
| 521 | finally: |
| 522 | # 尾包捕获 |
| 523 | if span is not None and span.is_recording() and count > 0: |
| 524 | span.add_event("last_chunk", {"time": last_time, "total_chunk": count}) |
| 525 | api_server_logger.debug(f"release: {connection_semaphore.status()}") |
| 526 | connection_semaphore.release() |
| 527 | else: |
| 528 | try: |
| 529 | async for chunk in original_generator: |
| 530 | yield chunk |
| 531 | finally: |
| 532 | api_server_logger.debug(f"release: {connection_semaphore.status()}") |
| 533 | connection_semaphore.release() |
| 534 | |
| 535 | return wrapped_generator |
| 536 |
no test coverage detected