src/nat/llm/utils/thinking.py (1)

32-41: Add proper return type annotation for the decorator function.

The decorator function returned by _thinking_injector is currently typed too generically as Callable[..., Any]. This loses type safety at call sites.

Apply this diff to add proper type hints using ParamSpec and TypeVar:

+from typing import ParamSpec
 from typing import TypeVar

 ModelType = TypeVar("ModelType")
 MessagesType = TypeVar("MessagesType")
+P = ParamSpec("P")
+R = TypeVar("R")

-def _thinking_injector(system_prompt_injector: Callable[[MessagesType], MessagesType], ) -> Callable[..., Any]:
+def _thinking_injector(system_prompt_injector: Callable[[MessagesType], MessagesType]) -> Callable[[Callable[P, R]], Callable[P, R]]:

packages/nvidia_nat_agno/src/nat/plugins/agno/llm.py (4)

17-17: Import missing symbols to support return annotations and generalized message types.

To add AsyncIterator return types and use Sequence in the injector signature, extend imports.

-from typing import TypeVar
+from typing import TypeVar, Any
+from collections.abc import AsyncIterator, Sequence

---

32-43: Docstring + generalize to Sequence and make idempotent to avoid duplicate system prompts.

Meet docstring requirements, avoid duplicate system messages, and accept any Sequence to match other providers and typing guidance.

-def agno_thinking_injector(client: ModelType, system_prompt: str) -> ModelType:
+def agno_thinking_injector(client: ModelType, system_prompt: str) -> ModelType:
+    """Patch an AGNO client to inject a system prompt as the first message."""
     from agno.models.message import Message
 
-    def injector(messages: list[Message]) -> list[Message]:
-        return [Message(role="system", content=system_prompt)] + messages
+    def injector(messages: Sequence[Message]) -> list[Message]:
+        if messages and getattr(messages[0], "role", None) == "system" and getattr(messages[0], "content", None) == system_prompt:
+            return list(messages)
+        return [Message(role="system", content=system_prompt), *messages]
 
     return patch_with_thinking(
         client,
         function_names=["invoke_stream", "invoke", "ainvoke", "ainvoke_stream"],
         system_prompt_injector=injector,
     )

---

46-47: Add return type and minimal docstring to public provider factory.

Public APIs must include return type hints and docstrings; annotate as AsyncIterator[Any] and add a concise summary.

-async def nim_agno(llm_config: NIMModelConfig, _builder: Builder):
+async def nim_agno(llm_config: NIMModelConfig, _builder: Builder) -> AsyncIterator[Any]:
+    """Yield an AGNO Nvidia client configured from nat, optionally patched for thinking and retries."""

---

89-90: Add return type and minimal docstring to public provider factory.

Apply the same return annotation and docstring pattern to OpenAI.

-async def openai_agno(llm_config: OpenAIModelConfig, _builder: Builder):
+async def openai_agno(llm_config: OpenAIModelConfig, _builder: Builder) -> AsyncIterator[Any]:
+    """Yield an AGNO OpenAI client configured from nat, optionally patched for thinking and retries."""

docs/source/extend/gated-fields.md (1)

80-84: Documentation formatting inconsistency.

The new ThinkingMixin documentation is indented differently from the previous entries. All other mixin entries use 2-space indentation for their bullet points, but this one uses inconsistent spacing.

Apply this formatting fix for consistency:

-- {py:class}`~nat.data_models.thinking_mixin.ThinkingMixin`
-  - Field: `thinking: bool | None`
-  - Default when supported: `None` (use model default)
-  - Only currently supported on Nemotron models
+- {py:class}`~nat.data_models.thinking_mixin.ThinkingMixin`
+  - Field: `thinking: bool | None`
+  - Default when supported: `None` (use model default)
+  - Only currently supported on Nemotron models

docs/source/extend/adding-an-llm-provider.md (1)

83-89: Clarify the thinking_system_prompt property description.

The documentation should be clearer about when thinking_system_prompt returns None versus a system prompt string.

Apply this clarification:

-- `ThinkingMixin`: adds a `thinking` field, with a default of `None` when supported by a model. If supported, the `thinking_system_prompt` property will return the system prompt to use for thinking.
+- `ThinkingMixin`: adds a `thinking` field, with a default of `None` when supported by a model. The `thinking_system_prompt` property returns the appropriate system prompt string when both the model supports thinking and the field is set, otherwise returns `None`.

src/nat/llm/utils/thinking.py (1)

107-119: Consolidate exception handling logic.

The exception messages use inconsistent formatting with special characters (‑ vs -) and the logic flow could be clearer.

Apply this refactor for consistency and clarity:

-        try:  # instance‑level first
+        try:  # instance-level first
             if not inspect.isclass(obj):
                 object.__setattr__(obj, name, types.MethodType(wrapped, obj))
                 continue
         except Exception as exc:
             logger.info(
-                "Instance‑level patch failed for %s.%s (%s); "
-                "falling back to class‑level patch.",
+                "Instance-level patch failed for %s.%s (%s); "
+                "falling back to class-level patch.",
                 cls_name,
                 name,
                 exc,
             )

-        try:  # class‑level fallback
+        try:  # class-level fallback
             setattr(cls, name, wrapped)

packages/nvidia_nat_llama_index/src/nat/plugins/llama_index/llm.py (3)

34-41: Add a concise docstring and make injection idempotent to avoid duplicate system prompts.

Prepend a Google-style docstring per guidelines and guard against double-injection when the first message already matches the same system prompt. This prevents duplicated system messages if the wrapper is patched multiple times or if upstream already injected one.

 def llama_index_thinking_injector(client: ModelType, system_prompt: str) -> ModelType:
+    """Patch a llama_index client to inject a system prompt as the first message."""
     from llama_index.core.base.llms.types import ChatMessage
 
-    def injector(messages: Sequence[ChatMessage]) -> Sequence[ChatMessage]:
-        msgs = list(messages)
-        msgs.insert(0, ChatMessage(role="system", content=system_prompt))
-        return msgs
+    def injector(messages: Sequence[ChatMessage]) -> Sequence[ChatMessage]:
+        # Idempotency: if first message is already the same system_prompt, do nothing.
+        if messages and getattr(messages[0], "role", None) == "system" and getattr(messages[0], "content", None) == system_prompt:
+            return list(messages)
+        msgs = list(messages)
+        msgs.insert(0, ChatMessage(role="system", content=system_prompt))
+        return msgs

