proscenium 0.0.1__py3-none-any.whl → 0.0.3__py3-none-any.whl

proscenium/__init__.py

@@ -0,0 +1,3 @@
+import logging
+
+logging.getLogger(__name__).addHandler(logging.NullHandler())

proscenium/admin/__init__.py

@@ -0,0 +1,37 @@
+from typing import Generator
+from typing import List
+from typing import Optional
+
+import logging
+
+from proscenium.core import Prop
+from proscenium.core import Character
+from rich.console import Console
+
+logging.getLogger(__name__).addHandler(logging.NullHandler())
+
+log = logging.getLogger(__name__)
+
+system_message = """
+You are an administrator of a chatbot.
+"""
+
+
+def props(console: Optional[Console]) -> List[Prop]:
+
+    return []
+
+
+class Admin(Character):
+
+    def __init__(self, admin_channel_id: str):
+        super().__init__(admin_channel_id)
+        self.admin_channel_id = admin_channel_id
+
+    def handle(
+        channel_id: str,
+        speaker_id: str,
+        question: str,
+    ) -> Generator[tuple[str, str], None, None]:
+
+        yield channel_id, "I am the administrator of this chat system."

proscenium/bin/bot.py

@@ -0,0 +1,142 @@
+#!/usr/bin/env python3
+
+import os
+import sys
+import logging
+import typer
+import importlib
+from rich.console import Console
+
+from proscenium.admin import Admin
+
+from proscenium.interfaces.slack import (
+    get_slack_auth,
+    channel_table,
+    bot_user_id,
+    places_table,
+    channel_maps,
+    make_slack_listener,
+    connect,
+    send_curtain_up,
+    listen,
+    send_curtain_down,
+    shutdown,
+)
+
+from proscenium.verbs.display import header
+
+logging.basicConfig(
+    stream=sys.stdout,
+    format="%(asctime)s  %(levelname)-8s %(name)s: %(message)s",
+    level=logging.WARNING,
+)
+
+logging.basicConfig(
+    stream=sys.stdout,
+    format="%(asctime)s  %(levelname)-8s %(name)s: %(message)s",
+    level=logging.WARNING,
+)
+
+app = typer.Typer(help="Proscenium Bot")
+
+log = logging.getLogger(__name__)
+
+
+@app.command(help="""Start the Proscenium Bot.""")
+def start(
+    verbose: bool = False,
+    production_module_name: str = typer.Option(
+        "demo.production",
+        "-p",
+        "--production",
+        help="The name of the python module in PYTHONPATH in which the variable production of type proscenium.core.Production is defined.",
+    ),
+    force_rebuild: bool = False,
+):
+
+    console = Console()
+    sub_console = None
+
+    if verbose:
+        log.setLevel(logging.INFO)
+        logging.getLogger("proscenium").setLevel(logging.INFO)
+        logging.getLogger("demo").setLevel(logging.INFO)
+        sub_console = console
+
+    console.print(header())
+
+    production_module = importlib.import_module(production_module_name, package=None)
+
+    slack_admin_channel_id = os.environ.get("SLACK_ADMIN_CHANNEL_ID")
+    # Note that the checking of the existence of the admin channel id is delayed
+    # until after the subscribed channels are shown.
+
+    production = production_module.make_production(slack_admin_channel_id, sub_console)
+
+    console.print("Preparing props...")
+    production.prepare_props()
+    console.print("Props are up-to-date.")
+
+    slack_app_token, slack_bot_token = get_slack_auth()
+
+    socket_mode_client = connect(slack_app_token, slack_bot_token)
+
+    user_id = bot_user_id(socket_mode_client, console)
+    console.print()
+
+    channels_by_id, channel_name_to_id = channel_maps(socket_mode_client)
+    console.print(channel_table(channels_by_id))
+    console.print()
+
+    if slack_admin_channel_id is None:
+        raise ValueError(
+            "SLACK_ADMIN_CHANNEL_ID environment variable not set. "
+            "Please set it to the channel ID of the Proscenium admin channel."
+        )
+    if slack_admin_channel_id not in channels_by_id:
+        raise ValueError(
+            f"Admin channel {slack_admin_channel_id} not found in subscribed channels."
+        )
+
+    admin = Admin(slack_admin_channel_id)
+    log.info("Admin handler started.")
+
+    log.info("Places, please!")
+    channel_id_to_character = production.places(channel_name_to_id)
+    channel_id_to_character[slack_admin_channel_id] = admin
+
+    console.print(places_table(channel_id_to_character, channels_by_id))
+    console.print()
+
+    slack_listener = make_slack_listener(
+        user_id,
+        slack_admin_channel_id,
+        channels_by_id,
+        channel_id_to_character,
+        console,
+    )
+
+    send_curtain_up(socket_mode_client, production, slack_admin_channel_id)
+
+    console.print("Starting the show. Listening for events...")
+    listen(
+        socket_mode_client,
+        slack_listener,
+        user_id,
+        console,
+    )
+
+    send_curtain_down(socket_mode_client, slack_admin_channel_id)
+
+    shutdown(
+        socket_mode_client,
+        slack_listener,
+        user_id,
+        production,
+        console,
+    )
+
+
+if __name__ == "__main__":
+
+    app()

