nvidia-nat-crewai 1.2.1rc1__py3-none-any.whl → 1.3.0__py3-none-any.whl

nat/plugins/crewai/crewai_callback_handler.py

@@ -41,6 +41,7 @@ class CrewAIProfilerHandler(BaseProfilerCallback):
     A callback manager/handler for CrewAI that intercepts calls to:
       - ToolUsage._use
       - LLM Calls
+
     to collect usage statistics (tokens, inputs, outputs, time intervals, etc.)
     and store them in NAT's usage_stats queue for subsequent analysis.
     """
@@ -94,7 +95,7 @@ class CrewAIProfilerHandler(BaseProfilerCallback):
                 if tool_info:
                     tool_name = tool_info.name
             except Exception as e:
-                logger.exception("Error getting tool name: %s", e, exc_info=True)
+                logger.exception("Error getting tool name: %s", e)
 
             try:
                 # Pre-call usage event
@@ -132,7 +133,7 @@ class CrewAIProfilerHandler(BaseProfilerCallback):
                 return result
 
             except Exception as e:
-                logger.exception("ToolUsage._use error: %s", e)
+                logger.error("ToolUsage._use error: %s", e)
                 raise
 
         return wrapped_tool_use
@@ -153,12 +154,15 @@ class CrewAIProfilerHandler(BaseProfilerCallback):
             seconds_between_calls = int(now - self.last_call_ts)
             model_name = kwargs.get('model', "")
 
-            model_input = ""
+            model_input = []
             try:
                 for message in kwargs.get('messages', []):
-                    model_input += message.get('content', "")
+                    content = message.get('content', "")
+                    model_input.append("" if content is None else str(content))
             except Exception as e:
-                logger.exception("Error getting model input: %s", e, exc_info=True)
+                logger.exception("Error getting model input: %s", e)
+
+            model_input = "".join(model_input)
 
             # Record the start event
             input_stats = IntermediateStepPayload(
@@ -176,13 +180,16 @@ class CrewAIProfilerHandler(BaseProfilerCallback):
             # Call the original litellm.completion(...)
             output = original_func(*args, **kwargs)
 
-            model_output = ""
+            model_output = []
             try:
                 for choice in output.choices:
                     msg = choice.model_extra["message"]
-                    model_output += msg.get('content', "")
+                    content = msg.get('content', "")
+                    model_output.append("" if content is None else str(content))
             except Exception as e:
-                logger.exception("Error getting model output: %s", e, exc_info=True)
+                logger.exception("Error getting model output: %s", e)
+
+            model_output = "".join(model_output)
 
             now = time.time()
             # Record the end event

nat/plugins/crewai/llm.py

@@ -14,67 +14,135 @@
 # limitations under the License.
 
 import os
+from typing import TypeVar
 
 from nat.builder.builder import Builder
 from nat.builder.framework_enum import LLMFrameworkEnum
 from nat.cli.register_workflow import register_llm_client
+from nat.data_models.llm import LLMBaseConfig
 from nat.data_models.retry_mixin import RetryMixin
+from nat.data_models.thinking_mixin import ThinkingMixin
+from nat.llm.azure_openai_llm import AzureOpenAIModelConfig
+from nat.llm.litellm_llm import LiteLlmModelConfig
 from nat.llm.nim_llm import NIMModelConfig
 from nat.llm.openai_llm import OpenAIModelConfig
+from nat.llm.utils.thinking import BaseThinkingInjector
+from nat.llm.utils.thinking import FunctionArgumentWrapper
+from nat.llm.utils.thinking import patch_with_thinking
 from nat.utils.exception_handlers.automatic_retries import patch_with_retry
+from nat.utils.type_utils import override
 
+ModelType = TypeVar("ModelType")
 
-@register_llm_client(config_type=NIMModelConfig, wrapper_type=LLMFrameworkEnum.CREWAI)
-async def nim_crewai(llm_config: NIMModelConfig, builder: Builder):
-
-    from crewai import LLM
-
-    config_obj = {
-        **llm_config.model_dump(exclude={"type"}, by_alias=True),
-        "model": f"nvidia_nim/{llm_config.model_name}",
-    }
 
-    # Because CrewAI uses a different environment variable for the API key, we need to set it here manually
-    if ("api_key" not in config_obj or config_obj["api_key"] is None):
-
-        if ("NVIDIA_NIM_API_KEY" in os.environ):
-            # Dont need to do anything. User has already set the correct key
-            pass
-        else:
-            nvidai_api_key = os.getenv("NVIDIA_API_KEY")
+def _patch_llm_based_on_config(client: ModelType, llm_config: LLMBaseConfig) -> ModelType:
 
-            if (nvidai_api_key is not None):
-                # Transfer the key to the correct environment variable for LiteLLM
-                os.environ["NVIDIA_NIM_API_KEY"] = nvidai_api_key
+    class CrewAIThinkingInjector(BaseThinkingInjector):
 
-    client = LLM(**config_obj)
+        @override
+        def inject(self, messages: list[dict[str, str]], *args, **kwargs) -> FunctionArgumentWrapper:
+            # Attempt to inject the system prompt into the first system message
+            for i, message in enumerate(messages):
+                if message["role"] == "system":
+                    if self.system_prompt not in message["content"]:
+                        messages = list(messages)
+                        messages[i] = {"role": "system", "content": f"{message['content']}\n{self.system_prompt}"}
+                    break
+            else:
+                messages = list(messages)
+                messages.insert(0, {"role": "system", "content": self.system_prompt})
+            return FunctionArgumentWrapper(messages, *args, **kwargs)
 
     if isinstance(llm_config, RetryMixin):
-
         client = patch_with_retry(client,
                                   retries=llm_config.num_retries,
                                   retry_codes=llm_config.retry_on_status_codes,
                                   retry_on_messages=llm_config.retry_on_errors)
 
-    yield client
+    if isinstance(llm_config, ThinkingMixin) and llm_config.thinking_system_prompt is not None:
+        client = patch_with_thinking(
+            client, CrewAIThinkingInjector(
+                system_prompt=llm_config.thinking_system_prompt,
+                function_names=["call"],
+            ))
+
+    return client
+
+
+@register_llm_client(config_type=AzureOpenAIModelConfig, wrapper_type=LLMFrameworkEnum.CREWAI)
+async def azure_openai_crewai(llm_config: AzureOpenAIModelConfig, _builder: Builder):
+
+    from crewai import LLM
+
+    # https://docs.crewai.com/en/concepts/llms#azure
+
+    api_key = llm_config.api_key or os.environ.get("AZURE_OPENAI_API_KEY") or os.environ.get("AZURE_API_KEY")
+    if api_key is None:
+        raise ValueError("Azure API key is not set")
+    os.environ["AZURE_API_KEY"] = api_key
+    api_base = (llm_config.azure_endpoint or os.environ.get("AZURE_OPENAI_ENDPOINT")
+                or os.environ.get("AZURE_API_BASE"))
+    if api_base is None:
+        raise ValueError("Azure endpoint is not set")
+    os.environ["AZURE_API_BASE"] = api_base
+
+    os.environ["AZURE_API_VERSION"] = llm_config.api_version
+    model = llm_config.azure_deployment or os.environ.get("AZURE_MODEL_DEPLOYMENT")
+    if model is None:
+        raise ValueError("Azure model deployment is not set")
+
+    client = LLM(
+        **llm_config.model_dump(
+            exclude={
+                "type",
+                "api_key",
+                "azure_endpoint",
+                "azure_deployment",
+                "thinking",
+            },
+            by_alias=True,
+            exclude_none=True,
+        ),
+        model=model,
+    )
+
+    yield _patch_llm_based_on_config(client, llm_config)
+
+
+@register_llm_client(config_type=NIMModelConfig, wrapper_type=LLMFrameworkEnum.CREWAI)
+async def nim_crewai(llm_config: NIMModelConfig, _builder: Builder):
+
+    from crewai import LLM
+
+    # Because CrewAI uses a different environment variable for the API key, we need to set it here manually
+    if llm_config.api_key is None and "NVIDIA_NIM_API_KEY" not in os.environ:
+        nvidia_api_key = os.getenv("NVIDIA_API_KEY")
+        if nvidia_api_key is not None:
+            os.environ["NVIDIA_NIM_API_KEY"] = nvidia_api_key
+
+    client = LLM(
+        **llm_config.model_dump(exclude={"type", "model_name", "thinking"}, by_alias=True, exclude_none=True),
+        model=f"nvidia_nim/{llm_config.model_name}",
+    )
+
+    yield _patch_llm_based_on_config(client, llm_config)
 
 
 @register_llm_client(config_type=OpenAIModelConfig, wrapper_type=LLMFrameworkEnum.CREWAI)
-async def openai_crewai(llm_config: OpenAIModelConfig, builder: Builder):
+async def openai_crewai(llm_config: OpenAIModelConfig, _builder: Builder):
 
     from crewai import LLM
 
-    config_obj = {
-        **llm_config.model_dump(exclude={"type"}, by_alias=True),
-    }
+    client = LLM(**llm_config.model_dump(exclude={"type", "thinking"}, by_alias=True, exclude_none=True))
 
-    client = LLM(**config_obj)
+    yield _patch_llm_based_on_config(client, llm_config)
 
-    if isinstance(llm_config, RetryMixin):
 
-        client = patch_with_retry(client,
-                                  retries=llm_config.num_retries,
-                                  retry_codes=llm_config.retry_on_status_codes,
-                                  retry_on_messages=llm_config.retry_on_errors)
+@register_llm_client(config_type=LiteLlmModelConfig, wrapper_type=LLMFrameworkEnum.CREWAI)
+async def litellm_crewai(llm_config: LiteLlmModelConfig, _builder: Builder):
+
+    from crewai import LLM
+
+    client = LLM(**llm_config.model_dump(exclude={"type", "thinking"}, by_alias=True, exclude_none=True))
 
-    yield client
+    yield _patch_llm_based_on_config(client, llm_config)

nat/plugins/crewai/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
 

{nvidia_nat_crewai-1.2.1rc1.dist-info → nvidia_nat_crewai-1.3.0.dist-info}/METADATA

@@ -1,13 +1,24 @@
 Metadata-Version: 2.4
 Name: nvidia-nat-crewai
-Version: 1.2.1rc1
+Version: 1.3.0
 Summary: Subpackage for CrewAI integration in NeMo Agent toolkit
+Author: NVIDIA Corporation
+Maintainer: NVIDIA Corporation
+License: Apache-2.0
+Project-URL: documentation, https://docs.nvidia.com/nemo/agent-toolkit/latest/
+Project-URL: source, https://github.com/NVIDIA/NeMo-Agent-Toolkit
 Keywords: ai,rag,agents
 Classifier: Programming Language :: Python
-Requires-Python: <3.13,>=3.11
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Python: <3.14,>=3.11
 Description-Content-Type: text/markdown
-Requires-Dist: nvidia-nat==v1.2.1-rc1
-Requires-Dist: crewai~=0.95.0
+License-File: LICENSE-3rd-party.txt
+License-File: LICENSE.md
+Requires-Dist: nvidia-nat[litellm]==v1.3.0
+Requires-Dist: crewai~=0.193.2
+Dynamic: license-file
 
 <!--
 SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.

nvidia_nat_crewai-1.3.0.dist-info/RECORD

@@ -0,0 +1,13 @@
+nat/meta/pypi.md,sha256=T68FnThRzDGFf1LR8u-okM-r11-skSnKqSyI6HOktQY,1107
+nat/plugins/crewai/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+nat/plugins/crewai/crewai_callback_handler.py,sha256=il537F5tD9pFL1P9Q38ReOZasD-GgcBrm8BX_w0-xdo,8582
+nat/plugins/crewai/llm.py,sha256=sfQsdSK6WGWHscald4tG4dyXEcn6zqdhkFjE-1zAtJE,5984
+nat/plugins/crewai/register.py,sha256=_R3bhGmz___696_NwyIcpw3koMBiWqIFoWEFJ0VAgXs,831
+nat/plugins/crewai/tool_wrapper.py,sha256=BNKEPQQCLKtXNzGDAKBLCdmGJXe9lBOVI1hObha8hoI,1569
+nvidia_nat_crewai-1.3.0.dist-info/licenses/LICENSE-3rd-party.txt,sha256=fOk5jMmCX9YoKWyYzTtfgl-SUy477audFC5hNY4oP7Q,284609
+nvidia_nat_crewai-1.3.0.dist-info/licenses/LICENSE.md,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
+nvidia_nat_crewai-1.3.0.dist-info/METADATA,sha256=73nP1c_6H0NHvZVzKlfZLNlPCysIQchATIPnFsXtH7w,1904
+nvidia_nat_crewai-1.3.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nvidia_nat_crewai-1.3.0.dist-info/entry_points.txt,sha256=YF5PUdQGr_OUDXB4TykElHJTsKT8yKkuE0bMX5n_RXs,58
+nvidia_nat_crewai-1.3.0.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
+nvidia_nat_crewai-1.3.0.dist-info/RECORD,,