test: add logprobs handling tests for chat completions

Author: NullPointerDepressiveDisorder

tests/unit/test_openai_compat.py

@@ -123,3 +123,86 @@ def test_generate_chat_reasoning_fallback() -> None:
         res = asyncio.run(backend.generate(prompt))
         assert res.text == "Just thinking"
         assert res.tokens == ["Just", "thinking"]
+
+
+def test_generate_chat_with_logprobs() -> None:
+    backend = OpenAICompatBackend(base_url="http://127.0.0.1:8000", model_id="dummy", chat=True)
+    prompt = Prompt(id="p1", text="Hello", max_tokens=10)
+
+    mock_response = httpx.Response(
+        200,
+        json={
+            "choices": [
+                {
+                    "message": {"content": "world"},
+                    "logprobs": {
+                        "content": [
+                            {
+                                "token": "world",
+                                "logprob": -0.1,
+                                "top_logprobs": [
+                                    {"token": "world", "logprob": -0.1},
+                                    {"token": "earth", "logprob": -2.0},
+                                ],
+                            }
+                        ]
+                    },
+                }
+            ],
+            "usage": {"completion_tokens": 1},
+        },
+        request=httpx.Request("POST", "http://127.0.0.1:8000/v1/chat/completions"),
+    )
+
+    with patch("httpx.AsyncClient.post", return_value=mock_response):
+        res = asyncio.run(backend.generate(prompt))
+        assert res.text == "world"
+        assert res.tokens == ["world"]
+        assert res.logprobs == [-0.1]
+        # top_logprobs are sorted by token: earth (-2.0), world (-0.1)
+        assert res.distributions == [[-2.0, -0.1]]
+        assert res.distribution_metadata == [{"id_0": "earth", "id_1": "world"}]
+
+
+def test_generate_chat_with_logprobs_nan_and_missing() -> None:
+    backend = OpenAICompatBackend(base_url="http://127.0.0.1:8000", model_id="dummy", chat=True)
+    prompt = Prompt(id="p1", text="Hello", max_tokens=10)
+
+    mock_response = httpx.Response(
+        200,
+        content=b"{}",
+        request=httpx.Request("POST", "http://127.0.0.1:8000/v1/chat/completions"),
+    )
+
+    # Use patch to return NaN for the raw logprob values to avoid JSON serialization issues in httpx.Response
+    with (
+        patch("httpx.AsyncClient.post", return_value=mock_response),
+        patch(
+            "httpx.Response.json",
+            return_value={
+                "choices": [
+                    {
+                        "message": {"content": "world"},
+                        "logprobs": {
+                            "content": [
+                                {
+                                    "token": "world",
+                                    "logprob": float("nan"),
+                                    "top_logprobs": [
+                                        {"token": "world", "logprob": float("nan")},
+                                        {"token": "earth"},
+                                    ],
+                                }
+                            ]
+                        },
+                    }
+                ],
+                "usage": {"completion_tokens": 1},
+            },
+        ),
+    ):
+        res = asyncio.run(backend.generate(prompt))
+        assert res.logprobs == [-9999.0]
+        # top_logprobs sorted by token: earth (-9999.0), world (-9999.0)
+        assert res.distributions == [[-9999.0, -9999.0]]
+        assert res.distribution_metadata == [{"id_0": "earth", "id_1": "world"}]