rasa-pro 3.11.0a2__py3-none-any.whl → 3.11.0a4.dev1__py3-none-any.whl

rasa/model_manager/runner_service.py

@@ -0,0 +1,279 @@
+import os
+import shutil
+from typing import Dict
+import aiohttp
+import structlog
+import subprocess
+from pydantic import BaseModel, ConfigDict
+from enum import Enum
+
+from rasa.exceptions import ModelNotFound
+from rasa.model_manager.utils import (
+    models_base_path,
+    subpath,
+    write_encoded_data_to_file,
+)
+from rasa.constants import MODEL_ARCHIVE_EXTENSION
+
+from rasa.model_manager import config
+from rasa.model_manager.utils import logs_path, ensure_base_directory_exists
+
+structlogger = structlog.get_logger()
+
+
+class BotSessionStatus(str, Enum):
+    """Enum for the bot status."""
+
+    QUEUED = "queued"
+    RUNNING = "running"
+    STOPPED = "stopped"
+
+
+class BotSession(BaseModel):
+    """Store information about a running bot."""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    deployment_id: str
+    status: BotSessionStatus
+    process: subprocess.Popen
+    url: str
+    internal_url: str
+    port: int
+
+    def is_alive(self) -> bool:
+        """Check if the bot is alive."""
+        return self.process.poll() is None
+
+    def is_status_indicating_alive(self) -> bool:
+        """Check if the status indicates that the bot is alive."""
+        return self.status in {BotSessionStatus.QUEUED, BotSessionStatus.RUNNING}
+
+    def has_died_recently(self) -> bool:
+        """Check if the bot has died recently.
+
+        Process will indicate that the bot exited,
+        but status is not yet updated.
+        """
+        return self.is_status_indicating_alive() and not self.is_alive()
+
+    async def completed_startup_recently(self) -> bool:
+        """Check if the bot has completed startup recently."""
+        return self.status == BotSessionStatus.QUEUED and await is_bot_startup_finished(
+            self
+        )
+
+
+def bot_path(deployment_id: str) -> str:
+    """Return the path to the bot directory for a given deployment id."""
+    return os.path.abspath(
+        f"{config.SERVER_BASE_WORKING_DIRECTORY}/bots/{deployment_id}"
+    )
+
+
+async def is_bot_startup_finished(bot: BotSession) -> bool:
+    """Send a request to the bot to see if the bot is up and running."""
+    health_timeout = aiohttp.ClientTimeout(total=5, sock_connect=2, sock_read=3)
+    try:
+        async with aiohttp.ClientSession(timeout=health_timeout) as session:
+            # can't use /status as by default the bot API is not enabled, only
+            # the input channel
+            async with session.get(f"{bot.internal_url}/license") as resp:
+                return resp.status == 200
+    except aiohttp.client_exceptions.ClientConnectorError:
+        structlogger.debug(
+            "model_runner.bot.not_running_yet", deployment_id=bot.deployment_id
+        )
+        return False
+
+
+def update_bot_to_stopped(bot: BotSession) -> None:
+    """Set a bots state to stopped."""
+    structlogger.info(
+        "model_runner.bot_stopped",
+        deployment_id=bot.deployment_id,
+        status=bot.process.returncode,
+    )
+    bot.status = BotSessionStatus.STOPPED
+
+
+def update_bot_to_running(bot: BotSession) -> None:
+    """Set a bots state to running."""
+    structlogger.info(
+        "model_runner.bot_running",
+        deployment_id=bot.deployment_id,
+    )
+    bot.status = BotSessionStatus.RUNNING
+
+
+def get_open_port() -> int:
+    """Get an open port on the system that is not in use yet."""
+    # from https://stackoverflow.com/questions/2838244/get-open-tcp-port-in-python/2838309#2838309
+    import socket
+
+    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    s.bind(("", 0))
+    s.listen(1)
+    port = s.getsockname()[1]
+    s.close()
+    return port
+
+
+def write_config_data_to_files(encoded_configs: Dict[str, str], base_path: str) -> None:
+    """Write the encoded config data to files."""
+    for key, value in encoded_configs.items():
+        write_encoded_data_to_file(value, subpath(base_path, f"{key}.yml"))
+
+
+def prepare_bot_directory(
+    bot_base_path: str,
+    model_name: str,
+    encoded_configs: Dict[str, str],
+) -> None:
+    """Prepare the bot directory for a new bot session."""
+    if not os.path.exists(bot_base_path):
+        os.makedirs(bot_base_path, exist_ok=True)
+    else:
+        shutil.rmtree(bot_base_path, ignore_errors=True)
+
+    model_file_name = f"{model_name}.{MODEL_ARCHIVE_EXTENSION}"
+    model_path = subpath(models_base_path(), model_file_name)
+
+    if config.SERVER_MODEL_REMOTE_STORAGE and not os.path.exists(model_path):
+        fetch_remote_model_to_dir(
+            model_file_name,
+            models_base_path(),
+            config.SERVER_MODEL_REMOTE_STORAGE,
+        )
+
+    if not os.path.exists(model_path):
+        raise ModelNotFound(f"Model '{model_file_name}' not found in '{model_path}'.")
+
+    os.makedirs(subpath(bot_base_path, "models"), exist_ok=True)
+    shutil.copy(
+        src=model_path,
+        dst=subpath(bot_base_path, "models"),
+    )
+
+    write_config_data_to_files(encoded_configs, bot_base_path)
+
+
+def fetch_remote_model_to_dir(
+    model_name: str, target_path: str, storage_type: str
+) -> str:
+    """Fetch the model from remote storage.
+
+    Returns the path to the model diretory.
+    """
+    from rasa.core.persistor import get_persistor
+
+    persistor = get_persistor(storage_type)
+
+    # we now there must be a persistor, because the config is set
+    # this is here to please the type checker for the call below
+    assert persistor is not None
+
+    try:
+        return persistor.retrieve(model_name=model_name, target_path=target_path)
+    except FileNotFoundError as e:
+        raise ModelNotFound() from e
+
+
+def start_bot_process(
+    deployment_id: str, bot_base_path: str, base_url_path: str
+) -> BotSession:
+    port = get_open_port()
+    log_path = logs_path(deployment_id)
+
+    ensure_base_directory_exists(log_path)
+
+    full_command = [
+        config.RASA_PYTHON_PATH,
+        "-m",
+        "rasa.__main__",
+        "run",
+        "--endpoints",
+        f"{bot_base_path}/endpoints.yml",
+        "--credentials",
+        f"{bot_base_path}/credentials.yml",
+        "--debug",
+        f"--port={port}",
+        "--model",
+        f"{bot_base_path}/models",
+    ]
+
+    structlogger.debug(
+        "model_runner.bot.starting_command",
+        deployment_id=deployment_id,
+        command=" ".join(full_command),
+    )
+
+    process = subprocess.Popen(
+        full_command,
+        cwd=bot_base_path,
+        stdout=open(log_path, "w"),
+        stderr=subprocess.STDOUT,
+        env=os.environ.copy(),
+    )
+
+    internal_bot_url = f"http://localhost:{port}"
+
+    structlogger.info(
+        "model_runner.bot.starting",
+        deployment_id=deployment_id,
+        log=log_path,
+        url=internal_bot_url,
+        port=port,
+        pid=process.pid,
+    )
+
+    return BotSession(
+        deployment_id=deployment_id,
+        status=BotSessionStatus.QUEUED,
+        process=process,
+        url=f"{base_url_path}?deployment_id={deployment_id}",
+        internal_url=internal_bot_url,
+        port=port,
+    )
+
+
+def run_bot(
+    deployment_id: str,
+    model_name: str,
+    base_url_path: str,
+    encoded_configs: Dict[str, str],
+) -> BotSession:
+    """Deploy a bot based on a given training id."""
+    bot_base_path = bot_path(deployment_id)
+    prepare_bot_directory(bot_base_path, model_name, encoded_configs)
+
+    return start_bot_process(deployment_id, bot_base_path, base_url_path)
+
+
+async def update_bot_status(bot: BotSession) -> None:
+    """Update the status of a bot based on the process return code."""
+    if bot.has_died_recently():
+        update_bot_to_stopped(bot)
+    elif await bot.completed_startup_recently():
+        update_bot_to_running(bot)
+
+
+def terminate_bot(bot: BotSession) -> None:
+    """Terminate the bot process."""
+    if not bot.is_status_indicating_alive():
+        # if the bot is not running, we don't need to terminate it
+        return
+
+    try:
+        bot.process.terminate()
+        structlogger.info(
+            "model_runner.stop_bot.stopped",
+            deployment_id=bot.deployment_id,
+            status=bot.process.returncode,
+        )
+        bot.status = BotSessionStatus.STOPPED
+    except ProcessLookupError:
+        structlogger.debug(
+            "model_runner.stop_bot.process_not_found",
+            deployment_id=bot.deployment_id,
+        )

