rasa-pro 3.11.0__py3-none-any.whl → 3.11.0a1__py3-none-any.whl

rasa/model_manager/runner_service.py

@@ -1,286 +0,0 @@
-import os
-import shutil
-from typing import Dict, Optional
-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
-from rasa.model_manager.warm_rasa_process import start_rasa_process
-
-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
-    log_id: str
-    returncode: Optional[int] = None
-
-    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:
-        return False
-
-
-def set_bot_status_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
-    bot.returncode = bot.process.returncode
-
-
-def set_bot_status_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_encoded_config_data_to_files(
-    encoded_configs: Dict[str, bytes], 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_encoded_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 fetch_size_of_remote_model(model_name: str, storage_type: str) -> int:
-    """Fetch the size of the model from remote storage."""
-    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
-
-    return persistor.size_of_persisted_model(model_name=model_name)
-
-
-def start_bot_process(
-    deployment_id: str, bot_base_path: str, base_url_path: str
-) -> BotSession:
-    port = get_open_port()
-
-    arguments = [
-        "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,
-        arguments=" ".join(arguments),
-    )
-
-    warm_process = start_rasa_process(cwd=bot_base_path, arguments=arguments)
-
-    internal_bot_url = f"http://localhost:{port}"
-
-    structlogger.info(
-        "model_runner.bot.starting",
-        deployment_id=deployment_id,
-        log=logs_path(warm_process.log_id),
-        url=internal_bot_url,
-        port=port,
-        pid=warm_process.process.pid,
-    )
-
-    return BotSession(
-        deployment_id=deployment_id,
-        status=BotSessionStatus.QUEUED,
-        process=warm_process.process,
-        url=f"{base_url_path}?deployment_id={deployment_id}",
-        internal_url=internal_bot_url,
-        port=port,
-        log_id=warm_process.log_id,
-    )
-
-
-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."""
-    with structlog.contextvars.bound_contextvars(model_name=model_name):
-        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():
-        set_bot_status_to_stopped(bot)
-    elif await bot.completed_startup_recently():
-        set_bot_status_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
-        bot.returncode = bot.process.returncode
-    except ProcessLookupError:
-        structlogger.debug(
-            "model_runner.stop_bot.process_not_found",
-            deployment_id=bot.deployment_id,
-        )

rasa/model_manager/socket_bridge.py

@@ -1,146 +0,0 @@
-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:
-        """Disconnect the bot connection."""
-        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 messages between user and bot.
-
-        Both incoming user messages to the bot_url and
-        bot responses sent back to the client 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

@@ -1,86 +0,0 @@
-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