---

31-32: Consider adding return type hints and minimal docstrings to the public provider factories.

Per repo guidelines, public APIs should include Python 3.11+ type hints and Google-style docstrings. Since these are async generators that yield a configured client, annotate with AsyncIterator[Any] and add one-line summaries.

Additional changes (outside the shown hunk):

# At top-level imports
from collections.abc import AsyncIterator
from typing import Any

# For each provider below, add return type + docstring, e.g.:
async def aws_bedrock_llama_index(llm_config: AWSBedrockModelConfig, _builder: Builder) -> AsyncIterator[Any]:
    """Yield a llama_index Bedrock client configured from nat, optionally patched for thinking and retries."""
    ...
    yield llm

async def azure_openai_llama_index(llm_config: AzureOpenAIModelConfig, _builder: Builder) -> AsyncIterator[Any]:
    """Yield a llama_index Azure OpenAI client configured from nat, optionally patched for thinking and retries."""
    ...
    yield llm

async def openai_llama_index(llm_config: OpenAIModelConfig, _builder: Builder) -> AsyncIterator[Any]:
    """Yield a llama_index OpenAI client configured from nat, optionally patched for thinking and retries."""
    ...
    yield llm

---

91-109: Parity: consider enabling thinking injection for llama_index NIM as well.

ThinkingMixin is supported by NIMModelConfig; wiring injection here would align behavior across providers and reduce surprises for users switching frameworks.

Additional changes (outside the shown hunk):

# Inside nim_llama_index after llm = NVIDIA(**kwargs)
if isinstance(llm_config, ThinkingMixin) and llm_config.thinking_system_prompt is not None:
    llm = llama_index_thinking_injector(llm, llm_config.thinking_system_prompt)

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

packages/nvidia_nat_agno/src/nat/plugins/agno/llm.py (3)

17-17: Import missing typing/ABC symbols to support required annotations.

To add return type hints (AsyncIterator[Any]) on public builders and to generalize the injector signature to Sequence, extend imports.

-from typing import TypeVar
+from typing import TypeVar, Any
+from collections.abc import AsyncIterator, Sequence

---

62-62: Add return type and Google-style docstring to public builder.

Public APIs must have return type hints (3.11+) and docstrings. This is an async generator.

-async def nim_agno(llm_config: NIMModelConfig, _builder: Builder):
+async def nim_agno(llm_config: NIMModelConfig, _builder: Builder) -> AsyncIterator[Any]:
+    """Yield an AGNO Nvidia client configured from nat.
+
+    Applies thinking prompt injection and retry policy if configured in llm_config.
+
+    Args:
+        llm_config: nat NIM provider configuration.
+        _builder: nat Builder instance (unused).
+
+    Yields:
+        Any: Configured AGNO Nvidia client.
+    """

---

104-105: Add return type and Google-style docstring to public builder.

Mirror nim_agno: annotate and document async generator return.

-async def openai_agno(llm_config: OpenAIModelConfig, _builder: Builder):
+async def openai_agno(llm_config: OpenAIModelConfig, _builder: Builder) -> AsyncIterator[Any]:
+    """Yield an AGNO OpenAI client configured from nat.
+
+    Applies thinking prompt injection and retry policy if configured in llm_config.
+
+    Args:
+        llm_config: nat OpenAI provider configuration.
+        _builder: nat Builder instance (unused).
+
+    Yields:
+        Any: Configured AGNO OpenAI client.
+    """

packages/nvidia_nat_crewai/src/nat/plugins/crewai/llm.py (1)

34-58: Add a docstring to _crewai_thinking_injector and broaden message typing; optionally patch call and avoid double-injection.

• Add a concise Google-style docstring to the outer helper (prior review also requested this on the public variant).
• Prefer collections.abc abstractions per guidelines (Sequence/Mapping) for messages.
• Optionally: include "call" in function_names for broader CrewAI compatibility and guard against injecting the same system message twice.

Apply within this hunk:

-def _crewai_thinking_injector(client: ModelType, system_prompt: str) -> ModelType:
+def _crewai_thinking_injector(client: ModelType, system_prompt: str) -> ModelType:
+    """
+    Inject a thinking system prompt into CrewAI LLM client calls.
+    
+    Prepends a system message {"role": "system", "content": system_prompt} to the first
+    positional messages argument for LLM.call/.__call__ invocations while preserving
+    the remaining args/kwargs. Returns the patched client.
+    """
 
-    def injector(messages: list[dict[str, str]], *args, **kwargs) -> FunctionArgumentWrapper:
+    def injector(messages: Sequence[Mapping[str, str]], *args, **kwargs) -> FunctionArgumentWrapper:
         """
         Inject a system prompt into the messages.
 
         The messages are the first (non-object) argument to the function.
         The rest of the arguments are passed through unchanged.
 
         Args:
             messages: The messages to inject the system prompt into.
             *args: The rest of the arguments to the function.
             **kwargs: The rest of the keyword arguments to the function.
 
         Returns:
             FunctionArgumentWrapper: An object that contains the transformed args and kwargs.
         """
-        new_messages = [{"role": "system", "content": system_prompt}] + messages
+        # Avoid double-injection if the first message already matches.
+        if messages and messages[0].get("role") == "system" and messages[0].get("content") == system_prompt:
+            return FunctionArgumentWrapper(messages, *args, **kwargs)
+        new_messages = [{"role": "system", "content": system_prompt}]
+        new_messages.extend(list(messages))
         return FunctionArgumentWrapper(new_messages, *args, **kwargs)
 
     return patch_with_thinking(
         client,
-        function_names=["call"],
+        function_names=["call", "__call__"],
         system_prompt_injector=injector,
     )

Add the needed imports near the other imports:

from collections.abc import Mapping, Sequence

packages/nvidia_nat_langchain/src/nat/plugins/langchain/llm.py (1)

16-18: Add AsyncIterator import and use it to annotate async generator wrappers’ return types.

Public APIs must have full type hints. The wrapper factories below are async generators (async def ... yield ...). Import AsyncIterator here so we can type their returns.

-from collections.abc import Sequence
+from collections.abc import AsyncIterator, Sequence

packages/nvidia_nat_agno/src/nat/plugins/agno/llm.py (3)

