rasa-pro 3.14.0.dev5__py3-none-any.whl → 3.14.0.dev7__py3-none-any.whl

rasa/shared/providers/llm/_base_litellm_client.py

@@ -21,7 +21,7 @@ from rasa.shared.providers._ssl_verification_utils import (
     ensure_ssl_certificates_for_litellm_non_openai_based_clients,
     ensure_ssl_certificates_for_litellm_openai_based_clients,
 )
-from rasa.shared.providers.llm.llm_response import LLMResponse, LLMUsage
+from rasa.shared.providers.llm.llm_response import LLMResponse, LLMToolCall, LLMUsage
 from rasa.shared.utils.io import resolve_environment_variables, suppress_logs
 
 structlogger = structlog.get_logger()
@@ -126,7 +126,9 @@ class _BaseLiteLLMClient:
             raise ProviderClientValidationError(event_info)
 
     @suppress_logs(log_level=logging.WARNING)
-    def completion(self, messages: Union[List[dict], List[str], str]) -> LLMResponse:
+    def completion(
+        self, messages: Union[List[dict], List[str], str], **kwargs: Any
+    ) -> LLMResponse:
         """Synchronously generate completions for given list of messages.
 
         Args:
@@ -138,6 +140,7 @@ class _BaseLiteLLMClient:
                 - a list of messages. Each message is a string and will be formatted
                     as a user message.
                 - a single message as a string which will be formatted as user message.
+            **kwargs: Additional parameters to pass to the completion call.
 
         Returns:
             List of message completions.
@@ -147,15 +150,19 @@ class _BaseLiteLLMClient:
         """
         try:
             formatted_messages = self._get_formatted_messages(messages)
-            arguments = resolve_environment_variables(self._completion_fn_args)
-            response = completion(messages=formatted_messages, **arguments)
+            arguments = cast(
+                Dict[str, Any], resolve_environment_variables(self._completion_fn_args)
+            )
+            response = completion(
+                messages=formatted_messages, **{**arguments, **kwargs}
+            )
             return self._format_response(response)
         except Exception as e:
             raise ProviderClientAPIException(e)
 
     @suppress_logs(log_level=logging.WARNING)
     async def acompletion(
-        self, messages: Union[List[dict], List[str], str]
+        self, messages: Union[List[dict], List[str], str], **kwargs: Any
     ) -> LLMResponse:
         """Asynchronously generate completions for given list of messages.
 
@@ -168,6 +175,7 @@ class _BaseLiteLLMClient:
                 - a list of messages. Each message is a string and will be formatted
                     as a user message.
                 - a single message as a string which will be formatted as user message.
+            **kwargs: Additional parameters to pass to the completion call.
 
         Returns:
             List of message completions.
@@ -177,8 +185,12 @@ class _BaseLiteLLMClient:
         """
         try:
             formatted_messages = self._get_formatted_messages(messages)
-            arguments = resolve_environment_variables(self._completion_fn_args)
-            response = await acompletion(messages=formatted_messages, **arguments)
+            arguments = cast(
+                Dict[str, Any], resolve_environment_variables(self._completion_fn_args)
+            )
+            response = await acompletion(
+                messages=formatted_messages, **{**arguments, **kwargs}
+            )
             return self._format_response(response)
         except Exception as e:
             message = ""
@@ -246,12 +258,32 @@ class _BaseLiteLLMClient:
                 else 0
             )
             formatted_response.usage = LLMUsage(prompt_tokens, completion_tokens)
+
+        # Extract tool calls from all choices
+        formatted_response.tool_calls = self._extract_tool_calls(response)
+
         structlogger.debug(
             "base_litellm_client.formatted_response",
             formatted_response=formatted_response.to_dict(),
         )
         return formatted_response
 
