| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | from starlette.requests import Request |
| 6 | |
| 7 | from vllm.entrypoints.generate.base.serving import GenerateBaseServing |
| 8 | from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest |
| 9 | from vllm.entrypoints.openai.completion.protocol import CompletionRequest |
| 10 | from vllm.entrypoints.openai.responses.protocol import ResponsesRequest |
| 11 | |
| 12 | |
| 13 | def _raw_request(headers: dict[str, str]) -> Request: |
| 14 | return Request( |
| 15 | { |
| 16 | "type": "http", |
| 17 | "method": "POST", |
| 18 | "path": "/v1/chat/completions", |
| 19 | "headers": [ |
| 20 | (key.lower().encode("latin-1"), value.encode("latin-1")) |
| 21 | for key, value in headers.items() |
| 22 | ], |
| 23 | } |
| 24 | ) |
| 25 | |
| 26 | |
| 27 | @pytest.mark.parametrize( |
| 28 | "openai_request", |
| 29 | [ |
| 30 | ChatCompletionRequest( |
| 31 | model="test-model", |
| 32 | messages=[{"role": "user", "content": "hi"}], |
| 33 | ), |
| 34 | CompletionRequest(model="test-model", prompt="hi"), |
| 35 | ResponsesRequest(model="test-model", input="hi"), |
| 36 | ], |
| 37 | ) |
| 38 | def test_get_session_id_accepts_body_field(openai_request): |
| 39 | openai_request.session_id = "body-session" |
| 40 | |
| 41 | session_id = GenerateBaseServing._get_session_id( |
| 42 | openai_request, |
| 43 | _raw_request({"X-Session-ID": "header-session"}), |
| 44 | ) |
| 45 | |
| 46 | assert session_id == "body-session" |
| 47 | |
| 48 | |
| 49 | def test_get_session_id_accepts_session_header(): |
| 50 | request = CompletionRequest(model="test-model", prompt="hi") |
| 51 | |
| 52 | session_id = GenerateBaseServing._get_session_id( |
| 53 | request, |
| 54 | _raw_request({"X-Session-ID": "header-session"}), |
| 55 | ) |
| 56 | |
| 57 | assert session_id == "header-session" |
| 58 | |
| 59 | |
| 60 | def test_get_session_id_ignores_correlation_header(): |
| 61 | request = CompletionRequest( |
| 62 | model="test-model", |
| 63 | prompt="hi", |
| 64 | vllm_xargs={"session_id": "xargs-session"}, |
| 65 | ) |
| 66 | |
| 67 | session_id = GenerateBaseServing._get_session_id( |
| 68 | request, |
| 69 | _raw_request({"X-Correlation-ID": "correlation-session"}), |
| 70 | ) |
| 71 | |
| 72 | assert session_id == "xargs-session" |
| 73 | |
| 74 | |
| 75 | def test_get_session_id_keeps_vllm_xargs_as_compatibility_fallback(): |
| 76 | request = CompletionRequest( |
| 77 | model="test-model", |
| 78 | prompt="hi", |
| 79 | vllm_xargs={"session_id": "xargs-session"}, |
| 80 | ) |
| 81 | |
| 82 | session_id = GenerateBaseServing._get_session_id(request, None) |
| 83 | |
| 84 | assert session_id == "xargs-session" |
| 85 | |
| 86 | |
| 87 | def test_get_session_id_ignores_empty_and_non_string_values(): |
| 88 | request = CompletionRequest( |
| 89 | model="test-model", |
| 90 | prompt="hi", |
| 91 | session_id="", |
| 92 | vllm_xargs={"session_id": 7}, |
| 93 | ) |
| 94 | |
| 95 | session_id = GenerateBaseServing._get_session_id(request, None) |
| 96 | |
| 97 | assert session_id is None |