33-59: Make thinking injection idempotent and accept Sequence[Message] instead of list.

• Prevents duplicate system prompts when upstream already inserted one.
• Broadens compatibility with any sequence type per repo guidelines.

-def _agno_thinking_injector(client: ModelType, system_prompt: str) -> ModelType:
+def _agno_thinking_injector(client: ModelType, system_prompt: str) -> ModelType:
+    """Patch an AGNO client to inject a thinking system prompt on every invocation."""
     from agno.models.message import Message
 
-    def injector(messages: list[Message], *args, **kwargs) -> FunctionArgumentWrapper:
+    def injector(messages: Sequence[Message], *args, **kwargs) -> FunctionArgumentWrapper:
         """
         Inject a system prompt into the messages.
@@
-        new_messages = [Message(role="system", content=system_prompt)] + messages
+        # Idempotency: if the first message already matches, do not duplicate.
+        if len(messages) > 0:
+            first = messages[0]
+            if getattr(first, "role", None) == "system" and getattr(first, "content", None) == system_prompt:
+                return FunctionArgumentWrapper(messages, *args, **kwargs)
+        new_messages = [Message(role="system", content=system_prompt)] + list(messages)
         return FunctionArgumentWrapper(new_messages, *args, **kwargs)
 
     return patch_with_thinking(
         client,
-        function_names=["invoke_stream", "invoke", "ainvoke", "ainvoke_stream"],
+        function_names=["invoke_stream", "invoke", "ainvoke", "ainvoke_stream"],
         system_prompt_injector=injector,
     )

---

30-30: Optional: factor method names into a single constant to avoid drift.

Keeps patch target names consistent across providers and eases future updates.

 ModelType = TypeVar("ModelType")
 
+AGNO_INVOKE_METHODS: tuple[str, ...] = ("invoke_stream", "invoke", "ainvoke", "ainvoke_stream")

And use it below:

-        function_names=["invoke_stream", "invoke", "ainvoke", "ainvoke_stream"],
+        function_names=list(AGNO_INVOKE_METHODS),

---

16-16: Optional: module docstring.

Add a brief module docstring per repo docs.

 import os
+"""AGNO builders and thinking/retry patching for nat LLM clients."""

packages/nvidia_nat_llama_index/src/nat/plugins/llama_index/llm.py (6)

72-74: Good placement: thinking before retries; reduce repetition with a helper.

The order (inject thinking, then wrap with retries) is correct. To avoid repeating the isinstance/None check across providers, factor a small helper and call it here.

Apply this local diff:

-    if isinstance(llm_config, ThinkingMixin) and llm_config.thinking_system_prompt is not None:
-        llm = _llama_index_thinking_injector(llm, llm_config.thinking_system_prompt)
+    llm = _maybe_inject_thinking(llm, llm_config)

And add this helper (place it near _llama_index_thinking_injector):

def _maybe_inject_thinking(llm: ModelType, llm_config: ThinkingMixin) -> ModelType:
    return (_llama_index_thinking_injector(llm, llm_config.thinking_system_prompt)
            if isinstance(llm_config, ThinkingMixin) and llm_config.thinking_system_prompt is not None
            else llm)

---

93-95: Same refactor applies here.

Call the shared helper for consistency and to keep these registration functions lean.

-    if isinstance(llm_config, ThinkingMixin) and llm_config.thinking_system_prompt is not None:
-        llm = _llama_index_thinking_injector(llm, llm_config.thinking_system_prompt)
+    llm = _maybe_inject_thinking(llm, llm_config)

---

117-119: Same refactor applies here (NIM).

Keeps all providers uniform and easier to maintain.

-    if isinstance(llm_config, ThinkingMixin) and llm_config.thinking_system_prompt is not None:
-        llm = _llama_index_thinking_injector(llm, llm_config.thinking_system_prompt)
+    llm = _maybe_inject_thinking(llm, llm_config)

---

141-143: Same refactor applies here (OpenAI).

One-liner helper simplifies future changes (e.g., adding another patched method).

-    if isinstance(llm_config, ThinkingMixin) and llm_config.thinking_system_prompt is not None:
-        llm = _llama_index_thinking_injector(llm, llm_config.thinking_system_prompt)
+    llm = _maybe_inject_thinking(llm, llm_config)

---

35-61: Add a quick test to lock behavior.

Recommend adding/expanding tests to assert:

• The first message is a system message equal to the computed thinking_system_prompt when thinking is enabled.
• No duplicate injection occurs if the first message is already the same system prompt.
• Retry wrapper still wraps patched methods (e.g., monkeypatch chat to count calls and ensure retry happens on transient errors).

I can draft pytest tests with a simple Mock LLM implementing chat/stream_chat/achat/astream_chat that records args, so we don’t bring in llama_index for unit tests. Want me to push those?

---

35-61: Harden the thinking injector and align with provider capabilities

• In llama_index 0.10.x and 0.11.x, ChatMessage’s role field is typed as MessageRole but Pydantic will coerce a matching string (e.g. "system") into the enum. Both syntaxes work, though using MessageRole.SYSTEM is recommended for clarity and typo safety. (docs.llamaindex.ai)
• The OpenAI, AzureOpenAI, and NVIDIA LLM wrappers all expose chat, stream_chat, achat, and astream_chat. (docs.llamaindex.ai)
• The Bedrock wrapper supports chat, stream_chat, and achat but does not implement astream_chat (it raises NotImplementedError in v0.10.x). If you intend to inject into async streaming on Bedrock, you’ll need to guard against or omit astream_chat. (aidoczh.com)
• Retain the duplicate-injection guard and add a concise injector docstring. Also consider adding a module-level docstring after the SPDX header.

Apply this refined diff in packages/nvidia_nat_llama_index/src/nat/plugins/llama_index/llm.py:

-def _llama_index_thinking_injector(client: ModelType, system_prompt: str) -> ModelType:
-    from llama_index.core.base.llms.types import ChatMessage
+def _llama_index_thinking_injector(client: ModelType, system_prompt: str) -> ModelType:
+    """Inject a LlamaIndex-compatible system prompt into chat calls before user messages."""
+    from llama_index.core.base.llms.types import ChatMessage, MessageRole
@@
     def injector(messages: Sequence[ChatMessage], *args, **kwargs) -> FunctionArgumentWrapper:
-        new_messages = [ChatMessage(role="system", content=system_prompt)] + list(messages)
+        # Avoid double-injecting if the first message already matches our system prompt.
+        if messages:
+            first = messages[0]
+            if getattr(first, "role", None) in (MessageRole.SYSTEM, "system") \
+               and getattr(first, "content", "").strip() == system_prompt:
+                return FunctionArgumentWrapper(messages, *args, **kwargs)
+        new_messages = [ChatMessage(role=MessageRole.SYSTEM, content=system_prompt)] + list(messages)
         return FunctionArgumentWrapper(new_messages, *args, **kwargs)

And add at top of module (just after the SPDX header):

"""LlamaIndex LLM client registrations for nat with optional thinking system-prompt injection."""

packages/nvidia_nat_crewai/src/nat/plugins/crewai/llm.py (1)

120-132: Minor: typo in variable name and comment clarity.

Rename nvidai_api_key -> nvidia_api_key; tighten comment to CrewAI context (not LiteLLM).

-            nvidai_api_key = os.getenv("NVIDIA_API_KEY")
+            nvidia_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 nvidia_api_key is not None:
+                # Transfer the key to the CrewAI env var expected by its NIM integration
+                os.environ["NVIDIA_NIM_API_KEY"] = nvidia_api_key

packages/nvidia_nat_langchain/src/nat/plugins/langchain/llm.py (1)

37-41: Use consistent, stable BaseMessage import path.

LangChain exposes BaseMessage from langchain_core.messages. Importing from .messages (not .messages.base) is more conventional and consistent with the adjacent imports.

-from langchain_core.messages.base import BaseMessage
+from langchain_core.messages import BaseMessage

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

packages/nvidia_nat_langchain/src/nat/plugins/langchain/llm.py (3)

80-84: Include batch/abatch to future-proof thinking injection across LangChain clients.

Some chat models expose batch/abatch; patch them now to avoid gaps if used later.

     return patch_with_thinking(
         client,
-        function_names=["invoke", "ainvoke", "stream", "astream"],
+        function_names=["invoke", "ainvoke", "stream", "astream", "batch", "abatch"],
         system_prompt_injector=injector,
     )

Run this to see if batch/abatch are currently used (helps prioritize urgency):

#!/bin/bash
# Audit for potential call sites that would bypass injection if not patched.
rg -nP -C2 '\b(a?batch|a?invoke|a?stream)\s*\(' --type=py

---

16-18: Add AsyncIterator import and annotate async-generator wrappers’ return types.

Public APIs must have return type hints; these four factory wrappers are async generators (they yield clients). Import AsyncIterator and annotate the four functions accordingly.

Apply this diff for the import:

-from collections.abc import Sequence
+from collections.abc import AsyncIterator, Sequence

Then update the wrapper signatures and add concise docstrings (example shown for all four):

# aws_bedrock_langchain (line ~88)
async def aws_bedrock_langchain(llm_config: AWSBedrockModelConfig, _builder: Builder) -> AsyncIterator[object]:
    """
    Build and yield a LangChain ChatBedrockConverse client with optional thinking and retry wrappers.

    Args:
        llm_config: AWS Bedrock model config.
        _builder: Workflow builder (unused).

    Yields:
        The configured ChatBedrockConverse client.
    """

# azure_openai_langchain (line ~107)
async def azure_openai_langchain(llm_config: AzureOpenAIModelConfig, _builder: Builder) -> AsyncIterator[object]:
    """
    Build and yield a LangChain AzureChatOpenAI client with optional thinking and retry wrappers.

    Args:
        llm_config: Azure OpenAI model config.
        _builder: Workflow builder (unused).

    Yields:
        The configured AzureChatOpenAI client.
    """

# nim_langchain (line ~126)
async def nim_langchain(llm_config: NIMModelConfig, _builder: Builder) -> AsyncIterator[object]:
    """
    Build and yield a LangChain ChatNVIDIA client with optional thinking and retry wrappers.

    Args:
        llm_config: NIM model config.
        _builder: Workflow builder (unused).

    Yields:
        The configured ChatNVIDIA client.
    """

# openai_langchain (line ~145)
async def openai_langchain(llm_config: OpenAIModelConfig, _builder: Builder) -> AsyncIterator[object]:
    """
    Build and yield a LangChain ChatOpenAI client with optional thinking and retry wrappers.

    Args:
        llm_config: OpenAI model config.
        _builder: Workflow builder (unused).

    Yields:
        The configured ChatOpenAI client.
    """

---

58-60: Raise on unsupported LanguageModelInput instead of silently passing through.

Docstring says ValueError is raised for invalid inputs, but the else branch currently passes them through unchanged, which can hide misuses and skip the thinking injection.

Apply this diff:

         Raises:
             ValueError: If the messages are not a valid type for LanguageModelInput.
         """