proscenium/core/__init__.py

@@ -0,0 +1,152 @@
+from typing import Generator
+from typing import Optional
+import logging
+from rich.console import Console
+
+logging.getLogger(__name__).addHandler(logging.NullHandler())
+
+log = logging.getLogger(__name__)
+
+
+class Prop:
+    """
+    A `Prop` is a resource available to the `Character`s in a `Scene`.
+    """
+
+    def __init__(
+        self,
+        console: Optional[Console] = None,
+    ):
+        self.console = console
+
+    def name(self) -> str:
+        return self.__class__.__name__
+
+    def description(self) -> str:
+        return self.__doc__ or ""
+
+    def curtain_up_message(self) -> str:
+        return f"- {self.name()}, {self.description().strip()}"
+
+    def already_built(self) -> bool:
+        return False
+
+    def build(self) -> None:
+        pass
+
+
+class Character:
+    """
+    A `Character` is a participant in a `Scene` that `handle`s utterances from the
+    scene by producing its own utterances."""
+
+    def __init__(self, admin_channel_id: str):
+        self.admin_channel_id = admin_channel_id
+
+    def name(self) -> str:
+        return self.__class__.__name__
+
+    def description(self) -> str:
+        return self.__doc__ or ""
+
+    def curtain_up_message(self) -> str:
+        return f"- {self.name()}, {self.description().strip()}"
+
+    def handle(
+        channel_id: str, speaker_id: str, utterance: str
+    ) -> Generator[tuple[str, str], None, None]:
+        pass
+
+
+class Scene:
+    """
+    A `Scene` is a setting in which `Character`s interact with each other and
+    with `Prop`s. It is a container for `Character`s and `Prop`s.
+    """
+
+    def __init__(self):
+        pass
+
+    def name(self) -> str:
+        return self.__class__.__name__
+
+    def description(self) -> str:
+        return self.__doc__ or ""
+
+    def curtain_up_message(self) -> str:
+
+        characters_msg = "\n".join(
+            [character.curtain_up_message() for character in self.characters()]
+        )
+
+        props_msg = "\n".join([prop.curtain_up_message() for prop in self.props()])
+
+        return f"""
+Scene: {self.name()}, {self.description().strip()}
+
+Characters:
+{characters_msg}
+
+Props:
+{props_msg}
+"""
+
+    def props(self) -> list[Prop]:
+        return []
+
+    def prepare_props(self, force_rebuild: bool = False) -> None:
+        for prop in self.props():
+            if force_rebuild:
+                prop.build()
+            elif not prop.already_built():
+                log.info("Prop %s not built. Building it now.", prop.name())
+                prop.build()
+
+    def characters(self) -> list[Character]:
+        return []
+
+    def places(self) -> dict[str, Character]:
+        pass
+
+    def curtain(self) -> None:
+        pass
+
+
+class Production:
+    """
+    A `Production` is a collection of `Scene`s."""
+
+    def __init__(self):
+        pass
+
+    def name(self) -> str:
+        return self.__class__.__name__
+
+    def description(self) -> str:
+        return self.__doc__ or ""
+
+    def prepare_props(self, force_rebuild: bool = False) -> None:
+        if force_rebuild:
+            log.info("Forcing rebuild of all props.")
+        else:
+            log.info("Building any missing props...")
+
+        for scene in self.scenes():
+            scene.prepare_props(force_rebuild=force_rebuild)
+
+    def curtain_up_message(self) -> str:
+
+        scenes_msg = "\n\n".join(
+            [scene.curtain_up_message() for scene in self.scenes()]
+        )
+
+        return f"""Production: {self.name()}, {self.description().strip()}
+
+{scenes_msg}"""
+
+    def scenes(self) -> list[Scene]:
+        return []
+
+    def curtain(self) -> None:
+        for scene in self.scenes():
+            scene.curtain()

proscenium/interfaces/__init__.py

@@ -0,0 +1,3 @@
+import logging
+
+logging.getLogger(__name__).addHandler(logging.NullHandler())

proscenium/interfaces/slack.py

