Refactor using janus queue

Decouple conversation tree from the slack bot
Reorganize data structures
Add comments

Author: Peter-Sh

pyproject.toml

@@ -10,7 +10,7 @@ authors = [
     {name = "Redis Team", email = "[email protected]"},
 ]
 readme = "README.md"
-requires-python = ">=3.8"
+requires-python = ">=3.9"
 classifiers = [
     "Development Status :: 3 - Alpha",
     "Intended Audience :: Developers",
@@ -35,6 +35,7 @@ dependencies = [
     "openai>=1.0.0",
     "PyJWT>=2.8.0",
     "cryptography>=41.0.0",
+    "janus>=2.0,<2.2",
 ]
 
 [project.optional-dependencies]

src/redis_release/bht/conversation_behaviours.py

@@ -9,7 +9,7 @@
 from slack_sdk import WebClient
 
 from ..config import Config
-from ..conversation_models import Command
+from ..conversation_models import Command, ConversationCockpit
 from ..models import ReleaseArgs, ReleaseType
 from .behaviours import ReleaseAction
 from .conversation_state import ConversationState
@@ -54,7 +54,11 @@ class CommandDetectionResult(BaseModel):
 
 class SimpleCommandClassifier(ReleaseAction):
     def __init__(
-        self, name: str, state: ConversationState, log_prefix: str = ""
+        self,
+        name: str,
+        state: ConversationState,
+        cockpit: ConversationCockpit,
+        log_prefix: str = "",
     ) -> None:
         self.state = state
         super().__init__(name=name, log_prefix=log_prefix)
@@ -85,12 +89,12 @@ class LLMCommandClassifier(ReleaseAction):
     def __init__(
         self,
         name: str,
-        llm: OpenAI,
         state: ConversationState,
+        cockpit: ConversationCockpit,
         log_prefix: str = "",
         confidence_threshold: float = 0.7,
     ) -> None:
-        self.llm = llm
+        self.llm = cockpit.llm
         self.state = state
         self.confidence_threshold = confidence_threshold
         super().__init__(name=name, log_prefix=log_prefix)
@@ -136,6 +140,7 @@ def update(self) -> Status:
             )
 
         try:
