nvidia-nat 1.4.0a20251015__py3-none-any.whl → 1.4.0a20251021__py3-none-any.whl

nat/front_ends/fastapi/message_handler.py

@@ -25,6 +25,7 @@ from pydantic import ValidationError
 from starlette.websockets import WebSocketDisconnect
 
 from nat.authentication.interfaces import FlowHandlerBase
+from nat.data_models.api_server import ChatRequest
 from nat.data_models.api_server import ChatResponse
 from nat.data_models.api_server import ChatResponseChunk
 from nat.data_models.api_server import Error
@@ -33,6 +34,8 @@ from nat.data_models.api_server import ResponsePayloadOutput
 from nat.data_models.api_server import ResponseSerializable
 from nat.data_models.api_server import SystemResponseContent
 from nat.data_models.api_server import TextContent
+from nat.data_models.api_server import UserMessageContentRoleType
+from nat.data_models.api_server import UserMessages
 from nat.data_models.api_server import WebSocketMessageStatus
 from nat.data_models.api_server import WebSocketMessageType
 from nat.data_models.api_server import WebSocketSystemInteractionMessage
@@ -64,12 +67,12 @@ class WebSocketMessageHandler:
         self._running_workflow_task: asyncio.Task | None = None
         self._message_parent_id: str = "default_id"
         self._conversation_id: str | None = None
-        self._workflow_schema_type: str = None
-        self._user_interaction_response: asyncio.Future[HumanResponse] | None = None
+        self._workflow_schema_type: str | None = None
+        self._user_interaction_response: asyncio.Future[TextContent] | None = None
 
         self._flow_handler: FlowHandlerBase | None = None
 
