nvidia-nat 1.4.0a20251102__py3-none-any.whl → 1.4.0a20251120__py3-none-any.whl

nat/retriever/milvus/register.py

@@ -48,6 +48,7 @@ class MilvusRetrieverConfig(RetrieverBaseConfig, name="milvus_retriever"):
     description: str | None = Field(default=None,
                                     description="If present it will be used as the tool description",
                                     alias="collection_description")
+    use_async_client: bool = Field(default=False, description="Use AsyncMilvusClient for async I/O operations. ")
 
 
 @register_retriever_provider(config_type=MilvusRetrieverConfig)
@@ -58,13 +59,20 @@ async def milvus_retriever(retriever_config: MilvusRetrieverConfig, builder: Bui
 
 @register_retriever_client(config_type=MilvusRetrieverConfig, wrapper_type=None)
 async def milvus_retriever_client(config: MilvusRetrieverConfig, builder: Builder):
-    from pymilvus import MilvusClient
-
     from nat.retriever.milvus.retriever import MilvusRetriever
 
     embedder = await builder.get_embedder(embedder_name=config.embedding_model, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
 
-    milvus_client = MilvusClient(uri=str(config.uri), **config.connection_args)
+    # Create Milvus client based on use_async_client flag
+    if config.use_async_client:
+        from pymilvus import AsyncMilvusClient
+
+        milvus_client = AsyncMilvusClient(uri=str(config.uri), **config.connection_args)
+    else:
+        from pymilvus import MilvusClient
+
+        milvus_client = MilvusClient(uri=str(config.uri), **config.connection_args)
+
     retriever = MilvusRetriever(
         client=milvus_client,
         embedder=embedder,

nat/retriever/milvus/retriever.py

@@ -13,13 +13,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import inspect
 import logging
 from functools import partial
+from typing import TYPE_CHECKING
 
 from langchain_core.embeddings import Embeddings
-from pymilvus import MilvusClient
 from pymilvus.client.abstract import Hit
 
+if TYPE_CHECKING:
+    from pymilvus import AsyncMilvusClient
+    from pymilvus import MilvusClient
+
 from nat.retriever.interface import Retriever
 from nat.retriever.models import Document
 from nat.retriever.models import RetrieverError
@@ -39,20 +44,27 @@ class MilvusRetriever(Retriever):
 
     def __init__(
         self,
-        client: MilvusClient,
+        client: "MilvusClient | AsyncMilvusClient",
         embedder: Embeddings,
         content_field: str = "text",
         use_iterator: bool = False,
     ) -> None:
         """
-        Initialize the Milvus Retriever using a preconfigured MilvusClient
+        Initialize the Milvus Retriever using a preconfigured MilvusClient or AsyncMilvusClient
 
         Args:
-           client (MilvusClient): Preinstantiate pymilvus.MilvusClient object.
         """
-        self._client = client
+        self._client: MilvusClient | AsyncMilvusClient = client
         self._embedder = embedder
 
+        # Detect if client is async by inspecting method capabilities
+        search_method = getattr(client, "search", None)
+        list_collections_method = getattr(client, "list_collections", None)
+        self._is_async = any(
+            inspect.iscoroutinefunction(method) for method in (search_method, list_collections_method)
+            if method is not None)
+        logger.info("Initialized Milvus Retriever with %s client", "async" if self._is_async else "sync")
+
         if use_iterator and "search_iterator" not in dir(self._client):
             raise ValueError("This version of the pymilvus.MilvusClient does not support the search iterator.")
 
@@ -60,7 +72,7 @@ class MilvusRetriever(Retriever):
         self._default_params = None
         self._bound_params = []
         self.content_field = content_field
-        logger.info("Mivlus Retriever using %s for search.", self._search_func.__name__)
+        logger.info("Milvus Retriever using %s for search.", self._search_func.__name__)
 
     def bind(self, **kwargs) -> None:
         """
@@ -81,8 +93,13 @@ class MilvusRetriever(Retriever):
         """
         return [param for param in ["query", "collection_name", "top_k", "filters"] if param not in self._bound_params]
 
-    def _validate_collection(self, collection_name: str) -> bool:
-        return collection_name in self._client.list_collections()
+    async def _validate_collection(self, collection_name: str) -> bool:
+        """Validate that a collection exists."""
+        if self._is_async:
+            collections = await self._client.list_collections()
+        else:
+            collections = self._client.list_collections()
+        return collection_name in collections
 
     async def search(self, query: str, **kwargs):
         return await self._search_func(query=query, **kwargs)
@@ -108,39 +125,64 @@ class MilvusRetriever(Retriever):
                      collection_name,
                      top_k)
 
-        if not self._validate_collection(collection_name):
+        if not await self._validate_collection(collection_name):
             raise CollectionNotFoundError(f"Collection: {collection_name} does not exist")
 
         # If no output fields are specified, return all of them
         if not output_fields:
-            collection_schema = self._client.describe_collection(collection_name)
+            if self._is_async:
+                collection_schema = await self._client.describe_collection(collection_name)
+            else:
+                collection_schema = self._client.describe_collection(collection_name)
             output_fields = [
                 field["name"] for field in collection_schema.get("fields") if field["name"] != vector_field_name
             ]
 
-        search_vector = self._embedder.embed_query(query)
-
-        search_iterator = self._client.search_iterator(
-            collection_name=collection_name,
-            data=[search_vector],
-            batch_size=kwargs.get("batch_size", 1000),
-            filter=filters,
-            limit=top_k,
-            output_fields=output_fields,
-            search_params=search_params if search_params else {"metric_type": "L2"},
-            timeout=timeout,
-            anns_field=vector_field_name,
-            round_decimal=kwargs.get("round_decimal", -1),
-            partition_names=kwargs.get("partition_names", None),
-        )
+        search_vector = await self._embedder.aembed_query(query)
+
+        # Create search iterator
+        if self._is_async:
+            search_iterator = await self._client.search_iterator(
+                collection_name=collection_name,
+                data=[search_vector],
+                batch_size=kwargs.get("batch_size", 1000),
+                filter=filters,
+                limit=top_k,
+                output_fields=output_fields,
+                search_params=search_params if search_params else {"metric_type": "L2"},
+                timeout=timeout,
+                anns_field=vector_field_name,
+                round_decimal=kwargs.get("round_decimal", -1),
+                partition_names=kwargs.get("partition_names", None),
+            )
+        else:
+            search_iterator = self._client.search_iterator(
+                collection_name=collection_name,
+                data=[search_vector],
+                batch_size=kwargs.get("batch_size", 1000),
+                filter=filters,
+                limit=top_k,
+                output_fields=output_fields,
+                search_params=search_params if search_params else {"metric_type": "L2"},
+                timeout=timeout,
+                anns_field=vector_field_name,
+                round_decimal=kwargs.get("round_decimal", -1),
+                partition_names=kwargs.get("partition_names", None),
+            )
 
         results = []
         try:
             while True:
-                _res = search_iterator.next()
+                if self._is_async:
+                    _res = await search_iterator.next()
+                else:
+                    _res = search_iterator.next()
                 res = _res.get_res()
                 if len(_res) == 0:
-                    search_iterator.close()
+                    if self._is_async:
+                        await search_iterator.close()
+                    else:
+                        search_iterator.close()
                     break
 
                 if distance_cutoff and res[0][-1].distance > distance_cutoff:
@@ -176,10 +218,16 @@ class MilvusRetriever(Retriever):
                      collection_name,
                      top_k)
 
-        if not self._validate_collection(collection_name):
+        if not await self._validate_collection(collection_name):
             raise CollectionNotFoundError(f"Collection: {collection_name} does not exist")
 
-        available_fields = [v.get("name") for v in self._client.describe_collection(collection_name).get("fields", {})]
+        # Get collection schema
+        if self._is_async:
+            collection_schema = await self._client.describe_collection(collection_name)
+        else:
+            collection_schema = self._client.describe_collection(collection_name)
+
+        available_fields = [v.get("name") for v in collection_schema.get("fields", [])]
 
         if self.content_field not in available_fields:
             raise ValueError(f"The specified content field: {self.content_field} is not part of the schema.")
@@ -194,17 +242,31 @@ class MilvusRetriever(Retriever):
         if self.content_field not in output_fields:
             output_fields.append(self.content_field)
 
-        search_vector = self._embedder.embed_query(query)
-        res = self._client.search(
-            collection_name=collection_name,
-            data=[search_vector],
-            filter=filters,
-            output_fields=output_fields,
-            search_params=search_params if search_params else {"metric_type": "L2"},
-            timeout=timeout,
-            anns_field=vector_field_name,
-            limit=top_k,
-        )
+        search_vector = await self._embedder.aembed_query(query)
+
+        # Perform search
+        if self._is_async:
+            res = await self._client.search(
+                collection_name=collection_name,
+                data=[search_vector],
+                filter=filters,
+                output_fields=output_fields,
+                search_params=search_params if search_params else {"metric_type": "L2"},
+                timeout=timeout,
+                anns_field=vector_field_name,
+                limit=top_k,
+            )
+        else:
+            res = self._client.search(
+                collection_name=collection_name,
+                data=[search_vector],
+                filter=filters,
+                output_fields=output_fields,
+                search_params=search_params if search_params else {"metric_type": "L2"},
+                timeout=timeout,
+                anns_field=vector_field_name,
+                limit=top_k,
+            )
 
         return _wrap_milvus_results(res[0], content_field=self.content_field)
 

nat/runtime/runner.py

@@ -26,6 +26,7 @@ from nat.data_models.intermediate_step import IntermediateStepType
 from nat.data_models.intermediate_step import StreamEventData
 from nat.data_models.intermediate_step import TraceMetadata
 from nat.data_models.invocation_node import InvocationNode
+from nat.data_models.runtime_enum import RuntimeTypeEnum
 from nat.observability.exporter_manager import ExporterManager
 from nat.utils.reactive.subject import Subject
 
@@ -53,7 +54,8 @@ class Runner:
                  input_message: typing.Any,
                  entry_fn: Function,
                  context_state: ContextState,
-                 exporter_manager: ExporterManager):
+                 exporter_manager: ExporterManager,
+                 runtime_type: RuntimeTypeEnum = RuntimeTypeEnum.RUN_OR_SERVE):
         """
         The Runner class is used to run a workflow. It handles converting input and output data types and running the
         workflow with the specified concurrency.
@@ -68,6 +70,8 @@ class Runner:
             The context state to use
         exporter_manager : ExporterManager
             The exporter manager to use
+        runtime_type : RuntimeTypeEnum
+            The runtime type (RUN_OR_SERVE, EVALUATE, OTHER)
         """
 
         if (entry_fn is None):
