├── .gitignore ├── media ├── tui.png ├── tui_demo.gif └── talking_to_user_and_retrying_nmap_scan.png ├── .env.example ├── requirements.txt ├── tools ├── __init__.py ├── user_io.py ├── searchtool.py ├── README.md ├── shodantool.py ├── webtool.py ├── shelltool.py └── file_io.py ├── prompt.py ├── README.md ├── main.py ├── tui.py └── task_manager.py /.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__/* 2 | __pycache__/ 3 | __pycache__ 4 | 5 | .env 6 | *.json -------------------------------------------------------------------------------- /media/tui.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AgeOfMarcus/1337GPT/HEAD/media/tui.png -------------------------------------------------------------------------------- /media/tui_demo.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AgeOfMarcus/1337GPT/HEAD/media/tui_demo.gif -------------------------------------------------------------------------------- /.env.example: -------------------------------------------------------------------------------- 1 | OPENAI_API_KEY= 2 | # the keys below are only needed for the tools Shodan, and WebReader 3 | SHODAN= 4 | EXTRACTOR= -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | langchain 2 | openai 3 | python-dotenv 4 | requests 5 | googlesearch.py 6 | duckduckgo-search 7 | beautifulsoup4 8 | rich -------------------------------------------------------------------------------- /media/talking_to_user_and_retrying_nmap_scan.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AgeOfMarcus/1337GPT/HEAD/media/talking_to_user_and_retrying_nmap_scan.png -------------------------------------------------------------------------------- /tools/__init__.py: -------------------------------------------------------------------------------- 1 | from .webtool import ScrapeTool, WebReadTool 2 | from .searchtool import GoogleSearchTool, DDGSearchTool 3 | from .user_io import TalkToUser 4 | from .file_io import WriteFileTool, ReadFileTool, ListDirTool 5 | from .shelltool import ShellTool 6 | from .shodantool import ShodanTool 7 | 8 | TOOLS = { 9 | 'WebScraper': ScrapeTool, 10 | 'WebReader': WebReadTool, 11 | 'GoogleSearch': GoogleSearchTool, 12 | 'DDGSearch': DDGSearchTool, 13 | 'AskUser': TalkToUser, 14 | 'ReadFile': ReadFileTool, 15 | 'WriteFile': WriteFileTool, 16 | 'ListDir': ListDirTool, 17 | 'Files': [ReadFileTool, WriteFileTool, ListDirTool], 18 | 'Shell': ShellTool, 19 | 'Shodan': ShodanTool, 20 | } -------------------------------------------------------------------------------- /tools/user_io.py: -------------------------------------------------------------------------------- 1 | from langchain.tools import BaseTool 2 | 3 | class TalkToUser(BaseTool): 4 | name = 'TalkToUser' 5 | description = ( 6 | 'Useful for asking the user for input.' 7 | 'Use this when you need to ask the user for more information that you cannot find on your own.' 8 | 'Do not ask the user something that you could find out yourself.' 9 | 'Accepts a single argument type string, which will be displayed to the user.' 10 | 'Returns a string response from the user.' 11 | ) 12 | 13 | def _run(self, message: str) -> str: 14 | print(f'[ai: question]: {message}') 15 | return input('[user response]: ') 16 | async def _arun(self, message): 17 | return self._run(message) -------------------------------------------------------------------------------- /tools/searchtool.py: -------------------------------------------------------------------------------- 1 | from langchain.tools import BaseTool 2 | from googlesearch_py import search 3 | from duckduckgo_search import ddg 4 | 5 | class GoogleSearchTool(BaseTool): 6 | name = "GoogleSearch" 7 | description = "searches google. Useful for when you need to get information about current events. The input to this tool should be a search query" 8 | 9 | def _run(self, query: str) -> list: 10 | return list(search(query)) 11 | #return [str(res) for res in search(query)] 12 | 13 | async def _arun(self, query: str) -> list: 14 | return self._run(query) 15 | 16 | class DDGSearchTool(BaseTool): 17 | name = 'DuckDuckGo' 18 | description = ( 19 | "Searches DuckDuckGo." 20 | "Useful for when you need to get information about current events." 21 | "The input to this tool should be a search query." 22 | "Returns a list of dicts containing a url, title, and description." 23 | ) 24 | 25 | def _run(self, query: str) -> list: 26 | return ddg(query) 27 | 28 | async def _arun(self, query: str) -> list: 29 | return self._run(query) -------------------------------------------------------------------------------- /tools/README.md: -------------------------------------------------------------------------------- 1 | # About 2 | 3 | These are tools I have written to extend the functionality of `1337GPT`. Here is a basic example of how to build a [langchain tool](https://python.langchain.com/en/latest/modules/agents/tools.html): 4 | 5 | ```python 6 | from langchain.tools import BaseTool 7 | 8 | class MyTool(BaseTool): 9 | name = 'ToolName' 10 | description = ( 11 | "Description of tool" 12 | "Useful for ..." 13 | "Use this when ..." 14 | "Accepts ... as input, returns ..." 15 | ) 16 | 17 | # gets called when ran synchronously 18 | def _run(self, argument): 19 | return 'thing' 20 | 21 | # gets called when ran asynchronously 22 | async def _arun(self, argument): 23 | # if you don't want to implement, do the following 24 | return self._run(argument) 25 | ``` 26 | 27 | # Tools with requirements 28 | 29 | * **GoogleSearch** requires `googlesearch.py` 30 | * **DDGSearch** requires `duckduckgo-search` 31 | * **Shodan** requires `shodan` 32 | 33 | # Tools requiring an API key 34 | 35 | * **Shodan** - If you are using a **free API key**, it helps to specify that in your goal for `1337GPT` -------------------------------------------------------------------------------- /tools/shodantool.py: -------------------------------------------------------------------------------- 1 | from langchain.tools import BaseTool 2 | from langchain.tools.base import Field 3 | import os 4 | try: 5 | from shodan import Shodan 6 | except ImportError: 7 | if os.getenv('REPL_SLUG'): 8 | # replit workaround 9 | os.system('pip install shodan') 10 | from shodan import Shodan 11 | else: 12 | exit('Please install the python shodan library using `pip install shodan`.') 13 | 14 | class ShodanTool(BaseTool): 15 | name = 'Shodan-io' 16 | description = ( 17 | "Useful for searching the shodan.io API for servers." 18 | "Use this for information gathering on a target." 19 | "Accepts a single string argument containing the full search query." 20 | "Returns a list of matches in dict format with the keys 'data', 'hostname', 'ip', 'port', 'os', etc." 21 | ) 22 | shodan_api_key: str = Field(default_factory=lambda: os.getenv('SHODAN')) 23 | 24 | def _run(self, query: str) -> list: 25 | shodan = Shodan(self.shodan_api_key) 26 | try: 27 | res = shodan.search(query) 28 | except Exception as e: 29 | return f'Error: {str(e)}' 30 | return res['matches'] 31 | async def _arun(self, query): 32 | return self._run(query) -------------------------------------------------------------------------------- /tools/webtool.py: -------------------------------------------------------------------------------- 1 | from langchain.tools import BaseTool 2 | from bs4 import BeautifulSoup 3 | import requests, os 4 | 5 | class ScrapeTool(BaseTool): 6 | name = "ScrapeTool" 7 | description = "scrapes the main text content of a website (does not include HTML). Useful for when you need to read the plaintext content of a website. The input to this tool should be a URL" 8 | 9 | def scrape_website(self, url: str) -> str: 10 | if not url.startswith('http'): 11 | url = f'https://{url}' 12 | r = requests.get(f'https://extractorapi.com/api/v1/extractor/?apikey={os.getenv("EXTRACTOR")}&url={url}') 13 | return r.json()['text'] 14 | 15 | def _run(self, url: str) -> str: 16 | return self.scrape_website(url) 17 | 18 | async def _arun(self, url) -> str: 19 | return self._run(url) 20 | 21 | class WebReadTool(BaseTool): 22 | name = 'WebsiteReader' 23 | description = ( 24 | "Reads the HTML content of a website excluding script and style tags." 25 | "Useful for reading the contents of a website URL." 26 | "The input to this tool should be a URL." 27 | ) 28 | 29 | def _run(self, url: str) -> str: 30 | headers = { 31 | 'User-Agent': 'Mozilla/5.0 (X11; Linux x86_64; rv:85.0) Gecko/20100101 Firefox/85.0' 32 | } 33 | r = requests.get(url, headers=headers) 34 | soup = BeautifulSoup(r.text, 'html.parser') 35 | 36 | for script in soup(['script', 'style']): 37 | script.extract() 38 | 39 | text = soup.get_text() 40 | lines = [line.strip() for line in text.splitlines()] 41 | chunks = [phrase.strip() for line in lines for phrase in line.split(' ')] 42 | text = '\n'.join(chunk for chunk in chunks if chunk) 43 | 44 | return text 45 | 46 | async def _arun(self, url: str) -> str: 47 | return self._run(url) -------------------------------------------------------------------------------- /tools/shelltool.py: -------------------------------------------------------------------------------- 1 | from langchain.tools import BaseTool 2 | from subprocess import Popen, PIPE 3 | from pydantic import Field 4 | import platform 5 | 6 | class ShellTool(BaseTool): 7 | name = 'RunCommand' 8 | description = ( 9 | 'Useful for running commands on the host system.' 10 | 'Use this for testing code, or to execute shell commands.' 11 | 'ONLY USE THIS WITH NON-BLOCKING COMMANDS. Always run commands with verbose mode when applicable.' 12 | f'Current host info: {str(platform.uname())}' 13 | 'Accepts a string as input which will be executed.' 14 | 'Returns a dict containing the keys "stdout" and "stderr".' 15 | ) 16 | confirm_before_exec: bool = Field(default=True) 17 | 18 | def _sh(self, cmd: str) -> dict: 19 | proc = Popen(cmd, shell=True, stdout=PIPE, stderr=PIPE) 20 | stdout, stderr = proc.communicate() 21 | return { 22 | 'stdout': stdout.decode(), 23 | 'stderr': stderr.decode() 24 | } 25 | 26 | def input(self, prompt: str) -> str: 27 | if str(type(self.verbose)) == "": # hack for tui 28 | return self.verbose(f'[red]{prompt}[/red]') 29 | else: 30 | return input(prompt) 31 | 32 | def _run(self, command: str) -> str: 33 | if self.confirm_before_exec: 34 | conf = self.input(f'[system] run the following command? `{command}`. [y]es/[N]o/[e]dit: ').lower() 35 | if conf.startswith('e'): 36 | new = self.input('enter new command: ') 37 | return self._sh(new) 38 | elif conf.startswith('y'): 39 | return self._sh(command) 40 | else: 41 | return { 42 | 'stdout':'', 43 | 'stderr':'User aborted the process before command was executed. Reformulate.' 44 | } 45 | async def _arun(self, command): 46 | return self._run(command) -------------------------------------------------------------------------------- /prompt.py: -------------------------------------------------------------------------------- 1 | AGENT_PREFIX = """Assistant is designed to be able to assist with a wide range of text, internet, and shell related tasks from answering simple questions to performing a full penetration test with multiple steps. Assistant is able to generate human-like text based on the input it receives, allowing it to engage in natural-sounding conversations and provide responses that are coherent and relevant to the topic at hand. 2 | Assistant is able to process and understand large amounts of text content. As a language model, Assistant can not directly search the web or run commands, but it has a list of tools to accomplish such tasks. When Assistant encounters information it does not fully understand, Assistant will determine an appropriate search query and use a tool to obtain relevant results. If a Google search for vulnerabilities doesn't return substantial results, Assistant will use a shell command to query a local exploitdb. When talking about targets (servers, IPs, etc), Assistant is very strict to the information it finds using tools, and never fabricates searches or command outputs. Assistant is able to use tools in a sequence, and is loyal to the tool observation outputs rather than faking the search content. 3 | Assistant is skilled at preforming static code analysis. When asked to analyse a github repository, Assistant will clone and traverse through files using shell tool and read file tool. 4 | Overall, Assistant is a powerful penetration testing assistant that can help with a wide range of tasks and execute advanced shell commands to exploit targets. 5 | TOOLS: 6 | ------ 7 | Assistant has access to the following tools:""" 8 | 9 | AGENT_FORMAT_INSTRUCTIONS = """To use a tool, please use the following format: 10 | ``` 11 | Thought: Do I need to use a tool? Yes 12 | Action: the action to take, should be one of [{tool_names}] 13 | Action Input: the input to the action 14 | Observation: the result of the action 15 | ``` 16 | When you have a response to say to the Human, or if you do not need to use a tool, you MUST use the format: 17 | ``` 18 | Thought: Do I need to use a tool? No 19 | {ai_prefix}: [your response here] 20 | ``` 21 | """ 22 | 23 | AGENT_SUFFIX = """You are very strict to the shell commands correctness and will never fake a command if it does not exist. 24 | Begin! 25 | Previous conversation history: 26 | {chat_history} 27 | New input: {input} 28 | Since Assistant is a text language model, Assistant must use tools to interact with computers rather than imagination. 29 | The thoughts and observations are only visible for Assistant, Assistant should remember to repeat important information in the final response for Human. 30 | Thought: Do I need to use a tool? {agent_scratchpad}""" 31 | -------------------------------------------------------------------------------- /tools/file_io.py: -------------------------------------------------------------------------------- 1 | from langchain.tools import BaseTool 2 | import json 3 | import os 4 | 5 | class WriteFileTool(BaseTool): 6 | name = 'WriteFileTool' 7 | description = ( 8 | 'Useful for writing text to files.' 9 | 'Use this when asked to save data to a file, or write a file.' 10 | 'Accepts a single argument in JSON format containing the keys "filename" and "content". The optional key "overwrite" can be set to True if you want to replace existing files, otherwise it is disabled by default.' 11 | 'Returns the filename if written successfully, or an error string.' 12 | ) 13 | 14 | def _run(self, *args, **kwargs): 15 | if args: 16 | if type(args[0]) == dict: 17 | args = args[0] 18 | else: 19 | try: 20 | args = json.loads(args[0]) 21 | except json.JSONDecodeError: 22 | return 'Error: invalid JSON argument. Make sure you are passing a single, string, argument containing VALID JSON with the keys "filename" containing a string, "content" containing a string, and the optional "overwrite" containing a boolean.' 23 | else: 24 | args = kwargs 25 | 26 | if not (filename := args.get('filename')): 27 | return 'Error: no "filename" specified.' 28 | if not (content := args.get('content')): 29 | return 'Error: no "content" provided.' 30 | 31 | if os.path.exists(filename) and not args.get('overwrite'): 32 | return 'Error: "overwrite" was not set to True, but the specified "filename" already exists.' 33 | 34 | with open(filename, 'w') as f: 35 | f.write(content) 36 | return filename 37 | 38 | async def _arun(self, args): 39 | return self._run(args) 40 | 41 | class ReadFileTool(BaseTool): 42 | name = 'ReadFileTool' 43 | description = ( 44 | 'Useful for reading file contents.' 45 | 'Use this to read code when debugging errors.' 46 | 'Accepts a filename as argument.' 47 | 'Returns the file contents as a string, or an error message beginning with "Error:".' 48 | ) 49 | 50 | def _run(self, filename: str) -> str: 51 | try: 52 | file = open(filename, 'r') 53 | except: 54 | return f'Error: cannot open file: {filename}' 55 | return file.read() 56 | async def _arun(self, filename): 57 | return self._run(filename) 58 | 59 | class ListDirTool(BaseTool): 60 | name = 'ListDirTool' 61 | description = ( 62 | 'Useful for listing directory contents.' 63 | 'Accepts a directory path string as argument (allows "." and such).' 64 | 'Returns a list of filenames, or an error message.' 65 | ) 66 | 67 | def _run(self, path: str) -> list: 68 | try: 69 | return os.listdir(path) 70 | except: 71 | return f'Error: cannot list directory: {path}' 72 | 73 | async def _arun(self, path): 74 | return self._run(path) -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # [1337GPT](https://marcus.hashnode.dev/1337gpt-yet-another-gpt-agent-for-penetration-testing) 2 | My attempt at making a GPT agent for pentesting 3 | 4 | This repository contains a copy of my [ai task manager](https://replit.com/@MarcusWeinberger/ai-task-manager) that has been modified to preform penetration tests. You can read a slightly rephrased version of this README [**on my hashnode blog**](https://Marcus.hashnode.dev/1337gpt-yet-another-gpt-agent-for-penetration-testing)! 5 | 6 | **Please note that due to the young age of libraries used, if you encounter an error, you may need to run `pip install --upgrade -r requirements.txt`.** If your error persists, please submit an issue. 7 | 8 | # **New Terminal UI** 9 | 10 | [![media/tui_demo.gif](media/tui.png)](media/tui_demo.gif) 11 | 12 | To use the new terminal user interface, run the program with the `--tui` argument! 13 | 14 | [*Click here to read my blog post about how I created this abomination of code*](https://marcus.hashnode.dev/creating-a-basic-split-panel-console-interface-in-python-with-io) 15 | 16 | ### Media of 1337GPT in use 17 | 18 | ![media/talking_to_user_and_retrying_nmap_scan.png](media/talking_to_user_and_retrying_nmap_scan.png) 19 | 20 | ### Stargazer history 21 | 22 | [![Star History Chart](https://api.star-history.com/svg?repos=AgeOfMarcus/1337GPT&type=Date)](https://star-history.com/#AgeOfMarcus/1337GPT) 23 | 24 | # Cool features 25 | 26 | * edit tasks before running, or skip entirely 27 | * persist sessions with `json` files 28 | * this allows you to edit previous tasks' outputs and more 29 | * built-in tools (using [langchain](https://python.langchain.com)): 30 | * Read/Write/List local files 31 | * Scrape text from websites (using `extractor API` - **requires API key**) 32 | * Read website html body with bs4 33 | * Search Google and DuckDuckGo (**no API key needed**) 34 | * Execute shell commands (by default, requires confirmation) 35 | * Search the shodan.io API (**requires API key**) 36 | * Ask the User a question 37 | 38 | You can now select which tools your agent can use with the `--tools` argument. By default, `DDGSearch,Shell` is set (yes, they have readable names now). 39 | 40 | # Usage 41 | 42 | 1. Clone this repo 43 | * `git clone https://github.com/AgeOfMarcus/1337GPT` 44 | 2. Change directory into the repo 45 | * `cd 1337GPT` 46 | 3. Install requirements 47 | * `pip install -r requirements.txt` 48 | 4. Make a `.env` file (using the provided template) 49 | * `cp .env.example .env` 50 | * edit the file and add your keys 51 | 5. Run the `main.py` file (use `--help` to see full list of arguments) 52 | * example: `python main.py --tui --goal "preform a pentest of localhost. start with an nmap scan" --tools "Shell,Files" --persist localhost.json` 53 | 54 | By default, `1337GPT` uses **GPT-4** for best results. However if you don't have access to the API, you can use `--model gpt-3.5-turbo` instead. 55 | 56 | ## Using tools 57 | 58 | Some tools that I have written to assist `1337GPT` can be found in the `tools/` directory. For more info, refer to [tools/README.md](tools/README.md). 59 | 60 | #### Tool groups 61 | 62 | Some tools that work closely have been grouped together to make it easier to use. You can still use each tool seperately, but if you want, you can refer to the group by one name to make it easier. 63 | 64 | **Groups:** 65 | 66 | * `Files` - contains `ReadFile`, `WriteFile`, and `ListDir` 67 | 68 | ## Best practices 69 | 70 | These are some prompt alignment techniques that you should use in your goal for the best results. 71 | 72 | * **Give it somewhere to start.** To avoid unpredicted workflows, tell the Assistant how it should begin, e.g. *"start by running an nmap scan"*. 73 | * **Tell it what tools you have installed.** While capable of installing tools itself, the Assistant generally will not do so without additional prompting. By telling it more details about your host system, `1337GPT `will have a better understanding of current limitations. 74 | * **If it's ignoring a tool, tell it to use it by name.** I found that often times, the Assistant would choose to use the `TalkToUser` tool to ask me a question, when it could have easily found an answer using the `GoogleSearch` tool. While I tried to fix this in the tool descriptions, if something like this happens, don't be afraid to tell it the name of the tool it should use for something. -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | from dotenv import load_dotenv; load_dotenv() 2 | 3 | from langchain.agents import initialize_agent, AgentType 4 | from langchain.memory import ConversationBufferMemory 5 | from langchain.chat_models import ChatOpenAI 6 | from langchain.llms import OpenAI 7 | # task_manager.py 8 | from task_manager import TaskManager, convert_langchain_tools, EmptyCallbackHandler 9 | # prompt.py - recycled 10 | from prompt import AGENT_PREFIX, AGENT_FORMAT_INSTRUCTIONS, AGENT_SUFFIX 11 | # tools/ 12 | from tools import TOOLS 13 | 14 | # parse arguments 15 | from argparse import ArgumentParser 16 | parser = ArgumentParser() 17 | parser.add_argument('--goal', '-g', help='Goal for task manager to complete.', required=True) 18 | parser.add_argument('--tui', help='Use the Terminal User Interface. Default False.', action='store_true', default=False) 19 | parser.add_argument('--persist', '-p', help='File to persist data to. If not set, persist will be disabled.', default=False) 20 | parser.add_argument('--repeat', '-r', help='Allow repeat tasks. Default False.', action='store_true', default=False) 21 | parser.add_argument('--model', '-m', help='Model to use for chat. Default gpt-4.', default='gpt-4') 22 | parser.add_argument('--temperature', help='Temperature for chat model. Default 0.', default=0) 23 | parser.add_argument('--tools', '-t', help=f'Comma separated list of tools to use (from: {", ".join(TOOLS.keys())}) . Default: DDGSearch,Shell', default='DDGSearch,Shell') 24 | parser.add_argument('--tool-args', help="A dictionary containing kwargs that will be passed to tools as they are initialized. Default: {'Shell': {'confirm_before_exec': True}}", type=dict, default={'Shell': {'confirm_before_exec': True}}) 25 | parser.add_argument('--use-smart-combine', help='Uses smart combination when formatting info for agents. Default False. Use this if your prompts are getting too large.', action='store_true', default=False) 26 | parser.add_argument('--include-completed-tasks', help='Include completed tasks in the prompt for agents. Default True. Turn this off for less tokens used.', action='store_false', default=True) 27 | args = parser.parse_args() 28 | 29 | # chat model for agent 30 | llm = ChatOpenAI( 31 | temperature=args.temperature, 32 | model_name=args.model 33 | ) 34 | # memory for agent 35 | memory = ConversationBufferMemory( 36 | memory_key="chat_history", 37 | output_key='output', 38 | return_messages=True 39 | ) 40 | # tools i wrote for the agent, can be used with any langchain program, from ./tools/ 41 | tools = [] 42 | for tool in args.tools.split(','): 43 | obj = TOOLS[tool] 44 | if type(obj) == list: 45 | tools += list(map(lambda x: x(**args.tool_args.get(tool, {})), obj)) 46 | else: 47 | tools.append(obj(**args.tool_args.get(tool, {}))) 48 | 49 | # create an instance of TaskManager 50 | taskman = TaskManager( 51 | args.goal, # goal 52 | convert_langchain_tools(tools), # list of dicts of tools 53 | OpenAI(temperature=0), # llm for taskmanager 54 | persist=args.persist, # i want to persist data 55 | allow_repeat_tasks=args.repeat, # so it doesn't get stuck in a loop 56 | ) 57 | 58 | # now the agent for doing tasks 59 | agent = initialize_agent( 60 | tools, # list of BaseTool objects 61 | llm, # our chat model 62 | memory=memory, # the memory we created 63 | agent = AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION, # const 64 | agent_kwargs={ 65 | 'prefix': AGENT_PREFIX, 66 | 'format_instructions': AGENT_FORMAT_INSTRUCTIONS, 67 | 'suffix': AGENT_SUFFIX 68 | }, 69 | callbacks=[EmptyCallbackHandler()], # so that we can dynamically add with taskman.init_agent 70 | verbose=True # so we can see when it runs each tool 71 | ) 72 | # this lets us save tool outputs 73 | taskman.init_agent(agent) 74 | 75 | # lightweight example task loop 76 | def main(): 77 | while not taskman.goal_completed: 78 | if not taskman.current_tasks: 79 | if taskman.ensure_goal_complete(): # makes sure we're done 80 | break 81 | continue # if not complete, more tasks will be added 82 | task = taskman.current_tasks.pop(0) # get first task 83 | print('\n\nNext task:', task) 84 | cont = input('Continue? [Y]es/[n]o/[e]dit: ').lower() 85 | if cont.startswith('n'): 86 | continue # skips task 87 | elif cont.startswith('e'): 88 | task = input('New task: ') 89 | res = agent.run(taskman.format_task_str( 90 | task, 91 | smart_combine=args.use_smart_combine, # i don't recommend using this lol 92 | include_completed_tasks=args.include_completed_tasks # this one, maybe so 93 | )) # task result 94 | taskman.refine(task, res) # refine process for taskmanager 95 | 96 | if __name__ == '__main__': 97 | if args.tui: 98 | from tui import main as tui_main 99 | tui_main(taskman, agent, args) 100 | else: 101 | main() # f it we ball 102 | -------------------------------------------------------------------------------- /tui.py: -------------------------------------------------------------------------------- 1 | from langchain.callbacks.base import BaseCallbackHandler 2 | # rich 3 | from rich.live import Live 4 | from rich.table import Table, Row 5 | from rich.json import JSON 6 | import json, os, math 7 | # task_manager.py 8 | from task_manager import TaskManager 9 | 10 | try: 11 | if os.name == 'nt': 12 | # windows 13 | from msvcrt import getch as _getch 14 | getch = lambda: _getch().decode() 15 | else: 16 | # linux based 17 | from getch import getch 18 | except ImportError as e: 19 | lib = str(e).split("'")[1] 20 | exit(f'This program requires the {lib} library. Please install it with `pip install {lib}`') 21 | 22 | # by god I will turn this Table into a TUI 23 | class TUI(object): 24 | def __init__(self): 25 | self.max_lines = os.get_terminal_size().lines - 12 26 | self.total_cols = os.get_terminal_size().columns - 6 27 | 28 | # if total_cols is odd, make it even 29 | if self.total_cols % 2 == 1: 30 | self.total_cols -= 1 31 | 32 | self.half_cols = self.total_cols // 2 33 | 34 | self.table = Table(title='1337GPT') 35 | self.table.add_column('Data') 36 | self.table.add_column('Console') 37 | 38 | for i in range(self.max_lines): 39 | self.table.add_row(' ' * self.half_cols, ' ' * self.half_cols) 40 | 41 | self.console = Console(self) 42 | 43 | def pad(self, text): 44 | return text.ljust(self.half_cols) 45 | 46 | class Console(object): 47 | def __init__(self, tui: TUI): 48 | self.tui = tui 49 | self.history = [] 50 | self._input = '' 51 | 52 | def getlines(self): 53 | return self.tui.table.columns[1]._cells 54 | 55 | def checklines(self, line): 56 | chars = len(line) 57 | if chars > self.tui.half_cols: 58 | over = chars / self.tui.half_cols 59 | return math.ceil(over) 60 | return 1 61 | 62 | def add_line(self, line): 63 | # pad line 64 | line = self.tui.pad(line) 65 | _lines = self.checklines(line) 66 | 67 | lines = self.getlines() 68 | for i in range(_lines): 69 | self.history.append(lines.pop(0)) 70 | lines.append(line) 71 | self.tui.table.columns[1]._cells = lines 72 | 73 | def print(self, *args): 74 | self.add_line(' '.join(map(repr, args))) 75 | 76 | def update_last(self, text): 77 | lines = self.getlines() 78 | lines[-1] = self.tui.pad(text) 79 | if (_ := self.checklines(text)) > 1: 80 | for i in range(_ - 1): 81 | self.history.append(lines.pop(0)) 82 | self.tui.table.columns[1]._cells = lines 83 | 84 | def input(self, prompt): 85 | self.add_line(prompt) 86 | while True: 87 | ch = getch() 88 | if ch == '\x03': 89 | raise KeyboardInterrupt 90 | elif ch in ('\n', '\r'): 91 | inp = self._input 92 | self._input = '' 93 | return inp 94 | elif (ord(ch) == 127) or (ch in ('\b', '\x08')): 95 | self._input = self._input[:-1] 96 | self.update_last(prompt + self._input) 97 | else: 98 | self._input += ch 99 | self.update_last(prompt + self._input) 100 | 101 | def update_data_from_taskman(taskman: TaskManager, tui: TUI): 102 | data = JSON(json.dumps({ 103 | 'Goal': taskman.final_goal, 104 | 'Next Task': taskman.current_tasks[0], 105 | 'Stored Info': taskman.stored_info, 106 | 'Final Result': taskman.final_result 107 | })).text 108 | for i, line in enumerate(data.split('\n')): 109 | try: 110 | tui.table.columns[0]._cells[i] = line 111 | except IndexError: 112 | tui.table.columns[0]._cells.append(line) # not gonna show but oh well its there 113 | 114 | 115 | def main(taskman: TaskManager, agent, args): 116 | tui = TUI() 117 | taskman.output_func = tui.console.print 118 | taskman.input_func = tui.console.input 119 | 120 | class TUIHandler(BaseCallbackHandler): 121 | def on_tool_start(self, tool, args, **kwargs): 122 | tui.console.add_line(f'Starting tool [green]{tool["name"]}[/green] with args [green]{args}[/green]') 123 | 124 | def on_tool_end(self, out, **kwargs): 125 | tui.console.add_line(f'Tool finished with output [darkgreen]{out}[/darkgreen]') 126 | 127 | cb_handler = TUIHandler() 128 | 129 | taskman.verbose = False # not needed as left pane shows info 130 | agent.verbose = False # cant redirect these outputs 131 | 132 | # patch input func 133 | for tool in agent.tools: 134 | if tool.name == 'RunCommand': 135 | tool.verbose = tui.console.input # yes this is dirty, but i cant create my own attr 136 | 137 | with Live(tui.table, refresh_per_second=30): 138 | 139 | # main loop 140 | while not taskman.goal_completed: 141 | update_data_from_taskman(taskman, tui) 142 | if not taskman.current_tasks: 143 | if taskman.ensure_goal_complete(): 144 | break 145 | continue 146 | 147 | task = taskman.current_tasks.pop(0) 148 | tui.console.add_line(f'[grey]Next task: {task}[/grey]') 149 | cont = tui.console.input('[blue]Continue?[/blue] \[Y]es/\[n]o/\[e]dit: ')[0].lower() 150 | if cont == 'n': 151 | continue 152 | elif cont == 'e': 153 | task = tui.console.input('[cyan]Enter task[/cyan]: ') 154 | 155 | res = agent.run(taskman.format_task_str( 156 | task, 157 | smart_combine=args.use_smart_combine, 158 | include_completed_tasks=args.include_completed_tasks 159 | ), callbacks=[cb_handler]) 160 | taskman.refine(task, res) -------------------------------------------------------------------------------- /task_manager.py: -------------------------------------------------------------------------------- 1 | from langchain.callbacks.base import BaseCallbackHandler 2 | from langchain.llms.base import BaseLLM 3 | from langchain.tools import BaseTool 4 | from code import InteractiveConsole 5 | import json 6 | import os 7 | 8 | class EmptyCallbackHandler(BaseCallbackHandler): 9 | pass 10 | 11 | def save_to_file(goal: str, result: dict): 12 | """Saves the results dict (second argument) to a file ending in '.result.txt' with the name set to the goal (first argument)""" 13 | fn = goal.replace(' ','_') + '.result.txt' 14 | print(f'saving final result to {fn}') 15 | with open(fn, 'w') as f: 16 | json.dump(result, f) 17 | 18 | def convert_langchain_tools(tools: list[BaseTool]) -> list[dict]: 19 | """Converts a list of BaseTools (used in langchain) to a list of dictionaries containing the keys: 'name', and 'description'.""" 20 | return [{'name': tool.name, 'description': tool.description} for tool in tools if type(tool) == BaseTool] 21 | 22 | class TaskManager(object): 23 | """Task Manager""" 24 | current_tasks: list = [] 25 | final_goal: str 26 | goal_completed: bool = False 27 | tools: list 28 | verbose: bool = True 29 | llm: BaseLLM 30 | final_result: dict = {} 31 | stored_info: dict = {} 32 | persist: str = None 33 | completed_tasks: dict = {} 34 | BASE_PROMPT: str = """ 35 | You are a task management system. Your job is to create, reformulate, and refine a set of tasks. The tasks must be focused on achieving your final goal. It is very important you keep the final goal in mind as you think. Your goal is a constant, throughout, and will never change. 36 | As tasks are completed, update your stored info with any info you will need to output at the end. As you go, add on to your final result. Your final result will be returned once, either, you cannot come up with any more reasonable tasks and all are complete, or your final result satisfies your final goal. 37 | The language models assigned to your tasks will have access to a list of tools available. As language models, you cannot interact with the internet, however the following tools have been made available so that the final goal can be met. As the tasks you create will be given to other agents, make sure to be specific with each tasks instructions. 38 | 39 | Tools 40 | ----- 41 | {tools_str} 42 | ----- 43 | 44 | Final Goal 45 | ---------- 46 | {final_goal} 47 | ---------- 48 | 49 | Current values 50 | -------------- 51 | current_tasks: {current_tasks} 52 | stored_info: {stored_info} 53 | final_result: {final_result} 54 | -------------- 55 | """ 56 | ENSURE_COMPLETE_PROMPT: str = ''' 57 | Based on your current values, assess whether you have completed your final_goal. Respond with a dictionary in valid JSON format with the following keys: 58 | "final_result" - dict - reformat your final result to better meet your final goal, 59 | "goal_complete" - bool - True if you have completed your final_goal, otherwise False if you need to continue 60 | "current_tasks" - list - list of strings containing tasks in natural language which you will need to complete to meet your final_goal. leave this empty if you set "goal_complete" to True. 61 | 62 | You always give your responses in VALID, JSON READABLE FORMAT. 63 | ''' 64 | REFINE_PROMPT: str = ''' 65 | Task Result 66 | ----------- 67 | task: {task} 68 | result: {result} 69 | ----------- 70 | 71 | Refine your current set of tasks based on the task result above. E.g., if information has already been gathered that satisfies the requests in a task, it is not needed anymore. However, if information gathered shows a new task needs to be added, include it. 72 | If the result included any info you may need to complete later tasks, add it to your stored_info. 73 | If the result included any info you may need to satisfy your final goal, add it to the final result. Format it as necessary, but make sure it includes all information needed. 74 | You always give your response in valid JSON format so that it can be parsed (in python) with `json.loads`. Return a dictionary with the keys: "current_tasks" a list of strings (your complete set of tasks, if you need to, add any new tasks and reorder as you see fit), "final_result" a dict (your final result to satisfy your final goal, add to this as you go), "stored_info" a dict (info you may need for later tasks), if you have any thoughts to output to the user, include them as a string with the key "thoughts", and lastly, the key "goal_complete" should contain a boolean value True or False indicating if the final goal has been reached. 75 | Make sure your list of tasks ends with a final task like "show results and terminate". 76 | ''' 77 | TASK_PROMPT = ''' 78 | You are one of many language models working on the same final goal: {final_goal}. 79 | 80 | Here is the list of tasks after yours needed to achieve this: {current_tasks}. Your job is to complete this one task: {task}. 81 | 82 | Here is some context from previous task results: {combined_info}. 83 | 84 | {task} 85 | ''' 86 | CREATE_PROMPT: str = 'Based on your end goal, come up with a list of tasks (in order) that you will need to take to achieve your goal.\nGive your response in valid JSON format so that it can be parsed (in python) with `json.loads`. Return a dictionary with the key "current_tasks" containing a list of strings. Make sure your list of tasks ends with a final step such as "show results and terminate".' 87 | FIX_JSON_PROMPT: str = """ 88 | Reformat the following JSON without losing content so that it can be loaded without errors in python using `json.loads`. The following output returned an error when trying to parse. Make sure your response doesn't contain things like: new lines, tabs. Make sure your response uses double quotes as according to the JSON spec. Your response must include an ending quote and ending bracket as needed. ONLY RETURN VALID JSON WITHOUT FORMATTING. 89 | 90 | Example of valid JSON: {example} 91 | 92 | Bad JSON: {bad_json} 93 | 94 | Error: {err} 95 | 96 | Good JSON: """ 97 | GOOD_JSON_EXAMPLE: str = '''{"current_tasks": ["Research Amjad Masad's career and background.", "Create a CSV called \"career.csv\" and write his careers to it."], "stored_info": {"username": "amasad"}, "thoughts": "I will research his career and background, and then save the results to \"career.csv\"."}''' 98 | 99 | def _make_tools_str(self, tools: list) -> str: 100 | """Tools should be a list of dictionaries with the keys: "name" and "description".""" 101 | return '-----\n'.join(['\n'.join([f'{k}: {v}' for k, v in tool.items()]) for tool in tools]) # the fn name has an _ so it doesn't have to be readable, right? 102 | 103 | def _load_persist(self): 104 | if os.path.exists(self.persist): 105 | with open(self.persist, 'r') as f: 106 | saved = json.load(f) 107 | self.output_func(f'[system] Loaded stored info from: {self.persist}') 108 | else: 109 | saved = {} 110 | self.output_func(f'[system] Could not read {self.persist}, assuming new file. It will be created later.') 111 | self.stored_info = saved.get('stored_info', {}) 112 | self.final_result = saved.get('final_result', {}) 113 | self.current_tasks = saved.get('current_tasks', []) 114 | self.completed_tasks = saved.get('completed_tasks', {}) 115 | def _save_persist(self): 116 | with open(self.persist, 'w') as f: 117 | json.dump({ 118 | 'stored_info': self.stored_info, 119 | 'final_result': self.final_result, 120 | 'current_tasks': self.current_tasks, 121 | 'completed_tasks': self.completed_tasks 122 | }, f) 123 | self.output_func(f'saved stored info to: {self.persist}') 124 | 125 | def __init__(self, goal: str, tools: list, llm: BaseLLM, verbose: bool = True, output_func: callable = print, complete_func: callable = save_to_file, input_func: callable = input, current_tasks: list = None, final_result: dict = None, allow_repeat_tasks: bool = True, completed_tasks: dict = None, persist: str = None, confirm_tool: bool = False): 126 | """ 127 | :param goal: str - final goal in natrual language 128 | :param tools: list - a list of tools (dicts) containing keys "name" and "description" 129 | :param llm: BaseLLM - LLM instance from langchain.llms 130 | 131 | :kwarg verbose: bool - defaults to True, if False, will not print updated info 132 | :kwarg allow_repeat_tasks: bool - defaults to True but you might want to disable, will not allow the bot to add tasks that have already been completed 133 | :kwarg output_func: callable - defaults to print, for verbose outout 134 | :kwarg input_func: callable - defaults to input, for user input 135 | :kwarg complete_func: callable - func to run when complete, accepts a goal (str) and results (dict), defaults to a func that saves to file 136 | :kwarg persist: str - defaults to None, but if set to a filepath, [stored_info, final_result, current_tasks] will be loaded and saved there 137 | :kwarg confirm_tool: bool - require user confirmation before running tools (default: False) 138 | :kwarg completed_tasks: dict - defaults to None for empty, already completed tasks for when allow_repeat_tasks=False (key = task name, value = task result), overwrites loaded tasks 139 | :kwarg current_tasks: list - defaults to None for empty, contains a list of (strings) tasks in natural language, overwrites loaded tasks 140 | :kwarg final_result: dict - defaults to None for empty, contains a dict of any results for the final goal, overwrites loaded result 141 | """ 142 | self.llm = llm 143 | self.final_goal = goal 144 | self.tools = tools 145 | self.output_func = output_func 146 | self.complete_func = complete_func 147 | self.input_func = input_func 148 | self.allow_repeat_tasks = allow_repeat_tasks 149 | self.confirm_tool = confirm_tool 150 | self.verbose = verbose 151 | if persist: # load from file 152 | self.persist = persist 153 | self._load_persist() 154 | # overwrite from kwargs 155 | if current_tasks: 156 | self.current_tasks = current_tasks 157 | if final_result: 158 | self.final_result = final_result 159 | if completed_tasks: 160 | self.completed_tasks = completed_tasks 161 | 162 | 163 | if not self.current_tasks: # if no loaded tasks 164 | self._create_initial_tasks() 165 | 166 | def init_agent(self, agent, on_tool_start: callable = None, on_tool_end: callable = None): 167 | agent.callbacks[0].on_tool_start = on_tool_start or self._on_tool_start 168 | agent.callbacks[0].on_tool_end = on_tool_end or self._on_tool_end 169 | 170 | def format_task_str(self, task: str, smart_combine: bool = False, include_completed_tasks: bool = False): 171 | """ 172 | Formats a task as a prompt which can be passed to an agent. 173 | 174 | :param task: str - task in natrual language 175 | :kwarg smart_combine: bool - defaults to False, I don't recommend using this, but if True will choose to include the larger out of final_result and stored_info 176 | :kwarg include_completed_tasks: bool - defaults to False, will include completed tasks (and results) if True, however this uses more tokens 177 | """ 178 | if smart_combine: 179 | # its really not so smart but hey it works for me 180 | if len(str(self.final_result)) > len(str(self.stored_info)): 181 | combined_info = self.final_result 182 | else: 183 | combined_info = self.stored_info 184 | else: 185 | combined_info = {'final_result': self.final_result, 'stored_info': self.stored_info} 186 | if include_completed_tasks: 187 | combined_info['completed_tasks'] = self.completed_tasks 188 | 189 | return self.TASK_PROMPT.format( 190 | task=task, # task for agent, the rest is context 191 | current_tasks=self.current_tasks, 192 | final_goal=self.final_goal, 193 | combined_info=combined_info 194 | ) 195 | 196 | def _base(self): 197 | return self.BASE_PROMPT.format( 198 | tools_str = self._make_tools_str(self.tools), 199 | final_goal = self.final_goal, 200 | current_tasks = self.current_tasks, 201 | final_result = self.final_result, 202 | stored_info=self.stored_info 203 | ) 204 | 205 | def fix_json(self, bad_json: str, err: Exception = None, retry: int = 1) -> dict: 206 | """ 207 | Uses the LLM to try fix JSON response. Prompt: `self.FIX_JSON_PROMPT` 208 | 209 | :param bad_json: str - invalid json 210 | :kwarg err: Exception - err that it caused when tryna load 211 | :kwarg retry: int - number of times to retry 212 | """ 213 | 214 | self.output_func(f'[system] fixing ai JSON output ({retry} retries left)...') 215 | resp = self.llm(self.FIX_JSON_PROMPT.format(bad_json=bad_json, err=err, example=self.GOOD_JSON_EXAMPLE)) 216 | try: 217 | return json.loads(resp.strip()) 218 | except json.JSONDecodeError as e: 219 | self.output_func('[system] cannot parse ai result as JSON: ' + str(e)) 220 | if retry > 0: 221 | return self.fix_json(resp, err=e, retry=(retry - 1)) 222 | else: 223 | console = InteractiveConsole(locals()) 224 | console.interact('dropping into debug shell, if you can fix it, set the variable "fixed" to the loaded json. data is in "resp". use ctrl+d to exit') 225 | return console.locals.get('fixed', {'error': 'could not parse json response'}) 226 | 227 | def _load_json(self, json_str: str): 228 | json_str = json_str.replace('\t', '').replace(' ', '').replace(' ', '').replace('\n', '').replace(' ', '').strip() 229 | try: 230 | return True, json.loads(json_str) 231 | except json.JSONDecodeError: 232 | pass 233 | if not json_str.endswith('}'): 234 | if not (json_str.endswith('"') or json_str.endswith("'")): 235 | json_str += '"' if '"' in json_str else "'" 236 | json_str += '}' 237 | try: 238 | return True, json.loads(json_str) 239 | except json.JSONDecodeError: 240 | return False, json_str 241 | 242 | def load_json(self, json_str: str, retry: int = 1) -> dict: 243 | """ 244 | Try loading a json_str, retrying 1 times by default. 245 | First tries manually, then uses LLM. 246 | """ 247 | try: 248 | ok, res = self._load_json(json_str) # try fix manually 249 | return res if ok else json.loads(res) # trigger err if not ok 250 | except json.JSONDecodeError as e: 251 | return self.fix_json(json_str, err=e, retry=retry) 252 | 253 | def _create_initial_tasks(self): 254 | """This gets called during __init__""" 255 | prompt = self._base() + self.CREATE_PROMPT 256 | resp = self.llm(prompt) 257 | res = self.load_json(resp) 258 | 259 | if self.verbose: 260 | self.output_func('[system] ai created task list: ' + ', '.join(res['current_tasks'])) 261 | self.current_tasks = res['current_tasks'] 262 | 263 | def _on_tool_start(self, tool, input_str, **kwargs): 264 | """Set the agent.callback_manager.on_tool_start to this to save tool inputs to self.stored_info['tools_used'].""" 265 | self.stored_info['tools_used'] = [*self.stored_info.get('tools_used', []), {'tool': tool, 'input': input_str}] 266 | def _on_tool_end(self, output, **kwargs): 267 | """Set the agent.callback_manager.on_tool_end to this to save tool outputs to self.stored_info['tool_used'].""" 268 | self.stored_info['tools_used'][-1]['output'] = output 269 | 270 | def refine(self, task_name: str, task_result: str): 271 | """ 272 | Use this after a task has been completed. This will update the current_tasks, final_result, stored_info, and completed_tasks - saving if persist is set. Uses base prompt plus `self.REFINE_PROMPT`. Returns True if goal has been met. 273 | 274 | :param task_name: str - task in natural language 275 | :param task_result: str - output from agent 276 | """ 277 | self.completed_tasks[task_name] = task_result 278 | needs_save = False 279 | prompt = self._base() + self.REFINE_PROMPT.format( 280 | task = task_name, 281 | result = task_result 282 | ) 283 | resp = self.llm(prompt) 284 | res = self.load_json(resp) 285 | 286 | if (err := res.get('error')): 287 | self.output_func(f'[system] skipping due to error: {err}') 288 | return 289 | 290 | if (thoughts := res.get('thoughts')): 291 | self.output_func(f'[ai] {thoughts}') 292 | if res.get('goal_complete'): 293 | self.current_tasks = [] # clear remaining tasks 294 | self.output_func('[system] goal complete') 295 | self.complete_func(self.final_goal, { 296 | 'final_result': {**self.final_result, **res['final_result']}, 297 | 'completed_tasks': self.completed_tasks, 298 | 'stored_info': self.stored_info, 299 | }) 300 | self.goal_completed = True 301 | 302 | if (current_tasks := res.get('current_tasks')): 303 | needs_save = True 304 | self.add_tasks(current_tasks) 305 | if (stored_info := res.get('stored_info')): 306 | needs_save = True 307 | if self.verbose: 308 | self.output_func(f'[system] new info: {stored_info}') 309 | self.stored_info.update(stored_info) 310 | if (final_result := res.get('final_result')): 311 | needs_save = True 312 | if self.verbose: 313 | self.output_func(f'[system] new final result: {final_result}') 314 | self.final_result.update(final_result) 315 | 316 | if needs_save and self.persist: 317 | self._save_persist() 318 | 319 | def add_tasks(self, current_tasks: list): 320 | if self.allow_repeat_tasks: 321 | if self.verbose: 322 | self.output_func(f'[system] new tasks: {current_tasks}') 323 | self.current_tasks = current_tasks 324 | else: 325 | new_current_tasks = [] 326 | for task in current_tasks: 327 | if not task in self.completed_tasks.keys(): 328 | new_current_tasks.append(task) 329 | self.output_func(f'[system] new tasks: {new_current_tasks} (skipped: {", ".join(t for t in new_current_tasks if not t in current_tasks)})') 330 | self.current_tasks = new_current_tasks 331 | 332 | def ensure_goal_complete(self): 333 | prompt = self._base() + self.ENSURE_COMPLETE_PROMPT 334 | resp = self.llm(prompt) 335 | res = self.load_json(resp) 336 | 337 | if (final_result := res.get('final_result')): 338 | if self.verbose: 339 | self.output_func(f'[system] new final result: {final_result}') 340 | self.final_result.update(final_result) 341 | if (current_tasks := res.get('current_tasks')): 342 | self.add_tasks(current_tasks) 343 | 344 | if res.get('goal_complete'): 345 | self.output_func('[system] Goal completed!') 346 | self.complete_func(self.final_goal, { 347 | 'final_result': self.final_result, 348 | 'completed_tasks': self.completed_tasks, 349 | 'stored_info': self.stored_info, 350 | }) 351 | self.current_tasks = [] 352 | self.goal_completed = True 353 | return True 354 | else: 355 | return False --------------------------------------------------------------------------------