[ { "path": ".devcontainer/devcontainer.json", "content": "{\n\t\"name\": \"Python 3\",\n\t\"image\": \"mcr.microsoft.com/devcontainers/python:0-3.9-bullseye\",\n\n\t\"customizations\": {\n\t\t\"vscode\": {\n\t\t\t\"settings\": {\n\t\t\t\t\"python.defaultInterpreterPath\": \"/usr/local/bin/python\",\n\t\t\t\t\"cSpell.words\": [\n\t\t\t\t\t\"OPENAI\",\n\t\t\t\t\t\"secho\",\n\t\t\t\t\t\"sgpt\",\n\t\t\t\t\t\"Typer\"\n\t\t\t\t],\n\t\t\t\t\"[python]\": {\n\t\t\t\t\t\"editor.defaultFormatter\": \"charliermarsh.ruff\"\n\t\t\t\t},\n\t\t\t\t\"files.exclude\": {\n \"**/.git/**\": true,\n \"**/.mypy_cache/**\": true,\n \"**/__pycache__/**\": true\n },\n \"files.watcherExclude\": {\n \"**/.git/**\": true,\n \"**/.mypy_cache/**\": true,\n \"**/.venv/**\": true,\n \"**/__pycache__/**\": true\n },\n\t\t\t\t\"launch\": {\n\t\t\t\t\t\"configurations\": [\n\t\t\t\t\t\t{\n\t\t\t\t\t\t\t\"name\": \"Python: Module\",\n\t\t\t\t\t\t\t\"type\": \"python\",\n\t\t\t\t\t\t\t\"request\": \"launch\",\n\t\t\t\t\t\t\t\"module\": \"sgpt\",\n\t\t\t\t\t\t\t\"justMyCode\": true,\n\t\t\t\t\t\t\t\"args\": [\"--chat\", \"init\", \"hello\"]\n\t\t\t\t\t\t}\n\t\t\t\t\t]\n\t\t\t\t}\n\t\t\t},\n\t\t\t\"extensions\": [\n\t\t\t\t\"GitHub.copilot\",\n\t\t\t\t\"charliermarsh.ruff\",\n\t\t\t\t\"ms-python.python\",\n\t\t\t\t\"ms-python.vscode-pylance\",\n\t\t\t\t\"ms-python.black-formatter\",\n \"ms-python.isort\",\n \"ms-python.mypy-type-checker\",\n \"ms-python.pylint\"\n\t\t\t]\n\t\t}\n\t},\n\n\t\"remoteUser\": \"vscode\",\n\t\"postCreateCommand\": \"echo __pycache__ > ~/.gitignore && git config --global core.excludesfile ~/.gitignore && pip3 install -e .'[dev,test]'\"\n}" }, { "path": ".github/FUNDING.yml", "content": "github: [ther1d]" }, { "path": ".github/workflows/codespell.yml", "content": "---\nname: Codespell\n\non:\n push:\n branches: [main]\n pull_request:\n branches: [main]\n\npermissions:\n contents: read\n\njobs:\n codespell:\n name: Check for spelling errors\n runs-on: ubuntu-latest\n\n steps:\n - name: Checkout\n uses: actions/checkout@v3\n - name: Codespell\n uses: codespell-project/actions-codespell@v2\n" }, { "path": ".github/workflows/docker.yml", "content": "name: Docker Image CI\n\non:\n push:\n tags:\n - '*'\n\njobs:\n build:\n runs-on: ubuntu-latest\n environment:\n name: Docker Image\n\n steps: \n - name: Checkout\n uses: actions/checkout@v3\n\n - name: Docker meta\n id: meta\n uses: docker/metadata-action@v4\n with:\n images: ghcr.io/${{ github.repository }}\n tags: type=ref,event=tag\n\n - name: Set up QEMU\n uses: docker/setup-qemu-action@v2\n \n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@v2 \n\n - name: Login to GitHub Container Registry\n uses: docker/login-action@v1\n with:\n registry: ghcr.io\n username: ${{ github.actor }}\n password: ${{ secrets.GITHUB_TOKEN }}\n\n - name: Build and push\n uses: docker/build-push-action@v4\n with:\n platforms: linux/amd64,linux/arm64\n push: true\n tags: ${{ steps.meta.outputs.tags }}\n labels: ${{ steps.meta.outputs.labels }}\n" }, { "path": ".github/workflows/lint_test.yml", "content": "name: Lint and Test\n\non:\n push:\n branches:\n - main\n pull_request:\n branches:\n - main\n\njobs:\n lint_test:\n runs-on: ubuntu-latest\n strategy:\n matrix:\n python-version: [\"3.10\", \"3.11\", \"3.12\"]\n steps:\n - uses: actions/checkout@v3\n - name: Set up Python ${{ matrix.python-version }}\n uses: actions/setup-python@v3\n with:\n python-version: ${{ matrix.python-version }}\n - name: Install dependencies\n run: |\n python -m pip install --upgrade pip\n pip install .\"[dev,test]\"\n - name: black\n run: black sgpt tests --check\n - name: isort\n run: isort sgpt tests scripts --check-only\n - name: ruff\n run: ruff sgpt tests scripts\n - name: mypy\n run: mypy sgpt --exclude llm_functions\n - name: tests\n run: |\n export OPENAI_API_KEY=test_api_key\n pytest tests/ -p no:warnings -v -s\n" }, { "path": ".github/workflows/release.yml", "content": "name: Publish to PyPI and release\n\non:\n push:\n branches:\n - main\n paths:\n - 'sgpt/__version__.py'\n\njobs:\n build:\n name: Build distribution\n runs-on: ubuntu-latest\n\n steps:\n - uses: actions/checkout@v4\n - name: Set up Python\n uses: actions/setup-python@v4\n with:\n python-version: '3.11'\n - name: Install pypa/build\n run: >-\n python3 -m\n pip install\n hatchling\n --user\n - name: Build a binary wheel and a source tarball\n run: python3 -m hatchling build\n - name: Store the distribution packages\n uses: actions/upload-artifact@v4\n with:\n name: python-package-distributions\n path: dist/\n\n publish-to-pypi:\n name: Publish to PyPI\n needs:\n - build\n runs-on: ubuntu-latest\n environment:\n name: PyPI\n url: https://pypi.org/p/shell-gpt\n permissions:\n id-token: write # IMPORTANT: mandatory for trusted publishing\n\n steps:\n - name: Download all the dists\n uses: actions/download-artifact@v4\n with:\n name: python-package-distributions\n path: dist/\n - name: Publish distribution to PyPI\n uses: pypa/gh-action-pypi-publish@release/v1\n\n github-release:\n name: Make release\n needs:\n - publish-to-pypi\n runs-on: ubuntu-latest\n\n permissions:\n contents: write # IMPORTANT: mandatory for making GitHub Releases\n id-token: write # IMPORTANT: mandatory for sigstore\n\n steps:\n - name: Download all the dists\n uses: actions/download-artifact@v4\n with:\n name: python-package-distributions\n path: dist/\n - name: Get ShellGPT version\n run: |\n echo \"SGPT_VERSION=$(find dist -type f -name '*.tar.gz' | grep -oP '\\d+.\\d+.\\d+')\" >> $GITHUB_ENV\n echo \"Release version $SGPT_VERSION\"\n - name: Sign the dists with Sigstore\n uses: sigstore/gh-action-sigstore-python@v3.0.0\n with:\n inputs: >-\n ./dist/*.tar.gz\n ./dist/*.whl\n - name: Create GitHub Release\n env:\n GITHUB_TOKEN: ${{ github.token }}\n run: >-\n gh release create\n \"$SGPT_VERSION\"\n --repo '${{ github.repository }}'\n --notes \"$SGPT_VERSION\"\n - name: Upload artifact signatures to GitHub Release\n env:\n GITHUB_TOKEN: ${{ github.token }}\n # Upload to GitHub Release using the `gh` CLI.\n # `dist/` contains the built packages, and the\n # sigstore-produced signatures and certificates.\n run: >-\n gh release upload\n \"$SGPT_VERSION\" dist/**\n --repo '${{ github.repository }}'\n" }, { "path": "CONTRIBUTING.md", "content": "# Contributing to ShellGPT\nThank you for considering contributing to ShellGPT! To ensure a smooth and enjoyable experience for everyone, please follow the steps outlined below.\n\n## Find an Issue to Work On\n- First, browse the existing issues to find one that interests you. If you find an issue you'd like to work on, assign it to yourself and leave a comment expressing your interest.\n- If you have a new feature idea that doesn't have an existing issue, please create a discussion in the \"ideas\" category using GitHub Discussions. Gather feedback from the community, and if you receive approval from at least a couple of people, create an issue and assign it to yourself.\n- If there is an urgent issue, such as a critical bug causing the app to crash, create a pull request immediately.\n\n## Development\nShellGPT is written with strict types, so you'll need to define types. The project uses several linting and testing tools: ruff, mypy, isort, black, and pytest.\n\n### Virtual Environment\nCreate and activate a virtual environment using Python venv:\n\n```shell\npython -m venv env && source ./env/bin/activate\n```\n\n### Install Dependencies\nInstall the necessary dependencies, including development and test dependencies:\n\n```shell\npip install -e .\"[dev,test]\"\n```\n\n### Start Coding\nWith your environment set up and the issue assigned, you can start working on your solution. Get to know the existing codebase and adhere to the project's coding style and conventions. Write clean, modular, and maintainable code to facilitate understanding and review. Commit your changes frequently to document your progress.\n\n### Testing\n**This is a crucial step.** Any changes that implement a new feature or modify existing features should include tests. **Unverified code will not be merged.** These tests should call `sgpt` with defined arguments, capture the output, and verify that the feature works as expected. Refer to the `tests` folder for examples.\n\n### Pull Request\nBefore creating a pull request, run `scripts/lint.sh` and `scripts/tests.sh` to ensure all linters and tests pass. In your pull request, provide a high-level description of your changes and detailed instructions for testing them, including any necessary commands.\n\n### Code Review\nAfter submitting your pull request, be patient and receptive to feedback from reviewers. Address any concerns they raise and collaborate to refine the code. Together, we can enhance the ShellGPT project.\n\nThank you once again for your contribution! We're excited to have you join us." }, { "path": "Dockerfile", "content": "FROM python:3-slim\n\nENV SHELL_INTERACTION=false\nENV PRETTIFY_MARKDOWN=false\nENV OS_NAME=auto\nENV SHELL_NAME=auto\n\nWORKDIR /app\nCOPY . /app\n\nRUN apt-get update && apt-get install -y gcc\nRUN pip install --no-cache /app && mkdir -p /tmp/shell_gpt\n\nVOLUME /tmp/shell_gpt\n\nENTRYPOINT [\"sgpt\"]\n" }, { "path": "LICENSE", "content": "MIT License\n\nCopyright (c) 2023 Farkhod Sadykov\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n" }, { "path": "README.md", "content": "# ShellGPT\nA command-line productivity tool powered by AI large language models (LLM). This command-line tool offers streamlined generation of **shell commands, code snippets, documentation**, eliminating the need for external resources (like Google search). Supports Linux, macOS, Windows and compatible with all major Shells like PowerShell, CMD, Bash, Zsh, etc.\n\nhttps://github.com/TheR1D/shell_gpt/assets/16740832/721ddb19-97e7-428f-a0ee-107d027ddd59\n\n## Installation\n```shell\npip install shell-gpt\n```\nBy default, ShellGPT uses OpenAI's API and GPT-4 model. You'll need an API key, you can generate one [here](https://platform.openai.com/api-keys). You will be prompted for your key which will then be stored in `~/.config/shell_gpt/.sgptrc`. OpenAI API is not free of charge, please refer to the [OpenAI pricing](https://openai.com/pricing) for more information.\n\n> [!TIP]\n> Alternatively, you can use locally hosted open source models which are available for free. To use local models, you will need to run your own LLM backend server such as [Ollama](https://github.com/ollama/ollama). To set up ShellGPT with Ollama, please follow this comprehensive [guide](https://github.com/TheR1D/shell_gpt/wiki/Ollama).\n>\n> **❗️Note that ShellGPT is not optimized for local models and may not work as expected.**\n\n## Usage\n**ShellGPT** is designed to quickly analyse and retrieve information. It's useful for straightforward requests ranging from technical configurations to general knowledge.\n```shell\nsgpt \"What is the fibonacci sequence\"\n# -> The Fibonacci sequence is a series of numbers where each number ...\n```\n\nShellGPT accepts prompt from both stdin and command line argument. Whether you prefer piping input through the terminal or specifying it directly as arguments, `sgpt` got you covered. For example, you can easily generate a git commit message based on a diff:\n```shell\ngit diff | sgpt \"Generate git commit message, for my changes\"\n# -> Added main feature details into README.md\n```\n\nYou can analyze logs from various sources by passing them using stdin, along with a prompt. For instance, we can use it to quickly analyze logs, identify errors and get suggestions for possible solutions:\n```shell\ndocker logs -n 20 my_app | sgpt \"check logs, find errors, provide possible solutions\"\n```\n```text\nError Detected: Connection timeout at line 7.\nPossible Solution: Check network connectivity and firewall settings.\nError Detected: Memory allocation failed at line 12.\nPossible Solution: Consider increasing memory allocation or optimizing application memory usage.\n```\n\nYou can also use all kind of redirection operators to pass input:\n```shell\nsgpt \"summarise\" < document.txt\n# -> The document discusses the impact...\nsgpt << EOF\nWhat is the best way to lear Golang?\nProvide simple hello world example.\nEOF\n# -> The best way to learn Golang...\nsgpt <<< \"What is the best way to learn shell redirects?\"\n# -> The best way to learn shell redirects is through...\n```\n\n\n### Shell commands\nHave you ever found yourself forgetting common shell commands, such as `find`, and needing to look up the syntax online? With `--shell` or shortcut `-s` option, you can quickly generate and execute the commands you need right in the terminal.\n```shell\nsgpt --shell \"find all json files in current folder\"\n# -> find . -type f -name \"*.json\"\n# -> [E]xecute, [D]escribe, [A]bort: e\n```\n\nShell GPT is aware of OS and `$SHELL` you are using, it will provide shell command for specific system you have. For instance, if you ask `sgpt` to update your system, it will return a command based on your OS. Here's an example using macOS:\n```shell\nsgpt -s \"update my system\"\n# -> sudo softwareupdate -i -a\n# -> [E]xecute, [D]escribe, [A]bort: e\n```\n\nThe same prompt, when used on Ubuntu, will generate a different suggestion:\n```shell\nsgpt -s \"update my system\"\n# -> sudo apt update && sudo apt upgrade -y\n# -> [E]xecute, [D]escribe, [A]bort: e\n```\n\nLet's try it with Docker:\n```shell\nsgpt -s \"start nginx container, mount ./index.html\"\n# -> docker run -d -p 80:80 -v $(pwd)/index.html:/usr/share/nginx/html/index.html nginx\n# -> [E]xecute, [D]escribe, [A]bort: e\n```\n\nWe can still use pipes to pass input to `sgpt` and generate shell commands:\n```shell\nsgpt -s \"POST localhost with\" < data.json\n# -> curl -X POST -H \"Content-Type: application/json\" -d '{\"a\": 1, \"b\": 2}' http://localhost\n# -> [E]xecute, [D]escribe, [A]bort: e\n```\n\nApplying additional shell magic in our prompt, in this example passing file names to `ffmpeg`:\n```shell\nls\n# -> 1.mp4 2.mp4 3.mp4\nsgpt -s \"ffmpeg combine $(ls -m) into one video file without audio.\"\n# -> ffmpeg -i 1.mp4 -i 2.mp4 -i 3.mp4 -filter_complex \"[0:v] [1:v] [2:v] concat=n=3:v=1 [v]\" -map \"[v]\" out.mp4\n# -> [E]xecute, [D]escribe, [A]bort: e\n```\n\nIf you would like to pass generated shell command using pipe, you can use `--no-interaction` option. This will disable interactive mode and will print generated command to stdout. In this example we are using `pbcopy` to copy generated command to clipboard:\n```shell\nsgpt -s \"find all json files in current folder\" --no-interaction | pbcopy\n```\n\n\n### Shell integration\nThis is a **very handy feature**, which allows you to use `sgpt` shell completions directly in your terminal, without the need to type `sgpt` with prompt and arguments. Shell integration enables the use of ShellGPT with hotkeys in your terminal, supported by both Bash and ZSH shells. This feature puts `sgpt` completions directly into terminal buffer (input line), allowing for immediate editing of suggested commands.\n\nhttps://github.com/TheR1D/shell_gpt/assets/16740832/bead0dab-0dd9-436d-88b7-6abfb2c556c1\n\nTo install shell integration, run `sgpt --install-integration` and restart your terminal to apply changes. This will add few lines to your `.bashrc` or `.zshrc` file. After that, you can use `Ctrl+l` (by default) to invoke ShellGPT. When you press `Ctrl+l` it will replace you current input line (buffer) with suggested command. You can then edit it and just press `Enter` to execute.\n\n### Generating code\nBy using the `--code` or `-c` parameter, you can specifically request pure code output, for instance:\n```shell\nsgpt --code \"solve fizz buzz problem using python\"\n```\n\n```python\nfor i in range(1, 101):\n if i % 3 == 0 and i % 5 == 0:\n print(\"FizzBuzz\")\n elif i % 3 == 0:\n print(\"Fizz\")\n elif i % 5 == 0:\n print(\"Buzz\")\n else:\n print(i)\n```\nSince it is valid python code, we can redirect the output to a file: \n```shell\nsgpt --code \"solve classic fizz buzz problem using Python\" > fizz_buzz.py\npython fizz_buzz.py\n# 1\n# 2\n# Fizz\n# 4\n# Buzz\n# ...\n```\n\nWe can also use pipes to pass input:\n```shell\ncat fizz_buzz.py | sgpt --code \"Generate comments for each line of my code\"\n```\n```python\n# Loop through numbers 1 to 100\nfor i in range(1, 101):\n # Check if number is divisible by both 3 and 5\n if i % 3 == 0 and i % 5 == 0:\n # Print \"FizzBuzz\" if number is divisible by both 3 and 5\n print(\"FizzBuzz\")\n # Check if number is divisible by 3\n elif i % 3 == 0:\n # Print \"Fizz\" if number is divisible by 3\n print(\"Fizz\")\n # Check if number is divisible by 5\n elif i % 5 == 0:\n # Print \"Buzz\" if number is divisible by 5\n print(\"Buzz\")\n # If number is not divisible by 3 or 5, print the number itself\n else:\n print(i)\n```\n\n### Chat Mode \nOften it is important to preserve and recall a conversation. `sgpt` creates conversational dialogue with each LLM completion requested. The dialogue can develop one-by-one (chat mode) or interactively, in a REPL loop (REPL mode). Both ways rely on the same underlying object, called a chat session. The session is located at the [configurable](#runtime-configuration-file) `CHAT_CACHE_PATH`.\n\nTo start a conversation, use the `--chat` option followed by a unique session name and a prompt.\n```shell\nsgpt --chat conversation_1 \"please remember my favorite number: 4\"\n# -> I will remember that your favorite number is 4.\nsgpt --chat conversation_1 \"what would be my favorite number + 4?\"\n# -> Your favorite number is 4, so if we add 4 to it, the result would be 8.\n```\n\nYou can use chat sessions to iteratively improve GPT suggestions by providing additional details. It is possible to use `--code` or `--shell` options to initiate `--chat`:\n```shell\nsgpt --chat conversation_2 --code \"make a request to localhost using python\"\n```\n```python\nimport requests\n\nresponse = requests.get('http://localhost')\nprint(response.text)\n```\n\nLet's ask LLM to add caching to our request:\n```shell\nsgpt --chat conversation_2 --code \"add caching\"\n```\n```python\nimport requests\nfrom cachecontrol import CacheControl\n\nsess = requests.session()\ncached_sess = CacheControl(sess)\n\nresponse = cached_sess.get('http://localhost')\nprint(response.text)\n```\n\nSame applies for shell commands:\n```shell\nsgpt --chat conversation_3 --shell \"what is in current folder\"\n# -> ls\nsgpt --chat conversation_3 \"Sort by name\"\n# -> ls | sort\nsgpt --chat conversation_3 \"Concatenate them using FFMPEG\"\n# -> ffmpeg -i \"concat:$(ls | sort | tr '\\n' '|')\" -codec copy output.mp4\nsgpt --chat conversation_3 \"Convert the resulting file into an MP3\"\n# -> ffmpeg -i output.mp4 -vn -acodec libmp3lame -ac 2 -ab 160k -ar 48000 final_output.mp3\n```\n\nTo list all the sessions from either conversational mode, use the `--list-chats` or `-lc` option: \n```shell\nsgpt --list-chats\n# .../shell_gpt/chat_cache/conversation_1 \n# .../shell_gpt/chat_cache/conversation_2\n```\n\nTo show all the messages related to a specific conversation, use the `--show-chat` option followed by the session name:\n```shell\nsgpt --show-chat conversation_1\n# user: please remember my favorite number: 4\n# assistant: I will remember that your favorite number is 4.\n# user: what would be my favorite number + 4?\n# assistant: Your favorite number is 4, so if we add 4 to it, the result would be 8.\n```\n\n### REPL Mode \nThere is very handy REPL (read–eval–print loop) mode, which allows you to interactively chat with GPT models. To start a chat session in REPL mode, use the `--repl` option followed by a unique session name. You can also use \"temp\" as a session name to start a temporary REPL session. Note that `--chat` and `--repl` are using same underlying object, so you can use `--chat` to start a chat session and then pick it up with `--repl` to continue the conversation in REPL mode.\n\n

\n \"gif\"\n

\n\n```text\nsgpt --repl temp\nEntering REPL mode, press Ctrl+C to exit.\n>>> What is REPL?\nREPL stands for Read-Eval-Print Loop. It is a programming environment ...\n>>> How can I use Python with REPL?\nTo use Python with REPL, you can simply open a terminal or command prompt ...\n```\n\nREPL mode can work with `--shell` and `--code` options, which makes it very handy for interactive shell commands and code generation:\n```text\nsgpt --repl temp --shell\nEntering shell REPL mode, type [e] to execute commands or press Ctrl+C to exit.\n>>> What is in current folder?\nls\n>>> Show file sizes\nls -lh\n>>> Sort them by file sizes\nls -lhS\n>>> e (enter just e to execute commands, or d to describe them)\n```\n\nTo provide multiline prompt use triple quotes `\"\"\"`:\n```text\nsgpt --repl temp\nEntering REPL mode, press Ctrl+C to exit.\n>>> \"\"\"\n... Explain following code:\n... import random\n... print(random.randint(1, 10))\n... \"\"\"\nIt is a Python script that uses the random module to generate and print a random integer.\n```\n\nYou can also enter REPL mode with initial prompt by passing it as an argument or stdin or even both:\n```shell\nsgpt --repl temp < my_app.py\n```\n```text\nEntering REPL mode, press Ctrl+C to exit.\n──────────────────────────────────── Input ────────────────────────────────────\nname = input(\"What is your name?\")\nprint(f\"Hello {name}\")\n───────────────────────────────────────────────────────────────────────────────\n>>> What is this code about?\nThe snippet of code you've provided is written in Python. It prompts the user...\n>>> Follow up questions...\n```\n\n### Function calling \n[Function calls](https://platform.openai.com/docs/guides/function-calling) is a powerful feature OpenAI provides. It allows LLM to execute functions in your system, which can be used to accomplish a variety of tasks. To install [default functions](https://github.com/TheR1D/shell_gpt/tree/main/sgpt/llm_functions/) run:\n```shell\nsgpt --install-functions\n```\n\nShellGPT has a convenient way to define functions and use them. In order to create your custom function, navigate to `~/.config/shell_gpt/functions` and create a new .py file with the function name. Inside this file, you can define your function using this [example](https://github.com/TheR1D/shell_gpt/blob/main/sgpt/llm_functions/common/execute_shell.py).\n\nThe docstring comment inside the class will be passed to OpenAI API as a description for the function, along with the `title` attribute and parameters descriptions. The `execute` function will be called if LLM decides to use your function. In this case we are allowing LLM to execute any Shell commands in our system. Since we are returning the output of the command, LLM will be able to analyze it and decide if it is a good fit for the prompt. Here is an example how the function might be executed by LLM:\n```shell\nsgpt \"What are the files in /tmp folder?\"\n# -> @FunctionCall execute_shell_command(shell_command=\"ls /tmp\")\n# -> The /tmp folder contains the following files and directories:\n# -> test.txt\n# -> test.json\n```\n\nNote that if for some reason the function (execute_shell_command) will return an error, LLM might try to accomplish the task based on the output. Let's say we don't have installed `jq` in our system, and we ask LLM to parse JSON file:\n```shell\nsgpt \"parse /tmp/test.json file using jq and return only email value\"\n# -> @FunctionCall execute_shell_command(shell_command=\"jq -r '.email' /tmp/test.json\")\n# -> It appears that jq is not installed on the system. Let me try to install it using brew.\n# -> @FunctionCall execute_shell_command(shell_command=\"brew install jq\")\n# -> jq has been successfully installed. Let me try to parse the file again.\n# -> @FunctionCall execute_shell_command(shell_command=\"jq -r '.email' /tmp/test.json\")\n# -> The email value in /tmp/test.json is johndoe@example.\n```\n\nIt is also possible to chain multiple function calls in the prompt:\n```shell\nsgpt \"Play music and open hacker news\"\n# -> @FunctionCall play_music()\n# -> @FunctionCall open_url(url=\"https://news.ycombinator.com\")\n# -> Music is now playing, and Hacker News has been opened in your browser. Enjoy!\n```\n\nThis is just a simple example of how you can use function calls. It is truly a powerful feature that can be used to accomplish a variety of complex tasks. We have dedicated [category](https://github.com/TheR1D/shell_gpt/discussions/categories/functions) in GitHub Discussions for sharing and discussing functions. \nLLM might execute destructive commands, so please use it at your own risk❗️\n\n### Roles\nShellGPT allows you to create custom roles, which can be utilized to generate code, shell commands, or to fulfill your specific needs. To create a new role, use the `--create-role` option followed by the role name. You will be prompted to provide a description for the role, along with other details. This will create a JSON file in `~/.config/shell_gpt/roles` with the role name. Inside this directory, you can also edit default `sgpt` roles, such as **shell**, **code**, and **default**. Use the `--list-roles` option to list all available roles, and the `--show-role` option to display the details of a specific role. Here's an example of a custom role:\n```shell\nsgpt --create-role json_generator\n# Enter role description: Provide only valid json as response.\nsgpt --role json_generator \"random: user, password, email, address\"\n```\n```json\n{\n \"user\": \"JohnDoe\",\n \"password\": \"p@ssw0rd\",\n \"email\": \"johndoe@example.com\",\n \"address\": {\n \"street\": \"123 Main St\",\n \"city\": \"Anytown\",\n \"state\": \"CA\",\n \"zip\": \"12345\"\n }\n}\n```\n\nIf the description of the role contains the words \"APPLY MARKDOWN\" (case sensitive), then chats will be displayed using markdown formatting unless it is explicitly turned off with `--no-md`.\n\n### Request cache\nControl cache using `--cache` (default) and `--no-cache` options. This caching applies for all `sgpt` requests to OpenAI API:\n```shell\nsgpt \"what are the colors of a rainbow\"\n# -> The colors of a rainbow are red, orange, yellow, green, blue, indigo, and violet.\n```\nNext time, same exact query will get results from local cache instantly. Note that `sgpt \"what are the colors of a rainbow\" --temperature 0.5` will make a new request, since we didn't provide `--temperature` (same applies to `--top-probability`) on previous request.\n\nThis is just some examples of what we can do using OpenAI GPT models, I'm sure you will find it useful for your specific use cases.\n\n### Runtime configuration file\nYou can setup some parameters in runtime configuration file `~/.config/shell_gpt/.sgptrc`:\n```text\n# API key, also it is possible to define OPENAI_API_KEY env.\nOPENAI_API_KEY=your_api_key\n# Base URL of the backend server. If \"default\" URL will be resolved based on --model.\nAPI_BASE_URL=default\n# Max amount of cached message per chat session.\nCHAT_CACHE_LENGTH=100\n# Chat cache folder.\nCHAT_CACHE_PATH=/tmp/shell_gpt/chat_cache\n# Request cache length (amount).\nCACHE_LENGTH=100\n# Request cache folder.\nCACHE_PATH=/tmp/shell_gpt/cache\n# Request timeout in seconds.\nREQUEST_TIMEOUT=60\n# Default OpenAI model to use.\nDEFAULT_MODEL=gpt-4o\n# Default color for shell and code completions.\nDEFAULT_COLOR=magenta\n# When in --shell mode, default to \"Y\" for no input.\nDEFAULT_EXECUTE_SHELL_CMD=false\n# Disable streaming of responses\nDISABLE_STREAMING=false\n# The pygment theme to view markdown (default/describe role).\nCODE_THEME=default\n# Path to a directory with functions.\nOPENAI_FUNCTIONS_PATH=/Users/user/.config/shell_gpt/functions\n# Print output of functions when LLM uses them.\nSHOW_FUNCTIONS_OUTPUT=false\n# Allows LLM to use functions.\nOPENAI_USE_FUNCTIONS=true\n# Enforce LiteLLM usage (for local LLMs).\nUSE_LITELLM=false\n```\nPossible options for `DEFAULT_COLOR`: black, red, green, yellow, blue, magenta, cyan, white, bright_black, bright_red, bright_green, bright_yellow, bright_blue, bright_magenta, bright_cyan, bright_white.\nPossible options for `CODE_THEME`: https://pygments.org/styles/\n\n### Full list of arguments\n```text\n╭─ Arguments ──────────────────────────────────────────────────────────────────────────────────────────────╮\n│ prompt [PROMPT] The prompt to generate completions for. │\n╰──────────────────────────────────────────────────────────────────────────────────────────────────────────╯\n╭─ Options ────────────────────────────────────────────────────────────────────────────────────────────────╮\n│ --model TEXT Large language model to use. [default: gpt-4o] │\n│ --temperature FLOAT RANGE [0.0<=x<=2.0] Randomness of generated output. [default: 0.0] │\n│ --top-p FLOAT RANGE [0.0<=x<=1.0] Limits highest probable tokens (words). [default: 1.0] │\n│ --md --no-md Prettify markdown output. [default: md] │\n│ --editor Open $EDITOR to provide a prompt. [default: no-editor] │\n│ --cache Cache completion results. [default: cache] │\n│ --version Show version. │\n│ --help Show this message and exit. │\n╰──────────────────────────────────────────────────────────────────────────────────────────────────────────╯\n╭─ Assistance Options ─────────────────────────────────────────────────────────────────────────────────────╮\n│ --shell -s Generate and execute shell commands. │\n│ --interaction --no-interaction Interactive mode for --shell option. [default: interaction] │\n│ --describe-shell -d Describe a shell command. │\n│ --code -c Generate only code. │\n│ --functions --no-functions Allow function calls. [default: functions] │\n╰──────────────────────────────────────────────────────────────────────────────────────────────────────────╯\n╭─ Chat Options ───────────────────────────────────────────────────────────────────────────────────────────╮\n│ --chat TEXT Follow conversation with id, use \"temp\" for quick session. [default: None] │\n│ --repl TEXT Start a REPL (Read–eval–print loop) session. [default: None] │\n│ --show-chat TEXT Show all messages from provided chat id. [default: None] │\n│ --list-chats -lc List all existing chat ids. │\n╰──────────────────────────────────────────────────────────────────────────────────────────────────────────╯\n╭─ Role Options ───────────────────────────────────────────────────────────────────────────────────────────╮\n│ --role TEXT System role for GPT model. [default: None] │\n│ --create-role TEXT Create role. [default: None] │\n│ --show-role TEXT Show role. [default: None] │\n│ --list-roles -lr List roles. │\n╰──────────────────────────────────────────────────────────────────────────────────────────────────────────╯\n```\n\n## Docker\nRun the container using the `OPENAI_API_KEY` environment variable, and a docker volume to store cache. Consider to set the environment variables `OS_NAME` and `SHELL_NAME` according to your preferences.\n```shell\ndocker run --rm \\\n --env OPENAI_API_KEY=api_key \\\n --env OS_NAME=$(uname -s) \\\n --env SHELL_NAME=$(echo $SHELL) \\\n --volume gpt-cache:/tmp/shell_gpt \\\n ghcr.io/ther1d/shell_gpt -s \"update my system\"\n```\n\nExample of a conversation, using an alias and the `OPENAI_API_KEY` environment variable:\n```shell\nalias sgpt=\"docker run --rm --volume gpt-cache:/tmp/shell_gpt --env OPENAI_API_KEY --env OS_NAME=$(uname -s) --env SHELL_NAME=$(echo $SHELL) ghcr.io/ther1d/shell_gpt\"\nexport OPENAI_API_KEY=\"your OPENAI API key\"\nsgpt --chat rainbow \"what are the colors of a rainbow\"\nsgpt --chat rainbow \"inverse the list of your last answer\"\nsgpt --chat rainbow \"translate your last answer in french\"\n```\n\nYou also can use the provided `Dockerfile` to build your own image:\n```shell\ndocker build -t sgpt .\n```\n\n### Docker + Ollama\n\nIf you want to send your requests to an Ollama instance and run ShellGPT inside a Docker container, you need to adjust the Dockerfile and build the container yourself: the litellm package is needed and env variables need to be set correctly.\n\nExample Dockerfile:\n```\nFROM python:3-slim\n\nENV DEFAULT_MODEL=ollama/mistral:7b-instruct-v0.2-q4_K_M\nENV API_BASE_URL=http://10.10.10.10:11434\nENV USE_LITELLM=true\nENV OPENAI_API_KEY=bad_key\nENV SHELL_INTERACTION=false\nENV PRETTIFY_MARKDOWN=false\nENV OS_NAME=\"Arch Linux\"\nENV SHELL_NAME=auto\n\nWORKDIR /app\nCOPY . /app\n\nRUN apt-get update && apt-get install -y gcc\nRUN pip install --no-cache /app[litellm] && mkdir -p /tmp/shell_gpt\n\nVOLUME /tmp/shell_gpt\n\nENTRYPOINT [\"sgpt\"]\n```\n\n\n## Additional documentation\n* [Azure integration](https://github.com/TheR1D/shell_gpt/wiki/Azure)\n* [Ollama integration](https://github.com/TheR1D/shell_gpt/wiki/Ollama)\n" }, { "path": "pyproject.toml", "content": "[build-system]\nrequires = [\"hatchling\"]\nbuild-backend = \"hatchling.build\"\n\n[project]\nname = \"shell_gpt\"\ndescription = \"A command-line productivity tool powered by large language models, will help you accomplish your tasks faster and more efficiently.\"\nkeywords = [\"shell\", \"gpt\", \"openai\", \"ollama\", \"cli\", \"productivity\", \"cheet-sheet\"]\nreadme = \"README.md\"\nlicense = \"MIT\"\nrequires-python = \">=3.10\"\nauthors = [{ name = \"Farkhod Sadykov\", email = \"farkhod@sadykov.dev\" }]\ndynamic = [\"version\"]\nclassifiers = [\n \"Operating System :: OS Independent\",\n \"Topic :: Software Development\",\n \"License :: OSI Approved :: MIT License\",\n \"Intended Audience :: Information Technology\",\n \"Intended Audience :: System Administrators\",\n \"Intended Audience :: Developers\",\n \"Programming Language :: Python :: 3.10\",\n \"Programming Language :: Python :: 3.11\",\n \"Programming Language :: Python :: 3.12\",\n \"Programming Language :: Python :: 3.13\",\n]\ndependencies = [\n \"openai >= 2.0.0, < 3.0.0\",\n \"typer >= 0.7.0, < 1.0.0\",\n \"rich >= 13.1.0, < 14.0.0\",\n \"distro >= 1.8.0, < 2.0.0\",\n 'pyreadline3 >= 3.4.1, < 4.0.0; sys_platform == \"win32\"',\n \"prompt_toolkit >= 3.0.51\",\n]\n\n[project.optional-dependencies]\nlitellm = [\n \"litellm == 1.42.5\"\n]\ntest = [\n \"pytest >= 7.2.2, < 8.0.0\",\n \"requests-mock[fixture] >= 1.10.0, < 2.0.0\",\n \"isort >= 5.12.0, < 6.0.0\",\n \"black == 23.1.0\",\n \"mypy == 1.1.1\",\n \"types-requests == 2.28.11.17\",\n \"codespell >= 2.2.5, < 3.0.0\"\n]\ndev = [\n \"ruff == 0.0.256\",\n \"pre-commit >= 3.1.1, < 4.0.0\",\n]\n\n[project.scripts]\nsgpt = \"sgpt:cli\"\n\n[project.urls]\nhomepage = \"https://github.com/ther1d/shell_gpt\"\nrepository = \"https://github.com/ther1d/shell_gpt\"\ndocumentation = \"https://github.com/TheR1D/shell_gpt/blob/main/README.md\"\n\n[tool.hatch.version]\npath = \"sgpt/__version__.py\"\n\n[tool.hatch.build.targets.wheel]\nonly-include = [\"sgpt\"]\n\n[tool.hatch.build.targets.sdist]\nonly-include = [\n \"sgpt\",\n \"tests\",\n \"README.md\",\n \"LICENSE\",\n \"pyproject.toml\",\n]\n\n[tool.isort]\nprofile = \"black\"\nskip = \"__init__.py\"\n\n[tool.mypy]\nstrict = true\nexclude = [\"llm_functions\"]\n\n[tool.ruff]\nselect = [\n \"E\", # pycodestyle errors.\n \"W\", # pycodestyle warnings.\n \"F\", # pyflakes.\n \"C\", # flake8-comprehensions.\n \"B\", # flake8-bugbear.\n]\nignore = [\n \"E501\", # line too long, handled by black.\n \"C901\", # too complex.\n \"B008\", # do not perform function calls in argument defaults.\n \"E731\", # do not assign a lambda expression, use a def.\n]\n\n[tool.codespell]\nskip = '.git,venv'\n" }, { "path": "scripts/format.sh", "content": "#!/bin/sh -e\nset -x\n\nruff sgpt tests scripts --fix\nblack sgpt tests scripts\nisort sgpt tests scripts\ncodespell --write-changes\n" }, { "path": "scripts/lint.sh", "content": "#!/usr/bin/env bash\n\nset -e\nset -x\n\nmypy sgpt\nruff sgpt tests scripts\nblack sgpt tests --check\nisort sgpt tests scripts --check-only\ncodespell\n" }, { "path": "scripts/test.sh", "content": "#!/usr/bin/env bash\n\nset -e\nset -x\n\n# shellcheck disable=SC2068\npytest tests ${@} -p no:warnings\n" }, { "path": "sgpt/__init__.py", "content": "from .app import main as main\nfrom .app import entry_point as cli # noqa: F401\n" }, { "path": "sgpt/__main__.py", "content": "from .app import entry_point\n\nentry_point()\n" }, { "path": "sgpt/__version__.py", "content": "__version__ = \"1.5.0\"\n" }, { "path": "sgpt/app.py", "content": "import os\n\n# To allow users to use arrow keys in the REPL.\nimport readline # noqa: F401\nimport sys\n\nimport typer\nfrom click import UsageError\nfrom click.types import Choice\nfrom prompt_toolkit import PromptSession\n\nfrom sgpt.config import cfg\nfrom sgpt.function import get_openai_schemas\nfrom sgpt.handlers.chat_handler import ChatHandler\nfrom sgpt.handlers.default_handler import DefaultHandler\nfrom sgpt.handlers.repl_handler import ReplHandler\nfrom sgpt.llm_functions.init_functions import install_functions as inst_funcs\nfrom sgpt.role import DefaultRoles, SystemRole\nfrom sgpt.utils import (\n get_edited_prompt,\n get_sgpt_version,\n install_shell_integration,\n run_command,\n)\n\n\ndef main(\n prompt: str = typer.Argument(\n \"\",\n show_default=False,\n help=\"The prompt to generate completions for.\",\n ),\n model: str = typer.Option(\n cfg.get(\"DEFAULT_MODEL\"),\n help=\"Large language model to use.\",\n ),\n temperature: float = typer.Option(\n 0.0,\n min=0.0,\n max=2.0,\n help=\"Randomness of generated output.\",\n ),\n top_p: float = typer.Option(\n 1.0,\n min=0.0,\n max=1.0,\n help=\"Limits highest probable tokens (words).\",\n ),\n md: bool = typer.Option(\n cfg.get(\"PRETTIFY_MARKDOWN\") == \"true\",\n help=\"Prettify markdown output.\",\n ),\n shell: bool = typer.Option(\n False,\n \"--shell\",\n \"-s\",\n help=\"Generate and execute shell commands.\",\n rich_help_panel=\"Assistance Options\",\n ),\n interaction: bool = typer.Option(\n cfg.get(\"SHELL_INTERACTION\") == \"true\",\n help=\"Interactive mode for --shell option.\",\n rich_help_panel=\"Assistance Options\",\n ),\n describe_shell: bool = typer.Option(\n False,\n \"--describe-shell\",\n \"-d\",\n help=\"Describe a shell command.\",\n rich_help_panel=\"Assistance Options\",\n ),\n code: bool = typer.Option(\n False,\n \"--code\",\n \"-c\",\n help=\"Generate only code.\",\n rich_help_panel=\"Assistance Options\",\n ),\n functions: bool = typer.Option(\n cfg.get(\"OPENAI_USE_FUNCTIONS\") == \"true\",\n help=\"Allow function calls.\",\n rich_help_panel=\"Assistance Options\",\n ),\n editor: bool = typer.Option(\n False,\n help=\"Open $EDITOR to provide a prompt.\",\n ),\n cache: bool = typer.Option(\n True,\n help=\"Cache completion results.\",\n ),\n version: bool = typer.Option(\n False,\n \"--version\",\n help=\"Show version.\",\n callback=get_sgpt_version,\n ),\n chat: str = typer.Option(\n None,\n help=\"Follow conversation with id, \" 'use \"temp\" for quick session.',\n rich_help_panel=\"Chat Options\",\n ),\n repl: str = typer.Option(\n None,\n help=\"Start a REPL (Read–eval–print loop) session.\",\n rich_help_panel=\"Chat Options\",\n ),\n show_chat: str = typer.Option(\n None,\n help=\"Show all messages from provided chat id.\",\n rich_help_panel=\"Chat Options\",\n ),\n list_chats: bool = typer.Option(\n False,\n \"--list-chats\",\n \"-lc\",\n help=\"List all existing chat ids.\",\n callback=ChatHandler.list_ids,\n rich_help_panel=\"Chat Options\",\n ),\n role: str = typer.Option(\n None,\n help=\"System role for GPT model.\",\n rich_help_panel=\"Role Options\",\n ),\n create_role: str = typer.Option(\n None,\n help=\"Create role.\",\n callback=SystemRole.create,\n rich_help_panel=\"Role Options\",\n ),\n show_role: str = typer.Option(\n None,\n help=\"Show role.\",\n callback=SystemRole.show,\n rich_help_panel=\"Role Options\",\n ),\n list_roles: bool = typer.Option(\n False,\n \"--list-roles\",\n \"-lr\",\n help=\"List roles.\",\n callback=SystemRole.list,\n rich_help_panel=\"Role Options\",\n ),\n install_integration: bool = typer.Option(\n False,\n help=\"Install shell integration (ZSH and Bash only)\",\n callback=install_shell_integration,\n hidden=True, # Hiding since should be used only once.\n ),\n install_functions: bool = typer.Option(\n False,\n help=\"Install default functions.\",\n callback=inst_funcs,\n hidden=True, # Hiding since should be used only once.\n ),\n) -> None:\n stdin_passed = not sys.stdin.isatty()\n\n if stdin_passed:\n stdin = \"\"\n # TODO: This is very hacky.\n # In some cases, we need to pass stdin along with inputs.\n # When we want part of stdin to be used as a init prompt,\n # but rest of the stdin to be used as a inputs. For example:\n # echo \"hello\\n__sgpt__eof__\\nThis is input\" | sgpt --repl temp\n # In this case, \"hello\" will be used as a init prompt, and\n # \"This is input\" will be used as \"interactive\" input to the REPL.\n # This is useful to test REPL with some initial context.\n for line in sys.stdin:\n if \"__sgpt__eof__\" in line:\n break\n stdin += line\n prompt = f\"{stdin}\\n\\n{prompt}\" if prompt else stdin\n try:\n # Switch to stdin for interactive input.\n if os.name == \"posix\":\n sys.stdin = open(\"/dev/tty\", \"r\")\n elif os.name == \"nt\":\n sys.stdin = open(\"CON\", \"r\")\n except OSError:\n # Non-interactive shell.\n pass\n\n if show_chat:\n ChatHandler.show_messages(show_chat, md)\n\n if sum((shell, describe_shell, code)) > 1:\n raise UsageError(\n \"Only one of --shell, --describe-shell, and --code options can be used at a time.\"\n )\n\n if chat and repl:\n raise UsageError(\"--chat and --repl options cannot be used together.\")\n\n if editor and stdin_passed:\n raise UsageError(\"--editor option cannot be used with stdin input.\")\n\n if editor:\n prompt = get_edited_prompt()\n\n role_class = (\n DefaultRoles.check_get(shell, describe_shell, code)\n if not role\n else SystemRole.get(role)\n )\n\n function_schemas = (get_openai_schemas() or None) if functions else None\n\n if repl:\n # Will be in infinite loop here until user exits with Ctrl+C.\n ReplHandler(repl, role_class, md).handle(\n init_prompt=prompt,\n model=model,\n temperature=temperature,\n top_p=top_p,\n caching=cache,\n functions=function_schemas,\n )\n\n if chat:\n full_completion = ChatHandler(chat, role_class, md).handle(\n prompt=prompt,\n model=model,\n temperature=temperature,\n top_p=top_p,\n caching=cache,\n functions=function_schemas,\n )\n else:\n full_completion = DefaultHandler(role_class, md).handle(\n prompt=prompt,\n model=model,\n temperature=temperature,\n top_p=top_p,\n caching=cache,\n functions=function_schemas,\n )\n\n session: PromptSession[str] = PromptSession()\n\n while shell and interaction:\n option = typer.prompt(\n text=\"[E]xecute, [M]odify, [D]escribe, [A]bort\",\n type=Choice((\"e\", \"m\", \"d\", \"a\", \"y\"), case_sensitive=False),\n default=\"e\" if cfg.get(\"DEFAULT_EXECUTE_SHELL_CMD\") == \"true\" else \"a\",\n show_choices=False,\n show_default=False,\n )\n\n if option in (\"e\", \"y\"):\n # \"y\" option is for keeping compatibility with old version.\n run_command(full_completion)\n elif option == \"m\":\n full_completion = session.prompt(\"\", default=full_completion)\n continue\n elif option == \"d\":\n DefaultHandler(DefaultRoles.DESCRIBE_SHELL.get_role(), md).handle(\n full_completion,\n model=model,\n temperature=temperature,\n top_p=top_p,\n caching=cache,\n functions=function_schemas,\n )\n continue\n break\n\n\ndef entry_point() -> None:\n typer.run(main)\n\n\nif __name__ == \"__main__\":\n entry_point()\n" }, { "path": "sgpt/cache.py", "content": "import json\nfrom hashlib import md5\nfrom pathlib import Path\nfrom typing import Any, Callable, Generator, no_type_check\n\n\nclass Cache:\n \"\"\"\n Decorator class that adds caching functionality to a function.\n \"\"\"\n\n def __init__(self, length: int, cache_path: Path) -> None:\n \"\"\"\n Initialize the Cache decorator.\n\n :param length: Integer, maximum number of cache files to keep.\n \"\"\"\n self.length = length\n self.cache_path = cache_path\n self.cache_path.mkdir(parents=True, exist_ok=True)\n\n def __call__(self, func: Callable[..., Any]) -> Callable[..., Any]:\n \"\"\"\n The Cache decorator.\n\n :param func: The function to cache.\n :return: Wrapped function with caching.\n \"\"\"\n\n def wrapper(*args: Any, **kwargs: Any) -> Generator[str, None, None]:\n key = md5(json.dumps((args[1:], kwargs)).encode(\"utf-8\")).hexdigest()\n file = self.cache_path / key\n if kwargs.pop(\"caching\") and file.exists():\n yield file.read_text()\n return\n result = \"\"\n for i in func(*args, **kwargs):\n result += i\n yield i\n if \"@FunctionCall\" not in result:\n file.write_text(result)\n self._delete_oldest_files(self.length) # type: ignore\n\n return wrapper\n\n @no_type_check\n def _delete_oldest_files(self, max_files: int) -> None:\n \"\"\"\n Class method to delete the oldest cached files in the CACHE_DIR folder.\n\n :param max_files: Integer, the maximum number of files to keep in the CACHE_DIR folder.\n \"\"\"\n # Get all files in the folder.\n files = self.cache_path.glob(\"*\")\n # Sort files by last modification time in ascending order.\n files = sorted(files, key=lambda f: f.stat().st_mtime)\n # Delete the oldest files if the number of files exceeds the limit.\n if len(files) > max_files:\n num_files_to_delete = len(files) - max_files\n for i in range(num_files_to_delete):\n files[i].unlink()\n" }, { "path": "sgpt/config.py", "content": "import os\nfrom getpass import getpass\nfrom pathlib import Path\nfrom tempfile import gettempdir\nfrom typing import Any\n\nfrom click import UsageError\n\nCONFIG_FOLDER = os.path.expanduser(\"~/.config\")\nSHELL_GPT_CONFIG_FOLDER = Path(CONFIG_FOLDER) / \"shell_gpt\"\nSHELL_GPT_CONFIG_PATH = SHELL_GPT_CONFIG_FOLDER / \".sgptrc\"\nROLE_STORAGE_PATH = SHELL_GPT_CONFIG_FOLDER / \"roles\"\nFUNCTIONS_PATH = SHELL_GPT_CONFIG_FOLDER / \"functions\"\nCHAT_CACHE_PATH = Path(gettempdir()) / \"chat_cache\"\nCACHE_PATH = Path(gettempdir()) / \"cache\"\n\n# TODO: Refactor ENV variables with SGPT_ prefix.\nDEFAULT_CONFIG = {\n # TODO: Refactor it to CHAT_STORAGE_PATH.\n \"CHAT_CACHE_PATH\": os.getenv(\"CHAT_CACHE_PATH\", str(CHAT_CACHE_PATH)),\n \"CACHE_PATH\": os.getenv(\"CACHE_PATH\", str(CACHE_PATH)),\n \"CHAT_CACHE_LENGTH\": int(os.getenv(\"CHAT_CACHE_LENGTH\", \"100\")),\n \"CACHE_LENGTH\": int(os.getenv(\"CHAT_CACHE_LENGTH\", \"100\")),\n \"REQUEST_TIMEOUT\": int(os.getenv(\"REQUEST_TIMEOUT\", \"60\")),\n \"DEFAULT_MODEL\": os.getenv(\"DEFAULT_MODEL\", \"gpt-4o\"),\n \"DEFAULT_COLOR\": os.getenv(\"DEFAULT_COLOR\", \"magenta\"),\n \"ROLE_STORAGE_PATH\": os.getenv(\"ROLE_STORAGE_PATH\", str(ROLE_STORAGE_PATH)),\n \"DEFAULT_EXECUTE_SHELL_CMD\": os.getenv(\"DEFAULT_EXECUTE_SHELL_CMD\", \"false\"),\n \"DISABLE_STREAMING\": os.getenv(\"DISABLE_STREAMING\", \"false\"),\n \"CODE_THEME\": os.getenv(\"CODE_THEME\", \"dracula\"),\n \"OPENAI_FUNCTIONS_PATH\": os.getenv(\"OPENAI_FUNCTIONS_PATH\", str(FUNCTIONS_PATH)),\n \"OPENAI_USE_FUNCTIONS\": os.getenv(\"OPENAI_USE_FUNCTIONS\", \"true\"),\n \"SHOW_FUNCTIONS_OUTPUT\": os.getenv(\"SHOW_FUNCTIONS_OUTPUT\", \"false\"),\n \"API_BASE_URL\": os.getenv(\"API_BASE_URL\", \"default\"),\n \"PRETTIFY_MARKDOWN\": os.getenv(\"PRETTIFY_MARKDOWN\", \"true\"),\n \"USE_LITELLM\": os.getenv(\"USE_LITELLM\", \"false\"),\n \"SHELL_INTERACTION\": os.getenv(\"SHELL_INTERACTION \", \"true\"),\n \"OS_NAME\": os.getenv(\"OS_NAME\", \"auto\"),\n \"SHELL_NAME\": os.getenv(\"SHELL_NAME\", \"auto\"),\n # New features might add their own config variables here.\n}\n\n\nclass Config(dict): # type: ignore\n def __init__(self, config_path: Path, **defaults: Any):\n self.config_path = config_path\n\n if self._exists:\n self._read()\n has_new_config = False\n for key, value in defaults.items():\n if key not in self:\n has_new_config = True\n self[key] = value\n if has_new_config:\n self._write()\n else:\n config_path.parent.mkdir(parents=True, exist_ok=True)\n # Don't write API key to config file if it is in the environment.\n if not defaults.get(\"OPENAI_API_KEY\") and not os.getenv(\"OPENAI_API_KEY\"):\n __api_key = getpass(prompt=\"Please enter your OpenAI API key: \")\n defaults[\"OPENAI_API_KEY\"] = __api_key\n super().__init__(**defaults)\n self._write()\n\n @property\n def _exists(self) -> bool:\n return self.config_path.exists()\n\n def _write(self) -> None:\n with open(self.config_path, \"w\", encoding=\"utf-8\") as file:\n string_config = \"\"\n for key, value in self.items():\n string_config += f\"{key}={value}\\n\"\n file.write(string_config)\n\n def _read(self) -> None:\n with open(self.config_path, \"r\", encoding=\"utf-8\") as file:\n for line in file:\n if line.strip() and not line.startswith(\"#\"):\n key, value = line.strip().split(\"=\", 1)\n self[key] = value\n\n def get(self, key: str) -> str: # type: ignore\n # Prioritize environment variables over config file.\n value = os.getenv(key) or super().get(key)\n if not value:\n raise UsageError(f\"Missing config key: {key}\")\n return value\n\n\ncfg = Config(SHELL_GPT_CONFIG_PATH, **DEFAULT_CONFIG)\n" }, { "path": "sgpt/function.py", "content": "import importlib.util\nimport sys\nfrom pathlib import Path\nfrom typing import Any, Callable, Dict, List\n\nfrom pydantic import BaseModel\n\nfrom .config import cfg\n\n\nclass Function:\n def __init__(self, path: str):\n module = self._read(path)\n self._function = module.Function.execute\n self._openai_schema = module.Function.openai_schema()\n self._name = self._openai_schema[\"function\"][\"name\"]\n\n @property\n def name(self) -> str:\n return self._name # type: ignore\n\n @property\n def openai_schema(self) -> dict[str, Any]:\n return self._openai_schema # type: ignore\n\n @property\n def execute(self) -> Callable[..., str]:\n return self._function # type: ignore\n\n @classmethod\n def _read(cls, path: str) -> Any:\n module_name = path.replace(\"/\", \".\").rstrip(\".py\")\n spec = importlib.util.spec_from_file_location(module_name, path)\n module = importlib.util.module_from_spec(spec) # type: ignore\n sys.modules[module_name] = module\n spec.loader.exec_module(module) # type: ignore\n\n if not issubclass(module.Function, BaseModel):\n raise TypeError(\n f\"Function {module_name} must be a subclass of pydantic.BaseModel\"\n )\n if not hasattr(module.Function, \"execute\"):\n raise TypeError(\n f\"Function {module_name} must have an 'execute' classmethod\"\n )\n if not hasattr(module.Function, \"openai_schema\"):\n raise TypeError(\n f\"Function {module_name} must have an 'openai_schema' classmethod\"\n )\n\n return module\n\n\nfunctions_folder = Path(cfg.get(\"OPENAI_FUNCTIONS_PATH\"))\nfunctions_folder.mkdir(parents=True, exist_ok=True)\nfunctions = [Function(str(path)) for path in functions_folder.glob(\"*.py\")]\n\n\ndef get_function(name: str) -> Callable[..., Any]:\n for function in functions:\n if function.name == name:\n return function.execute\n raise ValueError(f\"Function {name} not found\")\n\n\ndef get_openai_schemas() -> List[Dict[str, Any]]:\n return [function.openai_schema for function in functions]\n" }, { "path": "sgpt/handlers/__init__.py", "content": "" }, { "path": "sgpt/handlers/chat_handler.py", "content": "import json\nfrom pathlib import Path\nfrom typing import Any, Callable, Dict, Generator, List, Optional\n\nimport typer\nfrom click import BadParameter, UsageError\nfrom rich.console import Console\nfrom rich.markdown import Markdown\n\nfrom ..config import cfg\nfrom ..role import DefaultRoles, SystemRole\nfrom ..utils import option_callback\nfrom .handler import Handler\n\nCHAT_CACHE_LENGTH = int(cfg.get(\"CHAT_CACHE_LENGTH\"))\nCHAT_CACHE_PATH = Path(cfg.get(\"CHAT_CACHE_PATH\"))\n\n\nclass ChatSession:\n \"\"\"\n This class is used as a decorator for OpenAI chat API requests.\n The ChatSession class caches chat messages and keeps track of the\n conversation history. It is designed to store cached messages\n in a specified directory and in JSON format.\n \"\"\"\n\n def __init__(self, length: int, storage_path: Path):\n \"\"\"\n Initialize the ChatSession decorator.\n\n :param length: Integer, maximum number of cached messages to keep.\n \"\"\"\n self.length = length\n self.storage_path = storage_path\n self.storage_path.mkdir(parents=True, exist_ok=True)\n\n def __call__(self, func: Callable[..., Any]) -> Callable[..., Any]:\n \"\"\"\n The Cache decorator.\n\n :param func: The chat function to cache.\n :return: Wrapped function with chat caching.\n \"\"\"\n\n def wrapper(*args: Any, **kwargs: Any) -> Generator[str, None, None]:\n chat_id = kwargs.pop(\"chat_id\", None)\n if not kwargs.get(\"messages\"):\n return\n if not chat_id:\n yield from func(*args, **kwargs)\n return\n previous_messages = self._read(chat_id)\n for message in kwargs[\"messages\"]:\n previous_messages.append(message)\n kwargs[\"messages\"] = previous_messages\n response_text = \"\"\n for word in func(*args, **kwargs):\n response_text += word\n yield word\n previous_messages.append({\"role\": \"assistant\", \"content\": response_text})\n self._write(kwargs[\"messages\"], chat_id)\n\n return wrapper\n\n def _read(self, chat_id: str) -> List[Dict[str, str]]:\n file_path = self.storage_path / chat_id\n if not file_path.exists():\n return []\n parsed_cache = json.loads(file_path.read_text())\n return parsed_cache if isinstance(parsed_cache, list) else []\n\n def _write(self, messages: List[Dict[str, str]], chat_id: str) -> None:\n file_path = self.storage_path / chat_id\n # Retain the first message since it defines the role\n truncated_messages = (\n messages[:1] + messages[1 + max(0, len(messages) - self.length) :]\n )\n json.dump(truncated_messages, file_path.open(\"w\"))\n\n def invalidate(self, chat_id: str) -> None:\n file_path = self.storage_path / chat_id\n file_path.unlink(missing_ok=True)\n\n def get_messages(self, chat_id: str) -> List[str]:\n messages = self._read(chat_id)\n return [f\"{message['role']}: {message['content']}\" for message in messages]\n\n def exists(self, chat_id: Optional[str]) -> bool:\n return bool(chat_id and bool(self._read(chat_id)))\n\n def list(self) -> List[Path]:\n # Get all files in the folder.\n files = self.storage_path.glob(\"*\")\n # Sort files by last modification time in ascending order.\n return sorted(files, key=lambda f: f.stat().st_mtime)\n\n\nclass ChatHandler(Handler):\n chat_session = ChatSession(CHAT_CACHE_LENGTH, CHAT_CACHE_PATH)\n\n def __init__(self, chat_id: str, role: SystemRole, markdown: bool) -> None:\n super().__init__(role, markdown)\n self.chat_id = chat_id\n self.role = role\n\n if chat_id == \"temp\":\n # If the chat id is \"temp\", we don't want to save the chat session.\n self.chat_session.invalidate(chat_id)\n\n self.validate()\n\n @property\n def initiated(self) -> bool:\n return self.chat_session.exists(self.chat_id)\n\n @property\n def is_same_role(self) -> bool:\n # TODO: Should be optimized for REPL mode.\n return self.role.same_role(self.initial_message(self.chat_id))\n\n @classmethod\n def initial_message(cls, chat_id: str) -> str:\n chat_history = cls.chat_session.get_messages(chat_id)\n return chat_history[0] if chat_history else \"\"\n\n @classmethod\n @option_callback\n def list_ids(cls, value: str) -> None:\n # Prints all existing chat IDs to the console.\n for chat_id in cls.chat_session.list():\n typer.echo(chat_id)\n\n @classmethod\n def show_messages(cls, chat_id: str, markdown: bool) -> None:\n color = cfg.get(\"DEFAULT_COLOR\")\n if \"APPLY MARKDOWN\" in cls.initial_message(chat_id) and markdown:\n theme = cfg.get(\"CODE_THEME\")\n for message in cls.chat_session.get_messages(chat_id):\n if message.startswith(\"assistant:\"):\n Console().print(Markdown(message, code_theme=theme))\n else:\n typer.secho(message, fg=color)\n typer.echo()\n return\n\n for index, message in enumerate(cls.chat_session.get_messages(chat_id)):\n running_color = color if index % 2 == 0 else \"green\"\n typer.secho(message, fg=running_color)\n\n def validate(self) -> None:\n if self.initiated:\n chat_role_name = self.role.get_role_name(self.initial_message(self.chat_id))\n if not chat_role_name:\n raise BadParameter(f'Could not determine chat role of \"{self.chat_id}\"')\n if self.role.name == DefaultRoles.DEFAULT.value:\n # If user didn't pass chat mode, we will use the one that was used to initiate the chat.\n self.role = SystemRole.get(chat_role_name)\n else:\n if not self.is_same_role:\n raise UsageError(\n f'Cant change chat role to \"{self.role.name}\" '\n f'since it was initiated as \"{chat_role_name}\" chat.'\n )\n\n def make_messages(self, prompt: str) -> List[Dict[str, str]]:\n messages = []\n if not self.initiated:\n messages.append({\"role\": \"system\", \"content\": self.role.role})\n messages.append({\"role\": \"user\", \"content\": prompt})\n return messages\n\n @chat_session\n def get_completion(self, **kwargs: Any) -> Generator[str, None, None]:\n yield from super().get_completion(**kwargs)\n\n def handle(self, **kwargs: Any) -> str: # type: ignore[override]\n return super().handle(**kwargs, chat_id=self.chat_id)\n" }, { "path": "sgpt/handlers/default_handler.py", "content": "from pathlib import Path\nfrom typing import Dict, List\n\nfrom ..config import cfg\nfrom ..role import SystemRole\nfrom .handler import Handler\n\nCHAT_CACHE_LENGTH = int(cfg.get(\"CHAT_CACHE_LENGTH\"))\nCHAT_CACHE_PATH = Path(cfg.get(\"CHAT_CACHE_PATH\"))\n\n\nclass DefaultHandler(Handler):\n def __init__(self, role: SystemRole, markdown: bool) -> None:\n super().__init__(role, markdown)\n self.role = role\n\n def make_messages(self, prompt: str) -> List[Dict[str, str]]:\n messages = [\n {\"role\": \"system\", \"content\": self.role.role},\n {\"role\": \"user\", \"content\": prompt},\n ]\n return messages\n" }, { "path": "sgpt/handlers/handler.py", "content": "import json\nfrom pathlib import Path\nfrom typing import Any, Callable, Dict, Generator, List, Optional\n\nfrom ..cache import Cache\nfrom ..config import cfg\nfrom ..function import get_function\nfrom ..printer import MarkdownPrinter, Printer, TextPrinter\nfrom ..role import DefaultRoles, SystemRole\n\ncompletion: Callable[..., Any] = lambda *args, **kwargs: Generator[Any, None, None]\n\nbase_url = cfg.get(\"API_BASE_URL\")\nuse_litellm = cfg.get(\"USE_LITELLM\") == \"true\"\nadditional_kwargs = {\n \"timeout\": int(cfg.get(\"REQUEST_TIMEOUT\")),\n \"api_key\": cfg.get(\"OPENAI_API_KEY\"),\n \"base_url\": None if base_url == \"default\" else base_url,\n}\n\nif use_litellm:\n import litellm # type: ignore\n\n completion = litellm.completion\n litellm.suppress_debug_info = True\n additional_kwargs.pop(\"api_key\")\nelse:\n from openai import OpenAI\n\n client = OpenAI(**additional_kwargs) # type: ignore\n completion = client.chat.completions.create\n additional_kwargs = {}\n\n\nclass Handler:\n cache = Cache(int(cfg.get(\"CACHE_LENGTH\")), Path(cfg.get(\"CACHE_PATH\")))\n\n def __init__(self, role: SystemRole, markdown: bool) -> None:\n self.role = role\n\n api_base_url = cfg.get(\"API_BASE_URL\")\n self.base_url = None if api_base_url == \"default\" else api_base_url\n self.timeout = int(cfg.get(\"REQUEST_TIMEOUT\"))\n\n self.markdown = \"APPLY MARKDOWN\" in self.role.role and markdown\n self.code_theme, self.color = cfg.get(\"CODE_THEME\"), cfg.get(\"DEFAULT_COLOR\")\n\n @property\n def printer(self) -> Printer:\n return (\n MarkdownPrinter(self.code_theme)\n if self.markdown\n else TextPrinter(self.color)\n )\n\n def make_messages(self, prompt: str) -> List[Dict[str, str]]:\n raise NotImplementedError\n\n def handle_function_call(\n self,\n messages: List[dict[str, Any]],\n tool_call_id: str,\n name: str,\n arguments: str,\n ) -> Generator[str, None, None]:\n # Add assistant message with tool call\n messages.append(\n {\n \"role\": \"assistant\",\n \"content\": None,\n \"tool_calls\": [\n {\n \"id\": tool_call_id,\n \"type\": \"function\",\n \"function\": {\"name\": name, \"arguments\": arguments},\n }\n ],\n }\n )\n\n if messages and messages[-1][\"role\"] == \"assistant\":\n yield \"\\n\"\n\n dict_args = json.loads(arguments)\n joined_args = \", \".join(f'{k}=\"{v}\"' for k, v in dict_args.items())\n yield f\"> @FunctionCall `{name}({joined_args})` \\n\\n\"\n\n result = get_function(name)(**dict_args)\n if cfg.get(\"SHOW_FUNCTIONS_OUTPUT\") == \"true\":\n yield f\"```text\\n{result}\\n```\\n\"\n\n # Add tool response message\n messages.append(\n {\"role\": \"tool\", \"content\": result, \"tool_call_id\": tool_call_id}\n )\n\n @cache\n def get_completion(\n self,\n model: str,\n temperature: float,\n top_p: float,\n messages: List[Dict[str, Any]],\n functions: Optional[List[Dict[str, str]]],\n ) -> Generator[str, None, None]:\n tool_call_id = name = arguments = \"\"\n is_shell_role = self.role.name == DefaultRoles.SHELL.value\n is_code_role = self.role.name == DefaultRoles.CODE.value\n is_dsc_shell_role = self.role.name == DefaultRoles.DESCRIBE_SHELL.value\n if is_shell_role or is_code_role or is_dsc_shell_role:\n functions = None\n\n if functions:\n additional_kwargs[\"tool_choice\"] = \"auto\"\n additional_kwargs[\"tools\"] = functions\n additional_kwargs[\"parallel_tool_calls\"] = False\n\n response = completion(\n model=model,\n temperature=temperature,\n top_p=top_p,\n messages=messages,\n stream=True,\n **additional_kwargs,\n )\n\n try:\n for chunk in response:\n if not chunk.choices:\n continue\n delta = chunk.choices[0].delta\n\n # LiteLLM uses dict instead of Pydantic object like OpenAI does.\n tool_calls = (\n delta.get(\"tool_calls\") if use_litellm else delta.tool_calls\n )\n if tool_calls:\n for tool_call in tool_calls:\n if use_litellm:\n # TODO: test.\n tool_call_id = tool_call.get(\"id\") or tool_call_id\n name = tool_call.get(\"function\", {}).get(\"name\") or name\n arguments += tool_call.get(\"function\", {}).get(\n \"arguments\", \"\"\n )\n else:\n tool_call_id = tool_call.id or tool_call_id\n name = tool_call.function.name or name\n arguments += tool_call.function.arguments or \"\"\n if chunk.choices[0].finish_reason == \"tool_calls\":\n yield from self.handle_function_call(\n messages, tool_call_id, name, arguments\n )\n yield from self.get_completion(\n model=model,\n temperature=temperature,\n top_p=top_p,\n messages=messages,\n functions=functions,\n caching=False,\n )\n return\n\n yield delta.content or \"\"\n except KeyboardInterrupt:\n response.close()\n\n def handle(\n self,\n prompt: str,\n model: str,\n temperature: float,\n top_p: float,\n caching: bool,\n functions: Optional[List[Dict[str, str]]] = None,\n **kwargs: Any,\n ) -> str:\n disable_stream = cfg.get(\"DISABLE_STREAMING\") == \"true\"\n messages = self.make_messages(prompt.strip())\n generator = self.get_completion(\n model=model,\n temperature=temperature,\n top_p=top_p,\n messages=messages,\n functions=functions,\n caching=caching,\n **kwargs,\n )\n return self.printer(generator, not disable_stream)\n" }, { "path": "sgpt/handlers/repl_handler.py", "content": "from typing import Any\n\nimport typer\nfrom rich import print as rich_print\nfrom rich.rule import Rule\n\nfrom ..role import DefaultRoles, SystemRole\nfrom ..utils import run_command\nfrom .chat_handler import ChatHandler\nfrom .default_handler import DefaultHandler\n\n\nclass ReplHandler(ChatHandler):\n def __init__(self, chat_id: str, role: SystemRole, markdown: bool) -> None:\n super().__init__(chat_id, role, markdown)\n\n @classmethod\n def _get_multiline_input(cls) -> str:\n multiline_input = \"\"\n while (user_input := typer.prompt(\"...\", prompt_suffix=\"\")) != '\"\"\"':\n multiline_input += user_input + \"\\n\"\n return multiline_input\n\n def handle(self, init_prompt: str, **kwargs: Any) -> None: # type: ignore\n if self.initiated:\n rich_print(Rule(title=\"Chat History\", style=\"bold magenta\"))\n self.show_messages(self.chat_id, self.markdown)\n rich_print(Rule(style=\"bold magenta\"))\n\n info_message = (\n \"Entering REPL mode, press Ctrl+C to exit.\"\n if not self.role.name == DefaultRoles.SHELL.value\n else (\n \"Entering shell REPL mode, type [e] to execute commands \"\n \"or [d] to describe the commands, press Ctrl+C to exit.\"\n )\n )\n typer.secho(info_message, fg=\"yellow\")\n\n if init_prompt:\n rich_print(Rule(title=\"Input\", style=\"bold purple\"))\n typer.echo(init_prompt)\n rich_print(Rule(style=\"bold purple\"))\n\n full_completion = \"\"\n while True:\n # Infinite loop until user exits with Ctrl+C.\n prompt = typer.prompt(\">>>\", prompt_suffix=\" \")\n if prompt == '\"\"\"':\n prompt = self._get_multiline_input()\n if prompt == \"exit()\":\n raise typer.Exit()\n if init_prompt:\n prompt = f\"{init_prompt}\\n\\n\\n{prompt}\"\n init_prompt = \"\"\n if self.role.name == DefaultRoles.SHELL.value and prompt == \"e\":\n typer.echo()\n run_command(full_completion)\n typer.echo()\n rich_print(Rule(style=\"bold magenta\"))\n elif self.role.name == DefaultRoles.SHELL.value and prompt == \"d\":\n DefaultHandler(\n DefaultRoles.DESCRIBE_SHELL.get_role(), self.markdown\n ).handle(prompt=full_completion, **kwargs)\n else:\n full_completion = super().handle(prompt=prompt, **kwargs)\n" }, { "path": "sgpt/integration.py", "content": "bash_integration = \"\"\"\n# Shell-GPT integration BASH v0.2\n_sgpt_bash() {\nif [[ -n \"$READLINE_LINE\" ]]; then\n READLINE_LINE=$(sgpt --shell <<< \"$READLINE_LINE\" --no-interaction)\n READLINE_POINT=${#READLINE_LINE}\nfi\n}\nbind -x '\"\\\\C-l\": _sgpt_bash'\n# Shell-GPT integration BASH v0.2\n\"\"\"\n\nzsh_integration = \"\"\"\n# Shell-GPT integration ZSH v0.2\n_sgpt_zsh() {\nif [[ -n \"$BUFFER\" ]]; then\n _sgpt_prev_cmd=$BUFFER\n BUFFER+=\"⌛\"\n zle -I && zle redisplay\n BUFFER=$(sgpt --shell <<< \"$_sgpt_prev_cmd\" --no-interaction)\n zle end-of-line\nfi\n}\nzle -N _sgpt_zsh\nbindkey ^l _sgpt_zsh\n# Shell-GPT integration ZSH v0.2\n\"\"\"\n" }, { "path": "sgpt/llm_functions/__init__.py", "content": "" }, { "path": "sgpt/llm_functions/common/execute_shell.py", "content": "import subprocess\nfrom typing import Any, Dict\n\nfrom pydantic import BaseModel, Field\n\n\nclass Function(BaseModel):\n \"\"\"\n Executes a shell command and returns the output (result).\n \"\"\"\n\n shell_command: str = Field(\n ...,\n example=\"ls -la\",\n description=\"Shell command to execute.\",\n ) # type: ignore\n\n @classmethod\n def execute(cls, shell_command: str) -> str:\n process = subprocess.Popen(\n shell_command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT\n )\n output, _ = process.communicate()\n exit_code = process.returncode\n return f\"Exit code: {exit_code}, Output:\\n{output.decode()}\"\n\n @classmethod\n def openai_schema(cls) -> Dict[str, Any]:\n \"\"\"Generate OpenAI function schema from Pydantic model.\"\"\"\n schema = cls.model_json_schema()\n return {\n \"type\": \"function\",\n \"function\": {\n \"name\": \"execute_shell_command\",\n \"description\": cls.__doc__.strip() if cls.__doc__ else \"\",\n \"parameters\": {\n \"type\": \"object\",\n \"properties\": schema.get(\"properties\", {}),\n \"required\": schema.get(\"required\", []),\n },\n },\n }\n" }, { "path": "sgpt/llm_functions/init_functions.py", "content": "import os\nimport platform\nimport shutil\nfrom pathlib import Path\nfrom typing import Any\n\nfrom ..config import cfg\nfrom ..utils import option_callback\n\nFUNCTIONS_FOLDER = Path(cfg.get(\"OPENAI_FUNCTIONS_PATH\"))\n\n\n@option_callback\ndef install_functions(*_args: Any) -> None:\n current_folder = os.path.dirname(os.path.abspath(__file__))\n common_folder = Path(current_folder + \"/common\")\n common_files = [Path(path) for path in common_folder.glob(\"*.py\")]\n print(\"Installing default functions...\")\n\n for file in common_files:\n print(f\"Installed {FUNCTIONS_FOLDER}/{file.name}\")\n shutil.copy(file, FUNCTIONS_FOLDER, follow_symlinks=True)\n\n current_platform = platform.system()\n if current_platform == \"Linux\":\n print(\"Installing Linux functions...\")\n if current_platform == \"Windows\":\n print(\"Installing Windows functions...\")\n if current_platform == \"Darwin\":\n print(\"Installing Mac functions...\")\n mac_folder = Path(current_folder + \"/mac\")\n mac_files = [Path(path) for path in mac_folder.glob(\"*.py\")]\n for file in mac_files:\n print(f\"Installed {FUNCTIONS_FOLDER}/{file.name}\")\n shutil.copy(file, FUNCTIONS_FOLDER, follow_symlinks=True)\n" }, { "path": "sgpt/llm_functions/mac/apple_script.py", "content": "import subprocess\nfrom typing import Any, Dict\n\nfrom pydantic import BaseModel, Field\n\n\nclass Function(BaseModel):\n \"\"\"\n Executes Apple Script on macOS and returns the output (result).\n Can be used for actions like: draft (prepare) an email, show calendar events, create a note.\n \"\"\"\n\n apple_script: str = Field(\n default=...,\n example='tell application \"Finder\" to get the name of every disk',\n description=\"Apple Script to execute.\",\n ) # type: ignore\n\n @classmethod\n def execute(cls, apple_script):\n script_command = [\"osascript\", \"-e\", apple_script]\n try:\n process = subprocess.Popen(\n script_command, stdout=subprocess.PIPE, stderr=subprocess.PIPE\n )\n output, _ = process.communicate()\n output = output.decode(\"utf-8\").strip()\n return f\"Output: {output}\"\n except Exception as e:\n return f\"Error: {e}\"\n\n @classmethod\n def openai_schema(cls) -> Dict[str, Any]:\n \"\"\"Generate OpenAI function schema from Pydantic model.\"\"\"\n schema = cls.model_json_schema()\n return {\n \"type\": \"function\",\n \"function\": {\n \"name\": \"execute_apple_script\",\n \"description\": cls.__doc__.strip() if cls.__doc__ else \"\",\n \"parameters\": {\n \"type\": \"object\",\n \"properties\": schema.get(\"properties\", {}),\n \"required\": schema.get(\"required\", []),\n },\n },\n }\n" }, { "path": "sgpt/printer.py", "content": "from abc import ABC, abstractmethod\nfrom typing import Generator\n\nfrom rich.console import Console\nfrom rich.live import Live\nfrom rich.markdown import Markdown\nfrom typer import secho\n\n\nclass Printer(ABC):\n console = Console()\n\n @abstractmethod\n def live_print(self, chunks: Generator[str, None, None]) -> str:\n pass\n\n @abstractmethod\n def static_print(self, text: str) -> str:\n pass\n\n def __call__(self, chunks: Generator[str, None, None], live: bool = True) -> str:\n if live:\n return self.live_print(chunks)\n with self.console.status(\"[bold green]Loading...\"):\n full_completion = \"\".join(chunks)\n self.static_print(full_completion)\n return full_completion\n\n\nclass MarkdownPrinter(Printer):\n def __init__(self, theme: str) -> None:\n self.console = Console()\n self.theme = theme\n\n def live_print(self, chunks: Generator[str, None, None]) -> str:\n full_completion = \"\"\n with Live(console=self.console) as live:\n for chunk in chunks:\n full_completion += chunk\n markdown = Markdown(markup=full_completion, code_theme=self.theme)\n live.update(markdown, refresh=True)\n return full_completion\n\n def static_print(self, text: str) -> str:\n markdown = Markdown(markup=text, code_theme=self.theme)\n self.console.print(markdown)\n return text\n\n\nclass TextPrinter(Printer):\n def __init__(self, color: str) -> None:\n self.color = color\n\n def live_print(self, chunks: Generator[str, None, None]) -> str:\n full_text = \"\"\n for chunk in chunks:\n full_text += chunk\n secho(chunk, fg=self.color, nl=False)\n else:\n print() # Add new line after last chunk.\n return full_text\n\n def static_print(self, text: str) -> str:\n secho(text, fg=self.color)\n return text\n" }, { "path": "sgpt/role.py", "content": "import json\nimport platform\nfrom enum import Enum\nfrom os import getenv, pathsep\nfrom os.path import basename\nfrom pathlib import Path\nfrom typing import Dict, Optional\n\nimport typer\nfrom click import UsageError\nfrom distro import name as distro_name\n\nfrom .config import cfg\nfrom .utils import option_callback\n\nSHELL_ROLE = \"\"\"Provide only {shell} commands for {os} without any description.\nIf there is a lack of details, provide most logical solution.\nEnsure the output is a valid shell command.\nIf multiple steps required try to combine them together using &&.\nProvide only plain text without Markdown formatting.\nDo not provide markdown formatting such as ```.\n\"\"\"\n\nDESCRIBE_SHELL_ROLE = \"\"\"Provide a terse, single sentence description of the given shell command.\nDescribe each argument and option of the command.\nProvide short responses in about 80 words.\nAPPLY MARKDOWN formatting when possible.\"\"\"\n# Note that output for all roles containing \"APPLY MARKDOWN\" will be formatted as Markdown.\n\nCODE_ROLE = \"\"\"Provide only code as output without any description.\nProvide only code in plain text format without Markdown formatting.\nDo not include symbols such as ``` or ```python.\nIf there is a lack of details, provide most logical solution.\nYou are not allowed to ask for more details.\nFor example if the prompt is \"Hello world Python\", you should return \"print('Hello world')\".\"\"\"\n\nDEFAULT_ROLE = \"\"\"You are programming and system administration assistant.\nYou are managing {os} operating system with {shell} shell.\nProvide short responses in about 100 words, unless you are specifically asked for more details.\nIf you need to store any data, assume it will be stored in the conversation.\nAPPLY MARKDOWN formatting when possible.\"\"\"\n# Note that output for all roles containing \"APPLY MARKDOWN\" will be formatted as Markdown.\n\nROLE_TEMPLATE = \"You are {name}\\n{role}\"\n\n\nclass SystemRole:\n storage: Path = Path(cfg.get(\"ROLE_STORAGE_PATH\"))\n\n def __init__(\n self,\n name: str,\n role: str,\n variables: Optional[Dict[str, str]] = None,\n ) -> None:\n self.storage.mkdir(parents=True, exist_ok=True)\n self.name = name\n if variables:\n role = role.format(**variables)\n self.role = role\n\n @classmethod\n def create_defaults(cls) -> None:\n cls.storage.parent.mkdir(parents=True, exist_ok=True)\n variables = {\"shell\": cls._shell_name(), \"os\": cls._os_name()}\n for default_role in (\n SystemRole(\"ShellGPT\", DEFAULT_ROLE, variables),\n SystemRole(\"Shell Command Generator\", SHELL_ROLE, variables),\n SystemRole(\"Shell Command Descriptor\", DESCRIBE_SHELL_ROLE, variables),\n SystemRole(\"Code Generator\", CODE_ROLE),\n ):\n if not default_role._exists:\n default_role._save()\n\n @classmethod\n def get(cls, name: str) -> \"SystemRole\":\n file_path = cls.storage / f\"{name}.json\"\n if not file_path.exists():\n raise UsageError(f'Role \"{name}\" not found.')\n return cls(**json.loads(file_path.read_text()))\n\n @classmethod\n @option_callback\n def create(cls, name: str) -> None:\n role = typer.prompt(\"Enter role description\")\n role = cls(name, role)\n role._save()\n\n @classmethod\n @option_callback\n def list(cls, _value: str) -> None:\n if not cls.storage.exists():\n return\n # Get all files in the folder.\n files = cls.storage.glob(\"*\")\n # Sort files by last modification time in ascending order.\n for path in sorted(files, key=lambda f: f.stat().st_mtime):\n typer.echo(path)\n\n @classmethod\n @option_callback\n def show(cls, name: str) -> None:\n typer.echo(cls.get(name).role)\n\n @classmethod\n def get_role_name(cls, initial_message: str) -> Optional[str]:\n if not initial_message:\n return None\n message_lines = initial_message.splitlines()\n if \"You are\" in message_lines[0]:\n return message_lines[0].split(\"You are \")[1].strip()\n return None\n\n @classmethod\n def _os_name(cls) -> str:\n if cfg.get(\"OS_NAME\") != \"auto\":\n return cfg.get(\"OS_NAME\")\n current_platform = platform.system()\n if current_platform == \"Linux\":\n return \"Linux/\" + distro_name(pretty=True)\n if current_platform == \"Windows\":\n return \"Windows \" + platform.release()\n if current_platform == \"Darwin\":\n return \"Darwin/MacOS \" + platform.mac_ver()[0]\n return current_platform\n\n @classmethod\n def _shell_name(cls) -> str:\n if cfg.get(\"SHELL_NAME\") != \"auto\":\n return cfg.get(\"SHELL_NAME\")\n current_platform = platform.system()\n if current_platform in (\"Windows\", \"nt\"):\n is_powershell = len(getenv(\"PSModulePath\", \"\").split(pathsep)) >= 3\n return \"powershell.exe\" if is_powershell else \"cmd.exe\"\n return basename(getenv(\"SHELL\", \"/bin/sh\"))\n\n @property\n def _exists(self) -> bool:\n return self._file_path.exists()\n\n @property\n def _file_path(self) -> Path:\n return self.storage / f\"{self.name}.json\"\n\n def _save(self) -> None:\n if self._exists:\n typer.confirm(\n f'Role \"{self.name}\" already exists, overwrite it?',\n abort=True,\n )\n\n self.role = ROLE_TEMPLATE.format(name=self.name, role=self.role)\n self._file_path.write_text(json.dumps(self.__dict__), encoding=\"utf-8\")\n\n def delete(self) -> None:\n if self._exists:\n typer.confirm(\n f'Role \"{self.name}\" exist, delete it?',\n abort=True,\n )\n self._file_path.unlink()\n\n def same_role(self, initial_message: str) -> bool:\n if not initial_message:\n return False\n return True if f\"You are {self.name}\" in initial_message else False\n\n\nclass DefaultRoles(Enum):\n DEFAULT = \"ShellGPT\"\n SHELL = \"Shell Command Generator\"\n DESCRIBE_SHELL = \"Shell Command Descriptor\"\n CODE = \"Code Generator\"\n\n @classmethod\n def check_get(cls, shell: bool, describe_shell: bool, code: bool) -> SystemRole:\n if shell:\n return SystemRole.get(DefaultRoles.SHELL.value)\n if describe_shell:\n return SystemRole.get(DefaultRoles.DESCRIBE_SHELL.value)\n if code:\n return SystemRole.get(DefaultRoles.CODE.value)\n return SystemRole.get(DefaultRoles.DEFAULT.value)\n\n def get_role(self) -> SystemRole:\n return SystemRole.get(self.value)\n\n\nSystemRole.create_defaults()\n" }, { "path": "sgpt/utils.py", "content": "import os\nimport platform\nimport shlex\nfrom tempfile import NamedTemporaryFile\nfrom typing import Any, Callable\n\nimport typer\nfrom click import BadParameter, UsageError\n\nfrom sgpt.__version__ import __version__\nfrom sgpt.integration import bash_integration, zsh_integration\n\n\ndef get_edited_prompt() -> str:\n \"\"\"\n Opens the user's default editor to let them\n input a prompt, and returns the edited text.\n\n :return: String prompt.\n \"\"\"\n with NamedTemporaryFile(suffix=\".txt\", delete=False) as file:\n # Create file and store path.\n file_path = file.name\n editor = os.environ.get(\"EDITOR\", \"vim\")\n # This will write text to file using $EDITOR.\n os.system(f\"{editor} {file_path}\")\n # Read file when editor is closed.\n with open(file_path, \"r\", encoding=\"utf-8\") as file:\n output = file.read()\n os.remove(file_path)\n if not output:\n raise BadParameter(\"Couldn't get valid PROMPT from $EDITOR\")\n return output\n\n\ndef run_command(command: str) -> None:\n \"\"\"\n Runs a command in the user's shell.\n It is aware of the current user's $SHELL.\n :param command: A shell command to run.\n \"\"\"\n if platform.system() == \"Windows\":\n is_powershell = len(os.getenv(\"PSModulePath\", \"\").split(os.pathsep)) >= 3\n full_command = (\n f'powershell.exe -Command \"{command}\"'\n if is_powershell\n else f'cmd.exe /c \"{command}\"'\n )\n else:\n shell = os.environ.get(\"SHELL\", \"/bin/sh\")\n full_command = f\"{shell} -c {shlex.quote(command)}\"\n\n os.system(full_command)\n\n\ndef option_callback(func: Callable) -> Callable: # type: ignore\n def wrapper(cls: Any, value: str) -> None:\n if not value:\n return\n func(cls, value)\n raise typer.Exit()\n\n return wrapper\n\n\n@option_callback\ndef install_shell_integration(*_args: Any) -> None:\n \"\"\"\n Installs shell integration. Currently only supports ZSH and Bash.\n Allows user to get shell completions in terminal by using hotkey.\n Replaces current \"buffer\" of the shell with the completion.\n \"\"\"\n # TODO: Add support for Windows.\n # TODO: Implement updates.\n shell = os.getenv(\"SHELL\", \"\")\n if \"zsh\" in shell:\n typer.echo(\"Installing ZSH integration...\")\n with open(os.path.expanduser(\"~/.zshrc\"), \"a\", encoding=\"utf-8\") as file:\n file.write(zsh_integration)\n elif \"bash\" in shell:\n typer.echo(\"Installing Bash integration...\")\n with open(os.path.expanduser(\"~/.bashrc\"), \"a\", encoding=\"utf-8\") as file:\n file.write(bash_integration)\n else:\n raise UsageError(\"ShellGPT integrations only available for ZSH and Bash.\")\n\n typer.echo(\"Done! Restart your shell to apply changes.\")\n\n\n@option_callback\ndef get_sgpt_version(*_args: Any) -> None:\n \"\"\"\n Displays the current installed version of ShellGPT\n \"\"\"\n typer.echo(f\"ShellGPT {__version__}\")\n" }, { "path": "tests/__init__.py", "content": "" }, { "path": "tests/_integration.py", "content": "\"\"\"\nThis test module will execute real commands using shell.\nThis means it will call sgpt.py with command line arguments.\nMake sure you have your API key in place ~/.cfg/shell_gpt/.sgptrc\nor ENV variable OPENAI_API_KEY.\nIt is useful for quick tests, saves a bit time.\n\"\"\"\n\nimport json\nimport os\nimport subprocess\nfrom pathlib import Path\nfrom tempfile import NamedTemporaryFile\nfrom unittest import TestCase\nfrom unittest.mock import ANY, patch\nfrom uuid import uuid4\n\nimport typer\nfrom typer.testing import CliRunner\n\nfrom sgpt.__version__ import __version__\nfrom sgpt.app import main\nfrom sgpt.config import cfg\nfrom sgpt.handlers.handler import Handler\nfrom sgpt.role import SystemRole\n\nrunner = CliRunner()\napp = typer.Typer()\napp.command()(main)\n\n\nclass TestShellGpt(TestCase):\n @classmethod\n def setUpClass(cls):\n # Response streaming should be enabled for these tests.\n assert cfg.get(\"DISABLE_STREAMING\") == \"false\"\n # ShellGPT optimised and tested with gpt-4 turbo.\n assert cfg.get(\"DEFAULT_MODEL\") == \"gpt-4o\"\n # Make sure we will not call any functions.\n assert cfg.get(\"OPENAI_USE_FUNCTIONS\") == \"false\"\n\n @staticmethod\n def get_arguments(prompt, **kwargs):\n arguments = [prompt]\n for key, value in kwargs.items():\n arguments.append(key)\n if isinstance(value, bool):\n continue\n arguments.append(value)\n arguments.append(\"--no-cache\")\n return arguments\n\n def test_default(self):\n dict_arguments = {\n \"prompt\": \"What is the capital of the Czech Republic?\",\n }\n result = runner.invoke(app, self.get_arguments(**dict_arguments))\n assert result.exit_code == 0\n assert \"Prague\" in result.output\n\n def test_shell(self):\n dict_arguments = {\n \"prompt\": \"make a commit using git\",\n \"--shell\": True,\n }\n result = runner.invoke(app, self.get_arguments(**dict_arguments))\n assert result.exit_code == 0\n assert \"git commit\" in result.output\n\n def test_describe_shell(self):\n dict_arguments = {\n \"prompt\": \"ls\",\n \"--describe-shell\": True,\n }\n result = runner.invoke(app, self.get_arguments(**dict_arguments))\n assert result.exit_code == 0\n assert \"lists\" in result.output.lower()\n\n def test_code(self):\n \"\"\"\n This test will request from OpenAI API a python code to make CLI app,\n which will be written to a temp file, and then it will be executed\n in shell with two positional int arguments. As the output we are\n expecting the result of multiplying them.\n \"\"\"\n dict_arguments = {\n \"prompt\": (\n \"Create a command line application using Python that \"\n \"accepts two positional arguments \"\n \"and prints the result of multiplying them.\"\n ),\n \"--code\": True,\n }\n result = runner.invoke(app, self.get_arguments(**dict_arguments))\n assert result.exit_code == 0\n print(result.output)\n # Since output will be slightly different, there is no way how to test it precisely.\n assert \"print\" in result.output\n assert \"*\" in result.output\n with NamedTemporaryFile(\"w+\", delete=False) as file:\n try:\n compile(result.output, file.name, \"exec\")\n except SyntaxError:\n assert False, \"The output is not valid Python code.\" # noqa: B011\n file.seek(0)\n file.truncate()\n file.write(result.output)\n file_path = file.name\n number_a = number_b = 2\n # Execute output code in the shell with arguments.\n arguments = [\"python\", file.name, str(number_a), str(number_b)]\n script_output = subprocess.run(arguments, stdout=subprocess.PIPE, check=True)\n os.remove(file_path)\n assert script_output.stdout.decode().strip(), number_a * number_b\n\n def test_chat_default(self):\n chat_name = uuid4()\n dict_arguments = {\n \"prompt\": \"Remember my favorite number: 6\",\n \"--chat\": f\"test_{chat_name}\",\n \"--no-cache\": True,\n }\n runner.invoke(app, self.get_arguments(**dict_arguments))\n dict_arguments[\"prompt\"] = \"What is my favorite number + 2?\"\n result = runner.invoke(app, self.get_arguments(**dict_arguments))\n assert result.exit_code == 0\n assert \"8\" in result.output\n dict_arguments[\"--shell\"] = True\n result = runner.invoke(app, self.get_arguments(**dict_arguments))\n assert result.exit_code == 2\n dict_arguments[\"--code\"] = True\n result = runner.invoke(app, self.get_arguments(**dict_arguments))\n # If we have default chat, we cannot use --code or --shell.\n assert result.exit_code == 2\n\n def test_chat_shell(self):\n chat_name = uuid4()\n dict_arguments = {\n \"prompt\": \"Create nginx docker container, forward ports 80, \"\n \"mount current folder with index.html\",\n \"--chat\": f\"test_{chat_name}\",\n \"--shell\": True,\n }\n result = runner.invoke(app, self.get_arguments(**dict_arguments))\n assert result.exit_code == 0\n assert \"docker run\" in result.output\n assert \"-p 80:80\" in result.output\n assert \"nginx\" in result.output\n dict_arguments[\"prompt\"] = \"Also forward port 443.\"\n result = runner.invoke(app, self.get_arguments(**dict_arguments))\n assert result.exit_code == 0\n assert \"-p 80:80\" in result.output\n assert \"-p 443:443\" in result.output\n dict_arguments[\"--code\"] = True\n del dict_arguments[\"--shell\"]\n assert \"--shell\" not in dict_arguments\n result = runner.invoke(app, self.get_arguments(**dict_arguments))\n # If we are using --code, we cannot use --shell.\n assert result.exit_code == 2\n\n def test_chat_describe_shell(self):\n chat_name = uuid4()\n dict_arguments = {\n \"prompt\": \"git add\",\n \"--chat\": f\"test_{chat_name}\",\n \"--describe-shell\": True,\n }\n result = runner.invoke(app, self.get_arguments(**dict_arguments))\n assert result.exit_code == 0\n assert \"adds\" in result.output.lower() or \"stages\" in result.output.lower()\n dict_arguments[\"prompt\"] = \"'-A'\"\n result = runner.invoke(app, self.get_arguments(**dict_arguments))\n assert result.exit_code == 0\n assert \"all\" in result.output\n\n def test_chat_code(self):\n chat_name = uuid4()\n dict_arguments = {\n \"prompt\": \"Using python request localhost:80.\",\n \"--chat\": f\"test_{chat_name}\",\n \"--code\": True,\n }\n result = runner.invoke(app, self.get_arguments(**dict_arguments))\n assert result.exit_code == 0\n assert \"localhost:80\" in result.output\n dict_arguments[\"prompt\"] = \"Change port to 443.\"\n result = runner.invoke(app, self.get_arguments(**dict_arguments))\n assert result.exit_code == 0\n assert \"localhost:443\" in result.output\n del dict_arguments[\"--code\"]\n assert \"--code\" not in dict_arguments\n dict_arguments[\"--shell\"] = True\n result = runner.invoke(app, self.get_arguments(**dict_arguments))\n # If we have --code chat, we cannot use --shell.\n assert result.exit_code == 2\n\n def test_list_chat(self):\n result = runner.invoke(app, [\"--list-chats\"])\n assert result.exit_code == 0\n assert \"test_\" in result.output\n\n def test_show_chat(self):\n chat_name = uuid4()\n dict_arguments = {\n \"prompt\": \"Remember my favorite number: 6\",\n \"--chat\": f\"test_{chat_name}\",\n }\n runner.invoke(app, self.get_arguments(**dict_arguments))\n dict_arguments[\"prompt\"] = \"What is my favorite number + 2?\"\n runner.invoke(app, self.get_arguments(**dict_arguments))\n result = runner.invoke(app, [\"--show-chat\", f\"test_{chat_name}\"])\n assert result.exit_code == 0\n assert \"Remember my favorite number: 6\" in result.output\n assert \"What is my favorite number + 2?\" in result.output\n assert \"8\" in result.output\n\n def test_validation_code_shell(self):\n dict_arguments = {\n \"prompt\": \"What is the capital of the Czech Republic?\",\n \"--code\": True,\n \"--shell\": True,\n }\n result = runner.invoke(app, self.get_arguments(**dict_arguments))\n assert result.exit_code == 2\n assert \"Only one of --shell, --describe-shell, and --code\" in result.output\n\n def test_repl_default(\n self,\n ):\n dict_arguments = {\n \"prompt\": \"\",\n \"--repl\": \"temp\",\n }\n inputs = [\n \"Please remember my favorite number: 6\",\n \"What is my favorite number + 2?\",\n \"exit()\",\n ]\n result = runner.invoke(\n app, self.get_arguments(**dict_arguments), input=\"\\n\".join(inputs)\n )\n assert result.exit_code == 0\n assert \">>> Please remember my favorite number: 6\" in result.output\n assert \">>> What is my favorite number + 2?\" in result.output\n assert \"8\" in result.output\n\n def test_repl_multiline(\n self,\n ):\n dict_arguments = {\n \"prompt\": \"\",\n \"--repl\": \"temp\",\n }\n inputs = [\n '\"\"\"',\n \"Please remember my favorite number: 6\",\n \"What is my favorite number + 2?\",\n '\"\"\"',\n \"exit()\",\n ]\n result = runner.invoke(\n app, self.get_arguments(**dict_arguments), input=\"\\n\".join(inputs)\n )\n\n assert result.exit_code == 0\n assert '\"\"\"' in result.output\n assert \"Please remember my favorite number: 6\" in result.output\n assert \"What is my favorite number + 2?\" in result.output\n assert '\"\"\"' in result.output\n assert \"8\" in result.output\n\n def test_repl_shell(self):\n # Temp chat session from previous test should be overwritten.\n dict_arguments = {\n \"prompt\": \"\",\n \"--repl\": \"temp\",\n \"--shell\": True,\n }\n inputs = [\"What is in current folder?\", \"Simple sort by name\", \"exit()\"]\n result = runner.invoke(\n app, self.get_arguments(**dict_arguments), input=\"\\n\".join(inputs)\n )\n assert result.exit_code == 0\n assert \"type [e] to execute commands\" in result.output\n assert \">>> What is in current folder?\" in result.output\n assert \">>> Simple sort by name\" in result.output\n assert \"ls -la\" in result.output\n assert \"sort\" in result.output\n chat_storage = cfg.get(\"CHAT_CACHE_PATH\")\n tmp_chat = Path(chat_storage) / \"temp\"\n chat_messages = json.loads(tmp_chat.read_text())\n # TODO: Implement same check in chat mode tests.\n assert chat_messages[0][\"content\"].startswith(\"You are Shell Command Generator\")\n assert chat_messages[0][\"role\"] == \"system\"\n assert chat_messages[1][\"content\"].startswith(\"What is in current folder?\")\n assert chat_messages[1][\"role\"] == \"user\"\n assert chat_messages[2][\"content\"] == \"ls -la\"\n assert chat_messages[2][\"role\"] == \"assistant\"\n assert chat_messages[3][\"content\"] == \"Simple sort by name\"\n assert chat_messages[3][\"role\"] == \"user\"\n assert \"sort\" in chat_messages[4][\"content\"]\n assert chat_messages[4][\"role\"] == \"assistant\"\n\n def test_repl_describe_command(self):\n # Temp chat session from previous test should be overwritten.\n dict_arguments = {\n \"prompt\": \"\",\n \"--repl\": \"temp\",\n \"--describe-shell\": True,\n }\n inputs = [\"pacman -S\", \"-yu\", \"exit()\"]\n result = runner.invoke(\n app, self.get_arguments(**dict_arguments), input=\"\\n\".join(inputs)\n )\n assert result.exit_code == 0\n assert \"install\" in result.output.lower()\n assert \"upgrade\" in result.output.lower()\n\n chat_storage = cfg.get(\"CHAT_CACHE_PATH\")\n tmp_chat = Path(chat_storage) / \"temp\"\n chat_messages = json.loads(tmp_chat.read_text())\n assert chat_messages[0][\"content\"].startswith(\n \"You are Shell Command Descriptor\"\n )\n\n def test_repl_code(self):\n dict_arguments = {\n \"prompt\": \"\",\n \"--repl\": f\"test_{uuid4()}\",\n \"--code\": True,\n }\n inputs = (\n \"Using python make request to localhost:8080\",\n \"Change port to 443\",\n \"exit()\",\n )\n result = runner.invoke(\n app, self.get_arguments(**dict_arguments), input=\"\\n\".join(inputs)\n )\n assert result.exit_code == 0\n assert f\">>> {inputs[0]}\" in result.output\n assert \"requests.get\" in result.output\n assert \"localhost:8080\" in result.output\n assert f\">>> {inputs[1]}\" in result.output\n assert \"localhost:443\" in result.output\n\n chat_storage = cfg.get(\"CHAT_CACHE_PATH\")\n tmp_chat = Path(chat_storage) / dict_arguments[\"--repl\"]\n chat_messages = json.loads(tmp_chat.read_text())\n assert chat_messages[0][\"content\"].startswith(\"You are Code Generator\")\n assert chat_messages[0][\"role\"] == \"system\"\n\n # Coming back after exit.\n new_inputs = (\"Change port to 80\", \"exit()\")\n result = runner.invoke(\n app, self.get_arguments(**dict_arguments), input=\"\\n\".join(new_inputs)\n )\n # Should include previous chat history.\n assert \"Chat History\" in result.output\n assert f\"user: {inputs[1]}\" in result.output\n\n def test_zsh_command(self):\n \"\"\"\n The goal of this test is to verify that $SHELL\n specific commands are working as expected.\n In this case testing zsh specific \"print\" function.\n \"\"\"\n if os.getenv(\"SHELL\", \"\") != \"/bin/zsh\":\n return\n dict_arguments = {\n \"prompt\": 'Using zsh specific \"print\" function say hello world',\n \"--shell\": True,\n }\n result = runner.invoke(app, self.get_arguments(**dict_arguments), input=\"y\\n\")\n stdout = result.output.strip()\n print(stdout)\n # TODO: Fix this test.\n # Not sure how os.system pipes the output to stdout,\n # but it is not part of the result.output.\n # assert \"command not found\" not in result.output\n # assert \"hello world\" in stdout.split(\"\\n\")[-1]\n\n @patch(\"sgpt.handlers.handler.Handler.get_completion\")\n def test_model_option(self, mocked_get_completion):\n dict_arguments = {\n \"prompt\": \"What is the capital of the Czech Republic?\",\n \"--model\": \"gpt-4\",\n }\n result = runner.invoke(app, self.get_arguments(**dict_arguments))\n mocked_get_completion.assert_called_once_with(\n messages=ANY,\n model=\"gpt-4\",\n temperature=0.0,\n top_p=1.0,\n caching=False,\n functions=None,\n )\n assert result.exit_code == 0\n\n def test_color_output(self):\n color = cfg.get(\"DEFAULT_COLOR\")\n role = SystemRole.get(\"ShellGPT\")\n handler = Handler(role=role)\n assert handler.color == color\n os.environ[\"DEFAULT_COLOR\"] = \"red\"\n handler = Handler(role=role)\n assert handler.color == \"red\"\n\n def test_simple_stdin(self):\n result = runner.invoke(app, input=\"What is the capital of Germany?\\n\")\n assert \"Berlin\" in result.output\n\n def test_shell_stdin_with_prompt(self):\n dict_arguments = {\n \"prompt\": \"Sort by name\",\n \"--shell\": True,\n }\n stdin = \"What is in current folder\\n\"\n result = runner.invoke(app, self.get_arguments(**dict_arguments), input=stdin)\n assert \"ls\" in result.output\n assert \"sort\" in result.output\n\n def test_role(self):\n test_role = Path(cfg.get(\"ROLE_STORAGE_PATH\")) / \"json_generator.json\"\n test_role.unlink(missing_ok=True)\n dict_arguments = {\n \"prompt\": \"test\",\n \"--create-role\": \"json_generator\",\n }\n input = (\n \"Provide only valid plain JSON as response with valid field values. \"\n + \"Do not include any markdown formatting such as ```.\\n\"\n )\n result = runner.invoke(app, self.get_arguments(**dict_arguments), input=input)\n assert result.exit_code == 0\n\n dict_arguments = {\n \"prompt\": \"test\",\n \"--list-roles\": True,\n }\n result = runner.invoke(app, self.get_arguments(**dict_arguments))\n assert result.exit_code == 0\n assert \"json_generator\" in result.output\n\n dict_arguments = {\n \"prompt\": \"test\",\n \"--show-role\": \"json_generator\",\n }\n result = runner.invoke(app, self.get_arguments(**dict_arguments))\n assert result.exit_code == 0\n assert \"You are json_generator\" in result.output\n\n # Test with command line argument prompt.\n dict_arguments = {\n \"prompt\": \"random username, password, email\",\n \"--role\": \"json_generator\",\n }\n result = runner.invoke(app, self.get_arguments(**dict_arguments))\n assert result.exit_code == 0\n generated_json = json.loads(result.output)\n assert \"username\" in generated_json\n assert \"password\" in generated_json\n assert \"email\" in generated_json\n\n # Test with stdin prompt.\n dict_arguments = {\n \"prompt\": \"\",\n \"--role\": \"json_generator\",\n }\n stdin = \"random username, password, email\"\n result = runner.invoke(app, self.get_arguments(**dict_arguments), input=stdin)\n assert result.exit_code == 0\n generated_json = json.loads(result.output)\n assert \"username\" in generated_json\n assert \"password\" in generated_json\n assert \"email\" in generated_json\n test_role.unlink(missing_ok=True)\n\n def test_shell_command_run_description(self):\n dict_arguments = {\n \"prompt\": \"say hello\",\n \"--shell\": True,\n }\n result = runner.invoke(app, self.get_arguments(**dict_arguments), input=\"d\\n\")\n assert result.exit_code == 0\n # Can't really test it since stdin in disable for --shell flag.\n # for word in (\"prints\", \"hello\", \"console\"):\n # assert word in result.output\n\n def test_version(self):\n dict_arguments = {\n \"prompt\": \"\",\n \"--version\": True,\n }\n result = runner.invoke(app, self.get_arguments(**dict_arguments), input=\"d\\n\")\n assert __version__ in result.output\n\n # TODO: Implement function call tests.\n" }, { "path": "tests/conftest.py", "content": "import os\n\nimport pytest\n\n\n@pytest.fixture(autouse=True)\ndef mock_os_name(monkeypatch):\n monkeypatch.setattr(os, \"name\", \"test\")\n" }, { "path": "tests/test_code.py", "content": "from pathlib import Path\nfrom unittest.mock import patch\n\nfrom sgpt.config import cfg\nfrom sgpt.role import DefaultRoles, SystemRole\n\nfrom .utils import app, cmd_args, comp_args, mock_comp, runner\n\nrole = SystemRole.get(DefaultRoles.CODE.value)\n\n\n@patch(\"sgpt.handlers.handler.completion\")\ndef test_code_generation(completion):\n completion.return_value = mock_comp(\"print('Hello World')\")\n\n args = {\"prompt\": \"hello world python\", \"--code\": True}\n result = runner.invoke(app, cmd_args(**args))\n\n completion.assert_called_once_with(**comp_args(role, args[\"prompt\"]))\n assert result.exit_code == 0\n assert \"print('Hello World')\" in result.output\n\n\n@patch(\"sgpt.printer.TextPrinter.live_print\")\n@patch(\"sgpt.printer.MarkdownPrinter.live_print\")\n@patch(\"sgpt.handlers.handler.completion\")\ndef test_code_generation_no_markdown(completion, markdown_printer, text_printer):\n completion.return_value = mock_comp(\"print('Hello World')\")\n\n args = {\"prompt\": \"make a commit using git\", \"--code\": True, \"--md\": True}\n result = runner.invoke(app, cmd_args(**args))\n\n assert result.exit_code == 0\n # Should ignore --md for --code option and output code without markdown.\n markdown_printer.assert_not_called()\n text_printer.assert_called()\n\n\n@patch(\"sgpt.handlers.handler.completion\")\ndef test_code_generation_stdin(completion):\n completion.return_value = mock_comp(\"# Hello\\nprint('Hello')\")\n\n args = {\"prompt\": \"make comments for code\", \"--code\": True}\n stdin = \"print('Hello')\"\n result = runner.invoke(app, cmd_args(**args), input=stdin)\n\n expected_prompt = f\"{stdin}\\n\\n{args['prompt']}\"\n completion.assert_called_once_with(**comp_args(role, expected_prompt))\n assert result.exit_code == 0\n assert \"# Hello\" in result.output\n assert \"print('Hello')\" in result.output\n\n\n@patch(\"sgpt.handlers.handler.completion\")\ndef test_code_chat(completion):\n completion.side_effect = [\n mock_comp(\"print('hello')\"),\n mock_comp(\"print('hello')\\nprint('world')\"),\n ]\n chat_name = \"_test\"\n chat_path = Path(cfg.get(\"CHAT_CACHE_PATH\")) / chat_name\n chat_path.unlink(missing_ok=True)\n\n args = {\"prompt\": \"print hello\", \"--code\": True, \"--chat\": chat_name}\n result = runner.invoke(app, cmd_args(**args))\n assert result.exit_code == 0\n assert \"print('hello')\" in result.output\n assert chat_path.exists()\n\n args[\"prompt\"] = \"also print world\"\n result = runner.invoke(app, cmd_args(**args))\n assert result.exit_code == 0\n assert \"print('hello')\" in result.output\n assert \"print('world')\" in result.output\n\n expected_messages = [\n {\"role\": \"system\", \"content\": role.role},\n {\"role\": \"user\", \"content\": \"print hello\"},\n {\"role\": \"assistant\", \"content\": \"print('hello')\"},\n {\"role\": \"user\", \"content\": \"also print world\"},\n {\"role\": \"assistant\", \"content\": \"print('hello')\\nprint('world')\"},\n ]\n expected_args = comp_args(role, \"\", messages=expected_messages)\n completion.assert_called_with(**expected_args)\n assert completion.call_count == 2\n\n args[\"--shell\"] = True\n result = runner.invoke(app, cmd_args(**args))\n assert result.exit_code == 2\n assert \"Error\" in result.output\n chat_path.unlink()\n # TODO: Code chat can be recalled without --code option.\n\n\n@patch(\"sgpt.handlers.handler.completion\")\ndef test_code_repl(completion):\n completion.side_effect = [\n mock_comp(\"print('hello')\"),\n mock_comp(\"print('hello')\\nprint('world')\"),\n ]\n chat_name = \"_test\"\n chat_path = Path(cfg.get(\"CHAT_CACHE_PATH\")) / chat_name\n chat_path.unlink(missing_ok=True)\n\n args = {\"--repl\": chat_name, \"--code\": True}\n inputs = [\"__sgpt__eof__\", \"print hello\", \"also print world\", \"exit()\"]\n result = runner.invoke(app, cmd_args(**args), input=\"\\n\".join(inputs))\n\n expected_messages = [\n {\"role\": \"system\", \"content\": role.role},\n {\"role\": \"user\", \"content\": \"print hello\"},\n {\"role\": \"assistant\", \"content\": \"print('hello')\"},\n {\"role\": \"user\", \"content\": \"also print world\"},\n {\"role\": \"assistant\", \"content\": \"print('hello')\\nprint('world')\"},\n ]\n expected_args = comp_args(role, \"\", messages=expected_messages)\n completion.assert_called_with(**expected_args)\n assert completion.call_count == 2\n\n assert result.exit_code == 0\n assert \">>> print hello\" in result.output\n assert \"print('hello')\" in result.output\n assert \">>> also print world\" in result.output\n assert \"print('world')\" in result.output\n\n\n@patch(\"sgpt.handlers.handler.completion\")\ndef test_code_and_shell(completion):\n args = {\"--code\": True, \"--shell\": True}\n result = runner.invoke(app, cmd_args(**args))\n\n completion.assert_not_called()\n assert result.exit_code == 2\n assert \"Error\" in result.output\n\n\n@patch(\"sgpt.handlers.handler.completion\")\ndef test_code_and_describe_shell(completion):\n args = {\"--code\": True, \"--describe-shell\": True}\n result = runner.invoke(app, cmd_args(**args))\n\n completion.assert_not_called()\n assert result.exit_code == 2\n assert \"Error\" in result.output\n" }, { "path": "tests/test_default.py", "content": "from pathlib import Path\nfrom unittest.mock import patch\n\nimport typer\nfrom typer.testing import CliRunner\n\nfrom sgpt import config, main\nfrom sgpt.__version__ import __version__\nfrom sgpt.role import DefaultRoles, SystemRole\n\nfrom .utils import app, cmd_args, comp_args, mock_comp, runner\n\nrole = SystemRole.get(DefaultRoles.DEFAULT.value)\ncfg = config.cfg\n\n\n@patch(\"sgpt.handlers.handler.completion\")\ndef test_default(completion):\n completion.return_value = mock_comp(\"Prague\")\n\n args = {\"prompt\": \"capital of the Czech Republic?\"}\n result = runner.invoke(app, cmd_args(**args))\n\n completion.assert_called_once_with(**comp_args(role, **args))\n assert result.exit_code == 0\n assert \"Prague\" in result.output\n\n\n@patch(\"sgpt.handlers.handler.completion\")\ndef test_default_stdin(completion):\n completion.return_value = mock_comp(\"Prague\")\n\n stdin = \"capital of the Czech Republic?\"\n result = runner.invoke(app, cmd_args(), input=stdin)\n\n completion.assert_called_once_with(**comp_args(role, stdin))\n assert result.exit_code == 0\n assert \"Prague\" in result.output\n\n\n@patch(\"rich.console.Console.print\")\n@patch(\"sgpt.handlers.handler.completion\")\ndef test_show_chat_use_markdown(completion, console_print):\n completion.return_value = mock_comp(\"ok\")\n chat_name = \"_test\"\n chat_path = Path(cfg.get(\"CHAT_CACHE_PATH\")) / chat_name\n chat_path.unlink(missing_ok=True)\n\n args = {\"prompt\": \"my number is 2\", \"--chat\": chat_name}\n result = runner.invoke(app, cmd_args(**args))\n assert result.exit_code == 0\n assert chat_path.exists()\n\n result = runner.invoke(app, [\"--show-chat\", chat_name])\n assert result.exit_code == 0\n console_print.assert_called()\n\n\n@patch(\"rich.console.Console.print\")\n@patch(\"sgpt.handlers.handler.completion\")\ndef test_show_chat_no_use_markdown(completion, console_print):\n completion.return_value = mock_comp(\"ok\")\n chat_name = \"_test\"\n chat_path = Path(cfg.get(\"CHAT_CACHE_PATH\")) / chat_name\n chat_path.unlink(missing_ok=True)\n\n # Flag '--code' doesn't use markdown\n args = {\"prompt\": \"my number is 2\", \"--chat\": chat_name, \"--code\": True}\n result = runner.invoke(app, cmd_args(**args))\n assert result.exit_code == 0\n assert chat_path.exists()\n\n result = runner.invoke(app, [\"--show-chat\", chat_name, \"--no-md\"])\n assert result.exit_code == 0\n console_print.assert_not_called()\n\n\n@patch(\"sgpt.handlers.handler.completion\")\ndef test_default_chat(completion):\n completion.side_effect = [mock_comp(\"ok\"), mock_comp(\"4\")]\n chat_name = \"_test\"\n chat_path = Path(cfg.get(\"CHAT_CACHE_PATH\")) / chat_name\n chat_path.unlink(missing_ok=True)\n\n args = {\"prompt\": \"my number is 2\", \"--chat\": chat_name}\n result = runner.invoke(app, cmd_args(**args))\n assert result.exit_code == 0\n assert \"ok\" in result.output\n assert chat_path.exists()\n\n args[\"prompt\"] = \"my number + 2?\"\n result = runner.invoke(app, cmd_args(**args))\n assert result.exit_code == 0\n assert \"4\" in result.output\n\n expected_messages = [\n {\"role\": \"system\", \"content\": role.role},\n {\"role\": \"user\", \"content\": \"my number is 2\"},\n {\"role\": \"assistant\", \"content\": \"ok\"},\n {\"role\": \"user\", \"content\": \"my number + 2?\"},\n {\"role\": \"assistant\", \"content\": \"4\"},\n ]\n expected_args = comp_args(role, \"\", messages=expected_messages)\n completion.assert_called_with(**expected_args)\n assert completion.call_count == 2\n\n result = runner.invoke(app, [\"--list-chats\"])\n assert result.exit_code == 0\n assert \"_test\" in result.output\n\n result = runner.invoke(app, [\"--show-chat\", chat_name])\n assert result.exit_code == 0\n assert \"my number is 2\" in result.output\n assert \"ok\" in result.output\n assert \"my number + 2?\" in result.output\n assert \"4\" in result.output\n\n args[\"--shell\"] = True\n result = runner.invoke(app, cmd_args(**args))\n assert result.exit_code == 2\n assert \"Error\" in result.output\n\n args[\"--code\"] = True\n result = runner.invoke(app, cmd_args(**args))\n assert result.exit_code == 2\n assert \"Error\" in result.output\n chat_path.unlink()\n\n\n@patch(\"sgpt.handlers.handler.completion\")\ndef test_default_repl(completion):\n completion.side_effect = [mock_comp(\"ok\"), mock_comp(\"8\")]\n chat_name = \"_test\"\n chat_path = Path(cfg.get(\"CHAT_CACHE_PATH\")) / chat_name\n chat_path.unlink(missing_ok=True)\n\n args = {\"--repl\": chat_name}\n inputs = [\"__sgpt__eof__\", \"my number is 6\", \"my number + 2?\", \"exit()\"]\n result = runner.invoke(app, cmd_args(**args), input=\"\\n\".join(inputs))\n\n expected_messages = [\n {\"role\": \"system\", \"content\": role.role},\n {\"role\": \"user\", \"content\": \"my number is 6\"},\n {\"role\": \"assistant\", \"content\": \"ok\"},\n {\"role\": \"user\", \"content\": \"my number + 2?\"},\n {\"role\": \"assistant\", \"content\": \"8\"},\n ]\n expected_args = comp_args(role, \"\", messages=expected_messages)\n completion.assert_called_with(**expected_args)\n assert completion.call_count == 2\n\n assert result.exit_code == 0\n assert \">>> my number is 6\" in result.output\n assert \"ok\" in result.output\n assert \">>> my number + 2?\" in result.output\n assert \"8\" in result.output\n\n\n@patch(\"sgpt.handlers.handler.completion\")\ndef test_default_repl_stdin(completion):\n completion.side_effect = [mock_comp(\"ok init\"), mock_comp(\"ok another\")]\n chat_name = \"_test\"\n chat_path = Path(cfg.get(\"CHAT_CACHE_PATH\")) / chat_name\n chat_path.unlink(missing_ok=True)\n\n my_runner = CliRunner()\n my_app = typer.Typer()\n my_app.command()(main)\n\n args = {\"--repl\": chat_name}\n inputs = [\"this is stdin\", \"__sgpt__eof__\", \"prompt\", \"another\", \"exit()\"]\n result = my_runner.invoke(my_app, cmd_args(**args), input=\"\\n\".join(inputs))\n\n expected_messages = [\n {\"role\": \"system\", \"content\": role.role},\n {\"role\": \"user\", \"content\": \"this is stdin\\n\\n\\n\\nprompt\"},\n {\"role\": \"assistant\", \"content\": \"ok init\"},\n {\"role\": \"user\", \"content\": \"another\"},\n {\"role\": \"assistant\", \"content\": \"ok another\"},\n ]\n expected_args = comp_args(role, \"\", messages=expected_messages)\n completion.assert_called_with(**expected_args)\n assert completion.call_count == 2\n\n assert result.exit_code == 0\n assert \"this is stdin\" in result.output\n assert \">>> prompt\" in result.output\n assert \"ok init\" in result.output\n assert \">>> another\" in result.output\n assert \"ok another\" in result.output\n\n\n@patch(\"sgpt.handlers.handler.completion\")\ndef test_llm_options(completion):\n completion.return_value = mock_comp(\"Berlin\")\n\n args = {\n \"prompt\": \"capital of the Germany?\",\n \"--model\": \"gpt-4-test\",\n \"--temperature\": 0.5,\n \"--top-p\": 0.5,\n \"--no-functions\": True,\n }\n result = runner.invoke(app, cmd_args(**args))\n\n expected_args = comp_args(\n role=role,\n prompt=args[\"prompt\"],\n model=args[\"--model\"],\n temperature=args[\"--temperature\"],\n top_p=args[\"--top-p\"],\n )\n completion.assert_called_once_with(**expected_args)\n assert result.exit_code == 0\n assert \"Berlin\" in result.output\n\n\n@patch(\"sgpt.handlers.handler.completion\")\ndef test_version(completion):\n args = {\"--version\": True}\n result = runner.invoke(app, cmd_args(**args))\n\n completion.assert_not_called()\n assert __version__ in result.output\n\n\n@patch(\"sgpt.printer.TextPrinter.live_print\")\n@patch(\"sgpt.printer.MarkdownPrinter.live_print\")\n@patch(\"sgpt.handlers.handler.completion\")\ndef test_markdown(completion, markdown_printer, text_printer):\n completion.return_value = mock_comp(\"pong\")\n\n args = {\"prompt\": \"ping\", \"--md\": True}\n result = runner.invoke(app, cmd_args(**args))\n assert result.exit_code == 0\n markdown_printer.assert_called()\n text_printer.assert_not_called()\n\n\n@patch(\"sgpt.printer.TextPrinter.live_print\")\n@patch(\"sgpt.printer.MarkdownPrinter.live_print\")\n@patch(\"sgpt.handlers.handler.completion\")\ndef test_no_markdown(completion, markdown_printer, text_printer):\n completion.return_value = mock_comp(\"pong\")\n\n args = {\"prompt\": \"ping\", \"--no-md\": True}\n result = runner.invoke(app, cmd_args(**args))\n assert result.exit_code == 0\n markdown_printer.assert_not_called()\n text_printer.assert_called()\n" }, { "path": "tests/test_roles.py", "content": "import json\nfrom pathlib import Path\nfrom unittest.mock import patch\n\nfrom sgpt.config import cfg\nfrom sgpt.role import SystemRole\n\nfrom .utils import app, cmd_args, comp_args, mock_comp, runner\n\n\n@patch(\"sgpt.handlers.handler.completion\")\ndef test_role(completion):\n completion.return_value = mock_comp('{\"foo\": \"bar\"}')\n path = Path(cfg.get(\"ROLE_STORAGE_PATH\")) / \"json_gen_test.json\"\n path.unlink(missing_ok=True)\n args = {\"--create-role\": \"json_gen_test\"}\n stdin = \"you are a JSON generator\"\n result = runner.invoke(app, cmd_args(**args), input=stdin)\n completion.assert_not_called()\n assert result.exit_code == 0\n\n args = {\"--list-roles\": True}\n result = runner.invoke(app, cmd_args(**args))\n completion.assert_not_called()\n assert result.exit_code == 0\n assert \"json_gen_test\" in result.output\n\n args = {\"--show-role\": \"json_gen_test\"}\n result = runner.invoke(app, cmd_args(**args))\n completion.assert_not_called()\n assert result.exit_code == 0\n assert \"you are a JSON generator\" in result.output\n\n # Test with argument prompt.\n args = {\n \"prompt\": \"generate foo, bar\",\n \"--role\": \"json_gen_test\",\n }\n result = runner.invoke(app, cmd_args(**args))\n role = SystemRole.get(\"json_gen_test\")\n completion.assert_called_once_with(**comp_args(role, args[\"prompt\"]))\n assert result.exit_code == 0\n generated_json = json.loads(result.output)\n assert \"foo\" in generated_json\n\n # Test with stdin prompt.\n completion.return_value = mock_comp('{\"foo\": \"bar\"}')\n args = {\"--role\": \"json_gen_test\"}\n stdin = \"generate foo, bar\"\n result = runner.invoke(app, cmd_args(**args), input=stdin)\n completion.assert_called_with(**comp_args(role, stdin))\n assert result.exit_code == 0\n generated_json = json.loads(result.output)\n assert \"foo\" in generated_json\n path.unlink(missing_ok=True)\n" }, { "path": "tests/test_shell.py", "content": "import os\nfrom pathlib import Path\nfrom unittest.mock import patch\n\nfrom sgpt.config import cfg\nfrom sgpt.role import DefaultRoles, SystemRole\n\nfrom .utils import app, cmd_args, comp_args, mock_comp, runner\n\n\n@patch(\"sgpt.handlers.handler.completion\")\ndef test_shell(completion):\n role = SystemRole.get(DefaultRoles.SHELL.value)\n completion.return_value = mock_comp(\"git commit -m test\")\n\n args = {\"prompt\": \"make a commit using git\", \"--shell\": True}\n result = runner.invoke(app, cmd_args(**args))\n\n completion.assert_called_once_with(**comp_args(role, args[\"prompt\"]))\n assert \"git commit\" in result.output\n assert \"[E]xecute, [M]odify, [D]escribe, [A]bort:\" in result.output\n\n\n@patch(\"sgpt.printer.TextPrinter.live_print\")\n@patch(\"sgpt.printer.MarkdownPrinter.live_print\")\n@patch(\"sgpt.handlers.handler.completion\")\ndef test_shell_no_markdown(completion, markdown_printer, text_printer):\n completion.return_value = mock_comp(\"git commit -m test\")\n\n args = {\"prompt\": \"make a commit using git\", \"--shell\": True, \"--md\": True}\n runner.invoke(app, cmd_args(**args))\n\n # Should ignore --md for --shell option and output text without markdown.\n markdown_printer.assert_not_called()\n text_printer.assert_called()\n\n\n@patch(\"sgpt.handlers.handler.completion\")\ndef test_shell_stdin(completion):\n completion.return_value = mock_comp(\"ls -l | sort\")\n role = SystemRole.get(DefaultRoles.SHELL.value)\n\n args = {\"prompt\": \"Sort by name\", \"--shell\": True}\n stdin = \"What is in current folder\"\n result = runner.invoke(app, cmd_args(**args), input=stdin)\n\n expected_prompt = f\"{stdin}\\n\\n{args['prompt']}\"\n completion.assert_called_once_with(**comp_args(role, expected_prompt))\n assert \"ls -l | sort\" in result.output\n assert \"[E]xecute, [M]odify, [D]escribe, [A]bort:\" in result.output\n\n\n@patch(\"sgpt.handlers.handler.completion\")\ndef test_describe_shell(completion):\n completion.return_value = mock_comp(\"lists the contents of a folder\")\n role = SystemRole.get(DefaultRoles.DESCRIBE_SHELL.value)\n\n args = {\"prompt\": \"ls\", \"--describe-shell\": True}\n result = runner.invoke(app, cmd_args(**args))\n\n completion.assert_called_once_with(**comp_args(role, args[\"prompt\"]))\n assert result.exit_code == 0\n assert \"lists\" in result.output\n\n\n@patch(\"sgpt.handlers.handler.completion\")\ndef test_describe_shell_stdin(completion):\n completion.return_value = mock_comp(\"lists the contents of a folder\")\n role = SystemRole.get(DefaultRoles.DESCRIBE_SHELL.value)\n\n args = {\"--describe-shell\": True}\n stdin = \"What is in current folder\"\n result = runner.invoke(app, cmd_args(**args), input=stdin)\n\n expected_prompt = f\"{stdin}\"\n completion.assert_called_once_with(**comp_args(role, expected_prompt))\n assert result.exit_code == 0\n assert \"lists\" in result.output\n\n\n@patch(\"os.system\")\n@patch(\"sgpt.handlers.handler.completion\")\ndef test_shell_run_description(completion, system):\n completion.side_effect = [mock_comp(\"echo hello\"), mock_comp(\"prints hello\")]\n args = {\"prompt\": \"echo hello\", \"--shell\": True}\n inputs = \"__sgpt__eof__\\nd\\ne\\n\"\n result = runner.invoke(app, cmd_args(**args), input=inputs)\n shell = os.environ.get(\"SHELL\", \"/bin/sh\")\n system.assert_called_once_with(f\"{shell} -c 'echo hello'\")\n assert result.exit_code == 0\n assert \"echo hello\" in result.output\n assert \"prints hello\" in result.output\n\n\n@patch(\"sgpt.handlers.handler.completion\")\ndef test_shell_chat(completion):\n completion.side_effect = [mock_comp(\"ls\"), mock_comp(\"ls | sort\")]\n role = SystemRole.get(DefaultRoles.SHELL.value)\n chat_name = \"_test\"\n chat_path = Path(cfg.get(\"CHAT_CACHE_PATH\")) / chat_name\n chat_path.unlink(missing_ok=True)\n\n args = {\"prompt\": \"list folder\", \"--shell\": True, \"--chat\": chat_name}\n result = runner.invoke(app, cmd_args(**args))\n assert \"ls\" in result.output\n assert chat_path.exists()\n\n args[\"prompt\"] = \"sort by name\"\n result = runner.invoke(app, cmd_args(**args))\n assert \"ls | sort\" in result.output\n\n expected_messages = [\n {\"role\": \"system\", \"content\": role.role},\n {\"role\": \"user\", \"content\": \"list folder\"},\n {\"role\": \"assistant\", \"content\": \"ls\"},\n {\"role\": \"user\", \"content\": \"sort by name\"},\n {\"role\": \"assistant\", \"content\": \"ls | sort\"},\n ]\n expected_args = comp_args(role, \"\", messages=expected_messages)\n completion.assert_called_with(**expected_args)\n assert completion.call_count == 2\n\n args[\"--code\"] = True\n result = runner.invoke(app, cmd_args(**args))\n assert result.exit_code == 2\n assert \"Error\" in result.output\n chat_path.unlink()\n # TODO: Shell chat can be recalled without --shell option.\n\n\n@patch(\"os.system\")\n@patch(\"sgpt.handlers.handler.completion\")\ndef test_shell_repl(completion, mock_system):\n completion.side_effect = [mock_comp(\"ls\"), mock_comp(\"ls | sort\")]\n role = SystemRole.get(DefaultRoles.SHELL.value)\n chat_name = \"_test\"\n chat_path = Path(cfg.get(\"CHAT_CACHE_PATH\")) / chat_name\n chat_path.unlink(missing_ok=True)\n\n args = {\"--repl\": chat_name, \"--shell\": True}\n inputs = [\"__sgpt__eof__\", \"list folder\", \"sort by name\", \"e\", \"exit()\"]\n result = runner.invoke(app, cmd_args(**args), input=\"\\n\".join(inputs))\n shell = os.environ.get(\"SHELL\", \"/bin/sh\")\n mock_system.assert_called_once_with(f\"{shell} -c 'ls | sort'\")\n\n expected_messages = [\n {\"role\": \"system\", \"content\": role.role},\n {\"role\": \"user\", \"content\": \"list folder\"},\n {\"role\": \"assistant\", \"content\": \"ls\"},\n {\"role\": \"user\", \"content\": \"sort by name\"},\n {\"role\": \"assistant\", \"content\": \"ls | sort\"},\n ]\n expected_args = comp_args(role, \"\", messages=expected_messages)\n completion.assert_called_with(**expected_args)\n assert completion.call_count == 2\n\n assert result.exit_code == 0\n assert \">>> list folder\" in result.output\n assert \"ls\" in result.output\n assert \">>> sort by name\" in result.output\n assert \"ls | sort\" in result.output\n\n\n@patch(\"sgpt.handlers.handler.completion\")\ndef test_shell_and_describe_shell(completion):\n args = {\"prompt\": \"ls\", \"--describe-shell\": True, \"--shell\": True}\n result = runner.invoke(app, cmd_args(**args))\n\n completion.assert_not_called()\n assert result.exit_code == 2\n assert \"Error\" in result.output\n\n\n@patch(\"sgpt.handlers.handler.completion\")\ndef test_shell_no_interaction(completion):\n completion.return_value = mock_comp(\"git commit -m test\")\n role = SystemRole.get(DefaultRoles.SHELL.value)\n\n args = {\n \"prompt\": \"make a commit using git\",\n \"--shell\": True,\n \"--no-interaction\": True,\n }\n result = runner.invoke(app, cmd_args(**args))\n\n completion.assert_called_once_with(**comp_args(role, args[\"prompt\"]))\n assert result.exit_code == 0\n assert \"git commit\" in result.output\n assert \"[E]xecute\" not in result.output\n" }, { "path": "tests/utils.py", "content": "from datetime import datetime\n\nimport typer\nfrom openai.types.chat.chat_completion_chunk import ChatCompletionChunk\nfrom openai.types.chat.chat_completion_chunk import Choice as StreamChoice\nfrom openai.types.chat.chat_completion_chunk import ChoiceDelta\nfrom typer.testing import CliRunner\n\nfrom sgpt import main\nfrom sgpt.config import cfg\n\nrunner = CliRunner()\napp = typer.Typer()\napp.command()(main)\n\n\ndef mock_comp(tokens_string):\n return [\n ChatCompletionChunk(\n id=\"foo\",\n model=cfg.get(\"DEFAULT_MODEL\"),\n object=\"chat.completion.chunk\",\n choices=[\n StreamChoice(\n index=0,\n finish_reason=None,\n delta=ChoiceDelta(content=token, role=\"assistant\"),\n ),\n ],\n created=int(datetime.now().timestamp()),\n )\n for token in tokens_string\n ]\n\n\ndef cmd_args(prompt=\"\", **kwargs):\n arguments = [prompt]\n for key, value in kwargs.items():\n arguments.append(key)\n if isinstance(value, bool):\n continue\n arguments.append(value)\n arguments.append(\"--no-cache\")\n arguments.append(\"--no-functions\")\n return arguments\n\n\ndef comp_args(role, prompt, **kwargs):\n return {\n \"messages\": [\n {\"role\": \"system\", \"content\": role.role},\n {\"role\": \"user\", \"content\": prompt},\n ],\n \"model\": cfg.get(\"DEFAULT_MODEL\"),\n \"temperature\": 0.0,\n \"top_p\": 1.0,\n \"stream\": True,\n **kwargs,\n }\n" } ]