@@ -0,0 +1,265 @@
+from typing import Callable
+
+from typing import Generator
+import time
+import logging
+import os
+from rich.console import Console
+from rich.table import Table
+
+from slack_sdk.web import WebClient
+from slack_sdk.socket_mode import SocketModeClient
+from slack_sdk.socket_mode.request import SocketModeRequest
+from slack_sdk.socket_mode.response import SocketModeResponse
+from slack_sdk.socket_mode.listeners import SocketModeRequestListener
+
+from proscenium.core import Production
+from proscenium.core import Character
+
+log = logging.getLogger(__name__)
+
+
+def get_slack_auth() -> tuple[str, str]:
+
+    slack_app_token = os.environ.get("SLACK_APP_TOKEN")
+    if slack_app_token is None:
+        raise ValueError(
+            "SLACK_APP_TOKEN environment variable not set. "
+            "Please set it to the app token of the Proscenium Slack app."
+        )
+    slack_bot_token = os.environ.get("SLACK_BOT_TOKEN")
+    if slack_bot_token is None:
+        raise ValueError(
+            "SLACK_BOT_TOKEN environment variable not set. "
+            "Please set it to the bot token of the Proscenium Slack app."
+        )
+
+    return slack_app_token, slack_bot_token
+
+
+def connect(app_token: str, bot_token: str) -> SocketModeClient:
+
+    web_client = WebClient(token=bot_token)
+    socket_mode_client = SocketModeClient(app_token=app_token, web_client=web_client)
+
+    socket_mode_client.connect()
+    log.info("Connected to Slack.")
+
+    return socket_mode_client
+
+
+def make_slack_listener(
+    proscenium_user_id: str,
+    admin_channel_id: str,
+    channels_by_id: dict,
+    channel_id_to_handler: dict[
+        str, Callable[[str, str, str], Generator[tuple[str, str], None, None]]
+    ],
+    console: Console,
+):
+
+    def process(client: SocketModeClient, req: SocketModeRequest):
+
+        if req.type == "events_api":
+
+            event = req.payload["event"]
+
+            response = SocketModeResponse(envelope_id=req.envelope_id)
+            client.send_socket_mode_response(response)
+
+            if event.get("type") in [
+                "message",
+                "app_mention",
+            ]:
+                speaker_id = event.get("user")
+                if speaker_id == proscenium_user_id:
+                    return
+
+                text = event.get("text")
+                channel_id = event.get("channel")
+                console.print(f"{speaker_id} in {channel_id} said something")
+
+                channel = channels_by_id.get(channel_id, None)
+
+                if channel is None:
+
+                    # TODO: channels_by_id will get stale
+                    log.info("No handler for channel id %s", channel_id)
+
+                else:
+
+                    character = channel_id_to_handler[channel_id]
+                    log.info("Handler defined for channel id %s", channel_id)
+
+                    # TODO determine whether the handler has a good chance of being useful
+
+                    for receiving_channel_id, response in character.handle(
+                        channel_id, speaker_id, text
+                    ):
+                        response_response = client.web_client.chat_postMessage(
+                            channel=receiving_channel_id, text=response
+                        )
+                        log.info(
+                            "Response sent to channel %s",
+                            receiving_channel_id,
+                        )
+                        if receiving_channel_id == admin_channel_id:
+                            continue
+
+                        permalink = client.web_client.chat_getPermalink(
+                            channel=receiving_channel_id,
+                            message_ts=response_response["ts"],
+                        )["permalink"]
+                        log.info(
+                            "Response sent to channel %s link %s",
+                            receiving_channel_id,
+                            permalink,
+                        )
+                        client.web_client.chat_postMessage(
+                            channel=admin_channel_id,
+                            text=permalink,
+                        )
+
+        elif req.type == "interactive":
+            pass
+        elif req.type == "slash_commands":
+            pass
+        elif req.type == "app_home_opened":
+            pass
+        elif req.type == "block_actions":
+            pass
+        elif req.type == "message_actions":
+            pass
+
+    return process
+
+
+def channel_maps(
+    socket_mode_client: SocketModeClient,
+) -> tuple[dict[str, dict], dict[str, str]]:
+
+    subscribed_channels = socket_mode_client.web_client.users_conversations(
+        types="public_channel,private_channel,mpim,im",
+        limit=100,
+    )
+    log.info(
+        "Subscribed channels count: %s",
+        len(subscribed_channels["channels"]),
+    )
+
+    channels_by_id = {
+        channel["id"]: channel for channel in subscribed_channels["channels"]
+    }
+
+    channel_name_to_id = {
+        channel["name"]: channel["id"]
+        for channel in channels_by_id.values()
+        if channel.get("name")
+    }
+
+    return channels_by_id, channel_name_to_id
+
+
+def channel_table(channels_by_id) -> Table:
+    channel_table = Table(title="Subscribed channels")
+    channel_table.add_column("Channel ID", justify="left")
+    channel_table.add_column("Name", justify="left")
+    for channel_id, channel in channels_by_id.items():
+        channel_table.add_row(
+            channel_id,
+            channel.get("name", "-"),
+        )
+    return channel_table
+
+
+def bot_user_id(socket_mode_client: SocketModeClient, console: Console):
+
+    auth_response = socket_mode_client.web_client.auth_test()
+
+    console.print(auth_response["url"])
+    console.print()
+    console.print(f"Team '{auth_response["team"]}' ({auth_response["team_id"]})")
+    console.print(f"User '{auth_response["user"]}' ({auth_response["user_id"]})")
+
+    user_id = auth_response["user_id"]
+    console.print("Bot id", auth_response["bot_id"])
+
+    return user_id
+
+
+def places_table(
+    channel_id_to_character: dict[str, Character], channels_by_id: dict[str, dict]
+) -> Table:
+
+    table = Table(title="Characters in place")
+    table.add_column("Channel ID", justify="left")
+    table.add_column("Channel Name", justify="left")
+    table.add_column("Character", justify="left")
+    for channel_id, character in channel_id_to_character.items():
+        channel = channels_by_id[channel_id]
+        table.add_row(channel_id, channel["name"], character.name())
+
+    return table
+
+
+def send_curtain_up(
+    socket_mode_client: SocketModeClient,
+    production: Production,
+    slack_admin_channel_id: str,
+) -> None:
+
+    curtain_up_message = f"""
+Proscenium 🎭 https://the-ai-alliance.github.io/proscenium/
+
+```
+{production.curtain_up_message()}
+```
+
+Curtain up.
+"""
+
+    socket_mode_client.web_client.chat_postMessage(
+        channel=slack_admin_channel_id,
+        text=curtain_up_message,
+    )
+
+
+def listen(
+    socket_mode_client: SocketModeClient,
+    slack_listener: SocketModeRequestListener,
+    user_id: str,
+    console: Console,
+):
+    socket_mode_client.socket_mode_request_listeners.append(slack_listener)
+
+    try:
+        while True:
+            time.sleep(1)
+    except KeyboardInterrupt:
+        console.print("Exiting...")
+
+
+def send_curtain_down(
+    socket_mode_client: SocketModeClient, slack_admin_channel_id: str
+) -> None:
+    socket_mode_client.web_client.chat_postMessage(
+        channel=slack_admin_channel_id,
+        text="""Curtain down. We hope you enjoyed the show!""",
+    )
+
+
+def shutdown(
+    socket_mode_client: SocketModeClient,
+    slack_listener: SocketModeRequestListener,
+    user_id: str,
+    production: Production,
+    console: Console,
+):
+
+    socket_mode_client.socket_mode_request_listeners.remove(slack_listener)
+    socket_mode_client.disconnect()
+    console.print("Disconnected from Slack.")
+
+    production.curtain()
+
+    console.print("Handlers stopped.")