+            assert self.llm is not None
             # Call LLM with structured outputs
             response = self.llm.responses.parse(
                 model="gpt-4o-2024-08-06",
@@ -253,6 +258,7 @@ def __init__(
         self,
         name: str,
         state: ConversationState,
+        cockpit: ConversationCockpit,
         config: Config,
         log_prefix: str = "",
     ) -> None:
@@ -274,6 +280,18 @@ def update(self) -> Status:
         # Get release args
         release_args = self.state.release_args
 
+        # Check authorization
+        if (
+            self.state.authorized_users
+            and self.state.message
+            and self.state.message.user not in self.state.authorized_users
+        ):
+            logger.warning(
+                f"Unauthorized attempt by user {self.state.message.user}. Authorized users: {self.state.authorized_users}"
+            )
+            self.state.reply = "Sorry, you are not authorized to run releases. Please contact an administrator."
+            return Status.FAILURE
+
         self.logger.info(
             f"Starting release for tag {release_args.release_tag} in background thread"
         )
@@ -340,7 +358,11 @@ def run_release_in_thread() -> None:
 
 class IsLLMAvailable(ReleaseAction):
     def __init__(
-        self, name: str, state: ConversationState, log_prefix: str = ""
+        self,
+        name: str,
+        state: ConversationState,
+        cockpit: ConversationCockpit,
+        log_prefix: str = "",
     ) -> None:
         self.state = state
         super().__init__(name=name, log_prefix=log_prefix)
@@ -353,7 +375,11 @@ def update(self) -> Status:
 
 class HasReleaseArgs(ReleaseAction):
     def __init__(
-        self, name: str, state: ConversationState, log_prefix: str = ""
+        self,
+        name: str,
+        state: ConversationState,
+        cockpit: ConversationCockpit,
+        log_prefix: str = "",
     ) -> None:
         self.state = state
         super().__init__(name=name, log_prefix=log_prefix)
@@ -366,7 +392,11 @@ def update(self) -> Status:
 
 class IsCommandStarted(ReleaseAction):
     def __init__(
-        self, name: str, state: ConversationState, log_prefix: str = ""
+        self,
+        name: str,
+        state: ConversationState,
+        cockpit: ConversationCockpit,
+        log_prefix: str = "",
     ) -> None:
         self.state = state
         super().__init__(name=name, log_prefix=log_prefix)

src/redis_release/bht/conversation_state.py

@@ -2,15 +2,10 @@
 
 from pydantic import BaseModel, Field
 
-from ..conversation_models import Command
+from ..conversation_models import Command, InboxMessage
 from ..models import ReleaseArgs, SlackArgs, SlackFormat
 
 
-class InboxMessage(BaseModel):
-    message: str
-    context: List[str]
-
-
 class ConversationState(BaseModel):
     llm_available: bool = False
     message: Optional[InboxMessage] = None
@@ -20,3 +15,4 @@ class ConversationState(BaseModel):
     reply: Optional[str] = None
 
     slack_args: Optional[SlackArgs] = None
+    authorized_users: Optional[List[str]] = None

src/redis_release/bht/conversation_tree.py

@@ -1,12 +1,14 @@
-from typing import Optional, Tuple
+import logging
+from typing import List, Optional, Tuple
 
+from janus import SyncQueue
 from openai import OpenAI
 from py_trees.behaviour import Behaviour
 from py_trees.trees import BehaviourTree
 from py_trees.visitors import SnapshotVisitor
 
 from ..config import Config, load_config
-from ..conversation_models import ConversationArgs
+from ..conversation_models import ConversationArgs, ConversationCockpit, InboxMessage
 from ..models import SlackArgs
 from .backchain import create_PPA, latch_chain_to_chain
 from .conversation_behaviours import (
@@ -16,41 +18,45 @@
     RunCommand,
     SimpleCommandClassifier,
 )
-from .conversation_state import ConversationState, InboxMessage
+from .conversation_state import ConversationState
 from .tree import log_tree_state_with_markup
 
+logger = logging.getLogger(__name__)
+
 
 def create_conversation_root_node(
     input: InboxMessage,
     config: Config,
-    llm: Optional[OpenAI] = None,
+    cockpit: ConversationCockpit,
     slack_args: Optional[SlackArgs] = None,
+    authorized_users: Optional[List[str]] = None,
 ) -> Tuple[Behaviour, ConversationState]:
     state = ConversationState(
-        llm_available=llm is not None,
+        llm_available=cockpit.llm is not None,
         message=input,
         slack_args=slack_args,
+        authorized_users=authorized_users,
     )
     state.message = input
 
     # Use LLM classifier if available, otherwise use simple classifier
-    if llm is not None:
+    if cockpit.llm is not None:
         command_detector = create_PPA(
             "LLM Command Detector",
-            LLMCommandClassifier("LLM Command Detector", llm, state),
-            HasReleaseArgs("Has Release Args", state),
+            LLMCommandClassifier("LLM Command Detector", state, cockpit),
+            HasReleaseArgs("Has Release Args", state, cockpit),
         )
     else:
         command_detector = create_PPA(
             "Simple Command Detector",
-            SimpleCommandClassifier("Simple Command Classifier", state),
-            HasReleaseArgs("Has Release Args", state),
+            SimpleCommandClassifier("Simple Command Classifier", state, cockpit),
+            HasReleaseArgs("Has Release Args", state, cockpit),
         )
 
     run_command = create_PPA(
         "Run",
-        RunCommand("Run Command", state, config),
-        IsCommandStarted("Is Command Started", state),
+        RunCommand("Run Command", state, cockpit, config),
+        IsCommandStarted("Is Command Started", state, cockpit),
     )
 
     latch_chain_to_chain(run_command, command_detector)
@@ -61,6 +67,7 @@ def create_conversation_root_node(
 
 def initialize_conversation_tree(
     args: ConversationArgs,
+    reply_queue: Optional[SyncQueue] = None,
 ) -> Tuple[BehaviourTree, ConversationState]:
 
     # Load config
@@ -70,11 +77,19 @@ def initialize_conversation_tree(
     if args.openai_api_key:
         llm = OpenAI(api_key=args.openai_api_key)
 
+    cockpit = ConversationCockpit()
+    cockpit.llm = llm
+    cockpit.reply_queue = reply_queue
+
+    if not args.inbox:
+        raise ValueError("Inbox message is required")
+
     root, state = create_conversation_root_node(
-        InboxMessage(message=args.message, context=args.context or []),
+        args.inbox,
         config=config,
-        llm=llm,
+        cockpit=cockpit,
         slack_args=args.slack_args,
+        authorized_users=args.authorized_users,
     )
     tree = BehaviourTree(root)
     snapshot_visitor = SnapshotVisitor()
@@ -83,8 +98,24 @@ def initialize_conversation_tree(
     return tree, state
 
 
-def run_conversation_tree(tree: BehaviourTree) -> None:
+def run_conversation_tree(
+    tree: BehaviourTree, state: ConversationState, reply_queue: SyncQueue
+) -> None:
     """Abstacting away tree run
     Currently it's just a single tick, but it may change in future
     """
-    tree.tick()
+    try:
+        tree.tick()
+        try:
+            if state.reply:
+                reply_queue.put(state.reply)
+        except Exception as e:
+            logger.error(f"Error putting reply to queue: {e}", exc_info=True)
+    except Exception as e:
+        try:
+            reply_queue.put(f"Error running conversation tree: {str(e)}")
+        except Exception as e:
+            logger.error(f"Error putting error reply to queue: {e}", exc_info=True)
+    finally:
+        logger.debug("Shutting down reply queue")
+        reply_queue.shutdown(immediate=False)

src/redis_release/cli.py

@@ -18,7 +18,7 @@
 )
 from .bht.tree import TreeInspector, async_tick_tock, initialize_tree_and_state
 from .config import load_config
-from .conversation_models import ConversationArgs
+from .conversation_models import ConversationArgs, InboxMessage
 from .github_app_auth import GitHubAppAuth, load_private_key_from_file
 from .github_client_async import GitHubClientAsync
 from .logging_config import setup_logging
@@ -140,7 +140,9 @@ def release_print(
 def conversation_print() -> None:
     setup_logging()
     tree, state = initialize_conversation_tree(
-        ConversationArgs(message="test", openai_api_key="dummy")
+        ConversationArgs(
+            inbox=InboxMessage(message="test", context=[]), openai_api_key="dummy"
+        )
     )
     render_dot_tree(tree.root)
     print(unicode_tree(tree.root))
@@ -168,9 +170,11 @@ def conversation(
         openai_api_key = os.getenv("OPENAI_API_KEY")
 
     args = ConversationArgs(
-        message=message, openai_api_key=openai_api_key, config_path=config
+        inbox=InboxMessage(message=message, context=[]),
+        openai_api_key=openai_api_key,
+        config_path=config,
     )
-    tree = initialize_conversation_tree(args)
+    tree, _ = initialize_conversation_tree(args)
     tree.tick()
     print(unicode_tree(tree.root))
 

src/redis_release/conversation_models.py

@@ -1,6 +1,8 @@
 from enum import Enum
 from typing import List, Optional
 
+from janus import SyncQueue
+from openai import OpenAI
 from pydantic import BaseModel
 
 from .models import SlackArgs
@@ -15,10 +17,20 @@ class Command(str, Enum):
     HELP = "help"
 
 
-class ConversationArgs(BaseModel):
-    openai_api_key: Optional[str] = None
+class InboxMessage(BaseModel):
     message: str
-    context: Optional[List[str]] = None
-    config_path: Optional[str] = None
+    context: List[str]
+    user: Optional[str] = None
 
+
+class ConversationArgs(BaseModel):
+    inbox: Optional[InboxMessage]
+    config_path: Optional[str] = None
     slack_args: Optional[SlackArgs] = None
+    openai_api_key: Optional[str] = None
+    authorized_users: Optional[List[str]] = None
+
+
+class ConversationCockpit:
+    llm: Optional[OpenAI] = None
+    reply_queue: Optional[SyncQueue] = None

src/redis_release/github_client_async.py

@@ -242,7 +242,8 @@ async def trigger_workflow(
         logger.debug(f"[blue]Triggering workflow[/blue] {workflow_file} in {repo}")
         logger.debug(f"Inputs: {inputs}")
         logger.debug(f"Ref: {ref}")
-        logger.debug(f"Workflow UUID: [cyan]{inputs['workflow_uuid']}[/cyan]")
+        if "workflow_uuid" in inputs:
+            logger.debug(f"Workflow UUID: [cyan]{inputs['workflow_uuid']}[/cyan]")
 
         url = f"https://api.github.com/repos/{repo}/actions/workflows/{workflow_file}/dispatches"
         headers = {