@@ -86,6 +90,9 @@ class Runner:
 
         self._exporter_manager = exporter_manager
 
+        self._runtime_type = runtime_type
+        self._runtime_type_token = None
+
     @property
     def context(self) -> Context:
         return self._context
@@ -105,6 +112,8 @@ class Runner:
             function_id="root",
         ))
 
+        self._runtime_type_token = self._context_state.runtime_type.set(self._runtime_type)
+
         if (self._state == RunnerState.UNINITIALIZED):
             self._state = RunnerState.INITIALIZED
         else:
@@ -119,6 +128,8 @@ class Runner:
 
         self._context_state.input_message.reset(self._input_message_token)
 
+        self._context_state.runtime_type.reset(self._runtime_type_token)
+
         if (self._state not in (RunnerState.COMPLETED, RunnerState.FAILED)):
             raise ValueError("Cannot exit the context without completing the workflow")
 

nat/runtime/session.py

@@ -35,6 +35,7 @@ from nat.data_models.authentication import AuthProviderBaseConfig
 from nat.data_models.config import Config
 from nat.data_models.interactive import HumanResponse
 from nat.data_models.interactive import InteractionPrompt
+from nat.data_models.runtime_enum import RuntimeTypeEnum
 
 _T = typing.TypeVar("_T")
 