+    def _extract_tool_calls(self, response: Any) -> List[LLMToolCall]:
+        """Extract tool calls from response choices.
+
+        Args:
+            response: List of response choices from LiteLLM
+
+        Returns:
+            List of LLMToolCall objects, empty if no tool calls found
+        """
+        return [
+            LLMToolCall.from_litellm(tool_call)
+            for choice in response.choices
+            if choice.message.tool_calls
+            for tool_call in choice.message.tool_calls
+        ]
+
     def _format_text_completion_response(self, response: Any) -> LLMResponse:
         """Parses the LiteLLM text completion response to Rasa format."""
         formatted_response = LLMResponse(

rasa/shared/providers/llm/default_litellm_llm_client.py

@@ -39,7 +39,6 @@ class DefaultLiteLLMClient(_BaseLiteLLMClient):
 
     @classmethod
     def from_config(cls, config: Dict[str, Any]) -> DefaultLiteLLMClient:
-        """Creates a DefaultLiteLLMClient instance from a configuration dictionary."""
         default_config = DefaultLiteLLMClientConfig.from_dict(config)
         return cls(
             model=default_config.model,
@@ -111,7 +110,6 @@ class DefaultLiteLLMClient(_BaseLiteLLMClient):
                 self._litellm_model_name,
                 self._litellm_extra_parameters,
                 "default_litellm_llm_client",
-                provider=self.provider,
             )
         else:
             super().validate_client_setup()

rasa/shared/providers/llm/litellm_router_llm_client.py

@@ -122,7 +122,9 @@ class LiteLLMRouterLLMClient(_BaseLiteLLMRouterClient, _BaseLiteLLMClient):
             raise ProviderClientAPIException(e)
 
     @suppress_logs(log_level=logging.WARNING)
-    def completion(self, messages: Union[List[dict], List[str], str]) -> LLMResponse:
+    def completion(
+        self, messages: Union[List[dict], List[str], str], **kwargs: Any
+    ) -> LLMResponse:
         """
         Synchronously generate completions for given list of messages.
 
@@ -140,6 +142,7 @@ class LiteLLMRouterLLMClient(_BaseLiteLLMRouterClient, _BaseLiteLLMClient):
                 - a list of messages. Each message is a string and will be formatted
                     as a user message.
                 - a single message as a string which will be formatted as user message.
+            **kwargs: Additional parameters to pass to the completion call.
         Returns:
             List of message completions.
         Raises:
@@ -150,7 +153,7 @@ class LiteLLMRouterLLMClient(_BaseLiteLLMRouterClient, _BaseLiteLLMClient):
         try:
             formatted_messages = self._format_messages(messages)
             response = self.router_client.completion(
-                messages=formatted_messages, **self._completion_fn_args
+                messages=formatted_messages, **{**self._completion_fn_args, **kwargs}
             )
             return self._format_response(response)
         except Exception as e:
@@ -158,7 +161,7 @@ class LiteLLMRouterLLMClient(_BaseLiteLLMRouterClient, _BaseLiteLLMClient):
 
     @suppress_logs(log_level=logging.WARNING)
     async def acompletion(
-        self, messages: Union[List[dict], List[str], str]
+        self, messages: Union[List[dict], List[str], str], **kwargs: Any
     ) -> LLMResponse:
         """
         Asynchronously generate completions for given list of messages.
@@ -177,6 +180,7 @@ class LiteLLMRouterLLMClient(_BaseLiteLLMRouterClient, _BaseLiteLLMClient):
                 - a list of messages. Each message is a string and will be formatted
                     as a user message.
                 - a single message as a string which will be formatted as user message.
+            **kwargs: Additional parameters to pass to the completion call.
         Returns:
             List of message completions.
         Raises:
@@ -187,7 +191,7 @@ class LiteLLMRouterLLMClient(_BaseLiteLLMRouterClient, _BaseLiteLLMClient):
         try:
             formatted_messages = self._format_messages(messages)
             response = await self.router_client.acompletion(
-                messages=formatted_messages, **self._completion_fn_args
+                messages=formatted_messages, **{**self._completion_fn_args, **kwargs}
             )
             return self._format_response(response)
         except Exception as e:

rasa/shared/providers/llm/llm_client.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Dict, List, Protocol, Union, runtime_checkable
+from typing import Any, Dict, List, Protocol, Union, runtime_checkable
 
 from rasa.shared.providers.llm.llm_response import LLMResponse
 
@@ -32,7 +32,9 @@ class LLMClient(Protocol):
         """
         ...
 
-    def completion(self, messages: Union[List[dict], List[str], str]) -> LLMResponse:
+    def completion(
+        self, messages: Union[List[dict], List[str], str], **kwargs: Any
+    ) -> LLMResponse:
         """
         Synchronously generate completions for given list of messages.
 
@@ -48,13 +50,14 @@ class LLMClient(Protocol):
                 - a list of messages. Each message is a string and will be formatted
                     as a user message.
                 - a single message as a string which will be formatted as user message.
+            **kwargs: Additional parameters to pass to the completion call.
         Returns:
             LLMResponse
         """
         ...
 
     async def acompletion(
-        self, messages: Union[List[dict], List[str], str]
+        self, messages: Union[List[dict], List[str], str], **kwargs: Any
     ) -> LLMResponse:
         """
         Asynchronously generate completions for given list of messages.
@@ -71,6 +74,7 @@ class LLMClient(Protocol):
                 - a list of messages. Each message is a string and will be formatted
                     as a user message.
                 - a single message as a string which will be formatted as user message.
+            **kwargs: Additional parameters to pass to the completion call.
         Returns:
             LLMResponse
         """

rasa/shared/providers/llm/llm_response.py

@@ -1,9 +1,15 @@
 import functools
+import json
 import time
 from dataclasses import asdict, dataclass, field
 from typing import Any, Awaitable, Callable, Dict, List, Optional, Text, Union
 
 import structlog
+from litellm.utils import ChatCompletionMessageToolCall
+from pydantic import BaseModel
+
+from rasa.shared.constants import KEY_TOOL_CALLS
+from rasa.shared.exceptions import ProviderClientAPIException
 
 structlogger = structlog.get_logger()
 
@@ -38,6 +44,53 @@ class LLMUsage:
         return asdict(self)
 
 
+class LLMToolCall(BaseModel):
+    """A class representing a response from an LLM tool call."""
+
+    id: str
+    """The ID of the tool call."""
+
+    tool_name: str
+    """The name of the tool that was called."""
+
+    tool_args: Dict[str, Any]
+    """The arguments passed to the tool call."""
+
+    type: str = "function"
+    """The type of the tool call."""
+
+    @classmethod
+    def from_dict(cls, data: Dict[Text, Any]) -> "LLMToolCall":
+        """Creates an LLMToolResponse from a dictionary."""
+        return cls(**data)
+
+    @classmethod
+    def from_litellm(cls, data: ChatCompletionMessageToolCall) -> "LLMToolCall":
+        """Creates an LLMToolResponse from a dictionary."""
+        try:
+            tool_args = json.loads(data.function.arguments)
+        except json.JSONDecodeError as e:
+            structlogger.error(
+                "llm_response.litellm_tool_call.invalid_arguments",
+                tool_name=data.function.name,
+                tool_call=data.function.arguments,
+            )
+            raise ProviderClientAPIException(
+                original_exception=e,
+                message=(
+                    f"Invalid arguments for tool call `{data.function.name}`: "
+                    f"`{data.function.arguments}`"
+                ),
+            )
+
+        return cls(
+            id=data.id,
+            tool_name=data.function.name,
+            tool_args=tool_args,
+            type=data.type,
+        )
+
+
 @dataclass
 class LLMResponse:
     id: str
@@ -62,12 +115,22 @@ class LLMResponse:
     latency: Optional[float] = None
     """Optional field to store the latency of the LLM API call."""
 
+    tool_calls: Optional[List[LLMToolCall]] = None
+    """The list of tool calls the model generated for the input prompt."""
+
     @classmethod
     def from_dict(cls, data: Dict[Text, Any]) -> "LLMResponse":
         """Creates an LLMResponse from a dictionary."""
         usage_data = data.get("usage", {})
         usage_obj = LLMUsage.from_dict(usage_data) if usage_data else None
 
+        tool_calls_data = data.get(KEY_TOOL_CALLS, [])
+        tool_calls_obj = (
+            [LLMToolCall.from_dict(tool) for tool in tool_calls_data]
+            if tool_calls_data
+            else None
+        )
+
         return cls(
             id=data["id"],
             choices=data["choices"],
@@ -76,6 +139,7 @@ class LLMResponse:
             usage=usage_obj,
             additional_info=data.get("additional_info"),
             latency=data.get("latency"),
+            tool_calls=tool_calls_obj,
         )
 
     @classmethod
@@ -92,6 +156,8 @@ class LLMResponse:
         result = asdict(self)
         if self.usage:
             result["usage"] = self.usage.to_dict()
+        if self.tool_calls:
+            result[KEY_TOOL_CALLS] = [tool.model_dump() for tool in self.tool_calls]
         return result
 
 

rasa/shared/providers/llm/self_hosted_llm_client.py

@@ -237,7 +237,7 @@ class SelfHostedLLMClient(_BaseLiteLLMClient):
             raise ProviderClientAPIException(e)
 
     async def acompletion(
-        self, messages: Union[List[dict], List[str], str]
+        self, messages: Union[List[dict], List[str], str], **kwargs: Any
     ) -> LLMResponse:
         """Asynchronous completion of the model with the given messages.
 
@@ -255,15 +255,18 @@ class SelfHostedLLMClient(_BaseLiteLLMClient):
                 - a list of messages. Each message is a string and will be formatted
                     as a user message.
                 - a single message as a string which will be formatted as user message.
+            **kwargs: Additional parameters to pass to the completion call.
 
         Returns:
             The completion response.
         """
         if self._use_chat_completions_endpoint:
-            return await super().acompletion(messages)
+            return await super().acompletion(messages, **kwargs)
         return await self._atext_completion(messages)
 
-    def completion(self, messages: Union[List[dict], List[str], str]) -> LLMResponse:
+    def completion(
+        self, messages: Union[List[dict], List[str], str], **kwargs: Any
+    ) -> LLMResponse:
         """Completion of the model with the given messages.
 
         Method overrides the base class method to call the appropriate
@@ -273,12 +276,13 @@ class SelfHostedLLMClient(_BaseLiteLLMClient):
 
         Args:
             messages: The messages to be used for completion.
+            **kwargs: Additional parameters to pass to the completion call.
 
         Returns:
             The completion response.
         """
         if self._use_chat_completions_endpoint:
-            return super().completion(messages)
+            return super().completion(messages, **kwargs)
         return self._text_completion(messages)
 
     @staticmethod

rasa/shared/utils/health_check/health_check.py

@@ -1,7 +1,10 @@
 import os
-from typing import Any, Dict, Optional
+from typing import TYPE_CHECKING, Any, Dict, Optional
 
 from rasa.exceptions import HealthCheckError
+
+if TYPE_CHECKING:
+    pass
 from rasa.shared.constants import (
     LLM_API_HEALTH_CHECK_DEFAULT_VALUE,
     LLM_API_HEALTH_CHECK_ENV_VAR,
@@ -63,6 +66,7 @@ def perform_llm_health_check(
     log_source_component: str,
 ) -> None:
     """Try to instantiate the LLM Client to validate the provided config.
+
     If the LLM_API_HEALTH_CHECK environment variable is true, perform a test call
     to the LLM API. If config contains multiple models, perform a test call for each
     model in the model group.
@@ -125,6 +129,7 @@ def perform_embeddings_health_check(
     log_source_component: str,
 ) -> None:
     """Try to instantiate the Embedder to validate the provided config.
+
     If the LLM_API_HEALTH_CHECK environment variable is true, perform a test call
     to the Embeddings API. If config contains multiple models, perform a test call for
     each model in the model group.
@@ -240,8 +245,7 @@ def send_test_embeddings_api_request(
 
 
 def is_api_health_check_enabled() -> bool:
-    """Determines whether the API health check is enabled based on an environment
-    variable.
+    """Determines whether the API health check is enabled.
 
     Returns:
         bool: True if the API health check is enabled, False otherwise.

rasa/shared/utils/llm.py

@@ -49,7 +49,15 @@ from rasa.shared.constants import (
     RASA_PATTERN_INTERNAL_ERROR_USER_INPUT_TOO_LONG,
     ROUTER_CONFIG_KEY,
 )
-from rasa.shared.core.events import BotUttered, UserUttered
+from rasa.shared.core.events import (
+    AgentCancelled,
+    AgentCompleted,
+    AgentInterrupted,
+    AgentResumed,
+    AgentStarted,
+    BotUttered,
+    UserUttered,
+)
 from rasa.shared.core.slots import BooleanSlot, CategoricalSlot, Slot
 from rasa.shared.engine.caching import get_local_cache_location
 from rasa.shared.exceptions import (
@@ -112,7 +120,7 @@ DEPLOYMENT_CENTRIC_PROVIDERS = [AZURE_OPENAI_PROVIDER]
 
 # Placeholder messages used in the transcript for
 # instances where user input results in an error
-ERROR_PLACEHOLDER = {
+ERROR_PLACEHOLDER: Dict[str, str] = {
     RASA_PATTERN_INTERNAL_ERROR_USER_INPUT_TOO_LONG: "[User sent really long message]",
     RASA_PATTERN_INTERNAL_ERROR_USER_INPUT_EMPTY: "",
     "default": "[User input triggered an error]",
@@ -225,6 +233,7 @@ def tracker_as_readable_transcript(
     ai_prefix: str = AI,
     max_turns: Optional[int] = 20,
     turns_wrapper: Optional[Callable[[List[str]], List[str]]] = None,
+    highlight_agent_turns: bool = False,
 ) -> str:
     """Creates a readable dialogue from a tracker.
 
@@ -234,6 +243,7 @@ def tracker_as_readable_transcript(
         ai_prefix: the prefix to use for ai utterances
         max_turns: the maximum number of turns to include in the transcript
         turns_wrapper: optional function to wrap the turns in a custom way
+        highlight_agent_turns: whether to highlight agent turns in the transcript
 
     Example:
         >>> tracker = Tracker(
@@ -251,7 +261,9 @@ def tracker_as_readable_transcript(
     Returns:
     A string representing the transcript of the tracker
     """
-    transcript = []
+    transcript: List[str] = []
+
+    current_ai_prefix = ai_prefix
 
     # using `applied_events` rather than `events` means that only events after the
     # most recent `Restart` or `SessionStarted` are included in the transcript
@@ -266,9 +278,20 @@ def tracker_as_readable_transcript(
             else:
                 message = sanitize_message_for_prompt(event.text)
             transcript.append(f"{human_prefix}: {message}")
-
         elif isinstance(event, BotUttered):
-            transcript.append(f"{ai_prefix}: {sanitize_message_for_prompt(event.text)}")
+            transcript.append(
+                f"{current_ai_prefix}: {sanitize_message_for_prompt(event.text)}"
+            )
+
+        if highlight_agent_turns:
+            if isinstance(event, AgentStarted) or isinstance(event, AgentResumed):
+                current_ai_prefix = event.agent_id
+            elif (
+                isinstance(event, AgentCompleted)
+                or isinstance(event, AgentCancelled)
+                or isinstance(event, AgentInterrupted)
+            ):
+                current_ai_prefix = ai_prefix
 
     # turns_wrapper to count multiple utterances by bot/user as single turn
     if turns_wrapper:

rasa/shared/utils/mcp/server_connection.py

@@ -0,0 +1,186 @@
+"""MCP server connection utilities."""
+
+import asyncio
+import warnings
+from contextlib import AsyncExitStack
+from typing import Any, Dict, Optional
+
+import structlog
+from mcp import ClientSession
+from mcp.client.streamable_http import streamablehttp_client
+
+structlogger = structlog.get_logger()
+
+
+# Suppress RuntimeWarning about unawaited coroutines when MCP server is not reachable.
+warnings.filterwarnings(
+    "ignore",
+    message=".*BaseEventLoop.create_server.*was never awaited.*",
+    category=RuntimeWarning,
+)
+
+
+class MCPServerConnection:
+    """
+    Manages connection to an MCP server.
+
+    This class handles the lifecycle of connections to MCP servers,
+    including connection establishment, session management, and cleanup.
+    """
+
+    # Timeout for ping operations in seconds
+    PING_TIMEOUT_SECONDS = 3.0
+
+    def __init__(self, server_name: str, server_url: str, server_type: str):
+        """
+        Initialize the MCP server connection.
+
+        Args:
+            server_name: Server name to identify the server
+            server_url: Server URL
+            server_type: Server type (currently only 'http' is supported)
+        """
+        self.server_name = server_name
+        self.server_url = server_url
+        self.server_type = server_type
+        self.session: Optional[ClientSession] = None
+        self.exit_stack: Optional[AsyncExitStack] = None
+
+    @classmethod
+    def from_config(cls, server_config: Dict[str, Any]) -> "MCPServerConnection":
+        """Initialize the MCP server connection from a configuration dictionary."""
+        return cls(
+            server_config["name"],
+            server_config["url"],
+            server_config.get("type", "http"),
+        )
+
+    async def connect(self) -> None:
+        """Establish connection to the MCP server.
+
+        Raises:
+            ValueError: If the server type is not supported.
+            ConnectionError: If connection fails.
+        """
+        if self.server_type != "http":
+            raise ValueError(f"Unsupported server type: {self.server_type}")
+
+        # Create a new exit stack for this connection to avoid task boundary issues
+        self.exit_stack = AsyncExitStack()
+
+        try:
+            read_stream, write_stream, _ = await self.exit_stack.enter_async_context(
+                streamablehttp_client(self.server_url)
+            )
+            self.session = await self.exit_stack.enter_async_context(
+                ClientSession(read_stream, write_stream)
+            )
+            await self.session.initialize()
+        except asyncio.CancelledError as e:
+            event_info = f"Connection to MCP server `{self.server_name}` was cancelled."
+            structlogger.error(
+                "mcp_server_connection.connect.connection_cancelled",
+                event_info=event_info,
+                server_name=self.server_name,
+                server_url=self.server_url,
+            )
+            # Clean up on cancellation
+            await self._cleanup()
+            raise ConnectionError(e) from e
+
+        except Exception as e:
+            event_info = f"Failed to connect to MCP server `{self.server_name}`: {e}"
+            structlogger.error(
+                "mcp_server_connection.connect.connection_failed",
+                event_info=event_info,
+                server_name=self.server_name,
+                server_url=self.server_url,
+            )
+            # Clean up on error
+            await self._cleanup()
+            raise ConnectionError(e) from e
+
+    async def ensure_active_session(self) -> ClientSession:
+        """
+        Ensure an active session is available.
+
+        If no session exists or the current session is inactive,
+        a new connection will be established.
+
+        Returns:
+            Active ClientSession instance.
+        """
+        if self.session is None:
+            await self.connect()
+            structlogger.info(
+                "mcp_server_connection.ensure_active_session.no_session",
+                server_name=self.server_name,
+                server_url=self.server_url,
+                event_info=(
+                    "No session found, connecting to the server "
+                    f"`{self.server_name}` @ `{self.server_url}`"
+                ),
+            )
+        if self.session:
+            try:
+                # Add timeout to prevent hanging when MCP server is down
+                await asyncio.wait_for(
+                    self.session.send_ping(), timeout=self.PING_TIMEOUT_SECONDS
+                )
+            except asyncio.TimeoutError as e:
+                structlogger.error(
+                    "mcp_server_connection.ensure_active_session.ping_timeout",
+                    server_name=self.server_name,
+                    server_url=self.server_url,
+                    event_info=(
+                        "Ping timed out, Server not reachable - "
+                        f"`{self.server_name}` @ `{self.server_url}`"
+                    ),
+                )
+                raise e
+            except Exception as e:
+                structlogger.warning(
+                    "mcp_server_connection.ensure_active_session.ping_failed",
+                    error=str(e),
+                    server_name=self.server_name,
+                    server_url=self.server_url,
+                    event_info=(
+                        "Ping failed, trying to reconnect to the server "
+                        f"`{self.server_name}` @ `{self.server_url}`"
+                    ),
+                )
+                # Cleanup existing session
+                await self.close()
+                # Attempt to reconnect now
+                await self.connect()
+                structlogger.info(
+                    "mcp_server_connection.ensure_active_session.reconnected",
+                    server_name=self.server_name,
+                    server_url=self.server_url,
+                    event_info=(
+                        "Reconnected to the server "
+                        f"`{self.server_name}` @ `{self.server_url}`"
+                    ),
+                )
+        assert self.session is not None  # Ensures type for mypy
+        return self.session
+
+    async def close(self) -> None:
+        """Close the connection and clean up resources."""
+        await self._cleanup()
+
+    async def _cleanup(self) -> None:
+        """Internal cleanup method to safely close resources."""
+        if self.exit_stack:
+            try:
+                await self.exit_stack.aclose()
+            except Exception as e:
+                # Log cleanup errors but don't raise them
+                structlogger.warning(
+                    "mcp_server_connection.cleanup.failed",
+                    server_name=self.server_name,
+                    error=str(e),
+                )
+            finally:
+                self.exit_stack = None
+                self.session = None