proscenium/patterns/__init__.py

@@ -0,0 +1,3 @@
+import logging
+
+logging.getLogger(__name__).addHandler(logging.NullHandler())

proscenium/patterns/chunk_space.py

@@ -0,0 +1,51 @@
+from typing import Optional
+import logging
+from rich.console import Console
+from pymilvus import model
+
+from proscenium.verbs.read import load_file
+from proscenium.verbs.chunk import documents_to_chunks_by_characters
+from proscenium.verbs.display.milvus import collection_panel
+from proscenium.verbs.vector_database import vector_db
+from proscenium.verbs.vector_database import create_collection
+from proscenium.verbs.vector_database import add_chunks_to_vector_db
+
+log = logging.getLogger(__name__)
+
+
+def load_chunks_from_files(
+    data_files: list[str],
+    milvus_uri: str,
+    embedding_fn: model.dense.SentenceTransformerEmbeddingFunction,
+    collection_name: str,
+    console: Optional[Console] = None,
+) -> None:
+
+    vector_db_client = vector_db(milvus_uri)
+    log.info("Vector db stored at %s", milvus_uri)
+
+    for data_file in data_files:
+
+        log.info(
+            "Loading data file %s into vector db %s collection %s",
+            data_file,
+            milvus_uri,
+            collection_name,
+        )
+        create_collection(vector_db_client, embedding_fn, collection_name)
+
+        documents = load_file(data_file)
+        chunks = documents_to_chunks_by_characters(documents)
+        log.info("Data file %s has %s chunks", data_file, len(chunks))
+
+        info = add_chunks_to_vector_db(
+            vector_db_client,
+            embedding_fn,
+            chunks,
+            collection_name,
+        )
+        log.info("%s chunks inserted ", info["insert_count"])
+        if console is not None:
+            console.print(collection_panel(vector_db_client, collection_name))
+
+    vector_db_client.close()