@@ -45,7 +46,10 @@ class UserManagerBase:
 
 class SessionManager:
 
-    def __init__(self, workflow: Workflow, max_concurrency: int = 8):
+    def __init__(self,
+                 workflow: Workflow,
+                 max_concurrency: int = 8,
+                 runtime_type: RuntimeTypeEnum = RuntimeTypeEnum.RUN_OR_SERVE):
         """
         The SessionManager class is used to run and manage a user workflow session. It runs and manages the context,
         and configuration of a workflow with the specified concurrency.
@@ -56,6 +60,8 @@ class SessionManager:
             The workflow to run
         max_concurrency : int, optional
             The maximum number of simultaneous workflow invocations, by default 8
+        runtime_type : RuntimeTypeEnum, optional
+            The type of runtime the session manager is operating in, by default RuntimeTypeEnum.RUN_OR_SERVE
         """
 
         if (workflow is None):
@@ -66,6 +72,7 @@ class SessionManager:
         self._max_concurrency = max_concurrency
         self._context_state = ContextState.get()
         self._context = Context(self._context_state)
+        self._runtime_type = runtime_type
 
         # We save the context because Uvicorn spawns a new process
         # for each request, and we need to restore the context vars
@@ -128,7 +135,7 @@ class SessionManager:
                 self._context_state.user_auth_callback.reset(token_user_authentication)
 
     @asynccontextmanager
