rasa-pro 3.11.3a1.dev7__tar.gz → 3.12.0.dev2__tar.gz

{rasa_pro-3.11.3a1.dev7 → rasa_pro-3.12.0.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: rasa-pro
-Version: 3.11.3a1.dev7
+Version: 3.12.0.dev2
 Summary: State-of-the-art open-core Conversational AI framework for Enterprises that natively leverages generative AI for effortless assistant development.
 Home-page: https://rasa.com
 Keywords: nlp,machine-learning,machine-learning-library,bot,bots,botkit,rasa conversational-agents,conversational-ai,chatbot,chatbot-framework,bot-framework
@@ -32,6 +32,7 @@ Requires-Dist: aiogram (>=2.15,<2.26)
 Requires-Dist: aiohttp (>=3.9.4,<3.10)
 Requires-Dist: apscheduler (>=3.10,<3.11)
 Requires-Dist: attrs (>=23.1,<23.2)
+Requires-Dist: azure-identity (>=1.19.0,<2.0.0)
 Requires-Dist: azure-storage-blob (>=12.16.0,<12.17.0)
 Requires-Dist: boto3 (>=1.35.15,<1.36.0)
 Requires-Dist: certifi (>=2024.07.04)

{rasa_pro-3.11.3a1.dev7 → rasa_pro-3.12.0.dev2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "rasa-pro"
-version = "3.11.3a1.dev7"
+version = "3.12.0dev2"
 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>",]
@@ -193,6 +193,7 @@ numpy = "~1.26.4"
 scipy = "~1.13.1"
 scikit-learn = "~1.5.1"
 demoji = "^1.1.0"
+azure-identity = "^1.19.0"
 
 [[tool.poetry.dependencies.tensorflow-io-gcs-filesystem]]
 version = "==0.31" # newer versions don't provide wheels for Windows

{rasa_pro-3.11.3a1.dev7 → rasa_pro-3.12.0.dev2}/rasa/cli/arguments/default_arguments.py

@@ -13,7 +13,7 @@ from rasa.shared.constants import (
 
 
 def add_model_param(
-    parser: argparse.ArgumentParser,
+    parser: Union[argparse.ArgumentParser, argparse._ActionsContainer],
     model_name: Text = "Rasa",
     add_positional_arg: bool = True,
     default: Optional[Text] = DEFAULT_MODELS_PATH,

rasa_pro-3.12.0.dev2/rasa/cli/dialogue_understanding_test.py

@@ -0,0 +1,251 @@
+import argparse
+import asyncio
+import datetime
+import sys
+from typing import List
+
+import structlog
+
+import rasa.cli.utils
+from rasa.cli import SubParsersAction
+from rasa.cli.arguments.default_arguments import (
+    add_endpoint_param,
+    add_model_param,
+    add_remote_storage_param,
+)
+from rasa.core.exceptions import AgentNotReady
+from rasa.core.utils import AvailableEndpoints
+from rasa.dialogue_understanding_test.command_metric_calculation import (
+    calculate_command_metrics,
+)
+from rasa.dialogue_understanding_test.constants import (
+    DEFAULT_INPUT_TESTS_PATH,
+    KEY_STUB_CUSTOM_ACTIONS,
+)
+from rasa.dialogue_understanding_test.du_test_result import (
+    DialogueUnderstandingTestResult,
+)
+from rasa.dialogue_understanding_test.du_test_runner import (
+    DialogueUnderstandingTestRunner,
+)
+from rasa.dialogue_understanding_test.io import (
+    read_test_suite,
+    write_test_results_to_file,
+    print_test_results,
+)
+from rasa.dialogue_understanding_test.validation import (
+    validate_cli_arguments,
+    validate_test_cases,
+)
+from rasa.e2e_test.e2e_test_case import TestSuite
+from rasa.exceptions import RasaException
+from rasa.shared.constants import DEFAULT_ENDPOINTS_PATH
+from rasa.utils.beta import ensure_beta_feature_is_enabled
+from rasa.utils.endpoints import EndpointConfig
+
+RASA_PRO_BETA_DIALOGUE_UNDERSTANDING_TEST_ENV_VAR_NAME = (
+    "RASA_PRO_BETA_DIALOGUE_UNDERSTANDING_TEST"
+)
+
+structlogger = structlog.get_logger()
+
+
+def add_subparser(
+    subparsers: SubParsersAction, parents: List[argparse.ArgumentParser]
+) -> None:
+    """Add the dialogue understanding test subparser to `rasa test`.
+
+    Args:
+        subparsers: subparser we are going to attach to
+        parents: Parent parsers, needed to ensure tree structure in argparse
+    """
+    for subparser in subparsers.choices.values():
+        if subparser.prog == "rasa test":
+            du_test_subparser = create_du_test_subparser(parents)
+
+            for action in subparser._subparsers._actions:
+                if action.choices is not None:
+                    action.choices["du"] = du_test_subparser
+                    return
+
+    # If we get here, we couldn't hook the subparser to `rasa test`
+    raise RasaException(
+        "Hooking the dialogue understanding (du) test subparser to "
+        "`rasa test` command could not be completed. "
+        "Cannot run dialogue understanding testing."
+    )
+
+
+def create_du_test_subparser(
+    parents: List[argparse.ArgumentParser],
+) -> argparse.ArgumentParser:
+    """Create dialogue understanding test subparser."""
+    du_test_subparser = argparse.ArgumentParser(
+        prog="rasa test du",
+        parents=parents,
+        conflict_handler="resolve",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+        description="Runs dialogue understanding testing.",
+    )
+
+    du_test_subparser.set_defaults(func=execute_dialogue_understanding_tests)
+
+    add_du_test_arguments(du_test_subparser)
+    add_bot_arguments(du_test_subparser)
+
+    return du_test_subparser
+
+
+def add_bot_arguments(parser: argparse.ArgumentParser) -> None:
+    bot_arguments = parser.add_argument_group("Bot Settings")
+    add_model_param(bot_arguments, add_positional_arg=False)
+    add_endpoint_param(
+        bot_arguments,
+        help_text="Configuration file for the model server and the connectors as a "
+        "yml file.",
+    )
+    add_remote_storage_param(bot_arguments)
+
+
+def add_du_test_arguments(parser: argparse.ArgumentParser) -> None:
+    """Arguments for running dialogue understanding tests."""
+    du_arguments = parser.add_argument_group("Testing Settings")
+    du_arguments.add_argument(
+        "path-to-test-cases",
+        nargs="?",
+        type=str,
+        default=DEFAULT_INPUT_TESTS_PATH,
+        help="Input file or folder containing dialogue understanding test cases.",
+    )
+    du_arguments.add_argument(
+        "--output-file",
+        type=str,
+        default="dialogue_understanding_test_{date:%Y%m%d-%H%M%S}.yml".format(
+            date=datetime.datetime.now()
+        ),
+        help="Path to the output file to write the results to.",
+    )
+    du_arguments.add_argument(
+        "--no-output",
+        action="store_true",
+        help="If set, no output file will be written to disk.",
+    )
+    du_arguments.add_argument(
+        "--output-prompt",
+        action="store_true",
+        help="If set, the dialogue understanding test output will contain "
+        "prompts for each failure.",
+    )
+
+
+def execute_dialogue_understanding_tests(args: argparse.Namespace) -> None:
+    """Run the dialogue understanding tests.
+
+    Args:
+        args: Commandline arguments.
+    """
+    ensure_beta_feature_is_enabled(
+        "Dialogue Understanding (DU) Testing",
+        env_flag=RASA_PRO_BETA_DIALOGUE_UNDERSTANDING_TEST_ENV_VAR_NAME,
+    )
+
+    # basic validation of the passed CLI arguments
+    validate_cli_arguments(args)
+
+    # initialization of endpoints
+    endpoints = set_up_available_endpoints(args)
+
+    # read test cases from the given path
+    test_suite = get_valid_test_suite(args)
+
+    # setup stub custom actions if they are used
+    set_up_stub_custom_actions(test_suite, endpoints)
+
+    # set up the test runner, e.g. start the agent
+    try:
+        test_runner = DialogueUnderstandingTestRunner(
+            endpoints=endpoints,
+            model_path=args.model,
+            model_server=endpoints.model,
+            remote_storage=args.remote_storage,
+        )
+    except AgentNotReady as error:
+        structlogger.error(
+            "rasa.dialogue_understanding_test.agent_not_ready", message=error.message
+        )
+        sys.exit(1)
+
+    # run the actual test cases
+    test_results = asyncio.run(
+        test_runner.run_tests(
+            test_suite.test_cases, test_suite.fixtures, test_suite.metadata
+        )
+    )
+
+    # evaluate test results
+    failed_tests, passed_tests = split_test_results(test_results)
+    command_metrics = calculate_command_metrics(test_results)
+
+    # write results to console and file
+    print_test_results(failed_tests, passed_tests, command_metrics, args.output_prompt)
+    if not args.no_output:
+        write_test_results_to_file(
+            failed_tests,
+            passed_tests,
+            command_metrics,
+            args.output_file,
+            args.output_prompt,
+        )
+
+
+def get_valid_test_suite(args: argparse.Namespace) -> TestSuite:
+    """Read the test cases from the given test case path and validate them."""
+    path_to_test_cases = getattr(args, "path-to-test-cases", DEFAULT_INPUT_TESTS_PATH)
+    test_suite = read_test_suite(path_to_test_cases)
+    validate_test_cases(test_suite.test_cases)
+    return test_suite
+
+
+def set_up_available_endpoints(args: argparse.Namespace) -> AvailableEndpoints:
+    """Set up the available endpoints for the test runner."""
+    args.endpoints = rasa.cli.utils.get_validated_path(
+        args.endpoints, "endpoints", DEFAULT_ENDPOINTS_PATH, True
+    )
+    endpoints = AvailableEndpoints.get_instance(args.endpoints)
+
+    # Ignore all endpoints apart from action server, model, and nlu
+    # to ensure InMemoryTrackerStore is being used instead of production
+    # tracker store
+    endpoints.tracker_store = None
+    endpoints.lock_store = None
+    endpoints.event_broker = None
+
+    # disable nlg endpoint as we don't need it for dialogue understanding tests
+    endpoints.nlg = None
+
+    return endpoints
+
+
+def set_up_stub_custom_actions(
+    test_suite: TestSuite, endpoints: AvailableEndpoints
+) -> None:
+    """Set up the stub custom actions if they are used."""
+    if test_suite.stub_custom_actions:
+        if not endpoints.action:
+            endpoints.action = EndpointConfig()
+
+        endpoints.action.kwargs[KEY_STUB_CUSTOM_ACTIONS] = (
+            test_suite.stub_custom_actions
+        )
+
+
+def split_test_results(
+    results: List[DialogueUnderstandingTestResult],
+) -> tuple[
+    List[DialogueUnderstandingTestResult], List[DialogueUnderstandingTestResult]
+]:
+    """Split the test results into passed and failed test cases."""
+    passed_cases = [r for r in results if r.passed]
+    failed_cases = [r for r in results if not r.passed]
+
+    return passed_cases, failed_cases

{rasa_pro-3.11.3a1.dev7 → rasa_pro-3.12.0.dev2}/rasa/core/actions/action.py

@@ -17,15 +17,6 @@ from jsonschema import Draft202012Validator
 
 import rasa.core
 import rasa.shared.utils.io
-from rasa.shared.constants import (
-    TEXT,
-    ELEMENTS,
-    QUICK_REPLIES,
-    BUTTONS,
-    ATTACHMENT,
-    IMAGE,
-    CUSTOM,
-)
 from rasa.core.actions.custom_action_executor import (
     CustomActionExecutor,
     NoEndpointCustomActionExecutor,
@@ -264,18 +255,18 @@ def action_for_name_or_text(
 def create_bot_utterance(message: Dict[Text, Any]) -> BotUttered:
     """Create BotUttered event from message."""
     bot_message = BotUttered(
-        text=message.pop(TEXT, None),
+        text=message.pop("text", None),
         data={
-            ELEMENTS: message.pop(ELEMENTS, None),
-            QUICK_REPLIES: message.pop(QUICK_REPLIES, None),
-            BUTTONS: message.pop(BUTTONS, None),
+            "elements": message.pop("elements", None),
+            "quick_replies": message.pop("quick_replies", None),
+            "buttons": message.pop("buttons", None),
             # for legacy / compatibility reasons we need to set the image
             # to be the attachment if there is no other attachment (the
             # `.get` is intentional - no `pop` as we still need the image`
             # property to set it in the following line)
-            ATTACHMENT: message.pop(ATTACHMENT, None) or message.get(IMAGE, None),
-            IMAGE: message.pop(IMAGE, None),
-            CUSTOM: message.pop(CUSTOM, None),
+            "attachment": message.pop("attachment", None) or message.get("image", None),
+            "image": message.pop("image", None),
+            "custom": message.pop("custom", None),
         },
         metadata=message,
     )

{rasa_pro-3.11.3a1.dev7 → rasa_pro-3.12.0.dev2}/rasa/core/channels/__init__.py

@@ -31,7 +31,6 @@ from rasa.core.channels.vier_cvg import CVGInput
 from rasa.core.channels.voice_stream.twilio_media_streams import (
     TwilioMediaStreamsInputChannel,
 )
-from rasa.core.channels.studio_chat import StudioChatInput
 
 input_channel_classes: List[Type[InputChannel]] = [
     CmdlineInput,
@@ -54,7 +53,6 @@ input_channel_classes: List[Type[InputChannel]] = [
     JambonzVoiceReadyInput,
     TwilioMediaStreamsInputChannel,
     BrowserAudioInputChannel,
-    StudioChatInput,
 ]
 
 # Mapping from an input channel name to its class to allow name based lookup.

{rasa_pro-3.11.3a1.dev7 → rasa_pro-3.12.0.dev2}/rasa/core/channels/socketio.py

@@ -54,9 +54,31 @@ class SocketIOOutput(OutputChannel):
         super().__init__()
         self.sio = sio
         self.bot_message_evt = bot_message_evt
+        self.last_event_timestamp = (
+            -1
+        )  # Initialize with -1 to send all events on first message
+
+    def _get_new_events(self) -> List[Dict[Text, Any]]:
+        """Get events that are newer than the last sent event."""
+        events = self.tracker_state.get("events", []) if self.tracker_state else []
+        new_events = [
+            event for event in events if event["timestamp"] > self.last_event_timestamp
+        ]
+        if new_events:
+            self.last_event_timestamp = new_events[-1]["timestamp"]
+        return new_events
 
     async def _send_message(self, socket_id: Text, response: Any) -> None:
         """Sends a message to the recipient using the bot event."""
+        # send tracker state (contains stack, slots and more)
+        await self.sio.emit("tracker_state", self.tracker_state, room=socket_id)
+
+        # send new events
+        new_events = self._get_new_events()
+        if new_events:
+            await self.sio.emit("rasa_events", new_events, room=socket_id)
+
+        # send bot response
         await self.sio.emit(self.bot_message_evt, response, room=socket_id)
 
     async def send_text_message(
@@ -192,7 +214,7 @@ class SocketIOInput(InputChannel):
 
     def blueprint(
         self, on_new_message: Callable[[UserMessage], Awaitable[Any]]
-    ) -> SocketBlueprint:
+    ) -> Blueprint:
         """Defines a Sanic blueprint."""
         # Workaround so that socketio works with requests from other origins.
         # https://github.com/miguelgrinberg/python-socketio/issues/205#issuecomment-493769183

{rasa_pro-3.11.3a1.dev7 → rasa_pro-3.12.0.dev2}/rasa/core/nlg/contextual_response_rephraser.py

@@ -2,7 +2,6 @@ from typing import Any, Dict, Optional, Text
 
 import structlog
 from jinja2 import Template
-
 from rasa import telemetry
 from rasa.core.nlg.response import TemplatedNaturalLanguageGenerator
 from rasa.core.nlg.summarize import summarize_conversation
@@ -19,14 +18,6 @@ from rasa.shared.constants import (
 from rasa.shared.core.domain import KEY_RESPONSES_TEXT, Domain
 from rasa.shared.core.events import BotUttered, UserUttered
 from rasa.shared.core.trackers import DialogueStateTracker
-from rasa.shared.nlu.constants import (
-    PROMPTS,
-    KEY_USER_PROMPT,
-    KEY_LLM_RESPONSE_METADATA,
-    KEY_PROMPT_NAME,
-    KEY_COMPONENT_NAME,
-)
-from rasa.shared.providers.llm.llm_response import LLMResponse
 from rasa.shared.utils.health_check.llm_health_check_mixin import LLMHealthCheckMixin
 from rasa.shared.utils.llm import (
     DEFAULT_OPENAI_GENERATE_MODEL_NAME,
@@ -133,39 +124,6 @@ class ContextualResponseRephraser(
             ContextualResponseRephraser.__name__,
         )
 
-    @classmethod
-    def _add_prompt_and_llm_metadata_to_response(
-        cls,
-        response: Dict[str, Any],
-        prompt_name: str,
-        user_prompt: str,
-        llm_response: Optional["LLMResponse"] = None,
-    ) -> Dict[str, Any]:
-        """Stores the prompt and LLMResponse metadata to response.
-
-        Args:
-            response: The response to add the prompt and LLMResponse metadata to.
-            prompt_name: A name identifying prompt usage.
-            user_prompt: The user prompt that was sent to the LLM.
-            llm_response: The response object from the LLM (None if no response).
-        """
-        from rasa.dialogue_understanding.utils import record_commands_and_prompts
-
-        if not record_commands_and_prompts:
-            return response
-
-        prompt_data: Dict[Text, Any] = {
-            KEY_COMPONENT_NAME: cls.__name__,
-            KEY_PROMPT_NAME: prompt_name,
-            KEY_USER_PROMPT: user_prompt,
-            KEY_LLM_RESPONSE_METADATA: llm_response.to_dict() if llm_response else None,
-        }
-
-        prompts = response.get(PROMPTS, [])
-        prompts.append(prompt_data)
-        response[PROMPTS] = prompts
-        return response
-
     def _last_message_if_human(self, tracker: DialogueStateTracker) -> Optional[str]:
         """Returns the latest message from the tracker.
 
@@ -184,21 +142,20 @@ class ContextualResponseRephraser(
                 return None
         return None
 
-    async def _generate_llm_response(self, prompt: str) -> Optional[LLMResponse]:
-        """
-        Use LLM to generate a response, returning an LLMResponse object
-        containing both the generated text (choices) and metadata.
+    async def _generate_llm_response(self, prompt: str) -> Optional[str]:
+        """Use LLM to generate a response.
 
         Args:
-            prompt: The prompt to send to the LLM.
+            prompt: the prompt to send to the LLM
 
         Returns:
-            An LLMResponse object if successful, otherwise None.
+            generated text
         """
         llm = llm_factory(self.llm_config, DEFAULT_LLM_CONFIG)
 
         try:
-            return await llm.acompletion(prompt)
+            llm_response = await llm.acompletion(prompt)
+            return llm_response.choices[0]
         except Exception as e:
             # unfortunately, langchain does not wrap LLM exceptions which means
             # we have to catch all exceptions here
@@ -298,21 +255,11 @@ class ContextualResponseRephraser(
             or self.llm_property(MODEL_NAME_CONFIG_KEY),
             llm_model_group_id=self.llm_property(MODEL_GROUP_ID_CONFIG_KEY),
         )
-        llm_response = await self._generate_llm_response(prompt)
-        llm_response = LLMResponse.ensure_llm_response(llm_response)
-
-        response = self._add_prompt_and_llm_metadata_to_response(
-            response=response,
-            prompt_name="rephrase_prompt",
-            user_prompt=prompt,
-            llm_response=llm_response,
-        )
-
-        if not (llm_response and llm_response.choices and llm_response.choices[0]):
-            # If the LLM fails to generate a response, return the original response.
+        if not (updated_text := await self._generate_llm_response(prompt)):
+            # If the LLM fails to generate a response, we
+            # return the original response.
             return response
 
-        updated_text = llm_response.choices[0]
         structlogger.debug(
             "nlg.rewrite.complete",
             response_text=response_text,

{rasa_pro-3.11.3a1.dev7 → rasa_pro-3.12.0.dev2}/rasa/core/policies/enterprise_search_policy.py

@@ -2,7 +2,6 @@ import importlib.resources
 import json
 import re
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Text
-
 import dotenv
 import structlog
 from jinja2 import Template
@@ -64,19 +63,11 @@ from rasa.shared.core.events import Event, UserUttered, BotUttered
 from rasa.shared.core.generator import TrackerWithCachedStates
 from rasa.shared.core.trackers import DialogueStateTracker, EventVerbosity
 from rasa.shared.exceptions import RasaException, FileIOException
-from rasa.shared.nlu.constants import (
-    PROMPTS,
-    KEY_USER_PROMPT,
-    KEY_LLM_RESPONSE_METADATA,
-    KEY_PROMPT_NAME,
-    KEY_COMPONENT_NAME,
-)
 from rasa.shared.nlu.training_data.training_data import TrainingData
 from rasa.shared.providers.embedding._langchain_embedding_client_adapter import (
     _LangchainEmbeddingClientAdapter,
 )
 from rasa.shared.providers.llm.llm_client import LLMClient
-from rasa.shared.providers.llm.llm_response import LLMResponse
 from rasa.shared.utils.cli import print_error_and_exit
 from rasa.shared.utils.health_check.embeddings_health_check_mixin import (
     EmbeddingsHealthCheckMixin,
@@ -281,43 +272,6 @@ class EnterpriseSearchPolicy(LLMHealthCheckMixin, EmbeddingsHealthCheckMixin, Po
         # Wrap the embedding client in the adapter
         return _LangchainEmbeddingClientAdapter(client)
 
-    @classmethod
-    def _add_prompt_and_llm_response_to_latest_message(
-        cls,
-        tracker: DialogueStateTracker,
-        prompt_name: str,
-        user_prompt: str,
-        llm_response: Optional[LLMResponse] = None,
-    ) -> None:
-        """Stores the prompt and LLMResponse metadata in the tracker.
-
-        Args:
-            tracker: The DialogueStateTracker containing the current conversation state.
-            prompt_name: A name identifying prompt usage.
-            user_prompt: The user prompt that was sent to the LLM.
-            llm_response: The response object from the LLM (None if no response).
-        """
-        from rasa.dialogue_understanding.utils import record_commands_and_prompts
-
-        if not record_commands_and_prompts:
-            return
-
-        if not tracker.latest_message:
-            return
-
-        parse_data = tracker.latest_message.parse_data
-        if PROMPTS not in parse_data:
-            parse_data[PROMPTS] = []  # type: ignore[literal-required]
-
-        prompt_data: Dict[Text, Any] = {
-            KEY_COMPONENT_NAME: cls.__name__,
-            KEY_PROMPT_NAME: prompt_name,
-            KEY_USER_PROMPT: user_prompt,
-            KEY_LLM_RESPONSE_METADATA: llm_response.to_dict() if llm_response else None,
-        }
-
-        parse_data[PROMPTS].append(prompt_data)  # type: ignore[literal-required]
-
     def train(  # type: ignore[override]
         self,
         training_trackers: List[TrackerWithCachedStates],
@@ -544,27 +498,13 @@ class EnterpriseSearchPolicy(LLMHealthCheckMixin, EmbeddingsHealthCheckMixin, Po
 
         if self.use_llm:
             prompt = self._render_prompt(tracker, documents.results)
-            llm_response = await self._generate_llm_answer(llm, prompt)
-            llm_response = LLMResponse.ensure_llm_response(llm_response)
-
-            self._add_prompt_and_llm_response_to_latest_message(
-                tracker=tracker,
-                prompt_name="enterprise_search_prompt",
-                user_prompt=prompt,
-                llm_response=llm_response,
-            )
+            llm_answer = await self._generate_llm_answer(llm, prompt)
 
-            if llm_response is None or not llm_response.choices:
-                logger.debug(f"{logger_key}.no_llm_response")
-                response = None
-            else:
-                llm_answer = llm_response.choices[0]
+            if self.citation_enabled:
+                llm_answer = self.post_process_citations(llm_answer)
 
-                if self.citation_enabled:
-                    llm_answer = self.post_process_citations(llm_answer)
-
-                logger.debug(f"{logger_key}.llm_answer", llm_answer=llm_answer)
-                response = llm_answer
+            logger.debug(f"{logger_key}.llm_answer", llm_answer=llm_answer)
+            response = llm_answer
         else:
             response = documents.results[0].metadata.get("answer", None)
             if not response:
@@ -576,6 +516,7 @@ class EnterpriseSearchPolicy(LLMHealthCheckMixin, EmbeddingsHealthCheckMixin, Po
                 "enterprise_search_policy.predict_action_probabilities.no_llm",
                 search_results=documents,
             )
+
         if response is None:
             return self._create_prediction_internal_error(domain, tracker)
 
@@ -640,18 +581,10 @@ class EnterpriseSearchPolicy(LLMHealthCheckMixin, EmbeddingsHealthCheckMixin, Po
 
     async def _generate_llm_answer(
         self, llm: LLMClient, prompt: Text
-    ) -> Optional[LLMResponse]:
-        """Fetches an LLM completion for the provided prompt.
-
-        Args:
-            llm: The LLM client used to get the completion.
-            prompt: The prompt text to send to the model.
-
-        Returns:
-            An LLMResponse object, or None if the call fails.
-        """
+    ) -> Optional[Text]:
         try:
-            return await llm.acompletion(prompt)
+            llm_response = await llm.acompletion(prompt)
+            llm_answer = llm_response.choices[0]
         except Exception as e:
             # unfortunately, langchain does not wrap LLM exceptions which means
             # we have to catch all exceptions here
@@ -659,7 +592,9 @@ class EnterpriseSearchPolicy(LLMHealthCheckMixin, EmbeddingsHealthCheckMixin, Po
                 "enterprise_search_policy._generate_llm_answer.llm_error",
                 error=e,
             )
-            return None
+            llm_answer = None
+
+        return llm_answer
 
     def _create_prediction(
         self,