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

nat/registry_handlers/rest/rest_handler.py

@@ -42,15 +42,14 @@ logger = logging.getLogger(__name__)
 class RestRegistryHandler(AbstractRegistryHandler):
     """A registry handler for interactions with a remote REST registry."""
 
-    def __init__(  # pylint: disable=R0917
-            self,
-            endpoint: str,
-            token: str,
-            timeout: int = 30,
-            publish_route: str = "",
-            pull_route: str = "",
-            search_route: str = "",
-            remove_route: str = ""):
+    def __init__(self,
+                 endpoint: str,
+                 token: str,
+                 timeout: int = 30,
+                 publish_route: str = "",
+                 pull_route: str = "",
+                 search_route: str = "",
+                 remove_route: str = ""):
         super().__init__()
         self._endpoint = endpoint.rstrip("/")
         self._timeout = timeout
@@ -89,7 +88,7 @@ class RestRegistryHandler(AbstractRegistryHandler):
             validated_publish_response = PublishResponse(status={
                 "status": StatusEnum.ERROR, "message": msg, "action": ActionEnum.PUBLISH
             })
-            logger.exception(validated_publish_response.status.message, exc_info=True)
+            logger.exception(validated_publish_response.status.message)
 
             yield validated_publish_response
 
@@ -156,7 +155,7 @@ class RestRegistryHandler(AbstractRegistryHandler):
             validated_pull_response = PullResponse(status={
                 "status": StatusEnum.ERROR, "message": msg, "action": ActionEnum.PULL
             })
-            logger.exception(validated_pull_response.status.message, exc_info=True)
+            logger.exception(validated_pull_response.status.message)
 
             yield validated_pull_response
 
@@ -194,7 +193,7 @@ class RestRegistryHandler(AbstractRegistryHandler):
                                                            "message": msg,
                                                            "action": ActionEnum.SEARCH
                                                        })
-            logger.exception(validated_search_response.status.message, exc_info=True)
+            logger.exception(validated_search_response.status.message)
 
             yield validated_search_response
 
@@ -229,7 +228,7 @@ class RestRegistryHandler(AbstractRegistryHandler):
             validated_remove_response = RemoveResponse(status={
                 "status": StatusEnum.ERROR, "message": msg, "action": ActionEnum.REMOVE
             })
-            logger.exception(validated_remove_response.status.message, exc_info=True)
+            logger.exception(validated_remove_response.status.message)
 
             yield validated_remove_response
 

nat/retriever/milvus/retriever.py

@@ -154,7 +154,7 @@ class MilvusRetriever(Retriever):
                 return _wrap_milvus_results(results, content_field=self.content_field)
 
         except Exception as e:
-            logger.exception("Exception when retrieving results from milvus for query %s: %s", query, e)
+            logger.error("Exception when retrieving results from milvus for query %s: %s", query, e)
             raise RetrieverError(f"Error when retrieving documents from {collection_name} for query '{query}'") from e
 
     async def _search(self,
@@ -214,7 +214,7 @@ def _wrap_milvus_results(res: list[Hit], content_field: str):
 
 
 def _wrap_milvus_single_results(res: Hit | dict, content_field: str) -> Document:
-    if not isinstance(res, (Hit, dict)):
+    if not isinstance(res, Hit | dict):
         raise ValueError(f"Milvus search returned object of type {type(res)}. Expected 'Hit' or 'dict'.")
 
     if isinstance(res, Hit):

nat/retriever/nemo_retriever/retriever.py

@@ -143,7 +143,7 @@ class NemoRetriever(Retriever):
                 return _wrap_nemo_results(output=output, content_field="content")
 
         except Exception as e:
-            logger.exception("Encountered an error when retrieving results from Nemo Retriever: %s", e)
+            logger.error("Encountered an error when retrieving results from Nemo Retriever: %s", e)
             raise CollectionUnavailableError(
                 f"Error when retrieving documents from {collection_name} for query '{query}'") from e
 

nat/retriever/register.py

@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# pylint: disable=unused-import
 # flake8: noqa
 # isort:skip_file
 

nat/runtime/loader.py

@@ -114,7 +114,7 @@ async def load_workflow(config_file: StrPath, max_concurrency: int = -1):
     # Must yield the workflow function otherwise it cleans up
     async with WorkflowBuilder.from_config(config=config) as workflow:
 
-        yield SessionManager(workflow.build(), max_concurrency=max_concurrency)
+        yield SessionManager(await workflow.build(), max_concurrency=max_concurrency)
 
 
 @lru_cache
@@ -210,7 +210,7 @@ def discover_and_register_plugins(plugin_type: PluginTypes):
                 # Optionally, you can mark the plugin as unavailable or take other actions
 
             except Exception:
-                logger.exception("An error occurred while loading plugin '%s': {e}", entry_point.name, exc_info=True)
+                logger.exception("An error occurred while loading plugin '%s'", entry_point.name)
 
             finally:
                 count += 1

nat/runtime/runner.py

@@ -149,7 +149,8 @@ class Runner:
 
             return result
         except Exception as e:
-            logger.exception("Error running workflow: %s", e)
+            err_msg = f": {e}" if str(e).strip() else "."
+            logger.error("Error running workflow%s", err_msg)
             event_stream = self._context_state.event_stream.get()
             if event_stream:
                 event_stream.on_complete()
@@ -181,7 +182,7 @@ class Runner:
                     event_stream.on_complete()
 
         except Exception as e:
-            logger.exception("Error running workflow: %s", e)
+            logger.error("Error running workflow: %s", e)
             event_stream = self._context_state.event_stream.get()
             if event_stream:
                 event_stream.on_complete()

nat/runtime/session.py

@@ -21,7 +21,9 @@ from collections.abc import Callable
 from contextlib import asynccontextmanager
 from contextlib import nullcontext
 
+from fastapi import WebSocket
 from starlette.requests import HTTPConnection
+from starlette.requests import Request
 
 from nat.builder.context import Context
 from nat.builder.context import ContextState
@@ -89,7 +91,8 @@ class SessionManager:
     @asynccontextmanager
     async def session(self,
                       user_manager=None,
-                      request: HTTPConnection | None = None,
+                      http_connection: HTTPConnection | None = None,
+                      user_message_id: str | None = None,
                       conversation_id: str | None = None,
                       user_input_callback: Callable[[InteractionPrompt], Awaitable[HumanResponse]] = None,
                       user_authentication_callback: Callable[[AuthProviderBaseConfig, AuthFlowType],
@@ -107,10 +110,11 @@ class SessionManager:
         if user_authentication_callback is not None:
             token_user_authentication = self._context_state.user_auth_callback.set(user_authentication_callback)
 
-        if conversation_id is not None and request is None:
-            self._context_state.conversation_id.set(conversation_id)
+        if isinstance(http_connection, WebSocket):
+            self.set_metadata_from_websocket(http_connection, user_message_id, conversation_id)
 
-        self.set_metadata_from_http_request(request)
+        if isinstance(http_connection, Request):
+            self.set_metadata_from_http_request(http_connection)
 
         try:
             yield self
@@ -135,14 +139,11 @@ class SessionManager:
             async with self._workflow.run(message) as runner:
                 yield runner
 
-    def set_metadata_from_http_request(self, request: HTTPConnection | None) -> None:
+    def set_metadata_from_http_request(self, request: Request) -> None:
         """
         Extracts and sets user metadata request attributes from a HTTP request.
         If request is None, no attributes are set.
         """
-        if request is None:
-            return
-
         self._context.metadata._request.method = getattr(request, "method", None)
         self._context.metadata._request.url_path = request.url.path
         self._context.metadata._request.url_port = request.url.port
@@ -157,6 +158,40 @@ class SessionManager:
         if request.headers.get("conversation-id"):
             self._context_state.conversation_id.set(request.headers["conversation-id"])
 
+        if request.headers.get("user-message-id"):
+            self._context_state.user_message_id.set(request.headers["user-message-id"])
+
+    def set_metadata_from_websocket(self,
+                                    websocket: WebSocket,
+                                    user_message_id: str | None,
+                                    conversation_id: str | None) -> None:
+        """
+        Extracts and sets user metadata for Websocket connections.
+        """
+
+        # Extract cookies from WebSocket headers (similar to HTTP request)
+        if websocket and hasattr(websocket, 'scope') and 'headers' in websocket.scope:
+            cookies = {}
+            for header_name, header_value in websocket.scope.get('headers', []):
+                if header_name == b'cookie':
+                    cookie_header = header_value.decode('utf-8')
+                    # Parse cookie header: "name1=value1; name2=value2"
+                    for cookie in cookie_header.split(';'):
+                        cookie = cookie.strip()
+                        if '=' in cookie:
+                            name, value = cookie.split('=', 1)
+                            cookies[name.strip()] = value.strip()
+
+            # Set cookies in metadata (same as HTTP request)
+            self._context.metadata._request.cookies = cookies
+            self._context_state.metadata.set(self._context.metadata)
+
+        if conversation_id is not None:
+            self._context_state.conversation_id.set(conversation_id)
+
+        if user_message_id is not None:
+            self._context_state.user_message_id.set(user_message_id)
+
 
 # Compatibility aliases with previous releases
 AIQSessionManager = SessionManager

nat/settings/global_settings.py

@@ -47,6 +47,12 @@ class Settings(HashableBaseModel):
     # Registry Handeler Configuration
     channels: dict[str, RegistryHandlerBaseConfig] = {}
 
+    # Timezone fallback behavior
+    # Options:
+    #  - "utc": default to UTC
+    #  - "system": use the system's local timezone
+    fallback_timezone: typing.Literal["system", "utc"] = "utc"
+
     _configuration_directory: typing.ClassVar[str]
     _settings_changed_hooks: typing.ClassVar[list[Callable[[], None]]] = []
     _settings_changed_hooks_active: bool = True
@@ -118,8 +124,7 @@ class Settings(HashableBaseModel):
                 if (short_names[key.local_name] == 1):
                     type_list.append((key.local_name, key.config_type))
 
-            # pylint: disable=consider-alternative-union-syntax
-            return typing.Union[tuple(typing.Annotated[x_type, Tag(x_id)] for x_id, x_type in type_list)]
+            return typing.Union[*tuple(typing.Annotated[x_type, Tag(x_id)] for x_id, x_type in type_list)]
 
         RegistryHandlerAnnotation = dict[
             str,
@@ -164,8 +169,12 @@ class Settings(HashableBaseModel):
         if (not os.path.exists(configuration_file)):
             loaded_config = {}
         else:
-            with open(file_path, mode="r", encoding="utf-8") as f:
-                loaded_config = json.load(f)
+            with open(file_path, encoding="utf-8") as f:
+                try:
+                    loaded_config = json.load(f)
+                except Exception as e:
+                    logger.exception("Error loading configuration file %s: %s", file_path, e)
+                    loaded_config = {}
 
         settings = Settings(**loaded_config)
         settings.set_configuration_directory(configuration_directory)
@@ -214,13 +223,15 @@ class Settings(HashableBaseModel):
                 match field:
                     case "channels":
                         self.channels = validated_data.channels
+                    case "fallback_timezone":
+                        self.fallback_timezone = validated_data.fallback_timezone
                     case _:
                         raise ValueError(f"Encountered invalid model field: {field}")
 
             return True
 
         except Exception as e:
-            logger.exception("Unable to validate user settings configuration: %s", e, exc_info=True)
+            logger.exception("Unable to validate user settings configuration: %s", e)
             return False
 
     def print_channel_settings(self, channel_type: str | None = None) -> None:

nat/tool/chat_completion.py

@@ -44,7 +44,7 @@ async def register_chat_completion(config: ChatCompletionConfig, builder: Builde
     """Registers a chat completion function that can handle natural language queries."""
 
     # Get the LLM from the builder context using the configured LLM reference
-    # Use LangChain framework wrapper since we're using LangChain-based LLM
+    # Use LangChain/LangGraph framework wrapper since we're using LangChain/LangGraph-based LLM
     llm = await builder.get_llm(config.llm_name, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
 
     async def _chat_completion(query: str) -> str:
@@ -63,7 +63,10 @@ async def register_chat_completion(config: ChatCompletionConfig, builder: Builde
             # Generate response using the LLM
             response = await llm.ainvoke(prompt)
 
-            return response
+            if isinstance(response, str):
+                return response
+
+            return response.text()
 
         except Exception as e:
             # Fallback response if LLM call fails

nat/tool/code_execution/local_sandbox/local_sandbox_server.py

@@ -62,7 +62,7 @@ class CodeExecutionResponse(Response):
         super().__init__(status=status_code, mimetype="application/json", response=result.model_dump_json())
 
     @classmethod
-    def with_error(cls, status_code: int, error_message: str) -> 'CodeExecutionResponse':
+    def with_error(cls, status_code: int, error_message: str) -> CodeExecutionResponse:
         return cls(status_code,
                    CodeExecutionResult(process_status=CodeExecutionStatus.ERROR, stdout="", stderr=error_message))
 
@@ -121,13 +121,13 @@ def execute_code_subprocess(generated_code: str, queue):
         resource.setrlimit(resource.RLIMIT_AS, (limit, limit))
         resource.setrlimit(resource.RLIMIT_DATA, (limit, limit))
     except Exception as e:
-        logger.error("Failed to set resource limits, PID: %s, error: %s", os.getpid(), e)
+        logger.exception("Failed to set resource limits, PID: %s, error: %s", os.getpid(), e)
 
     stdout_capture = StringIO()
     stderr_capture = StringIO()
     try:
         with contextlib.redirect_stdout(stdout_capture), contextlib.redirect_stderr(stderr_capture):
-            exec(generated_code, {})  # pylint: disable=W0122
+            exec(generated_code, {})
         logger.debug("execute_code_subprocess finished, PID: %s", os.getpid())
         queue.put(CodeExecutionResult(stdout=stdout_capture.getvalue(), stderr=stderr_capture.getvalue()))
     except Exception as e:

nat/tool/datetime_tools.py

@@ -13,30 +13,70 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import datetime
+import zoneinfo
+
+from starlette.datastructures import Headers
+
 from nat.builder.builder import Builder
 from nat.builder.function_info import FunctionInfo
 from nat.cli.register_workflow import register_function
 from nat.data_models.function import FunctionBaseConfig
+from nat.settings.global_settings import GlobalSettings
 
 
 class CurrentTimeToolConfig(FunctionBaseConfig, name="current_datetime"):
     """
-    Simple tool which returns the current date and time in human readable format.
+    Simple tool which returns the current date and time in human readable format with timezone information. By default,
+    the timezone is in Etc/UTC. If the user provides a timezone in the header, we will use it. Timezone will be
+    provided in IANA zone name format. For example, "America/New_York" or "Etc/UTC".
     """
     pass
 
 
-@register_function(config_type=CurrentTimeToolConfig)
-async def current_datetime(config: CurrentTimeToolConfig, builder: Builder):
+def _get_timezone_obj(headers: Headers | None) -> zoneinfo.ZoneInfo | datetime.tzinfo:
+    # Default to UTC
+    timezone_obj = zoneinfo.ZoneInfo("Etc/UTC")
+
+    if headers:
+        # If user has provided a timezone in the header, we will prioritize on using it
+        timezone_header = headers.get("x-timezone")
+        if timezone_header:
+            try:
+                timezone_obj = zoneinfo.ZoneInfo(timezone_header)
+            except Exception:
+                pass
+    else:
+        # Only if a timezone is not in the header, we will determine default timezone based on global settings
+        fallback_tz = GlobalSettings.get().fallback_timezone
+
+        if fallback_tz == "system":
+            # Use the system's local timezone. Avoid requiring external deps.
+            timezone_obj = datetime.datetime.now().astimezone().tzinfo or zoneinfo.ZoneInfo("Etc/UTC")
+
+    return timezone_obj
 
-    import datetime
+
+@register_function(config_type=CurrentTimeToolConfig)
+async def current_datetime(_config: CurrentTimeToolConfig, _builder: Builder):
 
     async def _get_current_time(unused: str) -> str:
 
-        now = datetime.datetime.now()  # Get current time
-        now_human_readable = now.strftime(("%Y-%m-%d %H:%M:%S"))
+        del unused  # Unused parameter to avoid linting error
+
+        from nat.builder.context import Context
+        nat_context = Context.get()
+
+        headers: Headers | None = nat_context.metadata.headers
+
+        timezone_obj = _get_timezone_obj(headers)
+
+        now = datetime.datetime.now(timezone_obj)
+        now_machine_readable = now.strftime("%Y-%m-%d %H:%M:%S %z")
 
-        return f"The current time of day is {now_human_readable}"  # Format time in H:MM AM/PM format
+        # Returns the current time in machine readable format with timezone offset.
+        return f"The current time of day is {now_machine_readable}"
 
-    yield FunctionInfo.from_fn(_get_current_time,
-                               description="Returns the current date and time in human readable format.")
+    yield FunctionInfo.from_fn(
+        _get_current_time,
+        description="Returns the current date and time in human readable format with timezone information.")

nat/tool/document_search.py

@@ -53,7 +53,7 @@ async def document_search(config: MilvusDocumentSearchToolConfig, builder: Build
     from langchain_core.messages import HumanMessage
     from langchain_core.messages import SystemMessage
     from langchain_core.pydantic_v1 import BaseModel
-    from langchain_core.pydantic_v1 import Field  # pylint: disable=redefined-outer-name, reimported
+    from langchain_core.pydantic_v1 import Field
 
     # define collection store
     # create a list of tuples using enumerate()
@@ -119,7 +119,7 @@ Return only the name of the predicted collection."""
             if len(results["chunks"]) == 0:
                 return DocumentSearchOutput(collection_name=llm_pred.collection_name, documents="")
 
-            # parse docs from Langchain Document object to string
+            # parse docs from LangChain/LangGraph Document object to string
             parsed_docs = []
 
             # iterate over results and store parsed content