feat: add metrics for usages of prompt and completion tokens

Author: baixiac

app/api/routers/generative.py

@@ -6,6 +6,7 @@
 
 from typing import Union, Iterable, AsyncGenerator
 from typing_extensions import Annotated
+from functools import partial
 from fastapi import APIRouter, Depends, Request, Body, Query
 from fastapi.encoders import jsonable_encoder
 from fastapi.responses import PlainTextResponse, StreamingResponse, JSONResponse
@@ -15,6 +16,7 @@
 from app.utils import get_settings, get_prompt_from_messages
 from app.api.utils import get_rate_limiter
 from app.api.dependencies import validate_tracking_id
+from app.management.prometheus_metrics import cms_prompt_tokens, cms_completion_tokens, cms_total_tokens
 
 PATH_GENERATE = "/generate"
 PATH_GENERATE_ASYNC = "/stream/generate"
@@ -44,7 +46,7 @@ def generate_text(
     model_service: AbstractModelService = Depends(cms_globals.model_service_dep)
 ) -> PlainTextResponse:
     """
-    Generate text based on the prompt provided.
+    Generates text based on the prompt provided.
 
     Args:
         request (Request): The request object.
@@ -61,7 +63,12 @@ def generate_text(
     tracking_id = tracking_id or str(uuid.uuid4())
     if prompt:
         return PlainTextResponse(
-            model_service.generate(prompt, max_tokens=max_tokens, temperature=temperature),
+            model_service.generate(
+                prompt,
+                max_tokens=max_tokens,
+                temperature=temperature,
+                report_tokens=partial(_send_usage_metrics, handler=PATH_GENERATE),
+            ),
             headers={"x-cms-tracking-id": tracking_id},
             status_code=HTTP_200_OK,
         )
@@ -89,7 +96,7 @@ async def generate_text_stream(
     model_service: AbstractModelService = Depends(cms_globals.model_service_dep)
 ) -> StreamingResponse:
     """
-    Generate a stream of texts in near real-time.
+    Generates a stream of texts in near real-time.
 
     Args:
         request (Request): The request object.
