Skip to content

Commit fae43e7

Browse files
authored
[None][doc] add status labels to LLM class's api reference (#6899)
Signed-off-by: Superjomn <[email protected]>
1 parent c4535e6 commit fae43e7

File tree

2 files changed

+18
-7
lines changed

2 files changed

+18
-7
lines changed

tensorrt_llm/llmapi/utils.py

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -515,6 +515,8 @@ def generate_api_docs_as_docstring(model: Type[BaseModel],
515515

516516
# Format the argument documentation with 12 spaces indent for args
517517
arg_line = f"{indent} {field_name} ({type_str}): "
518+
if status := field_info.get("status", None):
519+
arg_line += f":tag:`{status}` "
518520
if field_description:
519521
arg_line += field_description.split('\n')[0] # First line with type
520522

@@ -554,20 +556,21 @@ class ApiParamTagger:
554556
'''
555557

556558
def __call__(self, cls: Type[BaseModel]) -> None:
557-
self.process_pydantic_model(cls)
559+
""" The main entry point to tag the api doc. """
560+
self._process_pydantic_model(cls)
558561

559-
def process_pydantic_model(self, cls: Type[BaseModel]) -> None:
562+
def _process_pydantic_model(self, cls: Type[BaseModel]) -> None:
560563
"""Process the Pydantic model to add tags to the fields.
561564
"""
562565
for field_name, field_info in cls.model_fields.items():
563566
if field_info.json_schema_extra and 'status' in field_info.json_schema_extra:
564567
status = field_info.json_schema_extra['status']
565-
self.amend_pydantic_field_description_with_tags(
568+
self._amend_pydantic_field_description_with_tags(
566569
cls, [field_name], status)
567570

568-
def amend_pydantic_field_description_with_tags(self, cls: Type[BaseModel],
569-
field_names: list[str],
570-
tag: str) -> None:
571+
def _amend_pydantic_field_description_with_tags(self, cls: Type[BaseModel],
572+
field_names: list[str],
573+
tag: str) -> None:
571574
"""Amend the description of the fields with tags.
572575
e.g. :tag:`beta` or :tag:`prototype`
573576
Args:

tests/unittest/llmapi/test_utils.py

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
1-
from tensorrt_llm.llmapi.utils import ApiStatusRegistry
1+
from tensorrt_llm.llmapi import LlmArgs
2+
from tensorrt_llm.llmapi.utils import (ApiStatusRegistry,
3+
generate_api_docs_as_docstring)
24

35

46
def test_api_status_registry():
@@ -24,3 +26,9 @@ def _my_method(self, *args, **kwargs):
2426
pass
2527

2628
assert ApiStatusRegistry.get_api_status(App._my_method) == "beta"
29+
30+
31+
def test_generate_api_docs_as_docstring():
32+
doc = generate_api_docs_as_docstring(LlmArgs)
33+
assert ":tag:`beta`" in doc, "the label is not generated"
34+
print(doc)

0 commit comments

Comments
 (0)