Mostly add slack integration. Slack tool exists, but is not necessary…

… for current usage.

README.md

@@ -17,6 +17,9 @@
 | GOOGLE_PROJECT_ID | Project ID for GCP |
 | GOOGLE_REGION | Region for GCP |
 | GOOGLE_CLUSTER_ID | Cluster Name for GCP |
+| SLACK_APP_TOKEN | Slack App Token for Slack API |
+| SLACK_BOT_TOKEN | Slack Bot Token for Slack API |
+| SLACK_CHANNEL_ID | Slack Channel ID for Slack API |
 
 This assumes that you are running locally and are able to authenticate by running:
 ```bash

agent/toolkits/base.py

@@ -37,5 +37,6 @@ def create_k8s_engineer_agent(
                          )
     tool_names = [tool.name for tool in tools]
     agent = ZeroShotAgent(llm_chain=llm_chain,
-                          allowed_tools=tool_names, **kwargs)
-    return AgentExecutor.from_agent_and_tools(agent=agent, tools=toolkit.get_tools(), verbose=verbose)
+                          allowed_tools=tool_names, return_intermediate_steps=True, **kwargs)
+    return AgentExecutor.from_agent_and_tools(agent=agent, tools=toolkit.get_tools(), verbose=verbose, return_intermediate_steps=True,)
+

agent/toolkits/k8s_explorer/prompt.py

@@ -10,8 +10,6 @@
 Second, determine which Namespace to act in. Some resources are not namespaced, so you may not need to do this.
 
 Third, determine the operation you are performing from the list of available operations.
-    If the operation is `list`, the final answer should be a list of the names of the objects.
-    If the operation is `logs`, the final answer should be the logs for the pod.
 
 Fourth, determine the name of the object you are interacting with. This is different than resource type.
 

agent/toolkits/prompt.py

@@ -6,6 +6,8 @@
 You may be asked to retrieve information from the cluster, or to interact with Infrastructure as Code (IaC) source code repositories.
 If the question does not seem related to Kubernetes, return I don't know. Do not make up an answer.
 
+If you are asked for logs, output the full text of the logs in your Final Answer.
+
 Be sure to always add an Action Input.  If no input makes sense, use None.
 """
 

agent/toolkits/toolkit.py

@@ -20,6 +20,7 @@
 from tools.git_integrator.tool import GitModel
 from tools.gitlab_integration.tool import GitlabModel
 from tools.k8s_explorer.tool import KubernetesOpsModel
+from tools.slack_integration.tool import SlackModel, SlackSendMessageTool
 
 
 class K8sEngineerToolkit(BaseToolkit):
@@ -28,6 +29,7 @@ class K8sEngineerToolkit(BaseToolkit):
     git_agent: AgentExecutor
     gitlab_agent: AgentExecutor
     k8s_explorer_agent: AgentExecutor
+    slack_tool: BaseTool
 
     def get_tools(self) -> List[BaseTool]:
         """Get the tools in the toolkit."""
@@ -46,7 +48,7 @@ def get_tools(self) -> List[BaseTool]:
             func=self.gitlab_agent.run,
             description=GITLAB_AGENT_DESCRIPTION,
         )
-        return [git_agent_tool, k8s_explorer_agent_tool, gitlab_agent_tool]
+        return [git_agent_tool, k8s_explorer_agent_tool, gitlab_agent_tool, self.slack_tool]
 
     @classmethod
     def from_llm(
@@ -55,6 +57,7 @@ def from_llm(
         k8s_model: KubernetesOpsModel,
         git_model: GitModel,
         gitlab_model: GitlabModel,
+        slack_model: SlackModel,
         verbose: bool = False,
         **kwargs: Any,
     ) -> K8sEngineerToolkit:
@@ -65,4 +68,5 @@ def from_llm(
             llm=llm, toolkit=K8sExplorerToolkit(model=k8s_model), verbose=verbose, **kwargs)
         gitlab_agent = create_git_integration_toolkit(
             llm=llm, toolkit=GitlabIntegrationToolkit(model=gitlab_model), verbose=verbose, **kwargs)
-        return cls(git_agent=git_agent, k8s_explorer_agent=k8s_explorer_agent, gitlab_agent=gitlab_agent, **kwargs)
+        slack_tool = SlackSendMessageTool(model = slack_model)
+        return cls(git_agent=git_agent, k8s_explorer_agent=k8s_explorer_agent, gitlab_agent=gitlab_agent, slack_tool=slack_tool, **kwargs)

listeners/agent_listener.py

@@ -0,0 +1,8 @@
+from langchain.agents.agent import AgentExecutor
+
+class AgentListener:
+    def __init__(self, agent: AgentExecutor):
+        self.agent = agent
+
+    def on_message(self, message: str) -> str:
+        return self.agent.run(message)

listeners/slack_listener.py

@@ -0,0 +1,125 @@
+import os
+import threading
+
+from slack_sdk.web import WebClient
+from slack_sdk.socket_mode import SocketModeClient
+from slack_sdk.socket_mode.response import SocketModeResponse
+from slack_sdk.socket_mode.request import SocketModeRequest
+from listeners.agent_listener import AgentListener
+from langchain.agents.agent import AgentExecutor
+
+
+class SlackListener(AgentListener):
+    def __init__(self, agent: AgentExecutor):
+        self.agent = agent
+        # Initialize SocketModeClient with an app-level token + WebClient
+        self.client = SocketModeClient(
+            # This app-level token will be used only for establishing a connection
+            app_token=os.environ.get("SLACK_APP_TOKEN"),  # xapp-A111-222-xyz
+            # You will be using this WebClient for performing Web API calls in listeners
+            web_client=WebClient(token=os.environ.get("SLACK_BOT_TOKEN"))  # xoxb-111-222-xyz
+
+        )
+
+    def on_message(self, message) -> str:
+        res = self.agent({"input": message})
+        return parse_intermediate_steps_into_slack_message(res["intermediate_steps"])
+
+
+    def listen(self, interrupt: threading.Event):
+        self.client.socket_mode_request_listeners.append(self.process)
+        self.client.connect()
+        interrupt.wait()
+
+    def start(self) -> threading.Event:
+    # run listen in a separate thread
+        interrupt = threading.Event()
+        thread = threading.Thread(target=self.listen, args=(interrupt,))
+        thread.start()
+        return interrupt
+
+    def process(self, client: SocketModeClient, req: SocketModeRequest):
+        if req.type == "events_api":
+            # Acknowledge the request anyway
+            response = SocketModeResponse(envelope_id=req.envelope_id)
+            client.send_socket_mode_response(response)
+
+            # Add a reaction to the message if it's a new message
+            if req.payload["event"]["type"] == "message" or req.payload["event"]["type"] == "app_mention"\
+                and req.payload["event"].get("subtype") is None:
+                # client.web_client.reactions_add(
+                #     name="eyes",
+                #     channel=req.payload["event"]["channel"],
+                #     timestamp=req.payload["event"]["ts"],
+                # )
+                # reply to the message with an acknowledgement
+
+                client.web_client.chat_postMessage(
+                    channel=req.payload["event"]["channel"],
+                    thread_ts=req.payload["event"]["ts"],
+                    text="Let me think about that...")
+
+                message = req.payload["event"]["text"]
+
+                # remove the bot mention
+                if req.payload["event"]["type"] == "app_mention":
+                    message = message.split(" ", 1)[1]
+
+                reply = self.on_message(message)
+                # the reply has newlines as literal \n, so we need to replace them with actual newlines
+                reply = reply.replace("\\n", "\n")
+                client.web_client.chat_postMessage(
+                    channel=req.payload["event"]["channel"],
+                    thread_ts=req.payload["event"]["ts"],
+                    text=reply)
+        if req.type == "interactive" \
+            and req.payload.get("type") == "shortcut":
+            if req.payload["callback_id"] == "hello-shortcut":
+                # Acknowledge the request
+                response = SocketModeResponse(envelope_id=req.envelope_id)
+                client.send_socket_mode_response(response)
+                # Open a welcome modal
+                client.web_client.views_open(
+                    trigger_id=req.payload["trigger_id"],
+                    view={
+                        "type": "modal",
+                        "callback_id": "hello-modal",
+                        "title": {
+                            "type": "plain_text",
+                            "text": "Greetings!"
+                        },
+                        "submit": {
+                            "type": "plain_text",
+                            "text": "Good Bye"
+                        },
+                        "blocks": [
+                            {
+                                "type": "section",
+                                "text": {
+                                    "type": "mrkdwn",
+                                    "text": "Hello!"
+                                }
+                            }
+                        ]
+                    }
+                )
+
+        if req.type == "interactive" \
+            and req.payload.get("type") == "view_submission":
+            if req.payload["view"]["callback_id"] == "hello-modal":
+                # Acknowledge the request and close the modal
+                response = SocketModeResponse(envelope_id=req.envelope_id)
+                client.send_socket_mode_response(response)
+
+
+
+def parse_intermediate_steps_into_slack_message(steps) -> str:
+    """steps is a NamedTuple with fields that looks like this:
+    [(AgentAction(tool='Search', tool_input='Leo DiCaprio girlfriend', log=' I should look up who Leo DiCaprio is dating\nAction: Search\nAction Input: "Leo DiCaprio girlfriend"'), 'Camila Morrone'), (AgentAction(tool='Search', tool_input='Camila Morrone age', log=' I should look up how old Camila Morrone is\nAction: Search\nAction Input: "Camila Morrone age"'), '25 years'), (AgentAction(tool='Calculator', tool_input='25^0.43', log=' I should calculate what 25 years raised to the 0.43 power is\nAction: Calculator\nAction Input: 25^0.43'), 'Answer: 3.991298452658078\n')]
+    """
+    message = ""
+    for step in steps:
+        message += f"{step[0].log}\n"
+        message += f"Answer: {step[1]}\n"
+    return message
+

load_env.sh

@@ -0,0 +1,2 @@
+#!/usr/bin/env bash
+export $(cat .env | xargs)

main.py

@@ -10,15 +10,18 @@
 import googleapiclient.discovery
 from tempfile import NamedTemporaryFile
 from kubernetes import  client
+from slack_sdk import WebClient
 from agent.toolkits.base import create_k8s_engineer_agent
 from agent.toolkits.git_integrator.toolkit import GitIntegratorToolkit
 from agent.toolkits.k8s_explorer.base import create_k8s_explorer_agent
 from agent.toolkits.k8s_explorer.toolkit import K8sExplorerToolkit
 from agent.toolkits.toolkit import K8sEngineerToolkit
+from listeners.slack_listener import SlackListener
 from tools.git_integrator.tool import GitModel
 from tools.gitlab_integration.tool import GitlabModel
 
 from tools.k8s_explorer.tool import KubernetesOpsModel
+from tools.slack_integration.tool import SlackModel
 
 llm = OpenAI(temperature=0, model_name="gpt-3.5-turbo", max_tokens=1024)
 
@@ -69,13 +72,23 @@ def get_cluster():
 gl = gitlab.Gitlab(url=gitlab_url, private_token=gitlab_private_token)
 gitlab_model = GitlabModel(gl=gl)
 
-k8s_engineer_toolkit = K8sEngineerToolkit.from_llm(llm=llm, k8s_model=k8s_model, git_model=git_model, gitlab_model=gitlab_model, verbose=True)
+slack_token = os.environ["SLACK_BOT_TOKEN"]
+slack_client = WebClient(token=slack_token)
+slack_channel = os.environ["SLACK_CHANNEL_ID"]
+slack_model = SlackModel(client=slack_client, channel=slack_channel)
+
+
+k8s_engineer_toolkit = K8sEngineerToolkit.from_llm(llm=llm, k8s_model=k8s_model, git_model=git_model, gitlab_model=gitlab_model, slack_model=slack_model,  verbose=True)
 k8s_engineer_agent = create_k8s_engineer_agent(llm=llm, toolkit=k8s_engineer_toolkit, verbose=True)
 
+slack_listener = SlackListener(k8s_engineer_agent)
+interrupt = slack_listener.start()
+# block until keyboard interrupt
+interrupt.wait()
 
 
-# k8s_agent.run("list all namespaces")
-# k8s_agent.run("list all services in the test-bed namespace")
-# k8s_agent.run("get the gitlab runner deployment in the gitlab-runner namespace")
-k8s_engineer_agent.run("get the logs for review-3 in test-bed")
+# # k8s_agent.run("list all namespaces")
+# # k8s_agent.run("list all services in the test-bed namespace")
+# # k8s_agent.run("get the gitlab runner deployment in the gitlab-runner namespace")
+# k8s_engineer_agent.run("get the logs for review-3 in test-bed")
 

requirements.txt

@@ -4,3 +4,4 @@ google-api-python-client
 google-api-client
 pygit2
 python-gitlab
+slack_sdk

tools/k8s_explorer/tool.py

@@ -417,11 +417,12 @@ class KubernetesGetPodLogsTool(BaseTool):
     """Tool for getting the logs of a pod."""
     name = "k8s_get_pod_logs"
     description = """
+    Can be used to get the logs of a pod.
     You should call the k8s_get_pod_name_like tool first to get the name of the pod.
     You should know the namespace and pod name before calling this tool.
     Executes a get in the specified namespace for the specified pod, with the specified name.
     Input should be a string containing the namespace and pod name, separated by commas.
-    Returns a yaml string containing the spec.
+    Returns the logs of the pod.
     """
     model: KubernetesOpsModel
 

tools/slack_integration/__init__.py

@@ -0,0 +1 @@
+"""Tools for integrating with Slack"""

tools/slack_integration/tool.py

@@ -0,0 +1,46 @@
+import os
+from slack_sdk import WebClient
+from langchain.tools.base import BaseTool
+from pydantic import BaseModel
+
+class SlackModel(BaseModel):
+    """Model for slack integration."""
+    client: WebClient
+    channel: str
+
+    class Config:
+        arbitrary_types_allowed = True
+
+    @classmethod
+    def from_client(cls, client: WebClient, channel: str):
+        """Create a model from a client."""
+        return cls(client=client, channel=channel)
+
+    def send_message(self, message: str) -> str:
+        """Send a message to slack."""
+        try: 
+            self.client.chat_postMessage(channel=self.channel, text=message)
+            return "Message sent"
+        except Exception as e:
+            return f"Error sending slack message: {e}"
+
+class SlackSendMessageTool(BaseTool):
+    """Tool for sending a message to slack."""
+    name = "slack_send_message"
+    description = """
+    Only use this tool if you are absolutely sure you want to send a message to slack.
+    Can be used to send a message in Slack.
+    Input should be a string.
+    """
+    model: SlackModel
+
+    def _run(self, tool_input: str) -> str:
+        """Send a message to slack."""
+        # the input has newlines as literal \n, so we need to replace them with actual newlines
+        tool_input = tool_input.replace("\\n", "\n")
+        return self.model.send_message(tool_input)
+
+    async def _arun(self, tool_input: str) -> str:
+        """Send a message to slack."""
+        return self.run(tool_input)
+