@@ -106,7 +113,12 @@ async def generate_text_stream(
     tracking_id = tracking_id or str(uuid.uuid4())
     if prompt:
         return StreamingResponse(
-            model_service.generate_async(prompt, max_tokens=max_tokens, temperature=temperature),
+            model_service.generate_async(
+                prompt,
+                max_tokens=max_tokens,
+                temperature=temperature,
+                report_tokens=partial(_send_usage_metrics, handler=PATH_GENERATE_ASYNC),
+            ),
             media_type="text/event-stream",
             headers={"x-cms-tracking-id": tracking_id},
             status_code=HTTP_200_OK,
@@ -121,7 +133,7 @@ async def generate_text_stream(
 
 
 @router.post(
-    "/v1/chat/completions",
+    PATH_OPENAI_COMPLETIONS,
     tags=[Tags.Generative.name],
     response_model=None,
     dependencies=[Depends(cms_globals.props.current_active_user)],
@@ -136,7 +148,7 @@ def generate_chat_completions(
     model_service: AbstractModelService = Depends(cms_globals.model_service_dep)
 ) -> Union[StreamingResponse, JSONResponse]:
     """
-    Generate chat response based on messages, mimicking OpenAI's /v1/chat/completions endpoint.
+    Generates chat response based on messages, mimicking OpenAI's /v1/chat/completions endpoint.
 
     Args:
         request (Request): The request object.
@@ -153,6 +165,7 @@ def generate_chat_completions(
     stream = request_data.stream
     max_tokens = request_data.max_tokens
     temperature = request_data.temperature
+    tracking_id = tracking_id or str(uuid.uuid4())
 
     if not messages:
         error_response = {
@@ -163,16 +176,25 @@ def generate_chat_completions(
                 "code": "missing_field",
             }
         }
-        return JSONResponse(content=error_response, status_code=HTTP_400_BAD_REQUEST)
+        return JSONResponse(
+            content=error_response,
+            status_code=HTTP_400_BAD_REQUEST,
+            headers={"x-cms-tracking-id": tracking_id},
+        )
 
-    async def _stream(p: str, mt: int, t: float) -> AsyncGenerator:
+    async def _stream(prompt: str, max_tokens: int, temperature: float) -> AsyncGenerator:
         data = {
-            "id": tracking_id or str(uuid.uuid4()),
+            "id": tracking_id,
             "object": "chat.completion.chunk",
             "choices": [{"delta": {"role": PromptRole.ASSISTANT.value}}],
         }
         yield f"data: {json.dumps(data)}\n\n"
-        async for chunk in model_service.generate_async(p, max_tokens=mt, temperature=t):
+        async for chunk in model_service.generate_async(
+            prompt,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            report_tokens=partial(_send_usage_metrics, handler=PATH_OPENAI_COMPLETIONS)
+        ):
             data = {
                 "choices": [
                     {
@@ -188,12 +210,18 @@ async def _stream(p: str, mt: int, t: float) -> AsyncGenerator:
     if stream:
         return StreamingResponse(
             _stream(prompt, max_tokens, temperature),
-            media_type="text/event-stream"
+            media_type="text/event-stream",
+            headers={"x-cms-tracking-id": tracking_id},
         )
     else:
-        generated_text = model_service.generate(prompt, max_tokens=max_tokens, temperature=temperature)
+        generated_text = model_service.generate(
+            prompt,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            send_metrics=partial(_send_usage_metrics, handler=PATH_OPENAI_COMPLETIONS),
+        )
         completion = OpenAIChatResponse(
-            id=str(uuid.uuid4()),
+            id=tracking_id,
             object="chat.completion",
             created=int(time.time()),
             model=model_service.model_name,
@@ -206,10 +234,19 @@ async def _stream(p: str, mt: int, t: float) -> AsyncGenerator:
                     ),
                     "finish_reason": "stop",
                 }
-            ]
+            ],
         )
-        return JSONResponse(content=jsonable_encoder(completion))
+        return JSONResponse(content=jsonable_encoder(completion), headers={"x-cms-tracking-id": tracking_id})
 
 
 def _empty_prompt_error() -> Iterable[str]:
     yield "ERROR: No prompt text provided\n"
+
+
+def _send_usage_metrics(handler: str, prompt_token_num: int, completion_token_num: int) -> None:
+    cms_prompt_tokens.labels(handler=handler).observe(prompt_token_num)
+    logger.debug(f"Sent prompt tokens usage: {prompt_token_num}")
+    cms_completion_tokens.labels(handler=handler).observe(completion_token_num)
+    logger.debug(f"Sent completion tokens usage: {completion_token_num}")
+    cms_total_tokens.labels(handler=handler).observe(prompt_token_num + completion_token_num)
+    logger.debug(f"Sent total tokens usage: {prompt_token_num + completion_token_num}")

app/domain.py

@@ -189,8 +189,8 @@ class OpenAIChatRequest(BaseModel):
 
 
 class OpenAIChatResponse(BaseModel):
-    id: str
-    object: str
-    created: int
-    model: str
-    choices: List
+    id: str = Field(..., description="The unique identifier for the chat completion request")
+    object: str = Field(..., description="The type of the response")
+    created: int = Field(..., description="The timestamp when the completion was generated")
+    model: str = Field(..., description="The name of the model used for generating the completion")
+    choices: List = Field(..., description="The generated messages and their metadata")

app/management/prometheus_metrics.py

@@ -34,3 +34,24 @@
     "Number of bulk-processed documents",
     ["handler"],
 )
+
+# The histogram metric to track the number of tokens in the messages of the input prompt
+cms_prompt_tokens = Histogram(
+    "cms_prompt_tokens",
+    "Number of tokens in the messages of the input prompt",
+    ["handler"],
+)
+
+# The histogram metric to track the number of tokens in the generated assistant reply
+cms_completion_tokens = Histogram(
+    "cms_completion_tokens",
+    "Number of tokens in the generated assistant reply",
+    ["handler"],
+)
+
+# The histogram metric to track the total number of tokens used in the prompt and the completion
+cms_total_tokens = Histogram(
+    "cms_total_tokens",
+    "Number of tokens used in the prompt and the completion",
+    ["handler"],
+)

app/model_services/huggingface_llm_model.py

@@ -2,7 +2,7 @@
 import logging
 import asyncio
 from concurrent.futures import ThreadPoolExecutor
-from typing import Dict, List, Optional, Tuple, Any, AsyncIterable
+from typing import Dict, List, Optional, Tuple, Any, AsyncIterable, Callable
 from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
@@ -198,6 +198,7 @@ def generate(
         prompt: str,
         max_tokens: int = 512,
         temperature: float = 0.7,
+        report_tokens: Optional[Callable[[str], None]] = None,
         **kwargs: Any
     ) -> str:
         """
@@ -207,6 +208,7 @@ def generate(
             prompt (str): The prompt for the text generation
             max_tokens (int): The maximum number of tokens to generate. Defaults to 512.
             temperature (float): The temperature for the text generation. Defaults to 0.7.
+            report_tokens (Optional[Callable[[str], None]]): The callback function to send metrics. Defaults to None.
             **kwargs (Any): Additional keyword arguments to be passed to this method.
 
         Returns:
@@ -230,17 +232,22 @@ def generate(
 
         outputs = self.model.generate(**generation_kwargs)
         generated_text = self.tokenizer.decode(outputs[0], skip_prompt=True, skip_special_tokens=True)
-
-
         logger.debug("Response generation completed")
 
+        if report_tokens:
+            report_tokens(
+                prompt_token_num=inputs.input_ids.shape[-1],    # type: ignore
+                completion_token_num=outputs[0].shape[-1],  # type: ignore
+            )
+
         return generated_text
 
     async def generate_async(
         self,
         prompt: str,
         max_tokens: int = 512,
         temperature: float = 0.7,
+        report_tokens: Optional[Callable[[str], None]] = None,
         **kwargs: Any
     ) -> AsyncIterable:
         """
@@ -250,6 +257,7 @@ async def generate_async(
             prompt (str): The prompt for the text generation.
             max_tokens (int): The maximum number of tokens to generate. Defaults to 512.
             temperature (float): The temperature for the text generation. Defaults to 0.7.
+            report_tokens (Optional[Callable[[str], None]]): The callback function to send metrics. Defaults to None.
             **kwargs (Any): Additional keyword arguments to be passed to the model loader.
 
         Returns:
@@ -279,9 +287,20 @@ async def generate_async(
 
         try:
             _ = self._text_generator.submit(self.model.generate, **generation_kwargs)
+            output = ""
             for content in streamer:
                 yield content
+                output += content
                 await asyncio.sleep(0.01)
+            if report_tokens:
+                report_tokens(
+                    prompt_token_num=inputs.input_ids.shape[-1],    # type: ignore
+                    completion_token_num=self.tokenizer(    # type: ignore
+                        output,
+                        add_special_tokens=False,
+                        return_tensors="pt"
+                    ).input_ids.shape[-1],
+                )
         except Exception as e:
             logger.error("An error occurred while generating the response")
             logger.exception(e)

app/utils.py

@@ -694,7 +694,7 @@ def dump_pydantic_object_to_dict(model: BaseModel) -> Dict:
     """
 
     if hasattr(model, "model_dump"):
-        return model.model_dump()    # type: ignore
+        return model.model_dump(mode="json")    # type: ignore
     elif hasattr(model, "dict"):
         return model.dict()    # type: ignore
     else:
@@ -835,3 +835,4 @@ def get_prompt_from_messages(tokenizer: PreTrainedTokenizer, messages: List[Prom
     "25624495": '© 2002-2020 International Health Terminology Standards Development Organisation (IHTSDO). All rights reserved. SNOMED CT®, was originally created by The College of American Pathologists. "SNOMED" and "SNOMED CT" are registered trademarks of the IHTSDO.',
     "55540447": "linkage concept"
 }
+

tests/app/model_services/test_huggingface_llm_model.py

@@ -46,6 +46,7 @@ def test_generate(huggingface_llm_model):
     huggingface_llm_model.init_model()
     huggingface_llm_model.model = MagicMock()
     huggingface_llm_model.tokenizer = MagicMock()
+    mock_send_metrics = MagicMock()
     inputs = MagicMock()
     inputs.input_ids = MagicMock(shape=[1, 2])
     inputs.attention_mask = MagicMock()
@@ -58,6 +59,7 @@ def test_generate(huggingface_llm_model):
         prompt="Alright?",
         max_tokens=128,
         temperature=0.5,
+        report_tokens=mock_send_metrics
     )
 
     huggingface_llm_model.tokenizer.assert_called_once_with(
@@ -78,13 +80,18 @@ def test_generate(huggingface_llm_model):
         skip_prompt=True,
         skip_special_tokens=True,
     )
+    mock_send_metrics.assert_called_once_with(
+        prompt_token_num=2,
+        completion_token_num=2,
+    )
     assert result == "Yeah."
 
 
 async def test_generate_async(huggingface_llm_model):
     huggingface_llm_model.init_model()
     huggingface_llm_model.model = MagicMock()
     huggingface_llm_model.tokenizer = MagicMock()
+    mock_send_metrics = MagicMock()
     inputs = MagicMock()
     inputs.input_ids = MagicMock(shape=[1, 2])
     inputs.attention_mask = MagicMock()
@@ -97,6 +104,7 @@ async def test_generate_async(huggingface_llm_model):
         prompt="Alright?",
         max_tokens=128,
         temperature=0.5,
+        report_tokens=mock_send_metrics
     )
 
     huggingface_llm_model.tokenizer.assert_called_once_with(
@@ -117,4 +125,8 @@ async def test_generate_async(huggingface_llm_model):
         skip_prompt=True,
         skip_special_tokens=True,
     )
+    mock_send_metrics.assert_called_once_with(
+        prompt_token_num=2,
+        completion_token_num=2,
+    )
     assert result == "Yeah."