-    async def run(self, message):
+    async def run(self, message, runtime_type: RuntimeTypeEnum = RuntimeTypeEnum.RUN_OR_SERVE):
         """
         Start a workflow run
         """
@@ -137,7 +144,7 @@ class SessionManager:
             for k, v in self._saved_context.items():
                 k.set(v)
 
-            async with self._workflow.run(message) as runner:
+            async with self._workflow.run(message, runtime_type=runtime_type) as runner:
                 yield runner
 
     def set_metadata_from_http_request(self, request: Request) -> None:

nat/tool/code_execution/code_sandbox.py

@@ -92,7 +92,9 @@ class Sandbox(abc.ABC):
             raise ValueError(f"Language {language} not supported")
 
         generated_code = generated_code.strip().strip("`")
-        code_to_execute = textwrap.dedent("""
+        # Use json.dumps to properly escape the generated_code instead of repr()
+        escaped_code = json.dumps(generated_code)
+        code_to_execute = textwrap.dedent(f"""
             import traceback
             import json
             import os
@@ -101,11 +103,6 @@ class Sandbox(abc.ABC):
             import io
             warnings.filterwarnings('ignore')
             os.environ['OPENBLAS_NUM_THREADS'] = '16'
-        """).strip()
-
-        # Use json.dumps to properly escape the generated_code instead of repr()
-        escaped_code = json.dumps(generated_code)
-        code_to_execute += textwrap.dedent(f"""
 
             generated_code = {escaped_code}
 
@@ -155,7 +152,7 @@ class LocalSandbox(Sandbox):
             output_json = output.json()
             assert isinstance(output_json, dict)
             return output_json
-        except json.JSONDecodeError as e:
+        except (requests.exceptions.JSONDecodeError, AssertionError) as e:
             logger.exception("Error parsing output: %s. %s", output.text, e)
             return {'process_status': 'error', 'stdout': '', 'stderr': f'Unknown error: {e} \"{output.text}\"'}
 

nat/tool/code_execution/local_sandbox/Dockerfile.sandbox

@@ -12,43 +12,26 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Use the base image with Python 3.10 and Flask
-FROM tiangolo/uwsgi-nginx-flask:python3.10
-
-# Install dependencies required for Lean 4 and other tools
-RUN apt-get update && \
-    apt-get install -y curl git && \
-    curl https://raw.githubusercontent.com/leanprover/elan/master/elan-init.sh -sSf | sh -s -- -y && \
-    /root/.elan/bin/elan toolchain install leanprover/lean4:v4.12.0 && \
-    /root/.elan/bin/elan default leanprover/lean4:v4.12.0 && \
-    /root/.elan/bin/elan self update
-
-# Set environment variables to include Lean and elan/lake in the PATH
-ENV PATH="/root/.elan/bin:$PATH"
-
-# Create Lean project directory and initialize a new Lean project with Mathlib4
-RUN mkdir -p /lean4 && cd /lean4 && \
-    /root/.elan/bin/lake new my_project && \
-    cd my_project && \
-    echo 'leanprover/lean4:v4.12.0' > lean-toolchain && \
-    echo 'require mathlib from git "https://github.com/leanprover-community/mathlib4" @ "v4.12.0"' >> lakefile.lean
-
-# Download and cache Mathlib4 to avoid recompiling, then build the project
-RUN cd /lean4/my_project && \
-    /root/.elan/bin/lake exe cache get && \
-    /root/.elan/bin/lake build
-
-# Set environment variables to include Lean project path
-ENV LEAN_PATH="/lean4/my_project"
-ENV PATH="/lean4/my_project:$PATH"
+# UWSGI_CHEAPER sets the number of initial uWSGI worker processes
+# UWSGI_PROCESSES sets the maximum number of uWSGI worker processes
+ARG UWSGI_CHEAPER=5
+ARG UWSGI_PROCESSES=10
+
+# Use the base image with Python 3.13
+FROM python:3.13-slim-bookworm
+
+
+RUN apt update && \
+    apt upgrade && \
+    apt install -y --no-install-recommends libexpat1 && \
+    apt clean && \
+    rm -rf /var/lib/apt/lists/*
 
 # Set up application code and install Python dependencies
 COPY sandbox.requirements.txt /app/requirements.txt
 RUN pip install --no-cache-dir -r /app/requirements.txt
 COPY local_sandbox_server.py /app/main.py
-
-# Set the working directory to /app
-WORKDIR /app
+RUN mkdir /workspace
 
 # Set Flask app environment variables and ports
 ARG UWSGI_CHEAPER
@@ -58,3 +41,7 @@ ARG UWSGI_PROCESSES
 ENV UWSGI_PROCESSES=$UWSGI_PROCESSES
 
 ENV LISTEN_PORT=6000
+EXPOSE 6000
+
+WORKDIR /app
+CMD uwsgi --http 0.0.0.0:${LISTEN_PORT} --master -p ${UWSGI_PROCESSES} --force-cwd /workspace -w main:app

nat/tool/code_execution/local_sandbox/local_sandbox_server.py

@@ -194,5 +194,10 @@ def execute():
     return do_execute(request)
 
 
+@app.route("/", methods=["GET"])
+def status() -> tuple[dict[str, str], int]:
+    return ({"status": "ok"}, 200)
+
+
 if __name__ == '__main__':
     app.run(port=6000)

nat/tool/code_execution/local_sandbox/sandbox.requirements.txt

@@ -1,6 +1,8 @@
+Flask==3.1
 numpy
 pandas
 scipy
 ipython
 plotly
 pydantic
+pyuwsgi==2.0.*

nat/tool/code_execution/local_sandbox/start_local_sandbox.sh

@@ -19,7 +19,11 @@
 
 DOCKER_COMMAND=${DOCKER_COMMAND:-"docker"}
 SANDBOX_NAME=${1:-'local-sandbox'}
-NUM_THREADS=10
+
+# UWSGI_CHEAPER sets the number of initial uWSGI worker processes
+# UWSGI_PROCESSES sets the maximum number of uWSGI worker processes
+UWSGI_CHEAPER=${UWSGI_CHEAPER:-5}
+UWSGI_PROCESSES=${UWSGI_PROCESSES:-10}
 
 # Get the output_data directory path for mounting
 # Priority: command line argument > environment variable > default path (current directory)
@@ -37,14 +41,16 @@ fi
 # Check if the Docker image already exists
 if ! ${DOCKER_COMMAND} images ${SANDBOX_NAME} | grep -q "${SANDBOX_NAME}"; then
     echo "Docker image not found locally. Building ${SANDBOX_NAME}..."
-    ${DOCKER_COMMAND} build --tag=${SANDBOX_NAME} --build-arg="UWSGI_PROCESSES=$((${NUM_THREADS} * 10))" --build-arg="UWSGI_CHEAPER=${NUM_THREADS}" -f Dockerfile.sandbox .
+    ${DOCKER_COMMAND} build --tag=${SANDBOX_NAME} \
+        --build-arg="UWSGI_PROCESSES=${UWSGI_PROCESSES}" \
+        --build-arg="UWSGI_CHEAPER=${UWSGI_CHEAPER}" \
+        -f Dockerfile.sandbox .
 else
     echo "Using existing Docker image: ${SANDBOX_NAME}"
 fi
 
 # Mount the output_data directory directly so files created in container appear in the local directory
-${DOCKER_COMMAND} run --rm --name=local-sandbox \
+${DOCKER_COMMAND} run --rm -ti --name=local-sandbox \
   --network=host \
   -v "${OUTPUT_DATA_PATH}:/workspace" \
-  -w /workspace \
   ${SANDBOX_NAME}

nat/tool/server_tools.py

@@ -32,14 +32,23 @@ class RequestAttributesTool(FunctionBaseConfig, name="current_request_attributes
 @register_function(config_type=RequestAttributesTool)
 async def current_request_attributes(config: RequestAttributesTool, builder: Builder):
 
+    from pydantic import RootModel
+    from pydantic.types import JsonValue
     from starlette.datastructures import Headers
     from starlette.datastructures import QueryParams
 
-    async def _get_request_attributes(unused: str) -> str:
+    class RequestBody(RootModel[JsonValue]):
+        """
+        Data model that accepts a request body of any valid JSON type.
+        """
+        root: JsonValue
+
+    async def _get_request_attributes(request_body: RequestBody) -> str:
 
         from nat.builder.context import Context
         nat_context = Context.get()
 
+        # Access request attributes from context
         method: str | None = nat_context.metadata.method
         url_path: str | None = nat_context.metadata.url_path
         url_scheme: str | None = nat_context.metadata.url_scheme
@@ -51,6 +60,9 @@ async def current_request_attributes(config: RequestAttributesTool, builder: Bui
         cookies: dict[str, str] | None = nat_context.metadata.cookies
         conversation_id: str | None = nat_context.conversation_id
 
+        # Access the request body data - can be any valid JSON type
+        request_body_data: JsonValue = request_body.root
+
         return (f"Method: {method}, "
                 f"URL Path: {url_path}, "
                 f"URL Scheme: {url_scheme}, "
@@ -60,7 +72,8 @@ async def current_request_attributes(config: RequestAttributesTool, builder: Bui
                 f"Client Host: {client_host}, "
                 f"Client Port: {client_port}, "
                 f"Cookies: {cookies}, "
-                f"Conversation Id: {conversation_id}")
+                f"Conversation Id: {conversation_id}, "
+                f"Request Body: {request_body_data}")
 
     yield FunctionInfo.from_fn(_get_request_attributes,
                                description="Returns the acquired user defined request attributes.")

nat/utils/__init__.py

@@ -29,7 +29,8 @@ async def run_workflow(*,
                        config: "Config | None" = None,
                        config_file: "StrPath | None" = None,
                        prompt: str,
-                       to_type: type[_T] = str) -> _T:
+                       to_type: type[_T] = str,
+                       session_kwargs: dict[str, typing.Any] | None = None) -> _T:
     """
     Wrapper to run a workflow given either a config or a config file path and a prompt, returning the result in the
     type specified by the `to_type`.
@@ -66,7 +67,10 @@ async def run_workflow(*,
 
         config = load_config(config_file)
 
+    session_kwargs = session_kwargs or {}
+
     async with WorkflowBuilder.from_config(config=config) as workflow_builder:
-        workflow = SessionManager(await workflow_builder.build())
-        async with workflow.run(prompt) as runner:
-            return await runner.result(to_type=to_type)
+        session_manager = SessionManager(await workflow_builder.build())
+        async with session_manager.session(**session_kwargs) as session:
+            async with session.run(prompt) as runner:
+                return await runner.result(to_type=to_type)