diff --git a/kairon/chat/converters/channels/whatsapp.py b/kairon/chat/converters/channels/whatsapp.py
index 1cd9135b8..1abbe35f2 100644
--- a/kairon/chat/converters/channels/whatsapp.py
+++ b/kairon/chat/converters/channels/whatsapp.py
@@ -71,6 +71,8 @@ def button_transformer(self, message):
             body_default = ElementTransformerOps.getChannelConfig(self.channel, "body_message")
             body_msg = {"text": body_default}
             for item in jsoniterator:
+                if item.get("type") == 'label':
+                    body_msg = {"text": item.get("text")}
                 if item.get("type") == ElementTypes.BUTTON.value:
                     title = ElementTransformerOps.json_generator(item.get("children"))
                     for titletext in title:
diff --git a/kairon/shared/llm/processor.py b/kairon/shared/llm/processor.py
index 5961cef95..26cc541d0 100644
--- a/kairon/shared/llm/processor.py
+++ b/kairon/shared/llm/processor.py
@@ -130,7 +130,7 @@ async def predict(self, query: Text, user, *args, **kwargs) -> Tuple:
 
             context = await self.__attach_similarity_prompt_if_enabled(query_embedding, context_prompt, **kwargs)
             answer = await self.__get_answer(query, system_prompt, context, user, invocation=invocation,llm_type = llm_type, **kwargs)
-            response = {"content": answer}
+            response = {"content": answer, "context": context}
         except Exception as e:
             logging.exception(e)
             if embeddings_created:
diff --git a/tests/integration_test/action_service_test.py b/tests/integration_test/action_service_test.py
index e5d0b42e8..cd36234b4 100644
--- a/tests/integration_test/action_service_test.py
+++ b/tests/integration_test/action_service_test.py
@@ -13623,6 +13623,8 @@ def test_prompt_action_dispatch_response_disabled(mock_get_embedding, aiorespons
     for event in events:
         if event.get('time_elapsed') is not None:
             del event['time_elapsed']
+        if event.get('llm_response_log'):
+            event['llm_response_log'].pop('context')
     assert events == [
         {'type': 'llm_response',
          'response': 'Python is dynamically typed, garbage-collected, high level, general purpose programming.',
@@ -13754,6 +13756,8 @@ def test_prompt_action_set_slots(mock_slot_set, mock_get_embedding, aioresponses
     for event in events:
         if event.get('time_elapsed') is not None:
             del event['time_elapsed']
+        if event.get('llm_response_log'):
+            event['llm_response_log'].pop('context')
     assert events == [
         {'type': 'llm_response',
          'response': '{"api_type": "filter", {"filter": {"must": [{"key": "Date Added", "match": {"value": 1673721000.0}}]}}}',
@@ -13902,9 +13906,12 @@ def test_prompt_action_response_action_slot_prompt(mock_get_embedding, aiorespon
     assert log["time_elapsed"]
     log.pop('time_elapsed')
     events = log.pop('events')
+    assert len(events[0]['llm_response_log']['context']) > 0
     for event in events:
         if event.get('time_elapsed') is not None:
             del event['time_elapsed']
+        if event.get('llm_response_log'):
+            event['llm_response_log'].pop('context')
     assert events == [
         {'type': 'llm_response',
          'response': 'Python is dynamically typed, garbage-collected, high level, general purpose programming.',
diff --git a/tests/unit_test/llm_test.py b/tests/unit_test/llm_test.py
index 288d94565..7909c551b 100644
--- a/tests/unit_test/llm_test.py
+++ b/tests/unit_test/llm_test.py
@@ -6,10 +6,7 @@
 
 import numpy as np
 import pytest
-import ujson as json
-from absl.logging import exception
 from aiohttp import ClientConnectionError
-from fastapi.dependencies.utils import request_body_to_args
 from mongoengine import connect
 
 from kairon.shared.rest_client import AioRestClient
@@ -17,12 +14,10 @@
 
 
 from kairon.exceptions import AppException
-from kairon.shared.admin.constants import BotSecretType
 from kairon.shared.admin.data_objects import BotSecrets, LLMSecret
 from kairon.shared.cognition.data_objects import CognitionData, CognitionSchema
 from kairon.shared.data.constant import DEFAULT_SYSTEM_PROMPT, DEFAULT_LLM
 from kairon.shared.llm.processor import LLMProcessor
-import litellm
 from deepdiff import DeepDiff