Merge pull request #697

Fix/langchain_nvidia_ai_endpoints patch

Author: drazvan

nemoguardrails/llm/providers/_langchain_nvidia_ai_endpoints_patch.py

@@ -14,66 +14,22 @@
 # limitations under the License.
 
 import logging
-from typing import Any, Dict, Iterator, List, Mapping, Optional, Sequence, Type
+from functools import wraps
+from typing import Any, List, Optional
 
-import pkg_resources
 from langchain_core.callbacks.manager import CallbackManagerForLLMRun
 from langchain_core.language_models.chat_models import generate_from_stream
-from langchain_core.messages import (
-    AIMessageChunk,
-    BaseMessage,
-    BaseMessageChunk,
-    ChatMessage,
-    ChatMessageChunk,
-    FunctionMessageChunk,
-    HumanMessageChunk,
-    SystemMessageChunk,
-    ToolMessageChunk,
-)
-from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
+from langchain_core.messages import BaseMessage
+from langchain_core.outputs import ChatResult
 from langchain_core.pydantic_v1 import Field
-from langchain_nvidia_ai_endpoints import ChatNVIDIA
-from packaging import version
+from langchain_nvidia_ai_endpoints import ChatNVIDIA as ChatNVIDIAOriginal
 
 log = logging.getLogger(__name__)
 
 
