[Score API][18132] return token usage in Score API response (#18381)

This commit is contained in:
Guy Stone
2026-03-03 13:45:35 -08:00
committed by GitHub
parent b0f26698f5
commit f749802402
7 changed files with 106 additions and 39 deletions
+25 -14
View File
@@ -164,7 +164,7 @@ class TestScoreAPI(CustomTestCase):
label_token_ids=label_token_ids,
apply_softmax=True,
item_first=case["item_first"],
)
).scores
# Get scores from HuggingFace using the same parameters
hf_scores = self.compute_hf_scores(
@@ -193,7 +193,7 @@ class TestScoreAPI(CustomTestCase):
items=texts,
label_token_ids=label_token_ids,
apply_softmax=True,
)
).scores
self.assertEqual(
len(scores),
@@ -245,7 +245,7 @@ class TestScoreAPI(CustomTestCase):
items=items,
label_token_ids=label_token_ids,
apply_softmax=True,
)
).scores
# Verify we got results
self.assertEqual(len(scores), len(items))
@@ -306,15 +306,18 @@ class TestScoreAPI(CustomTestCase):
label_token_ids = [9454, 2753] # "Yes" and "No" tokens
# Get scores using SGLang
scores = self.engine.score(
result = self.engine.score(
query=query,
items=items,
label_token_ids=label_token_ids,
apply_softmax=True,
)
scores = result.scores
prompt_tokens = result.prompt_tokens
# Verify we get the expected number of scores
self.assertEqual(len(scores), len(items), "Should get one score list per item")
self.assertGreater(prompt_tokens, 0, "Should have positive prompt_tokens")
# Verify each score list has the correct length
for i, score_list in enumerate(scores):
@@ -348,14 +351,14 @@ class TestScoreAPI(CustomTestCase):
items=items,
label_token_ids=label_token_ids,
apply_softmax=True,
)
).scores
scores2 = self.engine.score(
query=query,
items=items,
label_token_ids=label_token_ids,
apply_softmax=True,
)
).scores
# Results should be identical (deterministic)
self.assertEqual(len(scores1), len(scores2), "Should get same number of items")
@@ -391,7 +394,7 @@ class TestScoreAPI(CustomTestCase):
items=items,
label_token_ids=label_token_ids,
apply_softmax=True,
)
).scores
self.assertEqual(
len(scores), len(items), f"Should get {len(items)} score lists"
@@ -411,14 +414,19 @@ class TestScoreAPI(CustomTestCase):
items = []
label_token_ids = [1, 2]
scores = self.engine.score(
result = self.engine.score(
query=query,
items=items,
label_token_ids=label_token_ids,
apply_softmax=True,
)
scores = result.scores
prompt_tokens = result.prompt_tokens
self.assertEqual(len(scores), 0, "Should return empty list for empty items")
self.assertEqual(
prompt_tokens, 0, "Should return 0 prompt_tokens for empty items"
)
def test_multi_item_scoring_single_item(self):
"""Test multi-item scoring with single item (should work like regular scoring)."""
@@ -426,18 +434,21 @@ class TestScoreAPI(CustomTestCase):
items = ["Paris"]
label_token_ids = [1, 2, 3]
scores = self.engine.score(
result = self.engine.score(
query=query,
items=items,
label_token_ids=label_token_ids,
apply_softmax=True,
)
scores = result.scores
prompt_tokens = result.prompt_tokens
self.assertEqual(len(scores), 1, "Should get one score list")
self.assertEqual(
len(scores[0]), len(label_token_ids), "Should have correct number of scores"
)
self.assertAlmostEqual(sum(scores[0]), 1.0, places=6)
self.assertGreater(prompt_tokens, 0, "Should have positive prompt_tokens")
def test_multi_item_scoring_different_queries(self):
"""Test multi-item scoring with different types of queries."""
@@ -459,7 +470,7 @@ class TestScoreAPI(CustomTestCase):
items=items,
label_token_ids=label_token_ids,
apply_softmax=True,
)
).scores
self.assertEqual(
len(scores),
@@ -490,7 +501,7 @@ class TestScoreAPI(CustomTestCase):
items=items,
label_token_ids=label_token_ids,
apply_softmax=True,
)
).scores
self.assertEqual(len(scores), len(items))
@@ -513,7 +524,7 @@ class TestScoreAPI(CustomTestCase):
items=items,
label_token_ids=label_token_ids,
apply_softmax=False, # No softmax
)
).scores
self.assertEqual(len(scores), len(items))
@@ -537,7 +548,7 @@ class TestScoreAPI(CustomTestCase):
items=items,
label_token_ids=label_token_ids,
apply_softmax=True,
)
).scores
self.assertEqual(len(scores), len(items), "Should handle large batches")
@@ -556,7 +567,7 @@ class TestScoreAPI(CustomTestCase):
items=items,
label_token_ids=label_token_ids,
apply_softmax=True,
)
).scores
self.assertEqual(len(scores), len(items))
@@ -1038,6 +1038,18 @@ class TestOpenAIV1Score(CustomTestCase):
msg=f"Score {i} probabilities should sum to 1",
)
# Verify usage
self.assertIn("usage", response, "Response should have a 'usage' field")
self.assertGreater(response["usage"]["prompt_tokens"], 0)
self.assertEqual(
response["usage"]["prompt_tokens"], response["usage"]["total_tokens"]
)
self.assertEqual(
response["usage"]["completion_tokens"],
0,
"completion_tokens should be 0 for /v1/score",
)
def test_score_token_input(self):
"""Test scoring with token IDs input"""
query = "The capital of France is"
@@ -1088,6 +1100,18 @@ class TestOpenAIV1Score(CustomTestCase):
msg=f"Score {i} probabilities should sum to 1",
)
# Verify usage
self.assertIn("usage", response, "Response should have a 'usage' field")
self.assertGreater(response["usage"]["prompt_tokens"], 0)
self.assertEqual(
response["usage"]["prompt_tokens"], response["usage"]["total_tokens"]
)
self.assertEqual(
response["usage"]["completion_tokens"],
0,
"completion_tokens should be 0 for /v1/score",
)
def test_score_error_handling(self):
"""Test error handling for invalid inputs"""
query = "The capital of France is"
@@ -3,6 +3,7 @@ import unittest
from unittest.mock import Mock
from sglang.srt.entrypoints.openai.protocol import V1RerankReqInput
from sglang.srt.managers.tokenizer_manager_multiitem_mixin import ScoreResult
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
# Keep consistent with other openai_server/basic unit tests.
@@ -163,7 +164,7 @@ class TestOpenAIServingRerankUnit(unittest.TestCase):
# Return [p_yes, p_no] for each prompt
assert len(prompts) == 2
assert label_token_ids and len(label_token_ids) == 2
return [[0.9, 0.1], [0.2, 0.8]]
return ScoreResult(scores=[[0.9, 0.1], [0.2, 0.8]], prompt_tokens=42)
handler = OpenAIServingRerank(_TM())
req = V1RerankReqInput(query="q", documents=["d1", "d2"], return_documents=True)