@@
-        else:
-            return FunctionArgumentWrapper(messages, *args, **kwargs)
+        else:
+            raise ValueError(
+                f"Unsupported LanguageModelInput: {type(messages)}. "
+                "Expected str | BaseMessage | PromptValue | Sequence[BaseMessage]."
+            )

Also applies to: 77-78

packages/nvidia_nat_langchain/src/nat/plugins/langchain/llm.py (1)

35-43: Add a brief docstring to the injector for clarity.

Not required, but helpful for discoverability and to satisfy internal documentation checks.

 def _langchain_thinking_injector(client: ModelType, system_prompt: str) -> ModelType:
+    """
+    Patch a LangChain chat model client to prepend a thinking SystemMessage to inputs.
+
+    Args:
+        client: The LangChain chat model instance to patch.
+        system_prompt: The system prompt to inject (e.g., "/think" or "/no_think").
+
+    Returns:
+        The same client instance patched so invoke/ainvoke/stream/astream receive inputs
+        with the thinking SystemMessage prepended.
+    """

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

packages/nvidia_nat_crewai/src/nat/plugins/crewai/llm.py (1)

37-59: Add complete type hints to public function.

The function _patch_llm_based_on_config lacks a docstring and is missing complete type hints per coding guidelines. All public APIs require Google-style docstrings.

Apply this fix:

 def _patch_llm_based_on_config(client: ModelType, llm_config: LLMBaseConfig) -> ModelType:
+    """
+    Apply thinking and retry patches to a CrewAI LLM client based on config.
+    
+    Args:
+        client: The CrewAI LLM client to patch.
+        llm_config: Configuration containing thinking and retry settings.
+    
+    Returns:
+        The patched client with thinking and retry behavior applied.
+    """

packages/nvidia_nat_langchain/src/nat/plugins/langchain/llm.py (3)

38-39: Add docstring to public function.

The function _patch_llm_based_on_config requires a Google-style docstring per coding guidelines.

Apply this fix:

 def _patch_llm_based_on_config(client: ModelType, llm_config: LLMBaseConfig) -> ModelType:
+    """
+    Apply thinking and retry patches to a LangChain LLM client based on config.
+    
+    Args:
+        client: The LangChain LLM client to patch.
+        llm_config: Configuration containing thinking and retry settings.
+    
+    Returns:
+        The patched client with thinking and retry behavior applied.
+    """

---

91-96: Consider adding batch/abatch to function_names list.

Many LangChain chat models also expose batch and abatch methods. Consider including them for complete coverage of thinking injection.

                 function_names=[
                     "invoke",
                     "ainvoke",
                     "stream",
                     "astream",
+                    "batch",
+                    "abatch",
                 ],

---

84-84: Potential issue: Invalid inputs should raise an error.

The else branch returns the messages unchanged rather than raising an error for unsupported types. This could hide bugs where invalid input types are passed.

Apply this fix:

             else:
-                return FunctionArgumentWrapper(messages, *args, **kwargs)
+                raise ValueError(
+                    f"Unsupported LanguageModelInput: {type(messages)}. "
+                    "Expected str | BaseMessage | PromptValue | Sequence[BaseMessage]."
+                )

packages/nvidia_nat_agno/src/nat/plugins/agno/llm.py (3)

16-18: Add missing typing imports to support return annotations.

You’ll need AsyncIterator and Any for the provider factory return types. Recommend importing AsyncIterator from collections.abc per guidelines.

-from typing import TypeVar
+from typing import Any, TypeVar
+from collections.abc import AsyncIterator

---

70-74: Add return type and a minimal docstring to the public provider factory.

Public APIs require return hints and Google-style docstrings. Returning an async generator → AsyncIterator[Any].

-async def nim_agno(llm_config: NIMModelConfig, _builder: Builder):
+async def nim_agno(llm_config: NIMModelConfig, _builder: Builder) -> AsyncIterator[Any]:
+    """Yield an AGNO Nvidia client configured from nat, optionally patched for thinking and retries.
+
+    Yields:
+        The Nvidia client instance patched per llm_config.
+    """

---

103-106: Add return type and a minimal docstring to the OpenAI provider factory.

Mirror the NIM function’s annotations and docstring style.

-async def openai_agno(llm_config: OpenAIModelConfig, _builder: Builder):
+async def openai_agno(llm_config: OpenAIModelConfig, _builder: Builder) -> AsyncIterator[Any]:
+    """Yield an AGNO OpenAI client configured from nat, optionally patched for thinking and retries.
+
+    Yields:
+        The OpenAIChat client instance patched per llm_config.
+    """

packages/nvidia_nat_llama_index/src/nat/plugins/llama_index/llm.py (1)

42-48: Add docstring to the inject method.

The inject method overrides the base class method and should include a docstring explaining its specific behavior.

Apply this fix:

         @override
         def inject(self, messages: Sequence[ChatMessage], *args, **kwargs) -> FunctionArgumentWrapper:
+            """
+            Inject a system prompt as the first message in the sequence.
+            
+            Args:
+                messages: Sequence of ChatMessage objects to inject the prompt into.
+                *args: Additional positional arguments to pass through.
+                **kwargs: Additional keyword arguments to pass through.
+                
+            Returns:
+                FunctionArgumentWrapper containing the modified messages and other args.
+            """
             new_messages = [ChatMessage(role="system", content=self.system_prompt)] + list(messages)
             return FunctionArgumentWrapper(new_messages, *args, **kwargs)

src/nat/llm/utils/thinking.py (1)

97-100: Remove unnecessary return statement in generator.

The explicit return statement after yield from is unnecessary and can be removed.

Apply this fix:

         def _gen(obj: object, *call_args, **call_kwargs) -> Iterable[Any]:
             new_args = injector.inject(*call_args, **call_kwargs)
             yield from fn(obj, *new_args.args, **new_args.kwargs)
-            return

packages/nvidia_nat_agno/src/nat/plugins/agno/llm.py (2)

36-43: Add a concise docstring; remove redundant import in nested class.

• _patch_llm_based_on_config is a small public-ish helper used by two providers; a short docstring improves readability.
• Line 42 re-imports Message inside the class body; it’s already imported on Line 38. Drop the duplicate.

-def _patch_llm_based_on_config(client: ModelType, llm_config: LLMBaseConfig) -> ModelType:
+def _patch_llm_based_on_config(client: ModelType, llm_config: LLMBaseConfig) -> ModelType:
+    """Patch an AGNO client based on llm_config (thinking and retries) and return the patched client."""
@@
-    class AgnoThinkingInjector(BaseThinkingInjector):
-
-        from agno.models.message import Message
+    class AgnoThinkingInjector(BaseThinkingInjector):

---

86-88: Typo in local variable name.

nvidai_api_key → nvidia_api_key. This goes away if you remove the block above, but flagging for completeness.

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

packages/nvidia_nat_agno/src/nat/plugins/agno/llm.py (3)

87-99: Annotate return type and add a short docstring for the public factory.

Mirror the NIM path: add return type and docstring.

 @register_llm_client(config_type=OpenAIModelConfig, wrapper_type=LLMFrameworkEnum.AGNO)
-async def openai_agno(llm_config: OpenAIModelConfig, _builder: Builder):
+async def openai_agno(llm_config: OpenAIModelConfig, _builder: Builder) -> AsyncIterator[Any]:
+    """Create and yield an AGNO OpenAI client configured from NAT.
+
+    Applies thinking injection (if enabled) and retry logic.
+
+    Args:
+        llm_config: OpenAI model configuration.
+        _builder: Workflow builder (unused here).
+
+    Yields:
+        The configured and patched OpenAI client instance.
+    """

---

16-16: Add missing typing/ABC imports to support required annotations.

To add return type hints and use Sequence in signatures, extend imports accordingly.