-def _convert_delta_to_message_chunk(
-    _dict: Mapping[str, Any], default_class: Type[BaseMessageChunk]
-) -> BaseMessageChunk:
-    role = _dict.get("role")
-    content = _dict.get("content") or ""
-    additional_kwargs: Dict = {}
-    if _dict.get("function_call"):
-        function_call = dict(_dict["function_call"])
-        if "name" in function_call and function_call["name"] is None:
-            function_call["name"] = ""
-        additional_kwargs["function_call"] = function_call
-    if _dict.get("tool_calls"):
-        additional_kwargs["tool_calls"] = _dict["tool_calls"]
-
-    if role == "user" or default_class == HumanMessageChunk:
-        return HumanMessageChunk(content=content)
-    elif role == "assistant" or default_class == AIMessageChunk:
-        return AIMessageChunk(content=content, additional_kwargs=additional_kwargs)
-    elif role == "system" or default_class == SystemMessageChunk:
-        return SystemMessageChunk(content=content)
-    elif role == "function" or default_class == FunctionMessageChunk:
-        return FunctionMessageChunk(content=content, name=_dict["name"])
-    elif role == "tool" or default_class == ToolMessageChunk:
-        return ToolMessageChunk(content=content, tool_call_id=_dict["tool_call_id"])
-    elif role or default_class == ChatMessageChunk:
-        return ChatMessageChunk(content=content, role=role)  # type: ignore[arg-type]
-    else:
-        return default_class(content=content)  # type: ignore[call-arg]
-
-
-class PatchedChatNVIDIAV1(ChatNVIDIA):
-    streaming: bool = Field(
-        default=False, description="Whether to use streaming or not"
-    )
-
-    def _generate(
+def stream_decorator(func):
+    @wraps(func)
+    def wrapper(
         self,
         messages: List[BaseMessage],
         stop: Optional[List[str]] = None,
@@ -87,105 +43,30 @@ def _generate(
                 messages, stop=stop, run_manager=run_manager, **kwargs
             )
             return generate_from_stream(stream_iter)
-        inputs = self._custom_preprocess(messages)
-        payload = self._get_payload(inputs=inputs, stop=stop, stream=False, **kwargs)
-        response = self._client.client.get_req(payload=payload)
-        responses, _ = self._client.client.postprocess(response)
-        self._set_callback_out(responses, run_manager)
-        message = ChatMessage(**self._custom_postprocess(responses))
-        generation = ChatGeneration(message=message)
-        return ChatResult(generations=[generation], llm_output=responses)
+        else:
+            return func(self, messages, stop, run_manager, **kwargs)
 
-    def _stream(
-        self,
-        messages: List[BaseMessage],
-        stop: Optional[Sequence[str]] = None,
-        run_manager: Optional[CallbackManagerForLLMRun] = None,
-        **kwargs: Any,
-    ) -> Iterator[ChatGenerationChunk]:
-        """Allows streaming to model!"""
-        inputs = self._custom_preprocess(messages)
-        payload = self._get_payload(inputs=inputs, stop=stop, stream=True, **kwargs)
-        default_chunk_class = AIMessageChunk
-        for response in self._client.client.get_req_stream(payload=payload):
-            self._set_callback_out(response, run_manager)
-            chunk = _convert_delta_to_message_chunk(response, default_chunk_class)
-            default_chunk_class = chunk.__class__
-            cg_chunk = ChatGenerationChunk(message=chunk)
-            if run_manager:
-                run_manager.on_llm_new_token(cg_chunk.text, chunk=cg_chunk)
-            yield cg_chunk
+    return wrapper
 
 
-class PatchedChatNVIDIAV2(ChatNVIDIA):
+# NOTE: this needs to have the same name as the original class,
+#   otherwise, there's a check inside `langchain-nvidia-ai-endpoints` that will fail.
+class ChatNVIDIA(ChatNVIDIAOriginal):
     streaming: bool = Field(
         default=False, description="Whether to use streaming or not"
     )
 
+    @stream_decorator
     def _generate(
         self,
         messages: List[BaseMessage],
         stop: Optional[List[str]] = None,
         run_manager: Optional[CallbackManagerForLLMRun] = None,
-        stream: Optional[bool] = None,
         **kwargs: Any,
     ) -> ChatResult:
-        should_stream = stream if stream is not None else self.streaming
-        if should_stream:
-            stream_iter = self._stream(
-                messages, stop=stop, run_manager=run_manager, **kwargs
-            )
-            return generate_from_stream(stream_iter)
-        inputs = [
-            _nv_vlm_adjust_input(message)
-            for message in [convert_message_to_dict(message) for message in messages]
-        ]
-        payload = self._get_payload(inputs=inputs, stop=stop, stream=False, **kwargs)
-        response = self._client.client.get_req(payload=payload)
-        responses, _ = self._client.client.postprocess(response)
-        self._set_callback_out(responses, run_manager)
-        parsed_response = self._custom_postprocess(responses, streaming=False)
-        # for pre 0.2 compatibility w/ ChatMessage
-        # ChatMessage had a role property that was not present in AIMessage
-        parsed_response.update({"role": "assistant"})
-        generation = ChatGeneration(message=AIMessage(**parsed_response))
-        return ChatResult(generations=[generation], llm_output=responses)
-
-
-class ChatNVIDIAFactory:
-    RANGE1 = (version.parse("0.1.0"), version.parse("0.2.0"))
-    RANGE2 = (version.parse("0.2.0"), version.parse("0.3.0"))
-
-    @staticmethod
-    def get_package_version(package_name):
-        return version.parse(pkg_resources.get_distribution(package_name).version)
-
-    @staticmethod
-    def is_version_in_range(version, range):
-        return range[0] <= version < range[1]
-
-    @classmethod
-    def create(cls):
-        current_version = cls.get_package_version("langchain_nvidia_ai_endpoints")
-
-        if cls.is_version_in_range(current_version, cls.RANGE1):
-            log.debug(
-                f"Using pathed version of ChatNVIDIA for version {current_version}"
-            )
-            return PatchedChatNVIDIAV1
-        elif cls.is_version_in_range(current_version, cls.RANGE2):
-            log.debug(
-                f"Using pathed version of ChatNVIDIA for version {current_version}"
-            )
-            from langchain_community.adapters.openai import convert_message_to_dict
-            from langchain_nvidia_ai_endpoints.chat_models import _nv_vlm_adjust_input
-
-            return PatchedChatNVIDIAV2
-        else:
-            return ChatNVIDIA
-
-
-ChatNVIDIA = ChatNVIDIAFactory.create()
+        return super()._generate(
+            messages=messages, stop=stop, run_manager=run_manager, **kwargs
+        )
 
 
 __all__ = ["ChatNVIDIA"]

nemoguardrails/llm/providers/providers.py

@@ -20,8 +20,10 @@
 
 Additional providers can be registered using the `register_llm_provider` function.
 """
+
 import asyncio
 import logging
+from importlib.metadata import PackageNotFoundError, version
 from typing import Any, Dict, List, Optional, Type
 
 from langchain.base_language import BaseLanguageModel
@@ -33,6 +35,7 @@
 from langchain.schema.output import GenerationChunk
 from langchain_community import llms
 from langchain_community.llms import HuggingFacePipeline
+from packaging import version as pkg_version
 
 from nemoguardrails.rails.llm.config import Model
 
@@ -240,11 +243,18 @@ def get_llm_provider(model_config: Model) -> Type[BaseLanguageModel]:
             )
     elif model_config.engine == "nvidia_ai_endpoints" or model_config.engine == "nim":
         try:
-            from langchain_nvidia_ai_endpoints import ChatNVIDIA
-
             from ._langchain_nvidia_ai_endpoints_patch import ChatNVIDIA
 
+            # Check the version
+            package_version = version("langchain_nvidia_ai_endpoints")
+
+            if _parse_version(package_version) < (0, 2, 0):
+                raise ValueError(
+                    "langchain_nvidia_ai_endpoints version must be 0.2.0 or above."
+                    " Please upgrade it with `pip install langchain-nvidia-ai-endpoints --upgrade`."
+                )
             return ChatNVIDIA
+
         except ImportError:
             raise ImportError(
                 "Could not import langchain_nvidia_ai_endpoints, please install it with "
@@ -271,3 +281,7 @@ def get_llm_provider(model_config: Model) -> Type[BaseLanguageModel]:
 def get_llm_provider_names() -> List[str]:
     """Returns the list of supported LLM providers."""
     return list(sorted(list(_providers.keys())))
+
+
+def _parse_version(version_str):
+    return tuple(map(int, (version_str.split("."))))