-        self._schema_output_mapping: dict[str, type[BaseModel] | None] = {
+        self._schema_output_mapping: dict[str, type[BaseModel] | type[None]] = {
             WorkflowSchemaType.GENERATE: self._session_manager.workflow.single_output_schema,
             WorkflowSchemaType.CHAT: ChatResponse,
             WorkflowSchemaType.CHAT_STREAM: ChatResponseChunk,
@@ -114,36 +117,58 @@ class WebSocketMessageHandler:
                     pass
 
                 elif (isinstance(validated_message, WebSocketUserInteractionResponseMessage)):
-                    user_content = await self.process_user_message_content(validated_message)
+                    user_content = await self._process_websocket_user_interaction_response_message(validated_message)
+                    assert self._user_interaction_response is not None
                     self._user_interaction_response.set_result(user_content)
             except (asyncio.CancelledError, WebSocketDisconnect):
                 # TODO: Handle the disconnect
                 break
 
-    async def process_user_message_content(
-            self, user_content: WebSocketUserMessage | WebSocketUserInteractionResponseMessage) -> BaseModel | None:
+    def _extract_last_user_message_content(self, messages: list[UserMessages]) -> TextContent:
         """
-        Processes the contents of a user message.
+        Extracts the last user's TextContent from a list of messages.
 
-        :param user_content: Incoming content data model.
-        :return: A validated Pydantic user content model or None if not found.
-        """
+        Args:
+            messages: List of UserMessages.
 
-        for user_message in user_content.content.messages[::-1]:
-            if (user_message.role == "user"):
+        Returns:
+            TextContent object from the last user message.
 
+        Raises:
+            ValueError: If no user text content is found.
+        """
+        for user_message in messages[::-1]:
+            if user_message.role == UserMessageContentRoleType.USER:
                 for attachment in user_message.content:
-
                     if isinstance(attachment, TextContent):
                         return attachment
+        raise ValueError("No user text content found in messages.")
+
+    async def _process_websocket_user_interaction_response_message(
+            self, user_content: WebSocketUserInteractionResponseMessage) -> TextContent:
+        """
+        Processes a WebSocketUserInteractionResponseMessage.
+        """
+        return self._extract_last_user_message_content(user_content.content.messages)
 
-        return None
+    async def _process_websocket_user_message(self, user_content: WebSocketUserMessage) -> ChatRequest | str:
+        """
+        Processes a WebSocketUserMessage based on schema type.
+        """
+        if self._workflow_schema_type in [WorkflowSchemaType.CHAT, WorkflowSchemaType.CHAT_STREAM]:
+            return ChatRequest(**user_content.content.model_dump(include={"messages"}))
+
+        elif self._workflow_schema_type in [WorkflowSchemaType.GENERATE, WorkflowSchemaType.GENERATE_STREAM]:
+            return self._extract_last_user_message_content(user_content.content.messages).text
+
+        raise ValueError("Unsupported workflow schema type for WebSocketUserMessage")
 
     async def process_workflow_request(self, user_message_as_validated_type: WebSocketUserMessage) -> None:
         """
         Process user messages and routes them appropriately.
 
-        :param user_message_as_validated_type: A WebSocketUserMessage Data Model instance.
+        Args:
+            user_message_as_validated_type (WebSocketUserMessage): The validated user message to process.
         """
 
         try:
@@ -151,18 +176,15 @@ class WebSocketMessageHandler:
             self._workflow_schema_type = user_message_as_validated_type.schema_type
             self._conversation_id = user_message_as_validated_type.conversation_id
 
-            content: BaseModel | None = await self.process_user_message_content(user_message_as_validated_type)
-
-            if content is None:
-                raise ValueError(f"User message content could not be found: {user_message_as_validated_type}")
+            message_content: typing.Any = await self._process_websocket_user_message(user_message_as_validated_type)
 
-            if isinstance(content, TextContent) and (self._running_workflow_task is None):
+            if (self._running_workflow_task is None):
 
-                def _done_callback(task: asyncio.Task):
+                def _done_callback(_task: asyncio.Task):
                     self._running_workflow_task = None
 
                 self._running_workflow_task = asyncio.create_task(
-                    self._run_workflow(payload=content.text,
+                    self._run_workflow(payload=message_content,
                                        user_message_id=self._message_parent_id,
                                        conversation_id=self._conversation_id,
                                        result_type=self._schema_output_mapping[self._workflow_schema_type],
@@ -180,13 +202,14 @@ class WebSocketMessageHandler:
     async def create_websocket_message(self,
                                        data_model: BaseModel,
                                        message_type: str | None = None,
-                                       status: str = WebSocketMessageStatus.IN_PROGRESS) -> None:
+                                       status: WebSocketMessageStatus = WebSocketMessageStatus.IN_PROGRESS) -> None:
         """
         Creates a websocket message that will be ready for routing based on message type or data model.
 
-        :param data_model: Message content model.
-        :param message_type: Message content model.
-        :param status: Message content model.
+        Args:
+            data_model (BaseModel): Message content model.
+            message_type (str | None): Message content model.
+            status (WebSocketMessageStatus): Message content model.
         """
         try:
             message: BaseModel | None = None
@@ -196,8 +219,8 @@ class WebSocketMessageHandler:
 
             message_schema: type[BaseModel] = await self._message_validator.get_message_schema_by_type(message_type)
 
-            if 'id' in data_model.model_fields:
-                message_id: str = data_model.id
+            if hasattr(data_model, 'id'):
+                message_id: str = str(getattr(data_model, 'id'))
             else:
                 message_id = str(uuid.uuid4())
 
@@ -253,12 +276,15 @@ class WebSocketMessageHandler:
         Registered human interaction callback that processes human interactions and returns
         responses from websocket connection.
 
-        :param prompt: Incoming interaction content data model.
-        :return: A Text Content Base Pydantic model.
+        Args:
+            prompt: Incoming interaction content data model.
+
+        Returns:
+            A Text Content Base Pydantic model.
         """
 
         # First create a future from the loop for the human response
-        human_response_future: asyncio.Future[HumanResponse] = asyncio.get_running_loop().create_future()
+        human_response_future: asyncio.Future[TextContent] = asyncio.get_running_loop().create_future()
 
         # Then add the future to the outstanding human prompts dictionary
         self._user_interaction_response = human_response_future
@@ -274,10 +300,10 @@ class WebSocketMessageHandler:
                 return HumanResponseNotification()
 
             # Wait for the human response future to complete
-            interaction_response: HumanResponse = await human_response_future
+            text_content: TextContent = await human_response_future
 
             interaction_response: HumanResponse = await self._message_validator.convert_text_content_to_human_response(
-                interaction_response, prompt.content)
+                text_content, prompt.content)
 
             return interaction_response
 
@@ -293,13 +319,12 @@ class WebSocketMessageHandler:
                             output_type: type | None = None) -> None:
 
         try:
-            async with self._session_manager.session(
-                    user_message_id=user_message_id,
-                    conversation_id=conversation_id,
-                    http_connection=self._socket,
-                    user_input_callback=self.human_interaction_callback,
-                    user_authentication_callback=(self._flow_handler.authenticate
-                                                  if self._flow_handler else None)) as session:
+            auth_callback = self._flow_handler.authenticate if self._flow_handler else None
+            async with self._session_manager.session(user_message_id=user_message_id,
+                                                     conversation_id=conversation_id,
+                                                     http_connection=self._socket,
+                                                     user_input_callback=self.human_interaction_callback,
+                                                     user_authentication_callback=auth_callback) as session:
 
                 async for value in generate_streaming_response(payload,
                                                                session_manager=session,

nat/front_ends/fastapi/message_validator.py

@@ -240,8 +240,7 @@ class MessageValidator:
         thread_id: str = "default",
         parent_id: str = "default",
         conversation_id: str | None = None,
-        content: SystemResponseContent
-        | Error = SystemResponseContent(),
+        content: SystemResponseContent | Error = SystemResponseContent(),
         status: WebSocketMessageStatus = WebSocketMessageStatus.IN_PROGRESS,
         timestamp: str = str(datetime.datetime.now(datetime.UTC))
     ) -> WebSocketSystemResponseTokenMessage | None:

nat/front_ends/mcp/mcp_front_end_config.py

@@ -13,13 +13,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
 from typing import Literal
 
 from pydantic import Field
+from pydantic import model_validator
 
 from nat.authentication.oauth2.oauth2_resource_server_config import OAuth2ResourceServerConfig
 from nat.data_models.front_end import FrontEndBaseConfig
 
