rasa-pro 3.12.9__tar.gz → 3.12.10.dev1__tar.gz

{rasa_pro-3.12.9 → rasa_pro-3.12.10.dev1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: rasa-pro
-Version: 3.12.9
+Version: 3.12.10.dev1
 Summary: State-of-the-art open-core Conversational AI framework for Enterprises that natively leverages generative AI for effortless assistant development.
 Keywords: nlp,machine-learning,machine-learning-library,bot,bots,botkit,rasa conversational-agents,conversational-ai,chatbot,chatbot-framework,bot-framework
 Author: Rasa Technologies GmbH
@@ -63,6 +63,7 @@ Requires-Dist: keras (==2.14.0)
 Requires-Dist: langchain (>=0.2.17,<0.3.0)
 Requires-Dist: langchain-community (>=0.2.19,<0.3.0)
 Requires-Dist: langcodes (>=3.5.0,<4.0.0)
+Requires-Dist: langfuse (>=2.60.2,<2.61.0)
 Requires-Dist: litellm (>=1.68.0,<1.69.0)
 Requires-Dist: matplotlib (>=3.7,<3.8)
 Requires-Dist: mattermostwrapper (>=2.2,<2.3)

{rasa_pro-3.12.9 → rasa_pro-3.12.10.dev1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "rasa-pro"
-version = "3.12.9"
+version = "3.12.10.dev1"
 description = "State-of-the-art open-core Conversational AI framework for Enterprises that natively leverages generative AI for effortless assistant development."
 authors = [ "Rasa Technologies GmbH <hi@rasa.com>",]
 maintainers = [ "Tom Bocklisch <tom@rasa.com>",]
@@ -195,6 +195,7 @@ scikit-learn = "~1.5.1"
 demoji = "^1.1.0"
 azure-identity = "~1.19.0"
 langcodes = "^3.5.0"
+langfuse = "~2.60.2"
 
 [[tool.poetry.dependencies.tensorflow-io-gcs-filesystem]]
 version = "==0.31" # newer versions don't provide wheels for Windows
@@ -333,6 +334,7 @@ convention = "google"
 
 
 
+
 [tool.poetry.group.dev.dependencies]
 ruff = ">=0.3.5,<0.4.0"
 docker = "~7.1.0"

rasa_pro-3.12.10.dev1/rasa/__init__.py

@@ -0,0 +1,15 @@
+import logging
+
+from rasa import version
+
+# define the version before the other imports since these need it
+__version__ = version.__version__
+
+from litellm.integrations.langfuse.langfuse import LangFuseLogger
+
+from rasa.monkey_patches import litellm_langfuse_logger_init_fixed
+
+# Monkey-patch the init method as early as possible before the class is used
+LangFuseLogger.__init__ = litellm_langfuse_logger_init_fixed  # type: ignore
+
+logging.getLogger(__name__).addHandler(logging.NullHandler())

{rasa_pro-3.12.9 → rasa_pro-3.12.10.dev1}/rasa/core/channels/studio_chat.py

@@ -16,6 +16,7 @@ import structlog
 from sanic import Sanic
 
 from rasa.core.channels.socketio import SocketBlueprint, SocketIOInput
+from rasa.core.exceptions import AgentNotReady
 from rasa.hooks import hookimpl
 from rasa.plugin import plugin_manager
 from rasa.shared.core.constants import ACTION_LISTEN_NAME
@@ -149,8 +150,15 @@ class StudioChatInput(SocketIOInput):
         """
         await on_new_message(message)
 
-        if not self.agent:
+        if not self.agent or not self.agent.is_ready():
             structlogger.error("studio_chat.on_message_proxy.agent_not_initialized")
+            await self.emit_error(
+                "The Rasa Pro model could not be loaded. "
+                "Please check the training and deployment logs "
+                "for more information.",
+                message.sender_id,
+                AgentNotReady("The Rasa Pro model could not be loaded."),
+            )
             return
 
         tracker = await self.agent.tracker_store.retrieve(message.sender_id)
@@ -160,6 +168,17 @@ class StudioChatInput(SocketIOInput):
 
         await self.on_tracker_updated(tracker)
 
+    async def emit_error(self, message: str, room: str, e: Exception) -> None:
+        await self.emit(
+            "error",
+            {
+                "message": message,
+                "error": str(e),
+                "exception": str(type(e).__name__),
+            },
+            room=room,
+        )
+
     async def handle_tracker_update(self, sid: str, data: Dict) -> None:
         from rasa.shared.core.trackers import DialogueStateTracker
 
@@ -200,15 +219,12 @@ class StudioChatInput(SocketIOInput):
                     error=e,
                     sender_id=data["sender_id"],
                 )
-                await self.emit(
-                    "error",
-                    {
-                        "message": "An error occurred while updating the conversation.",
-                        "error": str(e),
-                        "exception": str(type(e).__name__),
-                    },
-                    room=sid,
+                await self.emit_error(
+                    "An error occurred while updating the conversation.",
+                    data["sender_id"],
+                    e,
                 )
+
         if not tracker:
             # in case the tracker couldn't be updated, we retrieve the prior
             # version and use that to populate the update

{rasa_pro-3.12.9 → rasa_pro-3.12.10.dev1}/rasa/core/policies/intentless_policy.py

@@ -715,7 +715,9 @@ class IntentlessPolicy(LLMHealthCheckMixin, EmbeddingsHealthCheckMixin, Policy):
                 final_response_examples.append(resp)
 
         llm_response = await self.generate_answer(
-            final_response_examples, conversation_samples, history
+            final_response_examples,
+            conversation_samples,
+            history,
         )
         if not llm_response:
             structlogger.debug("intentless_policy.prediction.skip_llm_fail")

{rasa_pro-3.12.9 → rasa_pro-3.12.10.dev1}/rasa/dialogue_understanding/coexistence/llm_based_router.py

@@ -166,7 +166,6 @@ class LLMBasedRouter(LLMHealthCheckMixin, GraphComponent):
         **kwargs: Any,
     ) -> "LLMBasedRouter":
         """Loads trained component (see parent class for full docstring)."""
-
         # Perform health check on the resolved LLM client config
         llm_config = resolve_model_client_config(config.get(LLM_CONFIG_KEY, {}))
         cls.perform_llm_health_check(

{rasa_pro-3.12.9 → rasa_pro-3.12.10.dev1}/rasa/dialogue_understanding/generator/llm_based_command_generator.py

@@ -1,6 +1,8 @@
 from abc import ABC, abstractmethod
+from asyncio import Lock
 from functools import lru_cache
 from typing import Any, Dict, List, Optional, Set, Text, Tuple, Union
+from uuid import UUID, uuid4
 
 import structlog
 from jinja2 import Environment, Template, select_autoescape
@@ -89,6 +91,9 @@ class LLMBasedCommandGenerator(
         else:
             self.flow_retrieval = None
 
+        self.sender_id_to_session_id_mapping: Dict[str, UUID] = {}
+        self._lock = Lock()
+
     ### Abstract methods
     @staticmethod
     @abstractmethod
@@ -225,8 +230,7 @@ class LLMBasedCommandGenerator(
 
     @lru_cache
     def compile_template(self, template: str) -> Template:
-        """
-        Compile the prompt template and register custom filters.
+        """Compile the prompt template and register custom filters.
         Compiling the template is an expensive operation,
         so we cache the result.
         """
@@ -328,7 +332,9 @@ class LLMBasedCommandGenerator(
 
     @measure_llm_latency
     async def invoke_llm(
-        self, prompt: Union[List[dict], List[str], str]
+        self,
+        prompt: Union[List[dict], List[str], str],
+        metadata: Optional[Dict[str, Any]] = None,
     ) -> Optional[LLMResponse]:
         """Use LLM to generate a response.
 
@@ -341,6 +347,7 @@ class LLMBasedCommandGenerator(
                 - 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.
+            metadata: Optional metadata to be passed to the LLM call.
 
         Returns:
             An LLMResponse object.
@@ -352,7 +359,7 @@ class LLMBasedCommandGenerator(
             self.config.get(LLM_CONFIG_KEY), self.get_default_llm_config()
         )
         try:
-            return await llm.acompletion(prompt)
+            return await llm.acompletion(prompt, metadata)
         except Exception as e:
             # unfortunately, langchain does not wrap LLM exceptions which means
             # we have to catch all exceptions here
@@ -655,3 +662,7 @@ class LLMBasedCommandGenerator(
     def get_default_llm_config() -> Dict[str, Any]:
         """Get the default LLM config for the command generator."""
         return DEFAULT_LLM_CONFIG
+
+    async def _get_or_create_session_id(self, sender_id: str) -> UUID:
+        async with self._lock:
+            return self.sender_id_to_session_id_mapping.setdefault(sender_id, uuid4())

{rasa_pro-3.12.9 → rasa_pro-3.12.10.dev1}/rasa/dialogue_understanding/generator/llm_command_generator.py

@@ -55,7 +55,9 @@ class LLMCommandGenerator(SingleStepLLMCommandGenerator):
         )
 
     async def invoke_llm(
-        self, prompt: Union[List[dict], List[str], str]
+        self,
+        prompt: Union[List[dict], List[str], str],
+        metadata: Optional[Dict[str, Any]] = None,
     ) -> Optional[LLMResponse]:
         try:
             return await super().invoke_llm(prompt)

{rasa_pro-3.12.9 → rasa_pro-3.12.10.dev1}/rasa/dialogue_understanding/generator/multi_step/multi_step_llm_command_generator.py

@@ -42,6 +42,9 @@ from rasa.engine.storage.resource import Resource
 from rasa.engine.storage.storage import ModelStorage
 from rasa.shared.constants import (
     EMBEDDINGS_CONFIG_KEY,
+    LANGFUSE_CUSTOM_METADATA_DICT,
+    LANGFUSE_METADATA_SESSION_ID,
+    LANGFUSE_TAGS,
     RASA_PATTERN_CANNOT_HANDLE_NOT_SUPPORTED,
     ROUTE_TO_CALM_SLOT,
 )
@@ -107,7 +110,7 @@ structlogger = structlog.get_logger()
 )
 @deprecated(
     reason=(
-        "The MultiStepLLMCommandGenerator is deprecated and will be removed in "
+        "The MultiStepLLMCommandGenerator is  deprecated and will be removed in "
         "Rasa `4.0.0`."
     )
 )
@@ -492,7 +495,20 @@ class MultiStepLLMCommandGenerator(LLMBasedCommandGenerator):
             prompt=prompt,
         )
 
-        response = await self.invoke_llm(prompt)
+        if tracker:
+            session_id = str(await self._get_or_create_session_id(tracker.sender_id))
+        else:
+            session_id = "unknown"
+        metadata = {
+            LANGFUSE_METADATA_SESSION_ID: session_id,
+            LANGFUSE_CUSTOM_METADATA_DICT: {
+                "component": self.__class__.__name__,
+                "function": "_predict_commands_for_active_flow",
+            },
+            LANGFUSE_TAGS: [self.__class__.__name__],
+        }
+
+        response = await self.invoke_llm(prompt, metadata)
         llm_response = LLMResponse.ensure_llm_response(response)
         actions = None
         if llm_response and llm_response.choices:
@@ -546,8 +562,20 @@ class MultiStepLLMCommandGenerator(LLMBasedCommandGenerator):
             ".prompt_rendered",
             prompt=prompt,
         )
+        if tracker:
+            session_id = str(await self._get_or_create_session_id(tracker.sender_id))
+        else:
+            session_id = "unknown"
+        metadata = {
+            LANGFUSE_METADATA_SESSION_ID: session_id,
+            LANGFUSE_CUSTOM_METADATA_DICT: {
+                "component": self.__class__.__name__,
+                "function": "_predict_commands_for_handling_flows",
+            },
+            LANGFUSE_TAGS: [self.__class__.__name__],
+        }
 
-        response = await self.invoke_llm(prompt)
+        response = await self.invoke_llm(prompt, metadata)
         llm_response = LLMResponse.ensure_llm_response(response)
         actions = None
         if llm_response and llm_response.choices:
@@ -636,8 +664,20 @@ class MultiStepLLMCommandGenerator(LLMBasedCommandGenerator):
             flow=newly_started_flow.id,
             prompt=prompt,
         )
+        if tracker:
+            session_id = str(await self._get_or_create_session_id(tracker.sender_id))
+        else:
+            session_id = "unknown"
+        metadata = {
+            LANGFUSE_METADATA_SESSION_ID: session_id,
+            LANGFUSE_CUSTOM_METADATA_DICT: {
+                "component": self.__class__.__name__,
+                "function": "_predict_commands_for_newly_started_flow",
+            },
+            LANGFUSE_TAGS: [self.__class__.__name__],
+        }
 
-        response = await self.invoke_llm(prompt)
+        response = await self.invoke_llm(prompt, metadata)
         llm_response = LLMResponse.ensure_llm_response(response)
         actions = None
         if llm_response and llm_response.choices:

{rasa_pro-3.12.9 → rasa_pro-3.12.10.dev1}/rasa/dialogue_understanding/generator/single_step/compact_llm_command_generator.py

@@ -47,6 +47,9 @@ from rasa.shared.constants import (
     AWS_BEDROCK_PROVIDER,
     AZURE_OPENAI_PROVIDER,
     EMBEDDINGS_CONFIG_KEY,
+    LANGFUSE_CUSTOM_METADATA_DICT,
+    LANGFUSE_METADATA_SESSION_ID,
+    LANGFUSE_TAGS,
     MAX_TOKENS_CONFIG_KEY,
     PROMPT_TEMPLATE_CONFIG_KEY,
     ROUTE_TO_CALM_SLOT,
@@ -366,7 +369,17 @@ class CompactLLMCommandGenerator(LLMBasedCommandGenerator):
             prompt=flow_prompt,
         )
 
-        response = await self.invoke_llm(flow_prompt)
+        if tracker:
+            session_id = str(await self._get_or_create_session_id(tracker.sender_id))
+        else:
+            session_id = "unknown"
+        metadata = {
+            LANGFUSE_METADATA_SESSION_ID: session_id,
+            LANGFUSE_CUSTOM_METADATA_DICT: {"component": self.__class__.__name__},
+            LANGFUSE_TAGS: [self.__class__.__name__],
+        }
+
+        response = await self.invoke_llm(flow_prompt, metadata)
         llm_response = LLMResponse.ensure_llm_response(response)
         # The check for 'None' maintains compatibility with older versions
         # of LLMCommandGenerator. In previous implementations, 'invoke_llm'

{rasa_pro-3.12.9 → rasa_pro-3.12.10.dev1}/rasa/hooks.py

@@ -1,8 +1,20 @@
 import argparse
 import logging
+import os
 from typing import TYPE_CHECKING, List, Optional, Text, Union
 
+import litellm
 import pluggy
+import structlog
+
+from rasa.shared.providers.constants import (
+    LANGFUSE_CALLBACK_NAME,
+    LANGFUSE_HOST_ENV_VAR,
+    LANGFUSE_PROJECT_ID_ENV_VAR,
+    LANGFUSE_PUBLIC_KEY_ENV_VAR,
+    LANGFUSE_SECRET_KEY_ENV_VAR,
+    RASA_LANGFUSE_INTEGRATION_ENABLED_ENV_VAR,
+)
 
 # IMPORTANT: do not import anything from rasa here - use scoped imports
 #  this avoids circular imports, as the hooks are used in different places
@@ -18,6 +30,7 @@ if TYPE_CHECKING:
 
 hookimpl = pluggy.HookimplMarker("rasa")
 logger = logging.getLogger(__name__)
+structlogger = structlog.get_logger()
 
 
 @hookimpl  # type: ignore[misc]
@@ -57,6 +70,8 @@ def configure_commandline(cmdline_arguments: argparse.Namespace) -> Optional[Tex
         config.configure_tracing(tracer_provider)
         config.configure_metrics(endpoints_file)
 
+    _init_langfuse_integration()
+
     return endpoints_file
 
 
@@ -115,3 +130,43 @@ def after_server_stop() -> None:
 
     if anon_pipeline is not None:
         anon_pipeline.stop()
+
+
+def _is_langfuse_integration_enabled() -> bool:
+    return (
+        os.environ.get(RASA_LANGFUSE_INTEGRATION_ENABLED_ENV_VAR, "false").lower()
+        == "true"
+    )
+
+
+def _init_langfuse_integration() -> None:
+    if not _is_langfuse_integration_enabled():
+        structlogger.info(
+            "hooks._init_langfuse_integration.disabled",
+            event_info="Langfuse integration is disabled.",
+        )
+        return
+
+    if (
+        not os.environ.get(LANGFUSE_HOST_ENV_VAR)
+        or not os.environ.get(LANGFUSE_PROJECT_ID_ENV_VAR)
+        or not os.environ.get(LANGFUSE_PUBLIC_KEY_ENV_VAR)
+        or not os.environ.get(LANGFUSE_SECRET_KEY_ENV_VAR)
+    ):
+        structlogger.warning(
+            "hooks._init_langfuse_integration.missing_langfuse_keys",
+            event_info=(
+                "Langfuse integration is enabled, but some environment variables "
+                "are missing. Please set LANGFUSE_HOST, LANGFUSE_PROJECT_ID, "
+                "LANGFUSE_PUBLIC_KEY and LANGFUSE_SECRET_KEY environment "
+                "variables to use Langfuse integration."
+            ),
+        )
+        return
+
+    litellm.success_callback = [LANGFUSE_CALLBACK_NAME]
+    litellm.failure_callback = [LANGFUSE_CALLBACK_NAME]
+    structlogger.info(
+        "hooks.langfuse_callbacks_initialized",
+        event_info="Langfuse integration initialized.",
+    )

rasa_pro-3.12.10.dev1/rasa/monkey_patches.py

@@ -0,0 +1,91 @@
+import os
+import traceback
+from typing import Any, Optional
+
+from litellm.secret_managers.main import str_to_bool
+from packaging.version import Version
+
+
+def litellm_langfuse_logger_init_fixed(
+    self: Any,  # we should not import LangfuseLogger class before we patch it
+    langfuse_public_key: Optional[str] = None,
+    langfuse_secret: Optional[str] = None,
+    langfuse_host: str = "https://cloud.langfuse.com",
+    flush_interval: int = 1,
+) -> None:
+    """Monkeypatched version of LangfuseLogger.__init__ from the LiteLLM library.
+
+    This patched version removes a call that fetched the `project_id` from
+    Langfuse Cloud even when it was already set via environment variables.
+    In the original implementation, this call was made *before* initializing
+    the LangfuseClient, which caused the application to freeze for up to 60 seconds.
+
+    By removing this premature call, the monkeypatch avoids the unnecessary network
+    request and prevents the timeout/freeze issue.
+
+    This workaround can be removed once the underlying bug is resolved in LiteLLM:
+    https://github.com/BerriAI/litellm/issues/7732
+    """
+    try:
+        import langfuse
+        from langfuse import Langfuse
+    except Exception as e:
+        raise Exception(
+            f"\033[91mLangfuse not installed, try running 'pip install langfuse' "
+            f"to fix this error: {e}\n{traceback.format_exc()}\033[0m"
+        )
+    # Instance variables
+    self.secret_key = langfuse_secret or os.getenv("LANGFUSE_SECRET_KEY", "")
+    self.public_key = langfuse_public_key or os.getenv("LANGFUSE_PUBLIC_KEY", "")
+
+    self.langfuse_host = langfuse_host or os.getenv(
+        "LANGFUSE_HOST", "https://cloud.langfuse.com"
+    )
+    self.langfuse_host.replace("http://", "https://")
+    if not self.langfuse_host.startswith("https://"):
+        self.langfuse_host = "https://" + self.langfuse_host
+
+    self.langfuse_release = os.getenv("LANGFUSE_RELEASE")
+    self.langfuse_debug = os.getenv("LANGFUSE_DEBUG")
+    self.langfuse_flush_interval = (
+        os.getenv("LANGFUSE_FLUSH_INTERVAL") or flush_interval
+    )
+
+    parameters = {
+        "public_key": self.public_key,
+        "secret_key": self.secret_key,
+        "host": self.langfuse_host,
+        "release": self.langfuse_release,
+        "debug": self.langfuse_debug,
+        "flush_interval": self.langfuse_flush_interval,  # flush interval in seconds
+    }
+
+    if Version(langfuse.version.__version__) >= Version("2.6.0"):
+        parameters["sdk_integration"] = "litellm"
+
+    self.Langfuse = Langfuse(**parameters)
+
+    if os.getenv("UPSTREAM_LANGFUSE_SECRET_KEY") is not None:
+        upstream_langfuse_debug = (
+            str_to_bool(self.upstream_langfuse_debug)
+            if self.upstream_langfuse_debug is not None
+            else None
+        )
+        self.upstream_langfuse_secret_key = os.getenv("UPSTREAM_LANGFUSE_SECRET_KEY")
+        self.upstream_langfuse_public_key = os.getenv("UPSTREAM_LANGFUSE_PUBLIC_KEY")
+        self.upstream_langfuse_host = os.getenv("UPSTREAM_LANGFUSE_HOST")
+        self.upstream_langfuse_release = os.getenv("UPSTREAM_LANGFUSE_RELEASE")
+        self.upstream_langfuse_debug = os.getenv("UPSTREAM_LANGFUSE_DEBUG")
+        self.upstream_langfuse = Langfuse(
+            public_key=self.upstream_langfuse_public_key,
+            secret_key=self.upstream_langfuse_secret_key,
+            host=self.upstream_langfuse_host,
+            release=self.upstream_langfuse_release,
+            debug=(
+                upstream_langfuse_debug
+                if upstream_langfuse_debug is not None
+                else False
+            ),
+        )
+    else:
+        self.upstream_langfuse = None

{rasa_pro-3.12.9 → rasa_pro-3.12.10.dev1}/rasa/shared/constants.py

@@ -339,3 +339,8 @@ ROLE_SYSTEM = "system"
 # Used for key values in ValidateSlotPatternFlowStackFrame
 REFILL_UTTER = "refill_utter"
 REJECTIONS = "rejections"
+
+LANGFUSE_METADATA_USER_ID = "trace_user_id"
+LANGFUSE_METADATA_SESSION_ID = "session_id"
+LANGFUSE_CUSTOM_METADATA_DICT = "trace_metadata"
+LANGFUSE_TAGS = "tags"

rasa_pro-3.12.10.dev1/rasa/shared/providers/constants.py

@@ -0,0 +1,15 @@
+DEFAULT_AZURE_API_KEY_NAME = "AZURE_API_KEY"
+LITE_LLM_API_BASE_FIELD = "api_base"
+LITE_LLM_API_KEY_FIELD = "api_key"
+LITE_LLM_API_VERSION_FIELD = "api_version"
+LITE_LLM_MODEL_FIELD = "model"
+LITE_LLM_AZURE_AD_TOKEN = "azure_ad_token"
+
+# Enable or disable Langfuse integration
+RASA_LANGFUSE_INTEGRATION_ENABLED_ENV_VAR = "RASA_LANGFUSE_INTEGRATION_ENABLED"
+# Langfuse configuration
+LANGFUSE_CALLBACK_NAME = "langfuse"
+LANGFUSE_HOST_ENV_VAR = "LANGFUSE_HOST"
+LANGFUSE_PROJECT_ID_ENV_VAR = "LANGFUSE_PROJECT_ID"
+LANGFUSE_PUBLIC_KEY_ENV_VAR = "LANGFUSE_PUBLIC_KEY"
+LANGFUSE_SECRET_KEY_ENV_VAR = "LANGFUSE_SECRET_KEY"

{rasa_pro-3.12.9 → rasa_pro-3.12.10.dev1}/rasa/shared/providers/llm/_base_litellm_client.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import logging
 from abc import abstractmethod
-from typing import Any, Dict, List, Union, cast
+from typing import Any, Dict, List, Optional, Union, cast
 
 import structlog
 from litellm import acompletion, completion, validate_environment
@@ -123,7 +123,11 @@ 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],
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> LLMResponse:
         """Synchronously generate completions for given list of messages.
 
         Args:
@@ -135,6 +139,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.
+            metadata: Optional metadata to be passed to the LLM call.
 
         Returns:
             List of message completions.
@@ -152,7 +157,9 @@ class _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],
+        metadata: Optional[Dict[str, Any]] = None,
     ) -> LLMResponse:
         """Asynchronously generate completions for given list of messages.
 
@@ -165,6 +172,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.
+            metadata: Optional metadata to be passed to the LLM call.
 
         Returns:
             List of message completions.
@@ -175,7 +183,9 @@ 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)
+            response = await acompletion(
+                messages=formatted_messages, metadata=metadata, **arguments
+            )
             return self._format_response(response)
         except Exception as e:
             message = ""

{rasa_pro-3.12.9 → rasa_pro-3.12.10.dev1}/rasa/shared/providers/llm/litellm_router_llm_client.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import logging
-from typing import Any, Dict, List, Union
+from typing import Any, Dict, List, Optional, Union
 
 import structlog
 
@@ -122,9 +122,12 @@ class LiteLLMRouterLLMClient(_BaseLiteLLMRouterClient, _BaseLiteLLMClient):
             raise ProviderClientAPIException(e)
 
     @suppress_logs(log_level=logging.WARNING)
-    def completion(self, messages: Union[List[dict], List[str], str]) -> LLMResponse:
-        """
-        Synchronously generate completions for given list of messages.
+    def completion(
+        self,
+        messages: Union[List[dict], List[str], str],
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> LLMResponse:
+        """Synchronously generate completions for given list of messages.
 
         Method overrides the base class method to call the appropriate
         completion method based on the configuration. If the chat completions
@@ -140,8 +143,11 @@ 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.
+            metadata: Optional metadata to be passed to the LLM call.
+
         Returns:
             List of message completions.
+
         Raises:
             ProviderClientAPIException: If the API request fails.
         """
@@ -158,10 +164,11 @@ 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],
+        metadata: Optional[Dict[str, Any]] = None,
     ) -> LLMResponse:
-        """
-        Asynchronously generate completions for given list of messages.
+        """Asynchronously generate completions for given list of messages.
 
         Method overrides the base class method to call the appropriate
         completion method based on the configuration. If the chat completions
@@ -177,8 +184,11 @@ 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.
+            metadata: Optional metadata to be passed to the LLM call.
+
         Returns:
             List of message completions.
+
         Raises:
             ProviderClientAPIException: If the API request fails.
         """