Model reload changes. (#1116)

* Model reload changes. * Model reload changes. * Model reload changes. --------- Co-authored-by: Nupur Khare <[email protected]> Co-authored-by: udit-pandey <[email protected]>
digiteinfotech · Jan 10, 2024 · 91775c1 · 91775c1
1 parent 8ffc87b
commit 91775c1
Show file tree

Hide file tree

Showing 6 changed files with 136 additions and 13 deletions.
diff --git a/kairon/chat/routers/web_client.py b/kairon/chat/routers/web_client.py
@@ -77,5 +77,5 @@ async def reload_model(
     """
     Retrieves chat client config of a bot.
     """
-    background_tasks.add_task(ChatUtils.reload, bot)
+    background_tasks.add_task(ChatUtils.reload, bot, current_user.get_user())
     return {"message": "Reloading Model!"}
diff --git a/kairon/chat/utils.py b/kairon/chat/utils.py
@@ -1,5 +1,6 @@
 import datetime
 import json
+import os
 from typing import Text, Dict
 
 from loguru import logger
@@ -11,7 +12,9 @@
 from .agent_processor import AgentProcessor
 from .. import Utility
 from ..live_agent.factory import LiveAgentFactory
+from ..shared.account.activity_log import UserActivityLogger
 from ..shared.actions.utils import ActionUtility
+from ..shared.constants import UserActivityType
 from ..shared.live_agent.processor import LiveAgentsProcessor
 from ..shared.metering.constants import MetricType
 from ..shared.metering.metering_processor import MeteringProcessor
@@ -29,8 +32,20 @@ async def chat(data: Text, account: int, bot: Text, user: Text, is_integration_u
         return chat_response
 
     @staticmethod
-    def reload(bot: Text):
-        AgentProcessor.reload(bot)
+    def reload(bot: Text, user: Text):
+        exc = None
+        status = "Success"
+        try:
+            AgentProcessor.reload(bot)
+        except Exception as e:
+            logger.error(e)
+            exc = str(e)
+            status = "Failed"
+        finally:
+            UserActivityLogger.add_log(a_type=UserActivityType.model_reload.value,
+                                       email=user, bot=bot,
+                                       data={"username": user, "process_id": os.getpid(), "exception": exc,
+                                             "status": status})
 
     @staticmethod
     def __attach_agent_handoff_metadata(account: int, bot: Text, sender_id: Text, bot_predictions, tracker):

diff --git a/kairon/shared/constants.py b/kairon/shared/constants.py
@@ -58,6 +58,7 @@ class UserActivityType(str, Enum):
     login_refresh_token = "login_refresh_token"
     invalid_login = 'invalid_login'
     template_creation = 'template_creation'
+    model_reload = "model_reload"
 
 
 class EventClass(str, Enum):

diff --git a/kairon/shared/utils.py b/kairon/shared/utils.py
@@ -62,8 +62,8 @@
 from websockets import connect
 
 from .actions.models import ActionParameterType
-from .constants import EventClass
-from .constants import MaskingStrategy, SYSTEM_TRIGGERED_UTTERANCES, ChannelTypes, PluginTypes
+from .constants import MaskingStrategy, SYSTEM_TRIGGERED_UTTERANCES, ChannelTypes, PluginTypes, UserActivityType, \
+    EventClass
 from .data.constant import TOKEN_TYPE, KAIRON_TWO_STAGE_FALLBACK, SLOT_TYPE
 from .data.dto import KaironStoryStep
 from .models import StoryStepType, LlmPromptType, LlmPromptSource
@@ -1056,14 +1056,27 @@ async def chat(data: Text, bot: Text, user: Text, email: Text):
 
     @staticmethod
     def reload_model(bot: Text, email: Text):
-        if Utility.environment.get('model') and Utility.environment['model']['agent'].get('url'):
-            from kairon.shared.auth import Authentication
-            agent_url = Utility.environment['model']['agent'].get('url')
-            token, _ = Authentication.generate_integration_token(bot, email, expiry=5, token_type=TOKEN_TYPE.CHANNEL.value)
-            response = Utility.http_request('get', urljoin(agent_url, f"/api/bot/{bot}/reload"), token, email)
-            return json.loads(response)
-        else:
-            raise AppException("Agent config not found!")
+        from kairon.shared.account.activity_log import UserActivityLogger
+
+        exc = None
+        status = "Initiated"
+        try:
+            if Utility.environment.get('model') and Utility.environment['model']['agent'].get('url'):
+                from kairon.shared.auth import Authentication
+                agent_url = Utility.environment['model']['agent'].get('url')
+                token, _ = Authentication.generate_integration_token(bot, email, expiry=5, token_type=TOKEN_TYPE.CHANNEL.value)
+                response = Utility.http_request('get', urljoin(agent_url, f"/api/bot/{bot}/reload"), token, email)
+                return json.loads(response)
+            else:
+                raise AppException("Agent config not found!")
+        except Exception as e:
+            logger.exception(e)
+            exc = str(e)
+            status = "Failed"
+            raise AppException(e)
+        finally:
+            UserActivityLogger.add_log(a_type=UserActivityType.model_reload.value,
+                                       email=email, bot=bot, data={"username": email, "exception": exc, "status": status})
 
     @staticmethod
     def validate_create_template_request(data: Dict):

diff --git a/tests/integration_test/chat_service_test.py b/tests/integration_test/chat_service_test.py
@@ -2,6 +2,7 @@
 import json
 import os
 import time
+from datetime import datetime, timedelta
 from unittest import mock
 from urllib.parse import urlencode, quote_plus
 
@@ -828,6 +829,63 @@ def test_reload(mock_store):
         assert header in actual_headers
 
 
+@patch('kairon.chat.utils.ChatUtils.reload')
+def test_reload_logging(mock_reload):
+    processor = MongoProcessor()
+    start_time = datetime.utcnow() - timedelta(days=1)
+    end_time = datetime.utcnow() + timedelta(days=1)
+    mock_reload.return_value = None
+    response = client.get(
+        f"/api/bot/{bot}/reload",
+        headers={
+            "Authorization": token_type + " " + token
+        },
+    )
+    actual = response.json()
+    assert actual["success"]
+    assert actual["error_code"] == 0
+    assert actual["data"] is None
+    assert actual["message"] == "Reloading Model!"
+    logs = processor.get_logs(bot, "audit_logs", start_time, end_time)
+    logs[0].pop('timestamp')
+    logs[0].pop('_id')
+    logs[0]['data'].pop('process_id')
+    assert logs[0] == {'attributes': [{'key': 'bot', 'value': bot}], 'user': '[email protected]', 'action': 'activity',
+                       'entity': 'model_reload',
+                       'data': {'message': None, 'username': '[email protected]', 'exception': None, 'status': 'Success'}}
+
+
+@mock.patch('kairon.chat.agent_processor.AgentProcessor.reload', autospec=True)
+def test_reload_event_exception(mock_reload):
+    processor = MongoProcessor()
+    start_time = datetime.utcnow() - timedelta(days=1)
+    end_time = datetime.utcnow() + timedelta(days=1)
+    def _reload(*args):
+        raise Exception('Simulated exception during model reload')
+
+    mock_reload.side_effect = _reload
+    response = client.get(
+        f"/api/bot/{bot}/reload",
+        headers={
+            "Authorization": token_type + " " + token
+        },
+    )
+    actual = response.json()
+    assert actual["success"]
+    assert actual["error_code"] == 0
+    assert actual["data"] is None
+    assert actual["message"] == "Reloading Model!"
+    logs = processor.get_logs(bot, "audit_logs", start_time, end_time)
+    logs[0].pop('timestamp')
+    logs[0].pop('_id')
+    logs[0]['data'].pop('process_id')
+    assert logs[0] == {'attributes': [{'key': 'bot', 'value': bot}], 'user': '[email protected]',
+                       'action': 'activity', 'entity': 'model_reload',
+                       'data': {'message': None, 'username': '[email protected]',
+                                'exception': 'Simulated exception during model reload', 'status': 'Failed'},
+                       }
+
+
 def test_reload_exception():
     response = client.get(
         f"/api/bot/{bot}/reload",

diff --git a/tests/integration_test/services_test.py b/tests/integration_test/services_test.py
@@ -5835,6 +5835,9 @@ def test_set_templates_insecure():
 
 @responses.activate
 def test_reload_model(monkeypatch):
+    processor = MongoProcessor()
+    start_time = datetime.utcnow() - timedelta(days=1)
+    end_time = datetime.utcnow() + timedelta(days=1)
     def mongo_store(*arge, **kwargs):
         return None
 
@@ -5859,6 +5862,39 @@ def mongo_store(*arge, **kwargs):
     assert actual['error_code'] == 0
     assert actual['message'] == "Reloading Model!"
     assert actual['success']
+    logs = processor.get_logs(pytest.bot, "audit_logs", start_time, end_time)
+    logs[0].pop('timestamp')
+    logs[0].pop('_id')
+    assert logs[0] == {'attributes': [{'key': 'bot', 'value': pytest.bot}], 'user': '[email protected]', 'action': 'activity', 'entity': 'model_reload', 'data': {'message': None, 'username': '[email protected]', 'exception': None, 'status': 'Initiated'}}
+
+
+@responses.activate
+def test_reload_model_exception():
+    processor = MongoProcessor()
+    start_time = datetime.utcnow() - timedelta(days=1)
+    end_time = datetime.utcnow() + timedelta(days=1)
+    # def mongo_store(*arge, **kwargs):
+    #     return None
+    #
+    # monkeypatch.setattr(Utility, "get_local_mongo_store", mongo_store)
+    # monkeypatch.setitem(Utility.environment['action'], "url", None)
+    # monkeypatch.setitem(Utility.environment['model']['agent'], "url", "http://localhost/")
+    response = client.get(
+        f"/api/bot/{pytest.bot}/model/reload",
+        headers={"Authorization": pytest.token_type + " " + pytest.access_token}
+    )
+    actual = response.json()
+    assert actual['data'] is None
+    assert not actual["success"]
+    assert actual["error_code"] == 422
+    assert actual['message'] == 'Agent config not found!'
+    logs = processor.get_logs(pytest.bot, "audit_logs", start_time, end_time)
+    logs[0].pop('timestamp')
+    logs[0].pop('_id')
+    assert logs[0] == {'attributes': [{'key': 'bot', 'value': pytest.bot}],
+                       'user': '[email protected]', 'action': 'activity', 'entity': 'model_reload',
+                       'data': {'message': None, 'username': '[email protected]',
+                                'exception': 'Agent config not found!', 'status': 'Failed'}}
 
 
 def test_get_config_templates():