-from typing import TypeVar
+from typing import TypeVar, Any
+from collections.abc import AsyncIterator, Sequence

---

69-84: Annotate return type and add a short docstring for the public factory.

Public APIs must have Python 3.11+ type hints and Google-style docstrings.

 @register_llm_client(config_type=NIMModelConfig, wrapper_type=LLMFrameworkEnum.AGNO)
-async def nim_agno(llm_config: NIMModelConfig, _builder: Builder):
+async def nim_agno(llm_config: NIMModelConfig, _builder: Builder) -> AsyncIterator[Any]:
+    """Create and yield an AGNO Nvidia client configured from NAT.
+
+    Applies thinking injection (if enabled) and retry logic.
+
+    Args:
+        llm_config: NIM model configuration.
+        _builder: Workflow builder (unused here).
+
+    Yields:
+        The configured and patched Nvidia client instance.
+    """

packages/nvidia_nat_crewai/src/nat/plugins/crewai/llm.py (3)

62-97: Add return type hints and Google-style docstring to the public Azure entrypoint.

Public APIs require explicit return types and docstrings. The function is an async generator yielding the client.

Example update (outside this hunk):

from collections.abc import AsyncIterator

@register_llm_client(config_type=AzureOpenAIModelConfig, wrapper_type=LLMFrameworkEnum.CREWAI)
async def azure_openai_crewai(llm_config: AzureOpenAIModelConfig, _builder: Builder) -> AsyncIterator[ModelType]:
    """
    Yield a CrewAI LLM client configured for Azure OpenAI and patched with thinking/retry behaviors.

    Args:
        llm_config: Azure OpenAI model configuration (API endpoint, deployment, version, etc.).
        _builder: NAT builder (unused here).

    Yields:
        The configured CrewAI LLM client instance.
    """
    ...
    yield _patch_llm_based_on_config(client, llm_config)

Minor sanity check: ensure llm_config.api_version is a non-empty str before os.environ["AZURE_API_VERSION"] = llm_config.api_version; otherwise raise a clear ValueError.

---

100-126: Add return type hints and docstring to the public NIM entrypoint.

Same requirement as Azure/OpenAI entrypoints.

Example (outside this hunk):

from collections.abc import AsyncIterator

@register_llm_client(config_type=NIMModelConfig, wrapper_type=LLMFrameworkEnum.CREWAI)
async def nim_crewai(llm_config: NIMModelConfig, _builder: Builder) -> AsyncIterator[ModelType]:
    """
    Yield a CrewAI LLM client configured for NVIDIA NIM and patched with thinking/retry behaviors.

    Args:
        llm_config: NIM model configuration (model name, keys, etc.).
        _builder: NAT builder (unused here).

    Yields:
        The configured CrewAI LLM client instance.
    """
    ...

---

128-139: Add return type hints and docstring to the public OpenAI entrypoint.

Mirror the Azure/NIM suggestions for consistency and to satisfy lint/type checks.

Example (outside this hunk):

from collections.abc import AsyncIterator

@register_llm_client(config_type=OpenAIModelConfig, wrapper_type=LLMFrameworkEnum.CREWAI)
async def openai_crewai(llm_config: OpenAIModelConfig, _builder: Builder) -> AsyncIterator[ModelType]:
    """
    Yield a CrewAI LLM client configured for OpenAI and patched with thinking/retry behaviors.

    Args:
        llm_config: OpenAI model configuration.
        _builder: NAT builder (unused here).

    Yields:
        The configured CrewAI LLM client instance.
    """
    ...

packages/nvidia_nat_agno/src/nat/plugins/agno/llm.py (5)

15-16: Add a concise module docstring per repo guidelines.

Public modules require a Google-style docstring with a one-line summary.

 # limitations under the License.
 
+"""AGNO plugin wrappers for NIM and OpenAI clients with optional thinking and retry patching."""
+
 from typing import TypeVar

---

35-47: Tidy injector: drop duplicate import and generalize parameter type to Sequence.

• Remove the redundant import inside the inner class (Line 41).
• Accept Sequence[Message] for broader compatibility and convert to list when constructing the new message list.

 def _patch_llm_based_on_config(client: ModelType, llm_config: LLMBaseConfig) -> ModelType:
 
     from agno.models.message import Message
 
     class AgnoThinkingInjector(BaseThinkingInjector):
-
-        from agno.models.message import Message
-
         @override
-        def inject(self, messages: list[Message], *args, **kwargs) -> FunctionArgumentWrapper:
-            new_messages = [Message(role="system", content=self.system_prompt)] + messages
+        def inject(self, messages: Sequence[Message], *args, **kwargs) -> FunctionArgumentWrapper:
+            new_messages = [Message(role="system", content=self.system_prompt)] + list(messages)
             return FunctionArgumentWrapper(new_messages, *args, **kwargs)

---

35-65: Document the central patch function for maintainability.

Add a short Google-style docstring clarifying ordering (thinking before retry) and parameters/return.

-def _patch_llm_based_on_config(client: ModelType, llm_config: LLMBaseConfig) -> ModelType:
+def _patch_llm_based_on_config(client: ModelType, llm_config: LLMBaseConfig) -> ModelType:
+    """Patch an AGNO client based on the NAT LLM config.
+
+    Applies thinking injection first (if configured) and then retry logic.
+
+    Args:
+        client: The AGNO client instance to patch.
+        llm_config: LLM configuration that may include thinking/retry options.
+
+    Returns:
+        The patched client instance.
+    """

---

74-76: Minor: avoid unnecessary f-string for id.

Model name is already a string.

-        "id": f"{llm_config.model_name}",
+        "id": llm_config.model_name,

---

91-94: Consider pruning None-valued fields in one pass.

Optional: instead of per-key pruning like base_url, drop all None from the dumped config. This reduces constructor noise and potential arg-type ignores.

-    config_obj = {
-        "id": llm_config.model_name,
-        **llm_config.model_dump(exclude={"type", "model_name"}, by_alias=True),
-    }
+    dumped = llm_config.model_dump(exclude={"type", "model_name"}, by_alias=True)
+    config_obj = {"id": llm_config.model_name, **{k: v for k, v in dumped.items() if v is not None}}

Note: verify against AGNO constructors so we don't inadvertently drop required-but-None defaults.

packages/nvidia_nat_crewai/src/nat/plugins/crewai/llm.py (3)

