Skip to content

Commit

Permalink
🎨 cleanup
Browse files Browse the repository at this point in the history
Signed-off-by: Joe Runde <[email protected]>
  • Loading branch information
joerunde committed Dec 11, 2024
1 parent d91700d commit a0ccd3c
Showing 1 changed file with 3 additions and 6 deletions.
9 changes: 3 additions & 6 deletions src/vllm_tgis_adapter/tgis_utils/logs.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@

from vllm import PromptType, RequestOutput, SamplingParams
from vllm.engine.protocol import EngineClient
from vllm.sequence import RequestMetrics

logger = init_logger(__name__)

Expand Down Expand Up @@ -97,7 +98,7 @@ async def generate_with_logging(*args, **kwargs) -> AsyncGenerator[RequestOutput
if last:
# Log the response
with suppress(BaseException):
_new_log_response(
_log_response(
request_id=request_id,
correlation_id=correlation_id,
response=last,
Expand All @@ -108,10 +109,6 @@ async def generate_with_logging(*args, **kwargs) -> AsyncGenerator[RequestOutput
engine.generate = generate_with_logging


if TYPE_CHECKING:
from vllm.sequence import RequestMetrics


def _log_error(request_id: str, correlation_id: str, exception_str: str) -> None:
logger.error(
"Request failed: request_id=%s correlation_id=%s error=%s",
Expand Down Expand Up @@ -143,7 +140,7 @@ def _log_request(
)


def _new_log_response(
def _log_response(
request_id: str,
correlation_id: str,
response: RequestOutput,
Expand Down

0 comments on commit a0ccd3c

Please sign in to comment.