commit f2306a3d11c3d6f812ca164f3373efeae4c2852e Author: wehub-resource-sync Date: Mon Jul 13 12:48:44 2026 +0800 chore: import upstream snapshot with attribution diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json new file mode 100644 index 0000000..5fb6e76 --- /dev/null +++ b/.devcontainer/devcontainer.json @@ -0,0 +1,57 @@ +{ + "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]'" +} \ No newline at end of file diff --git a/.github/FUNDING.yml b/.github/FUNDING.yml new file mode 100644 index 0000000..b12bc2c --- /dev/null +++ b/.github/FUNDING.yml @@ -0,0 +1 @@ +github: [ther1d] \ No newline at end of file diff --git a/.github/workflows/codespell.yml b/.github/workflows/codespell.yml new file mode 100644 index 0000000..3ebbf55 --- /dev/null +++ b/.github/workflows/codespell.yml @@ -0,0 +1,22 @@ +--- +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 diff --git a/.github/workflows/docker.yml b/.github/workflows/docker.yml new file mode 100644 index 0000000..a9c6e12 --- /dev/null +++ b/.github/workflows/docker.yml @@ -0,0 +1,44 @@ +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 }} diff --git a/.github/workflows/lint_test.yml b/.github/workflows/lint_test.yml new file mode 100644 index 0000000..e1f9c94 --- /dev/null +++ b/.github/workflows/lint_test.yml @@ -0,0 +1,38 @@ +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 diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml new file mode 100644 index 0000000..37c0836 --- /dev/null +++ b/.github/workflows/release.yml @@ -0,0 +1,98 @@ +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 }}' diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..e86dcb7 --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,38 @@ +# 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. \ No newline at end of file diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..2dc99b7 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,16 @@ +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"] diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..aa9d1c8 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +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. diff --git a/README.md b/README.md new file mode 100644 index 0000000..3573eb2 --- /dev/null +++ b/README.md @@ -0,0 +1,484 @@ +# 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 run open-source models locally for free. This requires setting up your own LLM backend, such as [Ollama](https://github.com/ollama/ollama). To get ShellGPT working with Ollama, follow this detailed [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-5.4-mini +# 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 +# Control how markdown live rendering handles overflow when output exceeds terminal height. +# Possible values: ellipsis, visible, crop +MARKDOWN_LIVE_VERTICAL_OVERFLOW=ellipsis +``` +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/ +Possible options for `MARKDOWN_LIVE_VERTICAL_OVERFLOW`: `ellipsis`, `visible`, `crop`. + +### Configuration Examples + +**Default behavior (ellipsis):** +```text +MARKDOWN_LIVE_VERTICAL_OVERFLOW=ellipsis +``` +When the markdown output exceeds the terminal height, only `...` is shown. This is the default and preserves backward compatibility. + +**Visible mode (recommended for REPL sessions):** +```text +MARKDOWN_LIVE_VERTICAL_OVERFLOW=visible +``` +All generated markdown content is visible in real-time. This is especially useful for long-running REPL interactions or agent workflows where you want to observe the model's reasoning process, tool calls, and intermediate outputs. + +```shell +sgpt --repl +``` +With `visible` mode, you can continuously observe generated markdown output, tool execution details, and progress updates instead of staring at `...` for several minutes. + +### Full list of arguments +```text +╭─ Arguments ──────────────────────────────────────────────────────────────────────────────────────────────╮ +│ prompt [PROMPT] The prompt to generate completions for. │ +╰──────────────────────────────────────────────────────────────────────────────────────────────────────────╯ +╭─ Options ────────────────────────────────────────────────────────────────────────────────────────────────╮ +│ --model TEXT Large language model to use. [default: gpt-5.4-mini] │ +│ --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 . +``` + +## Additional documentation +* [Azure integration](https://github.com/TheR1D/shell_gpt/wiki/Azure) +* [Ollama integration](https://github.com/TheR1D/shell_gpt/wiki/Ollama) diff --git a/README.wehub.md b/README.wehub.md new file mode 100644 index 0000000..0da74f7 --- /dev/null +++ b/README.wehub.md @@ -0,0 +1,7 @@ +# WeHub 来源说明 + +- 原始项目：`TheR1D/shell_gpt` +- 原始仓库：https://github.com/TheR1D/shell_gpt +- 导入方式：上游默认分支的最新快照 +- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准 +- 本文件仅用于记录来源，不代表 WeHub 是原项目作者 diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..2aa72f3 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,101 @@ +[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", + "click >= 8.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.83.4" +] +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' diff --git a/scripts/format.sh b/scripts/format.sh new file mode 100755 index 0000000..06c765e --- /dev/null +++ b/scripts/format.sh @@ -0,0 +1,7 @@ +#!/bin/sh -e +set -x + +ruff sgpt tests scripts --fix +black sgpt tests scripts +isort sgpt tests scripts +codespell --write-changes diff --git a/scripts/lint.sh b/scripts/lint.sh new file mode 100755 index 0000000..041b61c --- /dev/null +++ b/scripts/lint.sh @@ -0,0 +1,10 @@ +#!/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 diff --git a/scripts/test.sh b/scripts/test.sh new file mode 100755 index 0000000..151866a --- /dev/null +++ b/scripts/test.sh @@ -0,0 +1,7 @@ +#!/usr/bin/env bash + +set -e +set -x + +# shellcheck disable=SC2068 +pytest tests ${@} -p no:warnings diff --git a/sgpt/__init__.py b/sgpt/__init__.py new file mode 100644 index 0000000..49fc9a9 --- /dev/null +++ b/sgpt/__init__.py @@ -0,0 +1,2 @@ +from .app import main as main +from .app import entry_point as cli # noqa: F401 diff --git a/sgpt/__main__.py b/sgpt/__main__.py new file mode 100644 index 0000000..1c7bc11 --- /dev/null +++ b/sgpt/__main__.py @@ -0,0 +1,3 @@ +from .app import entry_point + +entry_point() diff --git a/sgpt/__version__.py b/sgpt/__version__.py new file mode 100644 index 0000000..0f228f2 --- /dev/null +++ b/sgpt/__version__.py @@ -0,0 +1 @@ +__version__ = "1.5.1" diff --git a/sgpt/app.py b/sgpt/app.py new file mode 100644 index 0000000..2778837 --- /dev/null +++ b/sgpt/app.py @@ -0,0 +1,276 @@ +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( + cfg.get("DEFAULT_TEMPERATURE"), + 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() diff --git a/sgpt/cache.py b/sgpt/cache.py new file mode 100644 index 0000000..c176fe2 --- /dev/null +++ b/sgpt/cache.py @@ -0,0 +1,61 @@ +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, encoding="utf-8") + 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() diff --git a/sgpt/config.py b/sgpt/config.py new file mode 100644 index 0000000..283e9c2 --- /dev/null +++ b/sgpt/config.py @@ -0,0 +1,97 @@ +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-5.4-mini"), + "DEFAULT_TEMPERATURE": os.getenv("DEFAULT_TEMPERATURE", 0.0), + "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"), + "MARKDOWN_LIVE_VERTICAL_OVERFLOW": os.getenv( + "MARKDOWN_LIVE_VERTICAL_OVERFLOW", "visible" + ), + "MARKDOWN_LIVE_REFRESH_INTERVAL": os.getenv("MARKDOWN_LIVE_REFRESH_INTERVAL", "0"), + "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) diff --git a/sgpt/function.py b/sgpt/function.py new file mode 100644 index 0000000..0e33196 --- /dev/null +++ b/sgpt/function.py @@ -0,0 +1,67 @@ +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] diff --git a/sgpt/handlers/__init__.py b/sgpt/handlers/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/sgpt/handlers/chat_handler.py b/sgpt/handlers/chat_handler.py new file mode 100644 index 0000000..30930d7 --- /dev/null +++ b/sgpt/handlers/chat_handler.py @@ -0,0 +1,177 @@ +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) diff --git a/sgpt/handlers/default_handler.py b/sgpt/handlers/default_handler.py new file mode 100644 index 0000000..e0fdad1 --- /dev/null +++ b/sgpt/handlers/default_handler.py @@ -0,0 +1,22 @@ +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 diff --git a/sgpt/handlers/handler.py b/sgpt/handlers/handler.py new file mode 100644 index 0000000..a3a34bb --- /dev/null +++ b/sgpt/handlers/handler.py @@ -0,0 +1,194 @@ +import json +from pathlib import Path +from typing import Any, Callable, Dict, Generator, List, Optional, cast + +from rich.live_render import VerticalOverflowMethod + +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 +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: + vertical_overflow = cast( + VerticalOverflowMethod, cfg.get("MARKDOWN_LIVE_VERTICAL_OVERFLOW") + ) + refresh_interval = float(cfg.get("MARKDOWN_LIVE_REFRESH_INTERVAL")) + return ( + MarkdownPrinter(self.code_theme, refresh_interval, vertical_overflow) + 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) diff --git a/sgpt/handlers/repl_handler.py b/sgpt/handlers/repl_handler.py new file mode 100644 index 0000000..04019fd --- /dev/null +++ b/sgpt/handlers/repl_handler.py @@ -0,0 +1,66 @@ +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) diff --git a/sgpt/integration.py b/sgpt/integration.py new file mode 100644 index 0000000..fd19fc6 --- /dev/null +++ b/sgpt/integration.py @@ -0,0 +1,27 @@ +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 +""" diff --git a/sgpt/llm_functions/__init__.py b/sgpt/llm_functions/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/sgpt/llm_functions/common/execute_shell.py b/sgpt/llm_functions/common/execute_shell.py new file mode 100644 index 0000000..3f5d290 --- /dev/null +++ b/sgpt/llm_functions/common/execute_shell.py @@ -0,0 +1,42 @@ +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", []), + }, + }, + } diff --git a/sgpt/llm_functions/init_functions.py b/sgpt/llm_functions/init_functions.py new file mode 100644 index 0000000..d0aca1c --- /dev/null +++ b/sgpt/llm_functions/init_functions.py @@ -0,0 +1,35 @@ +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) diff --git a/sgpt/llm_functions/mac/apple_script.py b/sgpt/llm_functions/mac/apple_script.py new file mode 100644 index 0000000..2277a63 --- /dev/null +++ b/sgpt/llm_functions/mac/apple_script.py @@ -0,0 +1,47 @@ +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", []), + }, + }, + } diff --git a/sgpt/printer.py b/sgpt/printer.py new file mode 100644 index 0000000..6bcd293 --- /dev/null +++ b/sgpt/printer.py @@ -0,0 +1,93 @@ +import time +from abc import ABC, abstractmethod +from typing import Generator + +from rich.console import Console +from rich.live import Live +from rich.live_render import VerticalOverflowMethod +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, + refresh_interval: float, + vertical_overflow: VerticalOverflowMethod, + ) -> None: + self.console = Console() + self.theme = theme + self.refresh_interval = refresh_interval + self.vertical_overflow: VerticalOverflowMethod = vertical_overflow + + def live_print(self, chunks: Generator[str, None, None]) -> str: + full_completion = "" + with Live( + console=self.console, + vertical_overflow=self.vertical_overflow, + auto_refresh=False, + ) as live: + last_refresh = time.monotonic() + for chunk in chunks: + full_completion += chunk + if ( + self.refresh_interval == 0 + or time.monotonic() - last_refresh >= self.refresh_interval + ): + live.update( + Markdown(markup=full_completion, code_theme=self.theme), + refresh=True, + ) + last_refresh = time.monotonic() + + # Ensure the complete output is always rendered when streaming finishes. + live.update( + Markdown(markup=full_completion, code_theme=self.theme), + 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 diff --git a/sgpt/role.py b/sgpt/role.py new file mode 100644 index 0000000..31371e4 --- /dev/null +++ b/sgpt/role.py @@ -0,0 +1,189 @@ +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() diff --git a/sgpt/utils.py b/sgpt/utils.py new file mode 100644 index 0000000..d49af9a --- /dev/null +++ b/sgpt/utils.py @@ -0,0 +1,95 @@ +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__}") diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/_integration.py b/tests/_integration.py new file mode 100644 index 0000000..958eb62 --- /dev/null +++ b/tests/_integration.py @@ -0,0 +1,501 @@ +""" +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 +from tests.utils import assert_usage_error + +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-5.4-mini. + assert cfg.get("DEFAULT_MODEL") == "gpt-5.4-mini" + # 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_usage_error(result) + 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_usage_error(result) + + 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_usage_error(result) + + 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_usage_error(result) + + 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_usage_error( + result, + "Only one of --shell, --describe-shell, and --code", + ) + + 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, markdown=False) + assert handler.color == color + os.environ["DEFAULT_COLOR"] = "red" + handler = Handler(role=role, markdown=False) + 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. diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..bf0903e --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,8 @@ +import os + +import pytest + + +@pytest.fixture(autouse=True) +def mock_os_name(monkeypatch): + monkeypatch.setattr(os, "name", "test") diff --git a/tests/test_code.py b/tests/test_code.py new file mode 100644 index 0000000..76aad75 --- /dev/null +++ b/tests/test_code.py @@ -0,0 +1,141 @@ +from pathlib import Path +from unittest.mock import patch + +from sgpt.config import cfg +from sgpt.role import DefaultRoles, SystemRole + +from .utils import app, assert_usage_error, 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_usage_error(result) + 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_usage_error(result) + + +@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_usage_error(result) diff --git a/tests/test_default.py b/tests/test_default.py new file mode 100644 index 0000000..6dc2054 --- /dev/null +++ b/tests/test_default.py @@ -0,0 +1,248 @@ +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, assert_usage_error, 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: dict[str, str | bool] = {"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_usage_error(result) + + args["--code"] = True + result = runner.invoke(app, cmd_args(**args)) + assert_usage_error(result) + 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() diff --git a/tests/test_roles.py b/tests/test_roles.py new file mode 100644 index 0000000..7a8553c --- /dev/null +++ b/tests/test_roles.py @@ -0,0 +1,55 @@ +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) diff --git a/tests/test_shell.py b/tests/test_shell.py new file mode 100644 index 0000000..331df6e --- /dev/null +++ b/tests/test_shell.py @@ -0,0 +1,187 @@ +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, assert_usage_error, 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_usage_error(result) + 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_usage_error(result) + + +@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 diff --git a/tests/utils.py b/tests/utils.py new file mode 100644 index 0000000..1be59ed --- /dev/null +++ b/tests/utils.py @@ -0,0 +1,69 @@ +from datetime import datetime +from typing import Any + +import typer +from click import UsageError +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 assert_usage_error(result, message=None): + assert result.exit_code == 1, result + assert isinstance(result.exception, UsageError), result.exception + if message is not None: + assert message in str(result.exception), result.exception + + +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(**kwargs: Any) -> list[str]: + prompt = kwargs.pop("prompt", "") + 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, + }