rasa/model_manager/socket_bridge.py

@@ -0,0 +1,143 @@
+from typing import Any, Dict, Optional
+
+from socketio import AsyncServer
+import structlog
+from socketio.asyncio_client import AsyncClient
+
+from rasa.model_manager.runner_service import BotSession
+from rasa.model_manager.studio_jwt_auth import (
+    UserToServiceAuthenticationError,
+    authenticate_user_to_service,
+)
+
+structlogger = structlog.get_logger()
+
+
+# A simple in-memory store for active chat connections to studio frontend
+socket_proxy_clients = {}
+
+
+async def socketio_websocket_traffic_wrapper(
+    sio: AsyncServer,
+    running_bots: Dict[str, BotSession],
+    sid: str,
+    auth: Optional[Dict],
+) -> bool:
+    """Wrapper for bridging the user chat websocket and the bot server."""
+    auth_token = auth.get("token") if auth else None
+
+    if auth_token is None:
+        structlogger.error("model_runner.user_no_token", sid=sid)
+        return False
+
+    try:
+        authenticate_user_to_service(auth_token)
+        structlogger.debug("model_runner.user_authenticated_successfully", sid=sid)
+    except UserToServiceAuthenticationError as error:
+        structlogger.error(
+            "model_runner.user_authentication_failed", sid=sid, error=str(error)
+        )
+        return False
+
+    deployment_id = auth.get("deployment_id") if auth else None
+
+    if deployment_id is None:
+        structlogger.error("model_runner.bot_no_deployment_id", sid=sid)
+        return False
+
+    bot = running_bots.get(deployment_id)
+    if bot is None:
+        structlogger.error("model_runner.bot_not_found", deployment_id=deployment_id)
+        return False
+
+    if not bot.is_alive():
+        structlogger.error("model_runner.bot_not_alive", deployment_id=deployment_id)
+        return False
+
+    client = await create_bridge_client(sio, bot.internal_url, sid, deployment_id)
+
+    if client.sid is not None:
+        structlogger.debug(
+            "model_runner.bot_connection_established", deployment_id=deployment_id
+        )
+        socket_proxy_clients[sid] = client
+        return True
+    else:
+        structlogger.error(
+            "model_runner.bot_connection_failed", deployment_id=deployment_id
+        )
+        return False
+
+
+def create_bridge_server(sio: AsyncServer, running_bots: Dict[str, BotSession]) -> None:
+    """Create handlers for the socket server side.
+
+    Forwards messages coming from the user to the bot.
+    """
+
+    @sio.on("connect")
+    async def socketio_websocket_traffic(
+        sid: str, environ: Dict, auth: Optional[Dict]
+    ) -> bool:
+        """Bridge websockets between user chat socket and bot server."""
+        return await socketio_websocket_traffic_wrapper(sio, running_bots, sid, auth)
+
+    @sio.on("disconnect")
+    async def disconnect(sid: str) -> None:
+        structlogger.debug("model_runner.bot_disconnect", sid=sid)
+        if sid in socket_proxy_clients:
+            await socket_proxy_clients[sid].disconnect()
+            del socket_proxy_clients[sid]
+
+    @sio.on("*")
+    async def handle_message(event: str, sid: str, data: Dict[str, Any]) -> None:
+        # bridge both, incoming messages to the bot_url but also
+        # send the response back to the client. both need to happen
+        # in parallel in an async way
+
+        client = socket_proxy_clients.get(sid)
+        if client is None:
+            structlogger.error("model_runner.bot_not_connected", sid=sid)
+            return
+
+        await client.emit(event, data)
+
+
+async def create_bridge_client(
+    sio: AsyncServer, url: str, sid: str, deployment_id: str
+) -> AsyncClient:
+    """Create a new socket bridge client.
+
+    Forwards messages comming from the bot to the user.
+    """
+    client = AsyncClient()
+
+    await client.connect(url)
+
+    @client.event  # type: ignore[misc]
+    async def session_confirm(data: Dict[str, Any]) -> None:
+        structlogger.debug(
+            "model_runner.bot_session_confirmed", deployment_id=deployment_id
+        )
+        await sio.emit("session_confirm", room=sid)
+
+    @client.event  # type: ignore[misc]
+    async def bot_message(data: Dict[str, Any]) -> None:
+        structlogger.debug("model_runner.bot_message", deployment_id=deployment_id)
+        await sio.emit("bot_message", data, room=sid)
+
+    @client.event  # type: ignore[misc]
+    async def disconnect() -> None:
+        structlogger.debug(
+            "model_runner.bot_connection_closed", deployment_id=deployment_id
+        )
+        await sio.emit("disconnect", room=sid)
+
+    @client.event  # type: ignore[misc]
+    async def connect_error() -> None:
+        structlogger.error(
+            "model_runner.bot_connection_error", deployment_id=deployment_id
+        )
+        await sio.emit("disconnect", room=sid)
+
+    return client

