nvidia-nat 1.3.0rc1__py3-none-any.whl → 1.3.0rc3__py3-none-any.whl

nat/cli/commands/workflow/workflow_commands.py

@@ -27,6 +27,50 @@ from jinja2 import FileSystemLoader
 logger = logging.getLogger(__name__)
 
 
+def _get_nat_version() -> str | None:
+    """
+    Get the current NAT version.
+
+    Returns:
+        str: The NAT version intended for use in a dependency string.
+        None: If the NAT version is not found.
+    """
+    from nat.cli.entrypoint import get_version
+
+    current_version = get_version()
+    if current_version == "unknown":
+        return None
+
+    version_parts = current_version.split(".")
+    if len(version_parts) < 3:
+        # If the version somehow doesn't have three parts, return the full version
+        return current_version
+
+    patch = version_parts[2]
+    try:
+        # If the patch is a number, keep only the major and minor parts
+        # Useful for stable releases and adheres to semantic versioning
+        _ = int(patch)
+        digits_to_keep = 2
+    except ValueError:
+        # If the patch is not a number, keep all three digits
+        # Useful for pre-release versions (and nightly builds)
+        digits_to_keep = 3
+
+    return ".".join(version_parts[:digits_to_keep])
+
+
+def _is_nat_version_prerelease() -> bool:
+    """
+    Check if the NAT version is a prerelease.
+    """
+    version = _get_nat_version()
+    if version is None:
+        return False
+
+    return len(version.split(".")) >= 3
+
+
 def _get_nat_dependency(versioned: bool = True) -> str:
     """
     Get the NAT dependency string with version.
@@ -44,16 +88,12 @@ def _get_nat_dependency(versioned: bool = True) -> str:
         logger.debug("Using unversioned NAT dependency: %s", dependency)
         return dependency
 
-    # Get the current NAT version
-    from nat.cli.entrypoint import get_version
-    current_version = get_version()
-    if current_version == "unknown":
-        logger.warning("Could not detect NAT version, using unversioned dependency")
+    version = _get_nat_version()
+    if version is None:
+        logger.debug("Could not detect NAT version, using unversioned dependency: %s", dependency)
         return dependency
 
-    # Extract major.minor (e.g., "1.2.3" -> "1.2")
-    major_minor = ".".join(current_version.split(".")[:2])
-    dependency += f"~={major_minor}"
+    dependency += f"~={version}"
     logger.debug("Using NAT dependency: %s", dependency)
     return dependency
 
@@ -219,12 +259,16 @@ def create_command(workflow_name: str, install: bool, workflow_dir: str, descrip
             install_cmd = ['uv', 'pip', 'install', '-e', str(new_workflow_dir)]
         else:
             install_cmd = ['pip', 'install', '-e', str(new_workflow_dir)]
+            if _is_nat_version_prerelease():
+                install_cmd.insert(2, "--pre")
+
+        python_safe_workflow_name = workflow_name.replace("-", "_")
 
         # List of templates and their destinations
         files_to_render = {
             'pyproject.toml.j2': new_workflow_dir / 'pyproject.toml',
             'register.py.j2': base_dir / 'register.py',
-            'workflow.py.j2': base_dir / f'{workflow_name}_function.py',
+            'workflow.py.j2': base_dir / f'{python_safe_workflow_name}.py',
             '__init__.py.j2': base_dir / '__init__.py',
             'config.yml.j2': configs_dir / 'config.yml',
         }
@@ -233,7 +277,7 @@ def create_command(workflow_name: str, install: bool, workflow_dir: str, descrip
         context = {
             'editable': editable,
             'workflow_name': workflow_name,
-            'python_safe_workflow_name': workflow_name.replace("-", "_"),
+            'python_safe_workflow_name': python_safe_workflow_name,
             'package_name': package_name,
             'rel_path_to_repo_root': rel_path_to_repo_root,
             'workflow_class_name': f"{_generate_valid_classname(workflow_name)}FunctionConfig",

nat/cli/entrypoint.py

@@ -29,6 +29,7 @@ import time
 
 import click
 import nest_asyncio
+from dotenv import load_dotenv
 
 from nat.utils.log_levels import LOG_LEVELS
 
@@ -45,6 +46,9 @@ from .commands.uninstall import uninstall_command
 from .commands.validate import validate_command
 from .commands.workflow.workflow import workflow_command
 
+# Load environment variables from .env file, if it exists
+load_dotenv()
+
 # Apply at the beginning of the file to avoid issues with asyncio
 nest_asyncio.apply()
 
@@ -52,7 +56,11 @@ nest_asyncio.apply()
 def setup_logging(log_level: str):
     """Configure logging with the specified level"""
     numeric_level = LOG_LEVELS.get(log_level.upper(), logging.INFO)
-    logging.basicConfig(level=numeric_level, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+    logging.basicConfig(
+        level=numeric_level,
+        format="%(asctime)s - %(levelname)-8s - %(name)s:%(lineno)d - %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
     return numeric_level
 
 

nat/cli/main.py

@@ -30,6 +30,9 @@ def run_cli():
     import os
     import sys
 
+    # Suppress warnings from transformers
+    os.environ["TRANSFORMERS_VERBOSITY"] = "error"
+
     parent_dir = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
 
     if (parent_dir not in sys.path):

nat/data_models/api_server.py

@@ -28,6 +28,7 @@ from pydantic import HttpUrl
 from pydantic import conlist
 from pydantic import field_serializer
 from pydantic import field_validator
+from pydantic import model_validator
 from pydantic_core.core_schema import ValidationInfo
 
 from nat.data_models.interactive import HumanPrompt
@@ -36,6 +37,15 @@ from nat.utils.type_converter import GlobalTypeConverter
 FINISH_REASONS = frozenset({'stop', 'length', 'tool_calls', 'content_filter', 'function_call'})
 
 
+class UserMessageContentRoleType(str, Enum):
+    """
+    Enum representing chat message roles in API requests and responses.
+    """
+    USER = "user"
+    ASSISTANT = "assistant"
+    SYSTEM = "system"
+
+
 class Request(BaseModel):
     """
     Request is a data model that represents HTTP request attributes.
