Repository: TheR1D/shell_gpt Branch: main Commit: 4ea2f834cff5 Files: 42 Total size: 126.2 KB Directory structure: gitextract_ni1c58zi/ ├── .devcontainer/ │ └── devcontainer.json ├── .github/ │ ├── FUNDING.yml │ └── workflows/ │ ├── codespell.yml │ ├── docker.yml │ ├── lint_test.yml │ └── release.yml ├── CONTRIBUTING.md ├── Dockerfile ├── LICENSE ├── README.md ├── pyproject.toml ├── scripts/ │ ├── format.sh │ ├── lint.sh │ └── test.sh ├── sgpt/ │ ├── __init__.py │ ├── __main__.py │ ├── __version__.py │ ├── app.py │ ├── cache.py │ ├── config.py │ ├── function.py │ ├── handlers/ │ │ ├── __init__.py │ │ ├── chat_handler.py │ │ ├── default_handler.py │ │ ├── handler.py │ │ └── repl_handler.py │ ├── integration.py │ ├── llm_functions/ │ │ ├── __init__.py │ │ ├── common/ │ │ │ └── execute_shell.py │ │ ├── init_functions.py │ │ └── mac/ │ │ └── apple_script.py │ ├── printer.py │ ├── role.py │ └── utils.py └── tests/ ├── __init__.py ├── _integration.py ├── conftest.py ├── test_code.py ├── test_default.py ├── test_roles.py ├── test_shell.py └── utils.py ================================================ FILE CONTENTS ================================================ ================================================ FILE: .devcontainer/devcontainer.json ================================================ { "name": "Python 3", "image": "mcr.microsoft.com/devcontainers/python:0-3.9-bullseye", "customizations": { "vscode": { "settings": { "python.defaultInterpreterPath": "/usr/local/bin/python", "cSpell.words": [ "OPENAI", "secho", "sgpt", "Typer" ], "[python]": { "editor.defaultFormatter": "charliermarsh.ruff" }, "files.exclude": { "**/.git/**": true, "**/.mypy_cache/**": true, "**/__pycache__/**": true }, "files.watcherExclude": { "**/.git/**": true, "**/.mypy_cache/**": true, "**/.venv/**": true, "**/__pycache__/**": true }, "launch": { "configurations": [ { "name": "Python: Module", "type": "python", "request": "launch", "module": "sgpt", "justMyCode": true, "args": ["--chat", "init", "hello"] } ] } }, "extensions": [ "GitHub.copilot", "charliermarsh.ruff", "ms-python.python", "ms-python.vscode-pylance", "ms-python.black-formatter", "ms-python.isort", "ms-python.mypy-type-checker", "ms-python.pylint" ] } }, "remoteUser": "vscode", "postCreateCommand": "echo __pycache__ > ~/.gitignore && git config --global core.excludesfile ~/.gitignore && pip3 install -e .'[dev,test]'" } ================================================ FILE: .github/FUNDING.yml ================================================ github: [ther1d] ================================================ FILE: .github/workflows/codespell.yml ================================================ --- name: Codespell on: push: branches: [main] pull_request: branches: [main] permissions: contents: read jobs: codespell: name: Check for spelling errors runs-on: ubuntu-latest steps: - name: Checkout uses: actions/checkout@v3 - name: Codespell uses: codespell-project/actions-codespell@v2 ================================================ FILE: .github/workflows/docker.yml ================================================ name: Docker Image CI on: push: tags: - '*' jobs: build: runs-on: ubuntu-latest environment: name: Docker Image steps: - name: Checkout uses: actions/checkout@v3 - name: Docker meta id: meta uses: docker/metadata-action@v4 with: images: ghcr.io/${{ github.repository }} tags: type=ref,event=tag - name: Set up QEMU uses: docker/setup-qemu-action@v2 - name: Set up Docker Buildx uses: docker/setup-buildx-action@v2 - name: Login to GitHub Container Registry uses: docker/login-action@v1 with: registry: ghcr.io username: ${{ github.actor }} password: ${{ secrets.GITHUB_TOKEN }} - name: Build and push uses: docker/build-push-action@v4 with: platforms: linux/amd64,linux/arm64 push: true tags: ${{ steps.meta.outputs.tags }} labels: ${{ steps.meta.outputs.labels }} ================================================ FILE: .github/workflows/lint_test.yml ================================================ name: Lint and Test on: push: branches: - main pull_request: branches: - main jobs: lint_test: runs-on: ubuntu-latest strategy: matrix: python-version: ["3.10", "3.11", "3.12"] steps: - uses: actions/checkout@v3 - name: Set up Python ${{ matrix.python-version }} uses: actions/setup-python@v3 with: python-version: ${{ matrix.python-version }} - name: Install dependencies run: | python -m pip install --upgrade pip pip install ."[dev,test]" - name: black run: black sgpt tests --check - name: isort run: isort sgpt tests scripts --check-only - name: ruff run: ruff sgpt tests scripts - name: mypy run: mypy sgpt --exclude llm_functions - name: tests run: | export OPENAI_API_KEY=test_api_key pytest tests/ -p no:warnings -v -s ================================================ FILE: .github/workflows/release.yml ================================================ name: Publish to PyPI and release on: push: branches: - main paths: - 'sgpt/__version__.py' jobs: build: name: Build distribution runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 - name: Set up Python uses: actions/setup-python@v4 with: python-version: '3.11' - name: Install pypa/build run: >- python3 -m pip install hatchling --user - name: Build a binary wheel and a source tarball run: python3 -m hatchling build - name: Store the distribution packages uses: actions/upload-artifact@v4 with: name: python-package-distributions path: dist/ publish-to-pypi: name: Publish to PyPI needs: - build runs-on: ubuntu-latest environment: name: PyPI url: https://pypi.org/p/shell-gpt permissions: id-token: write # IMPORTANT: mandatory for trusted publishing steps: - name: Download all the dists uses: actions/download-artifact@v4 with: name: python-package-distributions path: dist/ - name: Publish distribution to PyPI uses: pypa/gh-action-pypi-publish@release/v1 github-release: name: Make release needs: - publish-to-pypi runs-on: ubuntu-latest permissions: contents: write # IMPORTANT: mandatory for making GitHub Releases id-token: write # IMPORTANT: mandatory for sigstore steps: - name: Download all the dists uses: actions/download-artifact@v4 with: name: python-package-distributions path: dist/ - name: Get ShellGPT version run: | echo "SGPT_VERSION=$(find dist -type f -name '*.tar.gz' | grep -oP '\d+.\d+.\d+')" >> $GITHUB_ENV echo "Release version $SGPT_VERSION" - name: Sign the dists with Sigstore uses: sigstore/gh-action-sigstore-python@v3.0.0 with: inputs: >- ./dist/*.tar.gz ./dist/*.whl - name: Create GitHub Release env: GITHUB_TOKEN: ${{ github.token }} run: >- gh release create "$SGPT_VERSION" --repo '${{ github.repository }}' --notes "$SGPT_VERSION" - name: Upload artifact signatures to GitHub Release env: GITHUB_TOKEN: ${{ github.token }} # Upload to GitHub Release using the `gh` CLI. # `dist/` contains the built packages, and the # sigstore-produced signatures and certificates. run: >- gh release upload "$SGPT_VERSION" dist/** --repo '${{ github.repository }}' ================================================ FILE: CONTRIBUTING.md ================================================ # Contributing to ShellGPT Thank you for considering contributing to ShellGPT! To ensure a smooth and enjoyable experience for everyone, please follow the steps outlined below. ## Find an Issue to Work On - 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. - 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. - If there is an urgent issue, such as a critical bug causing the app to crash, create a pull request immediately. ## Development ShellGPT 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. ### Virtual Environment Create and activate a virtual environment using Python venv: ```shell python -m venv env && source ./env/bin/activate ``` ### Install Dependencies Install the necessary dependencies, including development and test dependencies: ```shell pip install -e ."[dev,test]" ``` ### Start Coding With 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. ### Testing **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. ### Pull Request Before 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. ### Code Review After 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. Thank you once again for your contribution! We're excited to have you join us. ================================================ FILE: Dockerfile ================================================ FROM python:3-slim ENV SHELL_INTERACTION=false ENV PRETTIFY_MARKDOWN=false ENV OS_NAME=auto ENV SHELL_NAME=auto WORKDIR /app COPY . /app RUN apt-get update && apt-get install -y gcc RUN pip install --no-cache /app && mkdir -p /tmp/shell_gpt VOLUME /tmp/shell_gpt ENTRYPOINT ["sgpt"] ================================================ FILE: LICENSE ================================================ MIT License Copyright (c) 2023 Farkhod Sadykov Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. ================================================ FILE: README.md ================================================ # ShellGPT A 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. https://github.com/TheR1D/shell_gpt/assets/16740832/721ddb19-97e7-428f-a0ee-107d027ddd59 ## Installation ```shell pip install shell-gpt ``` By 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. > [!TIP] > 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). > > **❗️Note that ShellGPT is not optimized for local models and may not work as expected.** ## Usage **ShellGPT** is designed to quickly analyse and retrieve information. It's useful for straightforward requests ranging from technical configurations to general knowledge. ```shell sgpt "What is the fibonacci sequence" # -> The Fibonacci sequence is a series of numbers where each number ... ``` ShellGPT 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: ```shell git diff | sgpt "Generate git commit message, for my changes" # -> Added main feature details into README.md ``` You 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: ```shell docker logs -n 20 my_app | sgpt "check logs, find errors, provide possible solutions" ``` ```text Error Detected: Connection timeout at line 7. Possible Solution: Check network connectivity and firewall settings. Error Detected: Memory allocation failed at line 12. Possible Solution: Consider increasing memory allocation or optimizing application memory usage. ``` You can also use all kind of redirection operators to pass input: ```shell sgpt "summarise" < document.txt # -> The document discusses the impact... sgpt << EOF What is the best way to lear Golang? Provide simple hello world example. EOF # -> The best way to learn Golang... sgpt <<< "What is the best way to learn shell redirects?" # -> The best way to learn shell redirects is through... ``` ### Shell commands Have 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. ```shell sgpt --shell "find all json files in current folder" # -> find . -type f -name "*.json" # -> [E]xecute, [D]escribe, [A]bort: e ``` Shell 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: ```shell sgpt -s "update my system" # -> sudo softwareupdate -i -a # -> [E]xecute, [D]escribe, [A]bort: e ``` The same prompt, when used on Ubuntu, will generate a different suggestion: ```shell sgpt -s "update my system" # -> sudo apt update && sudo apt upgrade -y # -> [E]xecute, [D]escribe, [A]bort: e ``` Let's try it with Docker: ```shell sgpt -s "start nginx container, mount ./index.html" # -> docker run -d -p 80:80 -v $(pwd)/index.html:/usr/share/nginx/html/index.html nginx # -> [E]xecute, [D]escribe, [A]bort: e ``` We can still use pipes to pass input to `sgpt` and generate shell commands: ```shell sgpt -s "POST localhost with" < data.json # -> curl -X POST -H "Content-Type: application/json" -d '{"a": 1, "b": 2}' http://localhost # -> [E]xecute, [D]escribe, [A]bort: e ``` Applying additional shell magic in our prompt, in this example passing file names to `ffmpeg`: ```shell ls # -> 1.mp4 2.mp4 3.mp4 sgpt -s "ffmpeg combine $(ls -m) into one video file without audio." # -> 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 # -> [E]xecute, [D]escribe, [A]bort: e ``` If 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: ```shell sgpt -s "find all json files in current folder" --no-interaction | pbcopy ``` ### Shell integration This 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. https://github.com/TheR1D/shell_gpt/assets/16740832/bead0dab-0dd9-436d-88b7-6abfb2c556c1 To 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. ### Generating code By using the `--code` or `-c` parameter, you can specifically request pure code output, for instance: ```shell sgpt --code "solve fizz buzz problem using python" ``` ```python for i in range(1, 101): if i % 3 == 0 and i % 5 == 0: print("FizzBuzz") elif i % 3 == 0: print("Fizz") elif i % 5 == 0: print("Buzz") else: print(i) ``` Since it is valid python code, we can redirect the output to a file: ```shell sgpt --code "solve classic fizz buzz problem using Python" > fizz_buzz.py python fizz_buzz.py # 1 # 2 # Fizz # 4 # Buzz # ... ``` We can also use pipes to pass input: ```shell cat fizz_buzz.py | sgpt --code "Generate comments for each line of my code" ``` ```python # Loop through numbers 1 to 100 for i in range(1, 101): # Check if number is divisible by both 3 and 5 if i % 3 == 0 and i % 5 == 0: # Print "FizzBuzz" if number is divisible by both 3 and 5 print("FizzBuzz") # Check if number is divisible by 3 elif i % 3 == 0: # Print "Fizz" if number is divisible by 3 print("Fizz") # Check if number is divisible by 5 elif i % 5 == 0: # Print "Buzz" if number is divisible by 5 print("Buzz") # If number is not divisible by 3 or 5, print the number itself else: print(i) ``` ### Chat Mode Often 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`. To start a conversation, use the `--chat` option followed by a unique session name and a prompt. ```shell sgpt --chat conversation_1 "please remember my favorite number: 4" # -> I will remember that your favorite number is 4. sgpt --chat conversation_1 "what would be my favorite number + 4?" # -> Your favorite number is 4, so if we add 4 to it, the result would be 8. ``` You 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`: ```shell sgpt --chat conversation_2 --code "make a request to localhost using python" ``` ```python import requests response = requests.get('http://localhost') print(response.text) ``` Let's ask LLM to add caching to our request: ```shell sgpt --chat conversation_2 --code "add caching" ``` ```python import requests from cachecontrol import CacheControl sess = requests.session() cached_sess = CacheControl(sess) response = cached_sess.get('http://localhost') print(response.text) ``` Same applies for shell commands: ```shell sgpt --chat conversation_3 --shell "what is in current folder" # -> ls sgpt --chat conversation_3 "Sort by name" # -> ls | sort sgpt --chat conversation_3 "Concatenate them using FFMPEG" # -> ffmpeg -i "concat:$(ls | sort | tr '\n' '|')" -codec copy output.mp4 sgpt --chat conversation_3 "Convert the resulting file into an MP3" # -> ffmpeg -i output.mp4 -vn -acodec libmp3lame -ac 2 -ab 160k -ar 48000 final_output.mp3 ``` To list all the sessions from either conversational mode, use the `--list-chats` or `-lc` option: ```shell sgpt --list-chats # .../shell_gpt/chat_cache/conversation_1 # .../shell_gpt/chat_cache/conversation_2 ``` To show all the messages related to a specific conversation, use the `--show-chat` option followed by the session name: ```shell sgpt --show-chat conversation_1 # user: please remember my favorite number: 4 # assistant: I will remember that your favorite number is 4. # user: what would be my favorite number + 4? # assistant: Your favorite number is 4, so if we add 4 to it, the result would be 8. ``` ### REPL Mode There 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.

gif

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