17-17: Broaden typing imports to prepare for collections.abc types.

Per guidelines, prefer typing/collections.abc abstractions (e.g., Sequence, Mapping). Let's import them here.

-from typing import TypeVar
+from typing import Mapping, Sequence, TypeVar

---

39-45: Generalize CrewAIThinkingInjector input type and avoid assuming a list.

CrewAI calls often accept sequences of message dicts. Using Sequence[Mapping[str, str]] improves interoperability and aligns with the guidelines. Cast to list on write.

-        def inject(self, messages: list[dict[str, str]], *args, **kwargs) -> FunctionArgumentWrapper:
-            new_messages = [{"role": "system", "content": self.system_prompt}] + messages
+        def inject(self, messages: Sequence[Mapping[str, str]], *args, **kwargs) -> FunctionArgumentWrapper:
+            new_messages = [{"role": "system", "content": self.system_prompt}] + list(messages)
             return FunctionArgumentWrapper(new_messages, *args, **kwargs)

---

117-122: Ensure NVIDIA_NIM_API_KEY is set even when config_obj includes api_key.

If CrewAI/LiteLLM reads the NIM key from env only, providing api_key in config but not exporting the env var could fail. Safest is to export env if not already set, preferring the explicit config value.

-    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:
-            nvidia_api_key = os.getenv("NVIDIA_API_KEY")
-
-            if (nvidia_api_key is not None):
-                # Transfer the key to the correct environment variable for LiteLLM
-                os.environ["NVIDIA_NIM_API_KEY"] = nvidia_api_key
+    # Ensure LiteLLM sees the correct API key for NIM
+    if "NVIDIA_NIM_API_KEY" not in os.environ:
+        if config_obj.get("api_key"):
+            os.environ["NVIDIA_NIM_API_KEY"] = str(config_obj["api_key"])
+        else:
+            nvidia_api_key = os.getenv("NVIDIA_API_KEY")
+            if nvidia_api_key:
+                os.environ["NVIDIA_NIM_API_KEY"] = nvidia_api_key

If CrewAI reliably honors the api_key argument directly, feel free to skip this change; otherwise, this makes the behavior robust.

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

packages/nvidia_nat_crewai/src/nat/plugins/crewai/llm.py (1)

37-45: Add docstrings for the patch helper and injector for maintainability.

Internal helpers still benefit from concise docstrings. This mirrors the earlier feedback about documenting the thinking injector.

Apply this diff:

 def _patch_llm_based_on_config(client: ModelType, llm_config: LLMBaseConfig) -> ModelType:
+    """
+    Patch a CrewAI LLM client with thinking and retry behavior based on the LLM config.
+    
+    Args:
+        client: The instantiated CrewAI LLM client.
+        llm_config: Provider config that may include ThinkingMixin and/or RetryMixin.
+    Returns:
+        The patched client (same instance), ready to be yielded.
+    """
 
     class CrewAIThinkingInjector(BaseThinkingInjector):
+        """CrewAI-specific thinking injector that prepends a system message."""
 
         @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)

packages/nvidia_nat_agno/src/nat/plugins/agno/llm.py (3)

16-16: Add missing typing imports for async generator annotations.

To add return type hints on the public factories, import AsyncIterator and Any. Prefer collections.abc for ABCs per guidelines.

-from typing import TypeVar
+from typing import TypeVar, Any
+from collections.abc import AsyncIterator, Sequence

---

69-84: Add return type and docstring to public factory (nim_agno).

Per repo guidelines, public APIs require explicit return type hints and Google-style docstrings. This is an async generator.

-async def nim_agno(llm_config: NIMModelConfig, _builder: Builder):
+async def nim_agno(llm_config: NIMModelConfig, _builder: Builder) -> AsyncIterator[Any]:
+    """Yield an AGNO Nvidia client configured from nat, patched for thinking and retries if enabled.
+
+    Yields:
+        Any: The configured AGNO Nvidia client.
+    """

---

87-98: Add return type and docstring to public factory (openai_agno).

Same requirement as above for the OpenAI path.

-async def openai_agno(llm_config: OpenAIModelConfig, _builder: Builder):
+async def openai_agno(llm_config: OpenAIModelConfig, _builder: Builder) -> AsyncIterator[Any]:
+    """Yield an AGNO OpenAI client configured from nat, patched for thinking and retries if enabled.
+
+    Yields:
+        Any: The configured AGNO OpenAI client.
+    """

packages/nvidia_nat_crewai/src/nat/plugins/crewai/llm.py (2)

41-45: Guard against duplicate system prompts and broaden input type slightly.

If upstream already prepends the same thinking directive, this will add it twice. Also, CrewAI messages can be a list of dicts but widening to Mapping/Sequence improves robustness.

Apply this diff:

-        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)
+        def inject(self, messages: list[dict[str, str]] | None, *args, **kwargs) -> FunctionArgumentWrapper:
+            msgs = messages or []
+            if msgs:
+                first = msgs[0]
+                if isinstance(first, dict) and first.get("role") == "system" and first.get("content") == self.system_prompt:
+                    # Already injected; return as-is.
+                    return FunctionArgumentWrapper(msgs, *args, **kwargs)
+            new_messages = [{"role": "system", "content": self.system_prompt}, *msgs]
+            return FunctionArgumentWrapper(new_messages, *args, **kwargs)

---

62-95: Add concise Google-style docstrings to public entrypoints.

The three registered LLM entrypoints are public APIs. Short docstrings describing config, env var behavior, and yielded object will aid users and satisfy the docs guideline.

I can draft these if helpful; let me know your preferred phrasing for env var precedence per provider.

Also applies to: 100-118, 121-133

packages/nvidia_nat_agno/src/nat/plugins/agno/llm.py (4)

15-16: Add a concise module docstring.

Public modules require a Google-style docstring. Place it after the license header.

 # limitations under the License.
 
+"""AGNO plugin: build and patch AGNO LLM clients with thinking and retry support driven by nat configs."""
+
 from typing import TypeVar

---

35-39: Document the central patcher for maintainability.

Add a short Google-style docstring to clarify what the helper composes and returns.

 def _patch_llm_based_on_config(client: ModelType, llm_config: LLMBaseConfig) -> ModelType:
-
+    """Apply thinking and retry patches to an AGNO client according to llm_config.
+
+    Args:
+        client: The AGNO client instance to patch.
+        llm_config: nat LLM config that may implement ThinkingMixin and/or RetryMixin.
+
+    Returns:
+        The patched client (same instance), with thinking/retry behavior applied.
+    """

---

41-46: Remove duplicate import and generalize injector signature to Sequence.

• The import of Message inside the class body (Line 41) is redundant because Message is already imported in the outer scope of this function.
• Accept a Sequence instead of list to be more flexible per guidelines, and avoid list-only concatenation.

-        from agno.models.message import Message
-
         @override
-        def inject(self, messages: list[Message], *args, **kwargs) -> FunctionArgumentWrapper:
-            new_messages = [Message(role="system", content=self.system_prompt)] + messages
+        def inject(self, messages: Sequence[Message], *args, **kwargs) -> FunctionArgumentWrapper:
+            new_messages = [Message(role="system", content=self.system_prompt)] + list(messages)
             return FunctionArgumentWrapper(new_messages, *args, **kwargs)

---

52-57: Minor: factor the AGNO method list into a constant.

Keeps the list DRY and discoverable.

-                                 function_names=[
-                                     "invoke_stream",
-                                     "invoke",
-                                     "ainvoke",
-                                     "ainvoke_stream",
-                                 ]))
+                                 function_names=AGNO_THINKING_METHODS))

Add near the top of the file (supporting snippet outside the selected lines):

AGNO_THINKING_METHODS: list[str] = ["invoke_stream", "invoke", "ainvoke", "ainvoke_stream"]

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

packages/nvidia_nat_crewai/src/nat/plugins/crewai/llm.py (1)

62-67: Add return type hints and brief docstrings for public async generators.

Per repo guidelines, public APIs need explicit return types and Google-style docstrings. This was flagged earlier and still applies.

Example updates (outside the changed ranges):

from typing import AsyncIterator

@register_llm_client(config_type=AzureOpenAIModelConfig, wrapper_type=LLMFrameworkEnum.CREWAI)
async def azure_openai_crewai(llm_config: AzureOpenAIModelConfig, _builder: Builder) -> AsyncIterator[ModelType]:
    """Yield a CrewAI LLM client configured for Azure OpenAI, patched with thinking/retry."""
    ...

@register_llm_client(config_type=NIMModelConfig, wrapper_type=LLMFrameworkEnum.CREWAI)
async def nim_crewai(llm_config: NIMModelConfig, _builder: Builder) -> AsyncIterator[ModelType]:
    """Yield a CrewAI LLM client configured for NVIDIA NIM, patched with thinking/retry."""
    ...

@register_llm_client(config_type=OpenAIModelConfig, wrapper_type=LLMFrameworkEnum.CREWAI)
async def openai_crewai(llm_config: OpenAIModelConfig, _builder: Builder) -> AsyncIterator[ModelType]:
    """Yield a CrewAI LLM client configured for OpenAI, patched with thinking/retry."""
    ...

To quickly verify in-repo:

#!/bin/bash
# Check public CrewAI wrappers have AsyncIterator return types and docstrings.
rg -nP -C1 $'@register_llm_client\\(config_type=.*CREWAI\\)\\s*\\nasync def\\s+\\w+\\(' packages/nvidia_nat_crewai/src |
  awk -F: '{print $1":"$2}' | while read -r defline; do
    file="${defline%%:*}"; line="${defline##*:}"
    sed -n "$((line)),+5p" "$file"
  done

Happy to push a small follow-up PR for the annotations and docstrings if you prefer.

Also applies to: 100-105, 121-126

packages/nvidia_nat_crewai/src/nat/plugins/crewai/llm.py (5)

37-39: Add a concise docstring and clarify patch order.

This helper is part of the public plugin surface; add a short Google-style docstring and state that thinking is applied before retry so each retry reuses the injected prompt.

 def _patch_llm_based_on_config(client: ModelType, llm_config: LLMBaseConfig) -> ModelType:
+    """
+    Patch a CrewAI LLM client based on the provided config.
+
+    Applies thinking injection first (if enabled) and then retry behavior, so
+    retries re-invoke the underlying call with the injected system prompt.
+
+    Args:
+        client: The CrewAI LLM client instance to patch.
+        llm_config: Provider-specific LLM configuration carrying thinking/retry knobs.
+
+    Returns:
+        The patched client with thinking and retry semantics applied.
+    """

---

41-44: Make the injector input-robust and align with typing/collections.abc guidance.

CrewAI might pass tuples or other Sequences; concatenating with a list will raise. Prefer Sequence[Mapping[str, str]] and build a new list. Also handle None gracefully.

-        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)
+        def inject(
+            self,
+            messages: Sequence[Mapping[str, str]] | None,
+            *args,
+            **kwargs,
+        ) -> FunctionArgumentWrapper:
+            system_message = {"role": "system", "content": self.system_prompt}
+            if messages is None:
+                new_messages = [system_message]
+            else:
+                new_messages = [system_message, *list(messages)]
+            return FunctionArgumentWrapper(new_messages, *args, **kwargs)

Add the missing imports in the typing block (see comment on Line 17).

---

78-87: Set both legacy and new Azure env var names for compatibility.

Some libs read AZURE_OPENAI_* while others read AZURE_API_*; set both to the resolved values to reduce surprises.

     os.environ["AZURE_API_KEY"] = api_key
+    # Keep both variants for compatibility across SDKs/plugins.
+    os.environ["AZURE_OPENAI_API_KEY"] = api_key
@@
     os.environ["AZURE_API_BASE"] = api_base
+    # Keep both variants for compatibility across SDKs/plugins.
+    os.environ["AZURE_OPENAI_ENDPOINT"] = api_base

---

89-93: Propagate deployment to ENV as well (compat fallback).

Some integrations look for an env var when the parameter isn’t threaded through; keeping it in ENV helps interop.

     config_obj["model"] = model
+    # Provide a compatibility env var for consumers that expect it.
+    os.environ["AZURE_MODEL_DEPLOYMENT"] = model

---

1-15: Consider adding a brief module docstring beneath the license header.

Keeps API docs cohesive and meets the “public module docstrings” guidance.

Example (outside the changed ranges):

"""CrewAI plugin LLM wrappers: build clients from nat configs and patch with thinking/retry semantics."""

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.