+logger = logging.getLogger(__name__)
+
 
 class MCPFrontEndConfig(FrontEndBaseConfig, name="mcp"):
     """MCP front end configuration.
@@ -56,3 +60,31 @@ class MCPFrontEndConfig(FrontEndBaseConfig, name="mcp"):
                                       le=50)
     memory_profile_log_level: str = Field(default="DEBUG",
                                           description="Log level for memory profiling output (default: DEBUG)")
+
+    @model_validator(mode="after")
+    def validate_security_configuration(self):
+        """Validate security configuration to prevent accidental misconfigurations."""
+        # Check if server is bound to a non-localhost interface without authentication
+        localhost_hosts = {"localhost", "127.0.0.1", "::1"}
+        if self.host not in localhost_hosts and self.server_auth is None:
+            logger.warning(
+                "MCP server is configured to bind to '%s' without authentication. "
+                "This may expose your server to unauthorized access. "
+                "Consider either: (1) binding to localhost for local-only access, "
+                "or (2) configuring server_auth for production deployments on public interfaces.",
+                self.host)
+
+        # Check if SSE transport is used (which doesn't support authentication)
+        if self.transport == "sse":
+            if self.server_auth is not None:
+                logger.warning("SSE transport does not support authentication. "
+                               "The configured server_auth will be ignored. "
+                               "For production use with authentication, use 'streamable-http' transport instead.")
+            elif self.host not in localhost_hosts:
+                logger.warning(
+                    "SSE transport does not support authentication and is bound to '%s'. "
+                    "This configuration is not recommended for production use. "
+                    "For production deployments, use 'streamable-http' transport with server_auth configured.",
+                    self.host)
+
+        return self

nat/front_ends/mcp/mcp_front_end_plugin.py

@@ -105,9 +105,12 @@ class MCPFrontEndPlugin(FrontEndBase[MCPFrontEndConfig]):
 
             # Start the MCP server with configurable transport
             # streamable-http is the default, but users can choose sse if preferred
-            if self.front_end_config.transport == "sse":
-                logger.info("Starting MCP server with SSE endpoint at /sse")
-                await mcp.run_sse_async()
-            else:  # streamable-http
-                logger.info("Starting MCP server with streamable-http endpoint at /mcp/")
-                await mcp.run_streamable_http_async()
+            try:
+                if self.front_end_config.transport == "sse":
+                    logger.info("Starting MCP server with SSE endpoint at /sse")
+                    await mcp.run_sse_async()
+                else:  # streamable-http
+                    logger.info("Starting MCP server with streamable-http endpoint at /mcp/")
+                    await mcp.run_streamable_http_async()
+            except KeyboardInterrupt:
+                logger.info("MCP server shutdown requested (Ctrl+C). Shutting down gracefully.")

nat/llm/aws_bedrock_llm.py

@@ -42,9 +42,9 @@ class AWSBedrockModelConfig(LLMBaseConfig,
     model_config = ConfigDict(protected_namespaces=(), extra="allow")
 
     # Completion parameters
-    model_name: str = Field(validation_alias=AliasChoices("model_name", "model"),
-                            serialization_alias="model",
-                            description="The model name for the hosted AWS Bedrock.")
+    model_name: str = OptimizableField(validation_alias=AliasChoices("model_name", "model"),
+                                       serialization_alias="model",
+                                       description="The model name for the hosted AWS Bedrock.")
     max_tokens: int = OptimizableField(default=300,
                                        description="Maximum number of tokens to generate.",
                                        space=SearchSpace(high=2176, low=128, step=512))

nat/llm/litellm_llm.py

@@ -23,6 +23,8 @@ from nat.builder.builder import Builder
 from nat.builder.llm import LLMProviderInfo
 from nat.cli.register_workflow import register_llm_provider
 from nat.data_models.llm import LLMBaseConfig
+from nat.data_models.optimizable import OptimizableField
+from nat.data_models.optimizable import OptimizableMixin
 from nat.data_models.retry_mixin import RetryMixin
 from nat.data_models.temperature_mixin import TemperatureMixin
 from nat.data_models.thinking_mixin import ThinkingMixin
@@ -31,6 +33,7 @@ from nat.data_models.top_p_mixin import TopPMixin
 
 class LiteLlmModelConfig(
         LLMBaseConfig,
+        OptimizableMixin,
         RetryMixin,
         TemperatureMixin,
         TopPMixin,
@@ -46,9 +49,9 @@ class LiteLlmModelConfig(
                                  description="Base url to the hosted model.",
                                  validation_alias=AliasChoices("base_url", "api_base"),
                                  serialization_alias="api_base")
-    model_name: str = Field(validation_alias=AliasChoices("model_name", "model"),
-                            serialization_alias="model",
-                            description="The LiteLlm hosted model name.")
+    model_name: str = OptimizableField(validation_alias=AliasChoices("model_name", "model"),
+                                       serialization_alias="model",
+                                       description="The LiteLlm hosted model name.")
     seed: int | None = Field(default=None, description="Random seed to set for generation.")
 
 

nat/llm/nim_llm.py

@@ -44,9 +44,9 @@ class NIMModelConfig(LLMBaseConfig,
 
     api_key: str | None = Field(default=None, description="NVIDIA API key to interact with hosted NIM.")
     base_url: str | None = Field(default=None, description="Base url to the hosted NIM.")
-    model_name: str = Field(validation_alias=AliasChoices("model_name", "model"),
-                            serialization_alias="model",
-                            description="The model name for the hosted NIM.")
+    model_name: str = OptimizableField(validation_alias=AliasChoices("model_name", "model"),
+                                       serialization_alias="model",
+                                       description="The model name for the hosted NIM.")
     max_tokens: PositiveInt = OptimizableField(default=300,
                                                description="Maximum number of tokens to generate.",
                                                space=SearchSpace(high=2176, low=128, step=512))

nat/llm/openai_llm.py

@@ -21,6 +21,7 @@ from nat.builder.builder import Builder
 from nat.builder.llm import LLMProviderInfo
 from nat.cli.register_workflow import register_llm_provider
 from nat.data_models.llm import LLMBaseConfig
+from nat.data_models.optimizable import OptimizableField
 from nat.data_models.optimizable import OptimizableMixin
 from nat.data_models.retry_mixin import RetryMixin
 from nat.data_models.temperature_mixin import TemperatureMixin
@@ -41,9 +42,9 @@ class OpenAIModelConfig(LLMBaseConfig,
 
     api_key: str | None = Field(default=None, description="OpenAI API key to interact with hosted model.")
     base_url: str | None = Field(default=None, description="Base url to the hosted model.")
-    model_name: str = Field(validation_alias=AliasChoices("model_name", "model"),
-                            serialization_alias="model",
-                            description="The OpenAI hosted model name.")
+    model_name: str = OptimizableField(validation_alias=AliasChoices("model_name", "model"),
+                                       serialization_alias="model",
+                                       description="The OpenAI hosted model name.")
     seed: int | None = Field(default=None, description="Random seed to set for generation.")
     max_retries: int = Field(default=10, description="The max number of retries for the request.")
 

nat/profiler/callbacks/langchain_callback_handler.py

@@ -33,6 +33,7 @@ from nat.builder.context import Context
 from nat.builder.framework_enum import LLMFrameworkEnum
 from nat.data_models.intermediate_step import IntermediateStepPayload
 from nat.data_models.intermediate_step import IntermediateStepType
+from nat.data_models.intermediate_step import ServerToolUseSchema
 from nat.data_models.intermediate_step import StreamEventData
 from nat.data_models.intermediate_step import ToolSchema
 from nat.data_models.intermediate_step import TraceMetadata
@@ -48,7 +49,14 @@ def _extract_tools_schema(invocation_params: dict) -> list:
     tools_schema = []
     if invocation_params is not None:
         for tool in invocation_params.get("tools", []):
-            tools_schema.append(ToolSchema(**tool))
+            try:
+                tools_schema.append(ToolSchema(**tool))
+            except Exception:
+                logger.debug(
+                    "Failed to parse tool schema from invocation params: %s. \n This "
+                    "can occur when the LLM server has native tools and can be ignored if "
+                    "using the responses API.",
+                    tool)
 
     return tools_schema
 
@@ -93,11 +101,15 @@ class LangchainProfilerHandler(AsyncCallbackHandler, BaseProfilerCallback):
             completion_tokens = usage_metadata.get("output_tokens", 0)
             total_tokens = usage_metadata.get("total_tokens", 0)
 
-            return TokenUsageBaseModel(
-                prompt_tokens=prompt_tokens,
-                completion_tokens=completion_tokens,
-                total_tokens=total_tokens,
-            )
+            cache_tokens = usage_metadata.get("input_token_details", {}).get("cache_read", 0)
+
+            reasoning_tokens = usage_metadata.get("output_token_details", {}).get("reasoning", 0)
+
+            return TokenUsageBaseModel(prompt_tokens=prompt_tokens,
+                                       completion_tokens=completion_tokens,
+                                       total_tokens=total_tokens,
+                                       cached_tokens=cache_tokens,
+                                       reasoning_tokens=reasoning_tokens)
         return TokenUsageBaseModel()
 
     async def on_llm_start(self, serialized: dict[str, Any], prompts: list[str], **kwargs: Any) -> None:
@@ -213,6 +225,7 @@ class LangchainProfilerHandler(AsyncCallbackHandler, BaseProfilerCallback):
         except IndexError:
             generation = None
 
+        message = None
         if isinstance(generation, ChatGeneration):
             try:
                 message = generation.message
@@ -232,6 +245,17 @@ class LangchainProfilerHandler(AsyncCallbackHandler, BaseProfilerCallback):
         else:
             llm_text_output = ""
 
+        tool_outputs_list = []
+        # Check if message.additional_kwargs as tool_outputs indicative of server side tool calling
+        if message and message.additional_kwargs and "tool_outputs" in message.additional_kwargs:
+            tools_outputs = message.additional_kwargs["tool_outputs"]
+            if isinstance(tools_outputs, list):
+                for tool in tools_outputs:
+                    try:
+                        tool_outputs_list.append(ServerToolUseSchema(**tool))
+                    except Exception:
+                        pass
+
         # update shared state behind lock
         with self._lock:
             usage_stat = IntermediateStepPayload(
@@ -243,7 +267,8 @@ class LangchainProfilerHandler(AsyncCallbackHandler, BaseProfilerCallback):
                 data=StreamEventData(input=self._run_id_to_llm_input.get(str(kwargs.get("run_id", "")), ""),
                                      output=llm_text_output),
                 usage_info=UsageInfo(token_usage=self._extract_token_base_model(usage_metadata)),
-                metadata=TraceMetadata(chat_responses=[generation] if generation else []))
+                metadata=TraceMetadata(chat_responses=[generation] if generation else [],
+                                       tool_outputs=tool_outputs_list if tool_outputs_list else []))
 
             self.step_manager.push_intermediate_step(usage_stat)
 

nat/profiler/callbacks/llama_index_callback_handler.py

@@ -30,6 +30,7 @@ from nat.builder.context import Context
 from nat.builder.framework_enum import LLMFrameworkEnum
 from nat.data_models.intermediate_step import IntermediateStepPayload
 from nat.data_models.intermediate_step import IntermediateStepType
+from nat.data_models.intermediate_step import ServerToolUseSchema
 from nat.data_models.intermediate_step import StreamEventData
 from nat.data_models.intermediate_step import TraceMetadata
 from nat.data_models.intermediate_step import UsageInfo
@@ -64,6 +65,26 @@ class LlamaIndexProfilerHandler(BaseCallbackHandler, BaseProfilerCallback):
         self._run_id_to_tool_input = {}
         self._run_id_to_timestamp = {}
 
+    @staticmethod
+    def _extract_token_usage(response: ChatResponse) -> TokenUsageBaseModel:
+        token_usage = TokenUsageBaseModel()
+        try:
+            if response and response.additional_kwargs and "usage" in response.additional_kwargs:
+                usage = response.additional_kwargs["usage"] if "usage" in response.additional_kwargs else {}
+                token_usage.prompt_tokens = usage.input_tokens if hasattr(usage, "input_tokens") else 0
+                token_usage.completion_tokens = usage.output_tokens if hasattr(usage, "output_tokens") else 0
+
+                if hasattr(usage, "input_tokens_details") and hasattr(usage.input_tokens_details, "cached_tokens"):
+                    token_usage.cached_tokens = usage.input_tokens_details.cached_tokens
+
+                if hasattr(usage, "output_tokens_details") and hasattr(usage.output_tokens_details, "reasoning_tokens"):
+                    token_usage.reasoning_tokens = usage.output_tokens_details.reasoning_tokens
+
+        except Exception as e:
+            logger.debug("Error extracting token usage: %s", e, exc_info=True)
+
+        return token_usage
+
     def on_event_start(
         self,
         event_type: CBEventType,
@@ -167,6 +188,18 @@ class LlamaIndexProfilerHandler(BaseCallbackHandler, BaseProfilerCallback):
                 except Exception as e:
                     logger.exception("Error getting model name: %s", e)
 
+                # Append usage data to NAT usage stats
+                tool_outputs_list = []
+                # Check if message.additional_kwargs as tool_outputs indicative of server side tool calling
+                if response and response.additional_kwargs and "built_in_tool_calls" in response.additional_kwargs:
+                    tools_outputs = response.additional_kwargs["built_in_tool_calls"]
+                    if isinstance(tools_outputs, list):
+                        for tool in tools_outputs:
+                            try:
+                                tool_outputs_list.append(ServerToolUseSchema(**tool.model_dump()))
+                            except Exception:
+                                pass
+
                 # Append usage data to NAT usage stats
                 with self._lock:
                     stats = IntermediateStepPayload(
@@ -176,8 +209,9 @@ class LlamaIndexProfilerHandler(BaseCallbackHandler, BaseProfilerCallback):
                         name=model_name,
                         UUID=event_id,
                         data=StreamEventData(input=self._run_id_to_llm_input.get(event_id), output=llm_text_output),
-                        metadata=TraceMetadata(chat_responses=response.message if response.message else None),
-                        usage_info=UsageInfo(token_usage=TokenUsageBaseModel(**response.additional_kwargs)))
+                        metadata=TraceMetadata(chat_responses=response.message if response.message else None,
+                                               tool_outputs=tool_outputs_list if tool_outputs_list else []),
+                        usage_info=UsageInfo(token_usage=self._extract_token_usage(response)))
                     self.step_manager.push_intermediate_step(stats)
 
         elif event_type == CBEventType.FUNCTION_CALL and payload:

nat/profiler/callbacks/token_usage_base_model.py

@@ -24,4 +24,6 @@ class TokenUsageBaseModel(BaseModel):
 
     prompt_tokens: int = Field(default=0, description="Number of tokens in the prompt.")
     completion_tokens: int = Field(default=0, description="Number of tokens in the completion.")
+    cached_tokens: int = Field(default=0, description="Number of tokens read from cache.")
+    reasoning_tokens: int = Field(default=0, description="Number of tokens used for reasoning.")
     total_tokens: int = Field(default=0, description="Number of tokens total.")