rasa/model_manager/studio_jwt_auth.py

@@ -0,0 +1,86 @@
+import os
+from functools import cache
+from http import HTTPStatus
+from typing import Any, Dict, Optional
+
+import jwt
+import requests
+import structlog
+from rasa.shared.exceptions import RasaException
+
+structlogger = structlog.get_logger()
+
+
+AUTH_URL = os.getenv("KEYCLOAK_URL", "http://localhost:8081/auth")
+
+AUTH_REALM = os.getenv("KEYCLOAK_REALM", "rasa-studio")
+
+
+class UserToServiceAuthenticationError(RasaException):
+    """Raised when the user authentication fails."""
+
+    def __init__(self, message: str) -> None:
+        self.message = message
+
+    def __str__(self) -> str:
+        return f"{self.__class__.__name__}: {self.message}"
+
+
+@cache
+def get_public_key_from_keycloak() -> Optional[str]:
+    """Fetch the public key from the keycloak server."""
+    realm_url = f"{AUTH_URL}/realms/{AUTH_REALM}"
+
+    try:
+        response = requests.get(realm_url)
+    except requests.RequestException as error:
+        structlogger.error("model_api.auth.keycloak_request_failed", error=str(error))
+        return None
+
+    if response.status_code != HTTPStatus.OK:
+        structlogger.error(
+            "model_api.auth.keycloak_public_key_fetch_failed",
+            status_code=response.status_code,
+            response=response.text,
+        )
+        return None
+
+    public_key = response.json().get("public_key")
+
+    if public_key is None:
+        structlogger.error(
+            "model_runner.keycloak_public_key_not_found",
+            response=response.text,
+        )
+        return None
+
+    public_key = f"-----BEGIN PUBLIC KEY-----\n{public_key}\n-----END PUBLIC KEY-----"
+    return public_key
+
+
+def authenticate_user_to_service(token: str) -> Dict[str, Any]:
+    """Authenticate the user to the model service."""
+    if not token:
+        structlogger.debug("model_api.auth.no_token_provided")
+        raise UserToServiceAuthenticationError("No token provided.")
+
+    public_key = get_public_key_from_keycloak()
+
+    if public_key is None:
+        raise UserToServiceAuthenticationError(
+            "Failed to fetch public key from keycloak."
+        )
+
+    try:
+        return jwt.decode(
+            token,
+            public_key,
+            algorithms=["RS256", "HS256", "HS512", "ES256"],
+            audience="account",
+        )
+    except jwt.InvalidKeyError as error:
+        structlogger.info("model_api.auth.invalid_jwt_key", error=str(error))
+        raise UserToServiceAuthenticationError("Invalid JWT key.")
+    except jwt.InvalidTokenError as error:
+        structlogger.info("model_api.auth.invalid_jwt_token", error=str(error))
+        raise UserToServiceAuthenticationError("Invalid JWT token.") from error