@@ -108,18 +118,10 @@ UserContent = typing.Annotated[TextContent | ImageContent | AudioContent, Discri
 
 class Message(BaseModel):
     content: str | list[UserContent]
-    role: str
-
-
-class ChatRequest(BaseModel):
-    """
-    ChatRequest is a data model that represents a request to the NAT chat API.
-    Fully compatible with OpenAI Chat Completions API specification.
-    """
+    role: UserMessageContentRoleType
 
-    # Required fields
-    messages: typing.Annotated[list[Message], conlist(Message, min_length=1)]
 
+class ChatRequestOptionals(BaseModel):
     # Optional fields (OpenAI Chat Completions API compatible)
     model: str | None = Field(default=None, description="name of the model to use")
     frequency_penalty: float | None = Field(default=0.0,
@@ -144,6 +146,16 @@ class ChatRequest(BaseModel):
     parallel_tool_calls: bool | None = Field(default=True, description="Whether to enable parallel function calling")
     user: str | None = Field(default=None, description="Unique identifier representing end-user")
 
+
+class ChatRequest(ChatRequestOptionals):
+    """
+    ChatRequest is a data model that represents a request to the NAT chat API.
+    Fully compatible with OpenAI Chat Completions API specification.
+    """
+
+    # Required fields
+    messages: typing.Annotated[list[Message], conlist(Message, min_length=1)]
+
     model_config = ConfigDict(extra="allow",
                               json_schema_extra={
                                   "example": {
@@ -164,7 +176,7 @@ class ChatRequest(BaseModel):
                     max_tokens: int | None = None,
                     top_p: float | None = None) -> "ChatRequest":
 
-        return ChatRequest(messages=[Message(content=data, role="user")],
+        return ChatRequest(messages=[Message(content=data, role=UserMessageContentRoleType.USER)],
                            model=model,
                            temperature=temperature,
                            max_tokens=max_tokens,
@@ -178,38 +190,85 @@ class ChatRequest(BaseModel):
                      max_tokens: int | None = None,
                      top_p: float | None = None) -> "ChatRequest":
 
-        return ChatRequest(messages=[Message(content=content, role="user")],
+        return ChatRequest(messages=[Message(content=content, role=UserMessageContentRoleType.USER)],
                            model=model,
                            temperature=temperature,
                            max_tokens=max_tokens,
                            top_p=top_p)
 
 
+class ChatRequestOrMessage(ChatRequestOptionals):
+    """
+    ChatRequestOrMessage is a data model that represents either a conversation or a string input.
+    This is useful for functions that can handle either type of input.
+
+    `messages` is compatible with the OpenAI Chat Completions API specification.
+
+    `input_string` is a string input that can be used for functions that do not require a conversation.
+    """
+
+    messages: typing.Annotated[list[Message] | None, conlist(Message, min_length=1)] = Field(
+        default=None, description="The conversation messages to process.")
+
+    input_string: str | None = Field(default=None, alias="input_message", description="The input message to process.")
+
+    @property
+    def is_string(self) -> bool:
+        return self.input_string is not None
+
+    @property
+    def is_conversation(self) -> bool:
+        return self.messages is not None
+
+    @model_validator(mode="after")
+    def validate_messages_or_input_string(self):
+        if self.messages is not None and self.input_string is not None:
+            raise ValueError("Either messages or input_message/input_string must be provided, not both")
+        if self.messages is None and self.input_string is None:
+            raise ValueError("Either messages or input_message/input_string must be provided")
+        if self.input_string is not None:
+            extra_fields = self.model_dump(exclude={"input_string"}, exclude_none=True, exclude_unset=True)
+            if len(extra_fields) > 0:
+                raise ValueError("no extra fields are permitted when input_message/input_string is provided")
+        return self
+
+
 class ChoiceMessage(BaseModel):
     content: str | None = None
-    role: str | None = None
+    role: UserMessageContentRoleType | None = None
 
 
 class ChoiceDelta(BaseModel):
     """Delta object for streaming responses (OpenAI-compatible)"""
     content: str | None = None
-    role: str | None = None
+    role: UserMessageContentRoleType | None = None
 
 
-class Choice(BaseModel):
+class ChoiceBase(BaseModel):
+    """Base choice model with common fields for both streaming and non-streaming responses"""
     model_config = ConfigDict(extra="allow")
-
-    message: ChoiceMessage | None = None
-    delta: ChoiceDelta | None = None
     finish_reason: typing.Literal['stop', 'length', 'tool_calls', 'content_filter', 'function_call'] | None = None
     index: int
-    # logprobs: ChoiceLogprobs | None = None
+
+
+class ChatResponseChoice(ChoiceBase):
+    """Choice model for non-streaming responses - contains message field"""
+    message: ChoiceMessage
+
+
+class ChatResponseChunkChoice(ChoiceBase):
+    """Choice model for streaming responses - contains delta field"""
+    delta: ChoiceDelta
+
+
+# Backward compatibility alias
+Choice = ChatResponseChoice
 
 
 class Usage(BaseModel):
-    prompt_tokens: int
-    completion_tokens: int
-    total_tokens: int
+    prompt_tokens: int | None = None
+    completion_tokens: int | None = None
+    total_tokens: int | None = None
 
 
 class ResponseSerializable(abc.ABC):
@@ -245,10 +304,10 @@ class ChatResponse(ResponseBaseModelOutput):
     model_config = ConfigDict(extra="allow")
     id: str
     object: str = "chat.completion"
-    model: str = ""
+    model: str = "unknown-model"
     created: datetime.datetime
-    choices: list[Choice]
-    usage: Usage | None = None
+    choices: list[ChatResponseChoice]
+    usage: Usage
     system_fingerprint: str | None = None
     service_tier: typing.Literal["scale", "default"] | None = None
 
@@ -264,14 +323,14 @@ class ChatResponse(ResponseBaseModelOutput):
                     object_: str | None = None,
                     model: str | None = None,
                     created: datetime.datetime | None = None,
-                    usage: Usage | None = None) -> "ChatResponse":
+                    usage: Usage) -> "ChatResponse":
 
         if id_ is None:
             id_ = str(uuid.uuid4())
         if object_ is None:
             object_ = "chat.completion"
         if model is None:
-            model = ""
+            model = "unknown-model"
         if created is None:
             created = datetime.datetime.now(datetime.UTC)
 
@@ -279,7 +338,12 @@ class ChatResponse(ResponseBaseModelOutput):
                             object=object_,
                             model=model,
                             created=created,
-                            choices=[Choice(index=0, message=ChoiceMessage(content=data), finish_reason="stop")],
+                            choices=[
+                                ChatResponseChoice(index=0,
+                                                   message=ChoiceMessage(content=data,
+                                                                         role=UserMessageContentRoleType.ASSISTANT),
+                                                   finish_reason="stop")
+                            ],
                             usage=usage)
 
 
@@ -293,9 +357,9 @@ class ChatResponseChunk(ResponseBaseModelOutput):
     model_config = ConfigDict(extra="allow")
 
     id: str
-    choices: list[Choice]
+    choices: list[ChatResponseChunkChoice]
     created: datetime.datetime
-    model: str = ""
+    model: str = "unknown-model"
     object: str = "chat.completion.chunk"
     system_fingerprint: str | None = None
     service_tier: typing.Literal["scale", "default"] | None = None
@@ -319,12 +383,18 @@ class ChatResponseChunk(ResponseBaseModelOutput):
         if created is None:
             created = datetime.datetime.now(datetime.UTC)
         if model is None:
-            model = ""
+            model = "unknown-model"
         if object_ is None:
             object_ = "chat.completion.chunk"
 
         return ChatResponseChunk(id=id_,
-                                 choices=[Choice(index=0, message=ChoiceMessage(content=data), finish_reason="stop")],
+                                 choices=[
+                                     ChatResponseChunkChoice(index=0,
+                                                             delta=ChoiceDelta(
+                                                                 content=data,
+                                                                 role=UserMessageContentRoleType.ASSISTANT),
+                                                             finish_reason="stop")
+                                 ],
                                  created=created,
                                  model=model,
                                  object=object_)
@@ -335,7 +405,7 @@ class ChatResponseChunk(ResponseBaseModelOutput):
                                id_: str | None = None,
                                created: datetime.datetime | None = None,
                                model: str | None = None,
-                               role: str | None = None,
+                               role: UserMessageContentRoleType | None = None,
                                finish_reason: str | None = None,
                                usage: Usage | None = None,
                                system_fingerprint: str | None = None) -> "ChatResponseChunk":
@@ -345,7 +415,7 @@ class ChatResponseChunk(ResponseBaseModelOutput):
         if created is None:
             created = datetime.datetime.now(datetime.UTC)
         if model is None:
-            model = ""
+            model = "unknown-model"
 
         delta = ChoiceDelta(content=content, role=role) if content is not None or role is not None else ChoiceDelta()
 
@@ -353,7 +423,14 @@ class ChatResponseChunk(ResponseBaseModelOutput):
 
         return ChatResponseChunk(
             id=id_,
-            choices=[Choice(index=0, message=None, delta=delta, finish_reason=final_finish_reason)],
+            choices=[
+                ChatResponseChunkChoice(
+                    index=0,
+                    delta=delta,
+                    finish_reason=typing.cast(
+                        typing.Literal['stop', 'length', 'tool_calls', 'content_filter', 'function_call'] | None,
+                        final_finish_reason))
+            ],
             created=created,
             model=model,
             object="chat.completion.chunk",
@@ -398,11 +475,6 @@ class GenerateResponse(BaseModel):
     value: str | None = "default"
 
 
-class UserMessageContentRoleType(str, Enum):
-    USER = "user"
-    ASSISTANT = "assistant"
-
-
 class WebSocketMessageType(str, Enum):
     """
     WebSocketMessageType is an Enum that represents WebSocket Message types.
@@ -622,12 +694,42 @@ GlobalTypeConverter.register_converter(_nat_chat_request_to_string)
 
 
 def _string_to_nat_chat_request(data: str) -> ChatRequest:
-    return ChatRequest.from_string(data, model="")
+    return ChatRequest.from_string(data, model="unknown-model")
 
 
 GlobalTypeConverter.register_converter(_string_to_nat_chat_request)
 
 
+def _chat_request_or_message_to_chat_request(data: ChatRequestOrMessage) -> ChatRequest:
+    if data.input_string is not None:
+        return _string_to_nat_chat_request(data.input_string)
+    return ChatRequest(**data.model_dump(exclude={"input_string"}))
+
+
+GlobalTypeConverter.register_converter(_chat_request_or_message_to_chat_request)
+
+
+def _chat_request_to_chat_request_or_message(data: ChatRequest) -> ChatRequestOrMessage:
+    return ChatRequestOrMessage(**data.model_dump(by_alias=True))
+
+
+GlobalTypeConverter.register_converter(_chat_request_to_chat_request_or_message)
+
+
+def _chat_request_or_message_to_string(data: ChatRequestOrMessage) -> str:
+    return data.input_string or ""
+
+
+GlobalTypeConverter.register_converter(_chat_request_or_message_to_string)
+
+
+def _string_to_chat_request_or_message(data: str) -> ChatRequestOrMessage:
+    return ChatRequestOrMessage(input_message=data)
+
+
+GlobalTypeConverter.register_converter(_string_to_chat_request_or_message)
+
+
 # ======== ChatResponse Converters ========
 def _nat_chat_response_to_string(data: ChatResponse) -> str:
     if data.choices and data.choices[0].message:
@@ -654,22 +756,12 @@ def _string_to_nat_chat_response(data: str) -> ChatResponse:
 GlobalTypeConverter.register_converter(_string_to_nat_chat_response)
 
 
-def _chat_response_to_chat_response_chunk(data: ChatResponse) -> ChatResponseChunk:
-    # Preserve original message structure for backward compatibility
-    return ChatResponseChunk(id=data.id, choices=data.choices, created=data.created, model=data.model)
-
-
-GlobalTypeConverter.register_converter(_chat_response_to_chat_response_chunk)
-
-
 # ======== ChatResponseChunk Converters ========
 def _chat_response_chunk_to_string(data: ChatResponseChunk) -> str:
     if data.choices and len(data.choices) > 0:
         choice = data.choices[0]
         if choice.delta and choice.delta.content:
             return choice.delta.content
-        if choice.message and choice.message.content:
-            return choice.message.content
     return ""
 
 
@@ -685,21 +777,6 @@ def _string_to_nat_chat_response_chunk(data: str) -> ChatResponseChunk:
 
 GlobalTypeConverter.register_converter(_string_to_nat_chat_response_chunk)
 
-
-# ======== AINodeMessageChunk Converters ========
-def _ai_message_chunk_to_nat_chat_response_chunk(data) -> ChatResponseChunk:
-    '''Converts LangChain/LangGraph AINodeMessageChunk to ChatResponseChunk'''
-    content = ""
-    if hasattr(data, 'content') and data.content is not None:
-        content = str(data.content)
-    elif hasattr(data, 'text') and data.text is not None:
-        content = str(data.text)
-    elif hasattr(data, 'message') and data.message is not None:
-        content = str(data.message)
-
-    return ChatResponseChunk.create_streaming_chunk(content=content, role="assistant", finish_reason=None)
-
-
 # Compatibility aliases with previous releases
 AIQChatRequest = ChatRequest
 AIQChoiceMessage = ChoiceMessage

nat/data_models/config.py

@@ -187,7 +187,7 @@ class TelemetryConfig(BaseModel):
 
 class GeneralConfig(BaseModel):
 
-    model_config = ConfigDict(protected_namespaces=())
+    model_config = ConfigDict(protected_namespaces=(), extra="forbid")
 
     use_uvloop: bool | None = Field(
         default=None,

nat/data_models/span.py

@@ -128,10 +128,48 @@ class SpanStatus(BaseModel):
     message: str | None = Field(default=None, description="The status message of the span.")
 
 
+def _generate_nonzero_trace_id() -> int:
+    """Generate a non-zero 128-bit trace ID."""
+    return uuid.uuid4().int
+
+
+def _generate_nonzero_span_id() -> int:
+    """Generate a non-zero 64-bit span ID."""
+    return uuid.uuid4().int >> 64
+
+
 class SpanContext(BaseModel):
-    trace_id: int = Field(default_factory=lambda: uuid.uuid4().int, description="The 128-bit trace ID of the span.")
-    span_id: int = Field(default_factory=lambda: uuid.uuid4().int & ((1 << 64) - 1),
-                         description="The 64-bit span ID of the span.")
+    trace_id: int = Field(default_factory=_generate_nonzero_trace_id,
+                          description="The OTel-syle 128-bit trace ID of the span.")
+    span_id: int = Field(default_factory=_generate_nonzero_span_id,
+                         description="The OTel-syle 64-bit span ID of the span.")
+
+    @field_validator("trace_id", mode="before")
+    @classmethod
+    def _validate_trace_id(cls, v: int | str | None) -> int:
+        """Regenerate if trace_id is None; raise an exception if trace_id is invalid;"""
+        if isinstance(v, str):
+            v = uuid.UUID(v).int
+        if isinstance(v, type(None)):
+            v = _generate_nonzero_trace_id()
+        if v <= 0 or v >> 128:
+            raise ValueError(f"Invalid trace_id: must be a non-zero 128-bit integer, got {v}")
+        return v
+
+    @field_validator("span_id", mode="before")
+    @classmethod
+    def _validate_span_id(cls, v: int | str | None) -> int:
+        """Regenerate if span_id is None; raise an exception if span_id is invalid;"""
+        if isinstance(v, str):
+            try:
+                v = int(v, 16)
+            except ValueError:
+                raise ValueError(f"span_id unable to be parsed: {v}")
+        if isinstance(v, type(None)):
+            v = _generate_nonzero_span_id()
+        if v <= 0 or v >> 64:
+            raise ValueError(f"Invalid span_id: must be a non-zero 64-bit integer, got {v}")
+        return v
 
 
 class Span(BaseModel):

nat/experimental/test_time_compute/functions/execute_score_select_function.py

@@ -46,7 +46,7 @@ async def execute_score_select_function(config: ExecuteScoreSelectFunctionConfig
 
     from pydantic import BaseModel
 
-    executable_fn: Function = builder.get_function(name=config.augmented_fn)
+    executable_fn: Function = await builder.get_function(name=config.augmented_fn)
 
     if config.scorer:
         scorer = await builder.get_ttc_strategy(strategy_name=config.scorer,

nat/experimental/test_time_compute/functions/ttc_tool_wrapper_function.py

@@ -98,8 +98,8 @@ async def register_ttc_tool_wrapper_function(
 
         augmented_function_desc = config.tool_description
 
-    fn_input_schema: BaseModel = augmented_function.input_schema
-    fn_output_schema: BaseModel = augmented_function.single_output_schema
+    fn_input_schema: type[BaseModel] = augmented_function.input_schema
+    fn_output_schema: type[BaseModel] | type[None] = augmented_function.single_output_schema
 
     runnable_llm = input_llm.with_structured_output(schema=fn_input_schema)
 

nat/front_ends/console/console_front_end_plugin.py

@@ -95,5 +95,14 @@ class ConsoleFrontEndPlugin(SimpleFrontEndPluginBase[ConsoleFrontEndConfig]):
         else:
             assert False, "Should not reach here. Should have been caught by pre_run"
 
-        # Print result
-        logger.info(f"\n{'-' * 50}\n{Fore.GREEN}Workflow Result:\n%s{Fore.RESET}\n{'-' * 50}", runner_outputs)
+        line = f"{'-' * 50}"
+        prefix = f"{line}\n{Fore.GREEN}Workflow Result:\n"
+        suffix = f"{Fore.RESET}\n{line}"
+
+        logger.info(f"{prefix}%s{suffix}", runner_outputs)
+
+        # (handler is a stream handler) => (level > INFO)
+        effective_level_too_high = all(
+            type(h) is not logging.StreamHandler or h.level > logging.INFO for h in logging.getLogger().handlers)
+        if effective_level_too_high:
+            print(f"{prefix}{runner_outputs}{suffix}")

nat/front_ends/fastapi/auth_flow_handlers/http_flow_handler.py

@@ -24,4 +24,4 @@ class HTTPAuthenticationFlowHandler(FlowHandlerBase):
     async def authenticate(self, config: AuthProviderBaseConfig, method: AuthFlowType) -> AuthenticatedContext:
 
         raise NotImplementedError(f"Authentication method '{method}' is not supported by the HTTP frontend."
-                                  f" Do you have Websockets enabled?")
+                                  f" Do you have WebSockets enabled?")