From d21c0291ec2833f6cb56ab5694c8dfdd6000df99 Mon Sep 17 00:00:00 2001 From: jjoneson Date: Wed, 29 Mar 2023 19:13:08 -0600 Subject: [PATCH] Mostly add slack integration. Slack tool exists, but is not necessary for current usage. --- README.md | 3 + agent/toolkits/base.py | 5 +- agent/toolkits/k8s_explorer/prompt.py | 2 - agent/toolkits/prompt.py | 2 + agent/toolkits/toolkit.py | 8 +- listeners/agent_listener.py | 8 ++ listeners/slack_listener.py | 125 ++++++++++++++++++++++++++ load_env.sh | 2 + main.py | 23 +++-- requirements.txt | 1 + tools/k8s_explorer/tool.py | 3 +- tools/slack_integration/__init__.py | 1 + tools/slack_integration/tool.py | 46 ++++++++++ 13 files changed, 217 insertions(+), 12 deletions(-) create mode 100644 listeners/agent_listener.py create mode 100644 listeners/slack_listener.py create mode 100644 load_env.sh create mode 100644 tools/slack_integration/__init__.py create mode 100644 tools/slack_integration/tool.py diff --git a/README.md b/README.md index 87fdb96..10219b6 100644 --- a/README.md +++ b/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 diff --git a/agent/toolkits/base.py b/agent/toolkits/base.py index 1200bf7..8a8b7df 100644 --- a/agent/toolkits/base.py +++ b/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) \ No newline at end of file + 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,) + \ No newline at end of file diff --git a/agent/toolkits/k8s_explorer/prompt.py b/agent/toolkits/k8s_explorer/prompt.py index 533fbe5..9bf153e 100644 --- a/agent/toolkits/k8s_explorer/prompt.py +++ b/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. diff --git a/agent/toolkits/prompt.py b/agent/toolkits/prompt.py index bdc23f4..b3d6ac9 100644 --- a/agent/toolkits/prompt.py +++ b/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. """ diff --git a/agent/toolkits/toolkit.py b/agent/toolkits/toolkit.py index da23d93..9d37919 100644 --- a/agent/toolkits/toolkit.py +++ b/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) diff --git a/listeners/agent_listener.py b/listeners/agent_listener.py new file mode 100644 index 0000000..a20e800 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/listeners/slack_listener.py b/listeners/slack_listener.py new file mode 100644 index 0000000..96ac5a3 --- /dev/null +++ b/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 + diff --git a/load_env.sh b/load_env.sh new file mode 100644 index 0000000..2c2c17a --- /dev/null +++ b/load_env.sh @@ -0,0 +1,2 @@ +#!/usr/bin/env bash +export $(cat .env | xargs) \ No newline at end of file diff --git a/main.py b/main.py index fa91edc..e8a811e 100644 --- a/main.py +++ b/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") diff --git a/requirements.txt b/requirements.txt index 2681b90..874707c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,3 +4,4 @@ google-api-python-client google-api-client pygit2 python-gitlab +slack_sdk diff --git a/tools/k8s_explorer/tool.py b/tools/k8s_explorer/tool.py index 4c17b3e..291ee45 100644 --- a/tools/k8s_explorer/tool.py +++ b/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 diff --git a/tools/slack_integration/__init__.py b/tools/slack_integration/__init__.py new file mode 100644 index 0000000..fb90f6a --- /dev/null +++ b/tools/slack_integration/__init__.py @@ -0,0 +1 @@ +"""Tools for integrating with Slack""" \ No newline at end of file diff --git a/tools/slack_integration/tool.py b/tools/slack_integration/tool.py new file mode 100644 index 0000000..130f9f1 --- /dev/null +++ b/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) + \ No newline at end of file