nvidia-nat-crewai 1.3.0a20250827__py3-none-any.whl → 1.3.0a20250828__py3-none-any.whl

nat/plugins/crewai/llm.py

@@ -14,15 +14,49 @@
 # 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.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")
+
+
+def _patch_llm_based_on_config(client: ModelType, llm_config: LLMBaseConfig) -> ModelType:
+
+    class CrewAIThinkingInjector(BaseThinkingInjector):
+
+        @override
+        def inject(self, messages: list[dict[str, str]], *args, **kwargs) -> FunctionArgumentWrapper:
+            new_messages = [{"role": "system", "content": self.system_prompt}] + messages
+            return FunctionArgumentWrapper(new_messages, *args, **kwargs)
+
+    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"],
+            ))
+
+    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)
+
+    return client
 
 
 @register_llm_client(config_type=AzureOpenAIModelConfig, wrapper_type=LLMFrameworkEnum.CREWAI)
@@ -41,33 +75,25 @@ async def azure_openai_crewai(llm_config: AzureOpenAIModelConfig, _builder: Buil
         }, by_alias=True),
     }
 
-    api_key = config_obj.get("api_key") or os.environ.get("AZURE_OPENAI_API_KEY") or os.environ.get("AZURE_API_KEY")
+    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 = (config_obj.get("azure_endpoint") or os.environ.get("AZURE_OPENAI_ENDPOINT")
+    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 = config_obj.get("azure_deployment") or os.environ.get("AZURE_MODEL_DEPLOYMENT")
+    model = llm_config.azure_deployment or os.environ.get("AZURE_MODEL_DEPLOYMENT")
     if model is None:
         raise ValueError("Azure model deployment is not set")
-
     config_obj["model"] = model
 
     client = LLM(**config_obj)
 
-    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
+    yield _patch_llm_based_on_config(client, llm_config)
 
 
 @register_llm_client(config_type=NIMModelConfig, wrapper_type=LLMFrameworkEnum.CREWAI)
@@ -81,28 +107,14 @@ async def nim_crewai(llm_config: NIMModelConfig, _builder: Builder):
     }
 
     # 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")
-
-            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
+    if config_obj.get("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(**config_obj)
 
-    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
+    yield _patch_llm_based_on_config(client, llm_config)
 
 
 @register_llm_client(config_type=OpenAIModelConfig, wrapper_type=LLMFrameworkEnum.CREWAI)
@@ -116,11 +128,4 @@ async def openai_crewai(llm_config: OpenAIModelConfig, _builder: Builder):
 
     client = LLM(**config_obj)
 
-    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
+    yield _patch_llm_based_on_config(client, llm_config)

{nvidia_nat_crewai-1.3.0a20250827.dist-info → nvidia_nat_crewai-1.3.0a20250828.dist-info}/METADATA

@@ -1,12 +1,12 @@
 Metadata-Version: 2.4
 Name: nvidia-nat-crewai
-Version: 1.3.0a20250827
+Version: 1.3.0a20250828
 Summary: Subpackage for CrewAI integration in NeMo Agent toolkit
 Keywords: ai,rag,agents
 Classifier: Programming Language :: Python
 Requires-Python: <3.13,>=3.11
 Description-Content-Type: text/markdown
-Requires-Dist: nvidia-nat==v1.3.0a20250827
+Requires-Dist: nvidia-nat==v1.3.0a20250828
 Requires-Dist: crewai~=0.95.0
 
 <!--

{nvidia_nat_crewai-1.3.0a20250827.dist-info → nvidia_nat_crewai-1.3.0a20250828.dist-info}/RECORD

@@ -1,11 +1,11 @@
 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=LDOctDQC9qdba1SVGoVkceCOSYuDj_mnl3HCuq2nIuQ,8382
-nat/plugins/crewai/llm.py,sha256=ErsUXsHpQ-chWqFhYjp-OhARAnCl-HdV_tkrrTZzIOU,4665
+nat/plugins/crewai/llm.py,sha256=DlkLN_fCXS7LUuNku4vvOUrMAgIA_om5LaQsKJOAAPk,5024
 nat/plugins/crewai/register.py,sha256=_R3bhGmz___696_NwyIcpw3koMBiWqIFoWEFJ0VAgXs,831
 nat/plugins/crewai/tool_wrapper.py,sha256=BNKEPQQCLKtXNzGDAKBLCdmGJXe9lBOVI1hObha8hoI,1569
-nvidia_nat_crewai-1.3.0a20250827.dist-info/METADATA,sha256=VEQudpIIMTEppwXQVv-6NPBtX_eTakcyVqPacygMySo,1453
-nvidia_nat_crewai-1.3.0a20250827.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-nvidia_nat_crewai-1.3.0a20250827.dist-info/entry_points.txt,sha256=YF5PUdQGr_OUDXB4TykElHJTsKT8yKkuE0bMX5n_RXs,58
-nvidia_nat_crewai-1.3.0a20250827.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
-nvidia_nat_crewai-1.3.0a20250827.dist-info/RECORD,,
+nvidia_nat_crewai-1.3.0a20250828.dist-info/METADATA,sha256=9B44N7vNvV2u57Py6bbSCurH3vn-5HUTZ9b_eJINOPM,1453
+nvidia_nat_crewai-1.3.0a20250828.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nvidia_nat_crewai-1.3.0a20250828.dist-info/entry_points.txt,sha256=YF5PUdQGr_OUDXB4TykElHJTsKT8yKkuE0bMX5n_RXs,58
+nvidia_nat_crewai-1.3.0a20250828.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
+nvidia_nat_crewai-1.3.0a20250828.dist-info/RECORD,,