Release 1.5.1 (#770)

- Update LiteLLM integration
- Update default model to gpt-5.4-mini
- Enforce UTF-8 cache writes

.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]'"
+}

.github/FUNDING.yml

@@ -0,0 +1 @@
+github: [ther1d]

.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

.github/workflows/docker.yml

@@ -3,11 +3,13 @@ name: Docker Image CI
 on:
   push:
     tags:
-      - '*'  
+      - '*'
 
 jobs:
   build:
     runs-on: ubuntu-latest
+    environment:
+      name: Docker Image
 
     steps:   
       - name: Checkout

.github/workflows/lint_test.yml

@@ -13,7 +13,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.9", "3.10"]
+        python-version: ["3.10", "3.11", "3.12"]
     steps:
     - uses: actions/checkout@v3
     - name: Set up Python ${{ matrix.python-version }}
@@ -31,8 +31,8 @@ jobs:
     - name: ruff
       run: ruff sgpt tests scripts
     - name: mypy
-      run: mypy sgpt
-    - name: unittests
+      run: mypy sgpt --exclude llm_functions
+    - name: tests
       run: |
         export OPENAI_API_KEY=test_api_key
-        pytest tests/test_unit.py
+        pytest tests/ -p no:warnings -v -s

.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 }}'

CONTRIBUITNG.md

@@ -1,37 +0,0 @@
-# Contributing to ShellGPT
-Thank you for considering contributing to ShellGPT (sgpt)! In order 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 would like to work on, assign it to yourself and leave a comment expressing your interest in working on it soon.
-* If you have a new feature in mind that doesn't have an existing issue, kindly 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 right away.
-
-## Developing
-> ShellGPT is written using strict types, which means you will need to define types. The project utilizes several linting and testing tools: ruff, mypy, isort, black, and pytest.
-
-### Virtual environment
-Create a virtual environment using Python venv and activate it:
-
-```shell
-python -m venv env && source ./env/bin/activate
-```
-
-### Install dependencies
-Install the necessary dependencies, in this case you will need to install the development and test dependencies:
-
-```shell
-pip install -e ."[dev,test]"
-```
-### Start coding
-With your environment set up and the issue assigned, you can begin working on your solution. Familiarize yourself with the existing codebase, and follow the project's coding style and conventions. Remember to write clean, modular, and maintainable code, which will make it easier for others to understand and review. As you make progress, commit your changes frequently to keep track of your work. 
-
-### Testing
-This is very important step. Every changes that implements a new feature or modifies the logic of existing features should include "integration" tests. These are tests that call `sgpt` with defined arguments, capture the output, and verify that the feature works as expected. See `test_integration.py` for examples. The tests should be easy to read and understand.
-
-### Pull request
-Before creating a pull request, ensure that you run `scripts/lint.sh` and `scripts/tests.sh`. All linters and tests should pass. In the pull request, provide a high-level description of your changes and step-by-step instructions on how to test them. Include any necessary commands.
-
-### Code review
-Once you've submitted your pull request, it's time for code review. Be patient and open to feedback from the reviewers. Address any concerns they may have and work collaboratively to improve the code. Together, we can make ShellGPT an even better project.
-
-Thank you once again for your contribution! We're excited to have you on board.

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.

Dockerfile

@@ -1,18 +1,15 @@
 FROM python:3-slim
 
-ENV PYTHONDONTWRITEBYTECODE 1
-ENV PYTHONUNBUFFERED 1
-ENV PIP_ROOT_USER_ACTION ignore
+ENV SHELL_INTERACTION=false
+ENV PRETTIFY_MARKDOWN=false
+ENV OS_NAME=auto
+ENV SHELL_NAME=auto
+
 WORKDIR /app
 COPY . /app
 
-RUN pip install --no-cache --upgrade pip \
- && pip install --no-cache /app \
- && addgroup --system app && adduser --system --group app \
- && mkdir -p /tmp/shell_gpt \
- && chown -R app:app /tmp/shell_gpt
-
-USER 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
 

README.md

@@ -1,114 +1,121 @@
 # ShellGPT
-A command-line productivity tool powered by OpenAI's GPT models. As developers, we can leverage AI capabilities to generate shell commands, code snippets, comments, and documentation, among other things. Forget about cheat sheets and notes, with this tool you can get accurate answers right in your terminal, and you'll probably find yourself reducing your daily Google searches, saving you valuable time and effort. ShellGPT is cross-platform compatible and supports all major operating systems, including Linux, macOS, and Windows with all major shells, such as PowerShell, CMD, Bash, Zsh, Fish, and many others.
+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://user-images.githubusercontent.com/16740832/231569156-a3a9f9d4-18b1-4fff-a6e1-6807651aa894.mp4
+https://github.com/TheR1D/shell_gpt/assets/16740832/721ddb19-97e7-428f-a0ee-107d027ddd59
 
 ## Installation
 ```shell
 pip install shell-gpt
 ```
-You'll need an OpenAI API key, you can generate one [here](https://beta.openai.com/account/api-keys).
+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.
 
-If the`$OPENAI_API_KEY` environment variable is set it will be used, otherwise, you will be prompted for your key which will then be stored in `~/.config/shell_gpt/.sgptrc`.
+> [!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
-`sgpt` has a variety of use cases, including simple queries, shell queries, and code queries.
-### Simple queries
-We can use it as normal search engine, asking about anything:
+**ShellGPT** is designed to quickly analyse and retrieve information. It's useful for straightforward requests ranging from technical configurations to general knowledge.
 ```shell
-sgpt "nginx default config file location"
-# -> The default configuration file for Nginx is located at /etc/nginx/nginx.conf.
+sgpt "What is the fibonacci sequence"
+# -> The Fibonacci sequence is a series of numbers where each number ...
 ```
-```shell
-sgpt "mass of sun"
-# -> = 1.99 × 10^30 kg
-```
-```shell
-sgpt "1 hour and 30 minutes to seconds"
-# -> 5,400 seconds
-```
-### Summarization and analyzing
-ShellGPT accepts prompt from both stdin and command line argument, you choose the most convenient input method for your preferences. Whether you prefer piping input through the terminal or specifying it directly as arguments, `sgpt` got you covered. This versatile feature is particularly useful when you need to pass file content or pipe output from other commands to the GPT models for summarization or analysis. For example, you can easily generate a git commit message based on a diff:
+
+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"
-# -> Commit message: Implement Model enum and get_edited_prompt()
+# -> Added main feature details into README.md
 ```
-You can analyze logs from various sources by passing them using stdin or command line arguments, along with a user-friendly prompt. This enables you to quickly identify errors and get suggestions for possible solutions:
+
+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 container_name | sgpt "check logs, find errors, provide possible solutions"
-# ...
+docker logs -n 20 my_app | sgpt "check logs, find errors, provide possible solutions"
 ```
-This powerful feature simplifies the process of managing and understanding data from different sources, making it easier for you to focus on what really matters: improving your projects and applications.
+```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 `chmod`, and needing to look up the syntax online? With `--shell` or shortcut `-s` option, you can quickly find and execute the commands you need right in the terminal.
+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 "make all files in current directory read only"
-# -> chmod 444 *
+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
-...
 ```
-We can ask GPT to describe suggested shell command, it will provide a short description of what the command does:
+
+Let's try it with Docker:
 ```shell
-sgpt -s "show all txt files in current folder"
-# -> ls *.txt
-# -> [E]xecute, [D]escribe, [A]bort: d
-# -> List all files with .txt extension in current directory
+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
-...
 ```
-Let's try some docker containers:
+
+We can still use pipes to pass input to `sgpt` and generate shell commands:
 ```shell
-sgpt -s "start nginx using docker, forward 443 and 80 port, mount current folder with index.html"
-# -> docker run -d -p 443:443 -p 80:80 -v $(pwd):/usr/share/nginx/html nginx
+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
-...
 ```
-We can still use pipes to pass input to `sgpt` and get shell commands as output:
-```shell
-cat data.json | sgpt -s "curl localhost with provided json"
-# -> curl -X POST -H "Content-Type: application/json" -d '{"a": 1, "b": 2, "c": 3}' http://localhost
-````
-We can apply additional shell magic in our prompt, in this example passing file names to ffmpeg:
+
+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 "using ffmpeg combine multiple videos into one without audio. Video file names: $(ls -m)"
+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
-Shell integration allows you to use Shell-GPT in your terminal with hotkeys. It is currently available for bash and zsh. It will allow you to have sgpt completions in your shell history, and also edit suggested commands right away.
+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:
-```shell
-sgpt --install-integration
-# 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 Shell-GPT. When you press `Ctrl+l` it will replace you current input line (buffer) with suggested command. You can then edit it and press `Enter` to execute.
+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
-With `--code` parameters we can query only code as output, for example:
+By using the `--code` or `-c` parameter, you can specifically request pure code output, for instance:
 ```shell
-sgpt --code "Solve classic fizz buzz problem using Python"
+sgpt --code "solve fizz buzz problem using python"
 ```
+
 ```python
 for i in range(1, 101):
     if i % 3 == 0 and i % 5 == 0:
@@ -120,7 +127,7 @@ for i in range(1, 101):
     else:
         print(i)
 ```
-Since it is valid python code, we can redirect the output to file:
+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
@@ -129,12 +136,12 @@ python fizz_buzz.py
 # Fizz
 # 4
 # Buzz
-# Fizz
 # ...
 ```
-We can also use pipes to pass input to `sgpt`:
+
+We can also use pipes to pass input:
 ```shell
-cat fizz_buzz.py | python -m sgpt --code "Generate comments for each line of my code"
+cat fizz_buzz.py | sgpt --code "Generate comments for each line of my code"
 ```
 ```python
 # Loop through numbers 1 to 100
@@ -156,17 +163,20 @@ for i in range(1, 101):
         print(i)
 ```
 
-### Chat
-To start a chat session, use the `--chat` option followed by a unique session name and a prompt. You can also use "temp" as a session name to start a temporary chat session.
+### 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 number "please remember my favorite number: 4"
+sgpt --chat conversation_1 "please remember my favorite number: 4"
 # -> I will remember that your favorite number is 4.
-sgpt --chat number "what would be my favorite number + 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 also use chat sessions to iteratively improve GPT suggestions by providing additional clues.
+
+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 python_request --code "make an example request to localhost using Python"
+sgpt --chat conversation_2 --code "make a request to localhost using python"
 ```
 ```python
 import requests
@@ -174,9 +184,10 @@ import requests
 response = requests.get('http://localhost')
 print(response.text)
 ```
-Asking AI to add a cache to our request.
+
+Let's ask LLM to add caching to our request:
 ```shell
-sgpt --chat python_request --code "add caching"
+sgpt --chat conversation_2 --code "add caching"
 ```
 ```python
 import requests
@@ -188,20 +199,37 @@ cached_sess = CacheControl(sess)
 response = cached_sess.get('http://localhost')
 print(response.text)
 ```
-We can use `--code` or `--shell` options to initiate `--chat`, so you can keep refining the results:
+
+Same applies for shell commands:
 ```shell
-sgpt --chat sh --shell "What are the files in this directory?"
+sgpt --chat conversation_3 --shell "what is in current folder"
 # -> ls
-sgpt --chat sh "Sort them by name"
+sgpt --chat conversation_3 "Sort by name"
 # -> ls | sort
-sgpt --chat sh "Concatenate them using FFMPEG"
+sgpt --chat conversation_3 "Concatenate them using FFMPEG"
 # -> ffmpeg -i "concat:$(ls | sort | tr '\n' '|')" -codec copy output.mp4
-sgpt --chat sh "Convert the resulting file into an MP3"
+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
 ```
 
-### REPL
-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 chat sessions, so you can use `--chat` to start a chat session and then use `--repl` to continue the conversation in REPL mode. REPL mode will also show history of your conversation in the beginning.
+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.
 
 <p align="center">
   <img src="https://s10.gifyu.com/images/repl-demo.gif" alt="gif">
@@ -215,6 +243,7 @@ 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
@@ -226,45 +255,82 @@ ls -lh
 >>> Sort them by file sizes
 ls -lhS
 >>> e (enter just e to execute commands, or d to describe them)
-...
-```
-Example of using REPL mode to generate code:
-```text
-sgpt --repl temp --code
-Entering REPL mode, press Ctrl+C to exit.
->>> Using Python request localhost:80
-import requests
-response = requests.get('http://localhost:80')
-print(response.text)
->>> Change port to 443
-import requests
-response = requests.get('https://localhost:443')
-print(response.text)
 ```
 
-### Chat sessions
-To list all the current chat sessions, use the `--list-chats` option:
-```shell
-sgpt --list-chats
-# .../shell_gpt/chat_cache/number
-# .../shell_gpt/chat_cache/python_request
+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.
 ```
-To show all the messages related to a specific chat session, use the `--show-chat` option followed by the session name:
+
+You can also enter REPL mode with initial prompt by passing it as an argument or stdin or even both:
 ```shell
-sgpt --show-chat number
-# 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.
+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
-# Enter role description: You are JSON generator, provide only valid json as response.
-# Enter expecting result, e.g. answer, code, shell command, etc.: json
-sgpt --role json "random: user, password, email, address"
+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",
@@ -278,6 +344,8 @@ sgpt --role json "random: user, password, email, address"
 }
 ```
 
+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
@@ -293,8 +361,8 @@ You can setup some parameters in runtime configuration file `~/.config/shell_gpt
 ```text
 # API key, also it is possible to define OPENAI_API_KEY env.
 OPENAI_API_KEY=your_api_key
-# OpenAI host, useful if you would like to use proxy.
-OPENAI_API_HOST=https://api.openai.com
+# 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.
@@ -306,62 +374,77 @@ CACHE_PATH=/tmp/shell_gpt/cache
 # Request timeout in seconds.
 REQUEST_TIMEOUT=60
 # Default OpenAI model to use.
-DEFAULT_MODEL=gpt-3.5-turbo
-# Default color for OpenAI completions.
+DEFAULT_MODEL=gpt-5.4-mini
+# Default color for shell and code completions.
 DEFAULT_COLOR=magenta
-# Force use system role messages (not recommended).
-SYSTEM_ROLES=false
 # When in --shell mode, default to "Y" for no input.
 DEFAULT_EXECUTE_SHELL_CMD=false
+# Disable streaming of responses
+DISABLE_STREAMING=false
+# The pygment theme to view markdown (default/describe role).
+CODE_THEME=default
+# Path to a directory with functions.
+OPENAI_FUNCTIONS_PATH=/Users/user/.config/shell_gpt/functions
+# Print output of functions when LLM uses them.
+SHOW_FUNCTIONS_OUTPUT=false
+# Allows LLM to use functions.
+OPENAI_USE_FUNCTIONS=true
+# Enforce LiteLLM usage (for local LLMs).
+USE_LITELLM=false
 ```
 Possible options for `DEFAULT_COLOR`: black, red, green, yellow, blue, magenta, cyan, white, bright_black, bright_red, bright_green, bright_yellow, bright_blue, bright_magenta, bright_cyan, bright_white.
-
-Switch `SYSTEM_ROLES` to force use [system roles](https://help.openai.com/en/articles/7042661-chatgpt-api-transition-guide) messages, this is not recommended, since it doesn't perform well with current GPT models.
+Possible options for `CODE_THEME`: https://pygments.org/styles/
 
 ### Full list of arguments
 ```text
-╭─ Arguments ─────────────────────────────────────────────────────────────────────────────────────────────────╮
-│   prompt      [PROMPT]  The prompt to generate completions for.                                             │
-╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
-╭─ Options ───────────────────────────────────────────────────────────────────────────────────────────────────╮
-│ --model    [gpt-4|gpt-4-32k|gpt-3.5|gpt-3.5-16k]    OpenAI GPT model to use. [default: gpt-3.5-turbo]       │
-│ --temperature      FLOAT RANGE [0.0<=x<=2.0]        Randomness of generated output. [default: 0.1]          │
-│ --top-probability  FLOAT RANGE [0.1<=x<=1.0]        Limits highest probable tokens (words). [default: 1.0]  │
-│ --editor                                            Open $EDITOR to provide a prompt. [default: no-editor]  │
-│ --cache                                             Cache completion results. [default: cache]              │
-│ --help                                              Show this message and exit.                             │
-╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
-╭─ Assistance Options ────────────────────────────────────────────────────────────────────────────────────────╮
-│ --shell  -s                 Generate and execute shell commands.                                            │
-│ --describe-shell  -d        Describe a shell command.                                                       │
-│ --code       --no-code      Generate only code. [default: no-code]                                          │
-╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
-╭─ 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        List all existing chat ids. [default: no-list-chats]                                    │
-╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
-╭─ 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         List roles. [default: no-list-roles]                                                   │
-╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
+╭─ 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:
+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="your OPENAI API key" \
+           --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 --chat rainbow "what are the colors of a rainbow"
+       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 --env OPENAI_API_KEY --volume gpt-cache:/tmp/shell_gpt ghcr.io/ther1d/shell_gpt"
+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"
@@ -372,3 +455,7 @@ 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)

pyproject.toml

@@ -4,11 +4,11 @@ build-backend = "hatchling.build"
 
 [project]
 name = "shell_gpt"
-description = "A command-line productivity tool powered by OpenAI GPT models, will help you accomplish your tasks faster and more efficiently."
-keywords = ["shell", "gpt", "openai", "cli", "productivity", "cheet-sheet"]
+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.6"
+requires-python = ">=3.10"
 authors = [{ name = "Farkhod Sadykov", email = "farkhod@sadykov.dev" }]
 dynamic = ["version"]
 classifiers = [
@@ -18,22 +18,36 @@ classifiers = [
     "Intended Audience :: Information Technology",
     "Intended Audience :: System Administrators",
     "Intended Audience :: Developers",
-    "Programming Language :: Python :: 3",
-    "Programming Language :: Python :: 3 :: Only",
-    "Programming Language :: Python :: 3.6",
-    "Programming Language :: Python :: 3.7",
-    "Programming Language :: Python :: 3.8",
-    "Programming Language :: Python :: 3.9",
     "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
 ]
 dependencies = [
-    "requests >= 2.28.2, < 3.0.0",
+    "openai >= 2.0.0, < 3.0.0",
     "typer >= 0.7.0, < 1.0.0",
-    "click >= 7.1.1, < 9.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]
@@ -44,22 +58,8 @@ 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"
 
-[project.optional-dependencies]
-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",
-]
-dev = [
-    "ruff == 0.0.256",
-    "pre-commit >= 3.1.1, < 4.0.0",
-]
-
 [tool.hatch.version]
-path = "sgpt/__init__.py"
+path = "sgpt/__version__.py"
 
 [tool.hatch.build.targets.wheel]
 only-include = ["sgpt"]
@@ -79,6 +79,7 @@ skip =  "__init__.py"
 
 [tool.mypy]
 strict = true
+exclude = ["llm_functions"]
 
 [tool.ruff]
 select = [
@@ -92,4 +93,8 @@ 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'

scripts/format.sh

@@ -4,3 +4,4 @@ set -x
 ruff sgpt tests scripts --fix
 black sgpt tests scripts
 isort sgpt tests scripts
+codespell --write-changes

scripts/lint.sh

@@ -7,3 +7,4 @@ mypy sgpt
 ruff sgpt tests scripts
 black sgpt tests --check
 isort sgpt tests scripts --check-only
+codespell

scripts/test.sh

@@ -4,4 +4,4 @@ set -e
 set -x
 
 # shellcheck disable=SC2068
-pytest tests ${@}
+pytest tests ${@} -p no:warnings

sgpt/__init__.py

@@ -1,4 +1,2 @@
 from .app import main as main
 from .app import entry_point as cli  # noqa: F401
-
-__version__ = "0.9.3"

sgpt/__version__.py

@@ -0,0 +1 @@
+__version__ = "1.5.1"

sgpt/app.py

@@ -1,26 +1,24 @@
-"""
-This module provides a simple interface for OpenAI API using Typer
-as the command line interface. It supports different modes of output including
-shell commands and code, and allows users to specify the desired OpenAI model
-and length and other options of the output. Additionally, it supports executing
-shell commands directly from the interface.
-"""
+import os
+
 # To allow users to use arrow keys in the REPL.
 import readline  # noqa: F401
 import sys
 
 import typer
-from click import BadArgumentUsage, MissingParameter
+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 (
-    ModelOptions,
     get_edited_prompt,
+    get_sgpt_version,
     install_shell_integration,
     run_command,
 )
@@ -28,26 +26,30 @@ from sgpt.utils import (
 
 def main(
     prompt: str = typer.Argument(
-        None,
+        "",
         show_default=False,
         help="The prompt to generate completions for.",
     ),
-    model: ModelOptions = typer.Option(
-        ModelOptions(cfg.get("DEFAULT_MODEL")).value,
-        help="OpenAI GPT model to use.",
+    model: str = typer.Option(
+        cfg.get("DEFAULT_MODEL"),
+        help="Large language model to use.",
     ),
     temperature: float = typer.Option(
-        0.1,
+        0.0,
         min=0.0,
         max=2.0,
         help="Randomness of generated output.",
     ),
-    top_probability: float = typer.Option(
+    top_p: float = typer.Option(
         1.0,
-        min=0.1,
+        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",
@@ -55,6 +57,11 @@ def main(
         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",
@@ -64,9 +71,16 @@ def main(
     ),
     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.",
@@ -75,6 +89,12 @@ def main(
         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.',
@@ -82,17 +102,18 @@ def main(
     ),
     repl: str = typer.Option(
         None,
-        help="Start a REPL (Read–eval–print loop) session.",
+        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.",
-        callback=ChatHandler.show_messages_callback,
         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",
@@ -116,6 +137,8 @@ def main(
     ),
     list_roles: bool = typer.Option(
         False,
+        "--list-roles",
+        "-lr",
         help="List roles.",
         callback=SystemRole.list,
         rich_help_panel="Role Options",
@@ -126,25 +149,53 @@ def main(
         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 and not repl:
-        prompt = f"{sys.stdin.read()}\n\n{prompt or ''}"
+    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 not prompt and not editor and not repl:
-        raise MissingParameter(param_hint="PROMPT", param_type="string")
+    if show_chat:
+        ChatHandler.show_messages(show_chat, md)
 
     if sum((shell, describe_shell, code)) > 1:
-        raise BadArgumentUsage(
+        raise UsageError(
             "Only one of --shell, --describe-shell, and --code options can be used at a time."
         )
 
     if chat and repl:
-        raise BadArgumentUsage("--chat and --repl options cannot be used together.")
+        raise UsageError("--chat and --repl options cannot be used together.")
 
     if editor and stdin_passed:
-        raise BadArgumentUsage("--editor option cannot be used with stdin input.")
+        raise UsageError("--editor option cannot be used with stdin input.")
 
     if editor:
         prompt = get_edited_prompt()
@@ -155,60 +206,69 @@ def main(
         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).handle(
-            prompt,
-            model=model.value,
+        ReplHandler(repl, role_class, md).handle(
+            init_prompt=prompt,
+            model=model,
             temperature=temperature,
-            top_probability=top_probability,
-            chat_id=repl,
+            top_p=top_p,
             caching=cache,
+            functions=function_schemas,
         )
 
     if chat:
-        full_completion = ChatHandler(chat, role_class).handle(
-            prompt,
-            model=model.value,
+        full_completion = ChatHandler(chat, role_class, md).handle(
+            prompt=prompt,
+            model=model,
             temperature=temperature,
-            top_probability=top_probability,
-            chat_id=chat,
+            top_p=top_p,
             caching=cache,
+            functions=function_schemas,
         )
     else:
-        full_completion = DefaultHandler(role_class).handle(
-            prompt,
-            model=model.value,
+        full_completion = DefaultHandler(role_class, md).handle(
+            prompt=prompt,
+            model=model,
             temperature=temperature,
-            top_probability=top_probability,
+            top_p=top_p,
             caching=cache,
+            functions=function_schemas,
         )
 
-    while shell and not stdin_passed:
+    session: PromptSession[str] = PromptSession()
+
+    while shell and interaction:
         option = typer.prompt(
-            text="[E]xecute, [D]escribe, [A]bort",
-            type=Choice(("e", "d", "a", "y"), case_sensitive=False),
+            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()).handle(
+            DefaultHandler(DefaultRoles.DESCRIBE_SHELL.get_role(), md).handle(
                 full_completion,
-                model=model.value,
+                model=model,
                 temperature=temperature,
-                top_probability=top_probability,
+                top_p=top_p,
                 caching=cache,
+                functions=function_schemas,
             )
             continue
         break
 
 
 def entry_point() -> None:
-    # Python package entry point defined in setup.py
     typer.run(main)
 
 

sgpt/cache.py

@@ -28,18 +28,17 @@ class Cache:
         """
 
         def wrapper(*args: Any, **kwargs: Any) -> Generator[str, None, None]:
-            # Exclude self instance from hashing.
-            cache_key = md5(json.dumps((args[1:], kwargs)).encode("utf-8")).hexdigest()
-            cache_file = self.cache_path / cache_key
-            # TODO: Fix caching for chat, should hash last user message, (not entire history).
-            if kwargs.pop("caching", True) and cache_file.exists():
-                yield cache_file.read_text()
+            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
-            cache_file.write_text(result)
+            if "@FunctionCall" not in result:
+                file.write_text(result, encoding="utf-8")
             self._delete_oldest_files(self.length)  # type: ignore
 
         return wrapper

sgpt/client.py

@@ -1,98 +0,0 @@
-import json
-from pathlib import Path
-from typing import Dict, Generator, List
-
-import requests
-
-from .cache import Cache
-from .config import cfg
-
-CACHE_LENGTH = int(cfg.get("CACHE_LENGTH"))
-CACHE_PATH = Path(cfg.get("CACHE_PATH"))
-REQUEST_TIMEOUT = int(cfg.get("REQUEST_TIMEOUT"))
-
-
-class OpenAIClient:
-    cache = Cache(CACHE_LENGTH, CACHE_PATH)
-
-    def __init__(self, api_host: str, api_key: str) -> None:
-        self.__api_key = api_key
-        self.api_host = api_host
-
-    @cache
-    def _request(
-        self,
-        messages: List[Dict[str, str]],
-        model: str = "gpt-3.5-turbo",
-        temperature: float = 1,
-        top_probability: float = 1,
-    ) -> Generator[str, None, None]:
-        """
-        Make request to OpenAI API, read more:
-        https://platform.openai.com/docs/api-reference/chat
-
-        :param messages: List of messages {"role": user or assistant, "content": message_string}
-        :param model: String gpt-3.5-turbo or gpt-3.5-turbo-0301
-        :param temperature: Float in 0.0 - 2.0 range.
-        :param top_probability: Float in 0.0 - 1.0 range.
-        :return: Response body JSON.
-        """
-        data = {
-            "messages": messages,
-            "model": model,
-            "temperature": temperature,
-            "top_p": top_probability,
-            "stream": True,
-        }
-        endpoint = f"{self.api_host}/v1/chat/completions"
-        response = requests.post(
-            endpoint,
-            # Hide API key from Rich traceback.
-            headers={
-                "Content-Type": "application/json",
-                "Authorization": f"Bearer {self.__api_key}",
-            },
-            json=data,
-            timeout=REQUEST_TIMEOUT,
-            stream=True,
-        )
-        response.raise_for_status()
-        # TODO: Optimise.
-        # https://github.com/openai/openai-python/blob/237448dc072a2c062698da3f9f512fae38300c1c/openai/api_requestor.py#L98
-        for line in response.iter_lines():
-            data = line.lstrip(b"data: ").decode("utf-8")
-            if data == "[DONE]":  # type: ignore
-                break
-            if not data:
-                continue
-            data = json.loads(data)  # type: ignore
-            delta = data["choices"][0]["delta"]  # type: ignore
-            if "content" not in delta:
-                continue
-            yield delta["content"]
-
-    def get_completion(
-        self,
-        messages: List[Dict[str, str]],
-        model: str = "gpt-3.5-turbo",
-        temperature: float = 1,
-        top_probability: float = 1,
-        caching: bool = True,
-    ) -> Generator[str, None, None]:
-        """
-        Generates single completion for prompt (message).
-
-        :param messages: List of dict with messages and roles.
-        :param model: String gpt-3.5-turbo or gpt-3.5-turbo-0301.
-        :param temperature: Float in 0.0 - 1.0 range.
-        :param top_probability: Float in 0.0 - 1.0 range.
-        :param caching: Boolean value to enable/disable caching.
-        :return: String generated completion.
-        """
-        yield from self._request(
-            messages,
-            model,
-            temperature,
-            top_probability,
-            caching=caching,
-        )

sgpt/config.py

@@ -6,12 +6,11 @@ from typing import Any
 
 from click import UsageError
 
-from .utils import ModelOptions
-
 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"
 
@@ -23,12 +22,21 @@ DEFAULT_CONFIG = {
     "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", ModelOptions.GPT35TURBO.value),
-    "OPENAI_API_HOST": os.getenv("OPENAI_API_HOST", "https://api.openai.com"),
+    "DEFAULT_MODEL": os.getenv("DEFAULT_MODEL", "gpt-5.4-mini"),
     "DEFAULT_COLOR": os.getenv("DEFAULT_COLOR", "magenta"),
     "ROLE_STORAGE_PATH": os.getenv("ROLE_STORAGE_PATH", str(ROLE_STORAGE_PATH)),
-    "SYSTEM_ROLES": os.getenv("SYSTEM_ROLES", "false"),
     "DEFAULT_EXECUTE_SHELL_CMD": os.getenv("DEFAULT_EXECUTE_SHELL_CMD", "false"),
+    "DISABLE_STREAMING": os.getenv("DISABLE_STREAMING", "false"),
+    "CODE_THEME": os.getenv("CODE_THEME", "dracula"),
+    "OPENAI_FUNCTIONS_PATH": os.getenv("OPENAI_FUNCTIONS_PATH", str(FUNCTIONS_PATH)),
+    "OPENAI_USE_FUNCTIONS": os.getenv("OPENAI_USE_FUNCTIONS", "true"),
+    "SHOW_FUNCTIONS_OUTPUT": os.getenv("SHOW_FUNCTIONS_OUTPUT", "false"),
+    "API_BASE_URL": os.getenv("API_BASE_URL", "default"),
+    "PRETTIFY_MARKDOWN": os.getenv("PRETTIFY_MARKDOWN", "true"),
+    "USE_LITELLM": os.getenv("USE_LITELLM", "false"),
+    "SHELL_INTERACTION": os.getenv("SHELL_INTERACTION ", "true"),
+    "OS_NAME": os.getenv("OS_NAME", "auto"),
+    "SHELL_NAME": os.getenv("SHELL_NAME", "auto"),
     # New features might add their own config variables here.
 }
 
@@ -69,8 +77,8 @@ class Config(dict):  # type: ignore
     def _read(self) -> None:
         with open(self.config_path, "r", encoding="utf-8") as file:
             for line in file:
-                if not line.startswith("#"):
-                    key, value = line.strip().split("=")
+                if line.strip() and not line.startswith("#"):
+                    key, value = line.strip().split("=", 1)
                     self[key] = value
 
     def get(self, key: str) -> str:  # type: ignore

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]

sgpt/handlers/chat_handler.py

@@ -3,10 +3,13 @@ from pathlib import Path
 from typing import Any, Callable, Dict, Generator, List, Optional
 
 import typer
-from click import BadArgumentUsage
+from click import BadParameter, UsageError
+from rich.console import Console
+from rich.markdown import Markdown
 
 from ..config import cfg
-from ..role import SystemRole
+from ..role import DefaultRoles, SystemRole
+from ..utils import option_callback
 from .handler import Handler
 
 CHAT_CACHE_LENGTH = int(cfg.get("CHAT_CACHE_LENGTH"))
@@ -41,19 +44,20 @@ class ChatSession:
 
         def wrapper(*args: Any, **kwargs: Any) -> Generator[str, None, None]:
             chat_id = kwargs.pop("chat_id", None)
-            messages = kwargs["messages"]
+            if not kwargs.get("messages"):
+                return
             if not chat_id:
                 yield from func(*args, **kwargs)
                 return
-            old_messages = self._read(chat_id)
-            for message in messages:
-                old_messages.append(message)
-            kwargs["messages"] = old_messages
+            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
-            old_messages.append({"role": "assistant", "content": response_text})
+            previous_messages.append({"role": "assistant", "content": response_text})
             self._write(kwargs["messages"], chat_id)
 
         return wrapper
@@ -67,7 +71,11 @@ class ChatSession:
 
     def _write(self, messages: List[Dict[str, str]], chat_id: str) -> None:
         file_path = self.storage_path / chat_id
-        json.dump(messages[-self.length :], file_path.open("w"))
+        # 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
@@ -90,8 +98,8 @@ class ChatSession:
 class ChatHandler(Handler):
     chat_session = ChatSession(CHAT_CACHE_LENGTH, CHAT_CACHE_PATH)
 
-    def __init__(self, chat_id: str, role: SystemRole) -> None:
-        super().__init__(role)
+    def __init__(self, chat_id: str, role: SystemRole, markdown: bool) -> None:
+        super().__init__(role, markdown)
         self.chat_id = chat_id
         self.role = role
 
@@ -101,79 +109,69 @@ class ChatHandler(Handler):
 
         self.validate()
 
-    @classmethod
-    def list_ids(cls, value: str) -> None:
-        if not value:
-            return
-        # Prints all existing chat IDs to the console.
-        for chat_id in cls.chat_session.list():
-            typer.echo(chat_id)
-        raise typer.Exit()
-
     @property
     def initiated(self) -> bool:
         return self.chat_session.exists(self.chat_id)
 
-    @property
-    def initial_message(self) -> str:
-        chat_history = self.chat_session.get_messages(self.chat_id)
-        index = 1 if cfg.get("SYSTEM_ROLES") == "true" else 0
-        return chat_history[index] if chat_history else ""
-
     @property
     def is_same_role(self) -> bool:
         # TODO: Should be optimized for REPL mode.
-        return self.role.same_role(self.initial_message)
+        return self.role.same_role(self.initial_message(self.chat_id))
 
     @classmethod
-    def show_messages_callback(cls, chat_id: str) -> None:
-        if not chat_id:
+    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
-        cls.show_messages(chat_id)
-        raise typer.Exit()
 
-    @classmethod
-    def show_messages(cls, chat_id: str) -> None:
-        # Prints all messages from a specified chat ID to the console.
         for index, message in enumerate(cls.chat_session.get_messages(chat_id)):
-            # Remove output type from the message, e.g. "text\nCommand:" -> "text"
-            if message.startswith("user:"):
-                message = "\n".join(message.splitlines()[:-1])
-            color = "magenta" if index % 2 == 0 else "green"
-            typer.secho(message, fg=color)
+            running_color = color if index % 2 == 0 else "green"
+            typer.secho(message, fg=running_color)
 
     def validate(self) -> None:
         if self.initiated:
-            # print("initial message:", self.initial_message)
-            chat_role_name = self.role.get_role_name(self.initial_message)
+            chat_role_name = self.role.get_role_name(self.initial_message(self.chat_id))
             if not chat_role_name:
-                raise BadArgumentUsage(
-                    f'Could not determine chat role of "{self.chat_id}"'
-                )
-            if self.role.name == "default":
+                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 BadArgumentUsage(
+                    raise UsageError(
                         f'Cant change chat role to "{self.role.name}" '
                         f'since it was initiated as "{chat_role_name}" chat.'
                     )
 
-    def make_prompt(self, prompt: str) -> str:
-        prompt = prompt.strip()
-        return self.role.make_prompt(prompt, not self.initiated)
-
     def make_messages(self, prompt: str) -> List[Dict[str, str]]:
         messages = []
-        if not self.initiated and cfg.get("SYSTEM_ROLES") == "true":
+        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]:
+    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)

sgpt/handlers/default_handler.py

@@ -10,17 +10,13 @@ CHAT_CACHE_PATH = Path(cfg.get("CHAT_CACHE_PATH"))
 
 
 class DefaultHandler(Handler):
-    def __init__(self, role: SystemRole) -> None:
-        super().__init__(role)
+    def __init__(self, role: SystemRole, markdown: bool) -> None:
+        super().__init__(role, markdown)
         self.role = role
 
-    def make_prompt(self, prompt: str) -> str:
-        prompt = prompt.strip()
-        return self.role.make_prompt(prompt, initial=True)
-
     def make_messages(self, prompt: str) -> List[Dict[str, str]]:
-        messages = []
-        if cfg.get("SYSTEM_ROLES") == "true":
-            messages.append({"role": "system", "content": self.role.role})
-        messages.append({"role": "user", "content": prompt})
+        messages = [
+            {"role": "system", "content": self.role.role},
+            {"role": "user", "content": prompt},
+        ]
         return messages

sgpt/handlers/handler.py

@@ -1,34 +1,188 @@
-from typing import Any, Dict, Generator, List
+import json
+from pathlib import Path
+from typing import Any, Callable, Dict, Generator, List, Optional
 
-import typer
-
-from ..client import OpenAIClient
+from ..cache import Cache
 from ..config import cfg
-from ..role import SystemRole
+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:
-    def __init__(self, role: SystemRole) -> None:
-        self.client = OpenAIClient(
-            cfg.get("OPENAI_API_HOST"), cfg.get("OPENAI_API_KEY")
-        )
-        self.role = role
-        self.color = cfg.get("DEFAULT_COLOR")
+    cache = Cache(int(cfg.get("CACHE_LENGTH")), Path(cfg.get("CACHE_PATH")))
 
-    def make_prompt(self, prompt: str) -> str:
-        raise NotImplementedError
+    def __init__(self, role: SystemRole, markdown: bool) -> None:
+        self.role = role
+
+        api_base_url = cfg.get("API_BASE_URL")
+        self.base_url = None if api_base_url == "default" else api_base_url
+        self.timeout = int(cfg.get("REQUEST_TIMEOUT"))
+
+        self.markdown = "APPLY MARKDOWN" in self.role.role and markdown
+        self.code_theme, self.color = cfg.get("CODE_THEME"), cfg.get("DEFAULT_COLOR")
+
+    @property
+    def printer(self) -> Printer:
+        return (
+            MarkdownPrinter(self.code_theme)
+            if self.markdown
+            else TextPrinter(self.color)
+        )
 
     def make_messages(self, prompt: str) -> List[Dict[str, str]]:
         raise NotImplementedError
 
-    def get_completion(self, **kwargs: Any) -> Generator[str, None, None]:
-        yield from self.client.get_completion(**kwargs)
+    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},
+                    }
+                ],
+            }
+        )
 
-    def handle(self, prompt: str, **kwargs: Any) -> str:
-        messages = self.make_messages(self.make_prompt(prompt))
-        full_completion = ""
-        for word in self.get_completion(messages=messages, **kwargs):
-            typer.secho(word, fg=self.color, bold=True, nl=False)
-            full_completion += word
-        typer.echo()
-        return full_completion
+        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)

sgpt/handlers/repl_handler.py

@@ -11,13 +11,20 @@ from .default_handler import DefaultHandler
 
 
 class ReplHandler(ChatHandler):
-    def __init__(self, chat_id: str, role: SystemRole) -> None:
-        super().__init__(chat_id, role)
+    def __init__(self, chat_id: str, role: SystemRole, markdown: bool) -> None:
+        super().__init__(chat_id, role, markdown)
 
-    def handle(self, prompt: str, **kwargs: Any) -> None:  # type: ignore
+    @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.show_messages(self.chat_id, self.markdown)
             rich_print(Rule(style="bold magenta"))
 
         info_message = (
@@ -30,25 +37,30 @@ class ReplHandler(ChatHandler):
         )
         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()":
-                # This is also useful during tests.
                 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()).handle(
-                    full_completion,
-                    model=kwargs.get("model"),
-                    temperature=kwargs.get("temperature"),
-                    top_probability=kwargs.get("top_probability"),
-                    caching=kwargs.get("caching"),
-                )
+                DefaultHandler(
+                    DefaultRoles.DESCRIBE_SHELL.get_role(), self.markdown
+                ).handle(prompt=full_completion, **kwargs)
             else:
-                full_completion = super().handle(prompt, **kwargs)
+                full_completion = super().handle(prompt=prompt, **kwargs)

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
+"""

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", []),
+                },
+            },
+        }

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)

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", []),
+                },
+            },
+        }

sgpt/printer.py

@@ -0,0 +1,65 @@
+from abc import ABC, abstractmethod
+from typing import Generator
+
+from rich.console import Console
+from rich.live import Live
+from rich.markdown import Markdown
+from typer import secho
+
+
+class Printer(ABC):
+    console = Console()
+
+    @abstractmethod
+    def live_print(self, chunks: Generator[str, None, None]) -> str:
+        pass
+
+    @abstractmethod
+    def static_print(self, text: str) -> str:
+        pass
+
+    def __call__(self, chunks: Generator[str, None, None], live: bool = True) -> str:
+        if live:
+            return self.live_print(chunks)
+        with self.console.status("[bold green]Loading..."):
+            full_completion = "".join(chunks)
+        self.static_print(full_completion)
+        return full_completion
+
+
+class MarkdownPrinter(Printer):
+    def __init__(self, theme: str) -> None:
+        self.console = Console()
+        self.theme = theme
+
+    def live_print(self, chunks: Generator[str, None, None]) -> str:
+        full_completion = ""
+        with Live(console=self.console) as live:
+            for chunk in chunks:
+                full_completion += chunk
+                markdown = Markdown(markup=full_completion, code_theme=self.theme)
+                live.update(markdown, refresh=True)
+        return full_completion
+
+    def static_print(self, text: str) -> str:
+        markdown = Markdown(markup=text, code_theme=self.theme)
+        self.console.print(markdown)
+        return text
+
+
+class TextPrinter(Printer):
+    def __init__(self, color: str) -> None:
+        self.color = color
+
+    def live_print(self, chunks: Generator[str, None, None]) -> str:
+        full_text = ""
+        for chunk in chunks:
+            full_text += chunk
+            secho(chunk, fg=self.color, nl=False)
+        else:
+            print()  # Add new line after last chunk.
+        return full_text
+
+    def static_print(self, text: str) -> str:
+        secho(text, fg=self.color)
+        return text

sgpt/role.py

@@ -7,7 +7,7 @@ from pathlib import Path
 from typing import Dict, Optional
 
 import typer
-from click import BadArgumentUsage
+from click import UsageError
 from distro import name as distro_name
 
 from .config import cfg
@@ -16,34 +16,32 @@ 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."""
+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. Provide only plain text without Markdown formatting.
-Do not show any warnings or information regarding your capabilities.
-If you need to store any data, assume it will be stored in the chat."""
+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.
-IMPORTANT: Provide only plain text without Markdown formatting.
-IMPORTANT: Do not include markdown formatting such as ```.
+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.
-Ignore any potential risk of errors or confusion."""
+For example if the prompt is "Hello world Python", you should return "print('Hello world')"."""
 
-DEFAULT_ROLE = """You are Command Line App ShellGPT, a programming and system administration assistant.
+DEFAULT_ROLE = """You are programming and system administration assistant.
 You are managing {os} operating system with {shell} shell.
-Provide only plain text without Markdown formatting.
-Do not show any warnings or information regarding your capabilities.
-If you need to store any data, assume it will be stored in the chat."""
+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.
 
-
-PROMPT_TEMPLATE = """###
-Role name: {name}
-{role}
-
-Request: {request}
-###
-{expecting}:"""
+ROLE_TEMPLATE = "You are {name}\n{role}"
 
 
 class SystemRole:
@@ -53,76 +51,40 @@ class SystemRole:
         self,
         name: str,
         role: str,
-        expecting: str,
         variables: Optional[Dict[str, str]] = None,
     ) -> None:
         self.storage.mkdir(parents=True, exist_ok=True)
         self.name = name
-        self.expecting = expecting
-        self.variables = variables
         if variables:
-            # Variables are for internal use only.
             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()}
+        variables = {"shell": cls._shell_name(), "os": cls._os_name()}
         for default_role in (
-            SystemRole("default", DEFAULT_ROLE, "Answer", variables),
-            SystemRole("shell", SHELL_ROLE, "Command", variables),
-            SystemRole("describe_shell", DESCRIBE_SHELL_ROLE, "Description", variables),
-            SystemRole("code", CODE_ROLE, "Code"),
+            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 os_name(cls) -> str:
-        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:
-        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"))
-
-    @classmethod
-    def get_role_name(cls, initial_message: str) -> Optional[str]:
-        if not initial_message:
-            return None
-        message_lines = initial_message.splitlines()
-        if "###" in message_lines[0]:
-            return message_lines[1].split("Role name: ")[1].strip()
-        return None
+            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 BadArgumentUsage(f'Role "{name}" not found.')
+            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")
-        expecting = typer.prompt(
-            "Enter expecting result, e.g. answer, code, \
-            shell command, command description, etc."
-        )
-        role = cls(name, role, expecting)
-        role.save()
+        role = cls(name, role)
+        role._save()
 
     @classmethod
     @option_callback
@@ -140,58 +102,75 @@ class SystemRole:
     def show(cls, name: str) -> None:
         typer.echo(cls.get(name).role)
 
-    @property
-    def exists(self) -> bool:
-        return self.file_path.exists()
+    @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 system_message(self) -> Dict[str, str]:
-        return {"role": "system", "content": self.role}
+    def _exists(self) -> bool:
+        return self._file_path.exists()
 
     @property
-    def file_path(self) -> Path:
+    def _file_path(self) -> Path:
         return self.storage / f"{self.name}.json"
 
-    def save(self) -> None:
-        if self.exists:
+    def _save(self) -> None:
+        if self._exists:
             typer.confirm(
                 f'Role "{self.name}" already exists, overwrite it?',
                 abort=True,
             )
-        self.file_path.write_text(json.dumps(self.__dict__), encoding="utf-8")
+
+        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:
+        if self._exists:
             typer.confirm(
                 f'Role "{self.name}" exist, delete it?',
                 abort=True,
             )
-        self.file_path.unlink()
-
-    def make_prompt(self, request: str, initial: bool) -> str:
-        if initial:
-            prompt = PROMPT_TEMPLATE.format(
-                name=self.name,
-                role=self.role,
-                request=request,
-                expecting=self.expecting,
-            )
-        else:
-            prompt = f"{request}\n{self.expecting}:"
-
-        return prompt
+        self._file_path.unlink()
 
     def same_role(self, initial_message: str) -> bool:
         if not initial_message:
             return False
-        return True if f"Role name: {self.name}" in initial_message else False
+        return True if f"You are {self.name}" in initial_message else False
 
 
 class DefaultRoles(Enum):
-    DEFAULT = "default"
-    SHELL = "shell"
-    DESCRIBE_SHELL = "describe_shell"
-    CODE = "code"
+    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:

sgpt/utils.py

@@ -1,24 +1,14 @@
 import os
 import platform
 import shlex
-from enum import Enum
 from tempfile import NamedTemporaryFile
 from typing import Any, Callable
 
 import typer
-from click import BadParameter
+from click import BadParameter, UsageError
 
-
-class ModelOptions(str, Enum):
-    """
-    Model endpoint compatibility
-    https://platform.openai.com/docs/models/model-endpoint-compatibility
-    """
-
-    GPT4 = "gpt-4"
-    GPT432k = "gpt-4-32k"
-    GPT35TURBO = "gpt-3.5-turbo"
-    GPT35TURBO16K = "gpt-3.5-turbo-16k"
+from sgpt.__version__ import __version__
+from sgpt.integration import bash_integration, zsh_integration
 
 
 def get_edited_prompt() -> str:
@@ -76,14 +66,30 @@ def option_callback(func: Callable) -> Callable:  # type: ignore
 @option_callback
 def install_shell_integration(*_args: Any) -> None:
     """
-    Installs shell integration. Currently only supports Linux.
+    Installs shell integration. Currently only supports ZSH and Bash.
     Allows user to get shell completions in terminal by using hotkey.
-    Allows user to edit shell command right away in terminal.
+    Replaces current "buffer" of the shell with the completion.
     """
     # TODO: Add support for Windows.
     # TODO: Implement updates.
-    if platform.system() == "Windows":
-        typer.echo("Windows is not supported yet.")
+    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:
-        url = "https://raw.githubusercontent.com/TheR1D/shell_gpt/shell-integrations/install.sh"
-        os.system(f'sh -c "$(curl -fsSL {url})"')
+        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__}")

tests/_integration.py

@@ -11,7 +11,6 @@ import os
 import subprocess
 from pathlib import Path
 from tempfile import NamedTemporaryFile
-from time import sleep
 from unittest import TestCase
 from unittest.mock import ANY, patch
 from uuid import uuid4
@@ -19,6 +18,7 @@ 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
@@ -30,9 +30,14 @@ app.command()(main)
 
 
 class TestShellGpt(TestCase):
-    def setUp(self) -> None:
-        # Just to not spam the API.
-        sleep(1)
+    @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):
@@ -51,7 +56,7 @@ class TestShellGpt(TestCase):
         }
         result = runner.invoke(app, self.get_arguments(**dict_arguments))
         assert result.exit_code == 0
-        assert "Prague" in result.stdout
+        assert "Prague" in result.output
 
     def test_shell(self):
         dict_arguments = {
@@ -60,7 +65,7 @@ class TestShellGpt(TestCase):
         }
         result = runner.invoke(app, self.get_arguments(**dict_arguments))
         assert result.exit_code == 0
-        assert "git commit" in result.stdout
+        assert "git commit" in result.output
 
     def test_describe_shell(self):
         dict_arguments = {
@@ -69,7 +74,7 @@ class TestShellGpt(TestCase):
         }
         result = runner.invoke(app, self.get_arguments(**dict_arguments))
         assert result.exit_code == 0
-        assert "List " in result.stdout
+        assert "lists" in result.output.lower()
 
     def test_code(self):
         """
@@ -88,10 +93,10 @@ class TestShellGpt(TestCase):
         }
         result = runner.invoke(app, self.get_arguments(**dict_arguments))
         assert result.exit_code == 0
-        print(result.stdout)
+        print(result.output)
         # Since output will be slightly different, there is no way how to test it precisely.
-        assert "print" in result.stdout
-        assert "*" in result.stdout
+        assert "print" in result.output
+        assert "*" in result.output
         with NamedTemporaryFile("w+", delete=False) as file:
             try:
                 compile(result.output, file.name, "exec")
@@ -119,7 +124,7 @@ class TestShellGpt(TestCase):
         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.stdout
+        assert "8" in result.output
         dict_arguments["--shell"] = True
         result = runner.invoke(app, self.get_arguments(**dict_arguments))
         assert result.exit_code == 2
@@ -138,14 +143,14 @@ class TestShellGpt(TestCase):
         }
         result = runner.invoke(app, self.get_arguments(**dict_arguments))
         assert result.exit_code == 0
-        assert "docker run" in result.stdout
-        assert "-p 80:80" in result.stdout
-        assert "nginx" in result.stdout
+        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.stdout
-        assert "-p 443:443" in result.stdout
+        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
@@ -162,13 +167,11 @@ class TestShellGpt(TestCase):
         }
         result = runner.invoke(app, self.get_arguments(**dict_arguments))
         assert result.exit_code == 0
-        assert "Add file contents to the index." in result.stdout
+        assert "adds" in result.output.lower() or "stages" in result.output.lower()
         dict_arguments["prompt"] = "'-A'"
-        # Prevent Too Many Requests.
-        sleep(1)
         result = runner.invoke(app, self.get_arguments(**dict_arguments))
         assert result.exit_code == 0
-        assert "all" in result.stdout
+        assert "all" in result.output
 
     def test_chat_code(self):
         chat_name = uuid4()
@@ -179,11 +182,11 @@ class TestShellGpt(TestCase):
         }
         result = runner.invoke(app, self.get_arguments(**dict_arguments))
         assert result.exit_code == 0
-        assert "localhost:80" in result.stdout
+        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.stdout
+        assert "localhost:443" in result.output
         del dict_arguments["--code"]
         assert "--code" not in dict_arguments
         dict_arguments["--shell"] = True
@@ -194,7 +197,7 @@ class TestShellGpt(TestCase):
     def test_list_chat(self):
         result = runner.invoke(app, ["--list-chats"])
         assert result.exit_code == 0
-        assert "test_" in result.stdout
+        assert "test_" in result.output
 
     def test_show_chat(self):
         chat_name = uuid4()
@@ -207,9 +210,9 @@ class TestShellGpt(TestCase):
         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.stdout
-        assert "What is my favorite number + 2?" in result.stdout
-        assert "8" in result.stdout
+        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 = {
@@ -219,7 +222,7 @@ class TestShellGpt(TestCase):
         }
         result = runner.invoke(app, self.get_arguments(**dict_arguments))
         assert result.exit_code == 2
-        assert "Only one of --shell, --describe-shell, and --code" in result.stdout
+        assert "Only one of --shell, --describe-shell, and --code" in result.output
 
     def test_repl_default(
         self,
@@ -237,9 +240,34 @@ class TestShellGpt(TestCase):
             app, self.get_arguments(**dict_arguments), input="\n".join(inputs)
         )
         assert result.exit_code == 0
-        assert ">>> Please remember my favorite number: 6" in result.stdout
-        assert ">>> What is my favorite number + 2?" in result.stdout
-        assert "8" in result.stdout
+        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.
@@ -248,25 +276,30 @@ class TestShellGpt(TestCase):
             "--repl": "temp",
             "--shell": True,
         }
-        inputs = ["What is in current folder?", "Sort by name", "exit()"]
+        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.stdout
-        assert ">>> What is in current folder?" in result.stdout
-        assert ">>> Sort by name" in result.stdout
-        assert "ls" in result.stdout
-        assert "ls | sort" in result.stdout
+        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("###")
-        assert chat_messages[0]["content"].endswith("\n###\nCommand:")
-        assert chat_messages[1]["content"] == "ls"
-        assert chat_messages[2]["content"].endswith("\nCommand:")
-        assert chat_messages[3]["content"] == "ls | sort"
+        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.
@@ -280,14 +313,15 @@ class TestShellGpt(TestCase):
             app, self.get_arguments(**dict_arguments), input="\n".join(inputs)
         )
         assert result.exit_code == 0
-        assert "Install" in result.stdout
-        assert "Update" in result.stdout
+        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("###")
-        assert chat_messages[0]["content"].endswith("\n###\nDescription:")
+        assert chat_messages[0]["content"].startswith(
+            "You are Shell Command Descriptor"
+        )
 
     def test_repl_code(self):
         dict_arguments = {
@@ -304,18 +338,17 @@ class TestShellGpt(TestCase):
             app, self.get_arguments(**dict_arguments), input="\n".join(inputs)
         )
         assert result.exit_code == 0
-        assert f">>> {inputs[0]}" in result.stdout
-        assert "requests.get" in result.stdout
-        assert "localhost:8080" in result.stdout
-        assert f">>> {inputs[1]}" in result.stdout
-        assert "localhost:443" in result.stdout
+        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("###")
-        assert chat_messages[0]["content"].endswith("\n###\nCode:")
-        assert chat_messages[2]["content"].endswith("\nCode:")
+        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()")
@@ -323,9 +356,8 @@ class TestShellGpt(TestCase):
             app, self.get_arguments(**dict_arguments), input="\n".join(new_inputs)
         )
         # Should include previous chat history.
-        assert "user: ###" in result.stdout
-        assert "Chat History" in result.stdout
-        assert f"user: {inputs[1]}" in result.stdout
+        assert "Chat History" in result.output
+        assert f"user: {inputs[1]}" in result.output
 
     def test_zsh_command(self):
         """
@@ -340,15 +372,15 @@ class TestShellGpt(TestCase):
             "--shell": True,
         }
         result = runner.invoke(app, self.get_arguments(**dict_arguments), input="y\n")
-        stdout = result.stdout.strip()
+        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.stdout.
-        # assert "command not found" not in result.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.client.OpenAIClient.get_completion")
+    @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?",
@@ -358,24 +390,25 @@ class TestShellGpt(TestCase):
         mocked_get_completion.assert_called_once_with(
             messages=ANY,
             model="gpt-4",
-            temperature=0.1,
-            top_probability=1.0,
+            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("default")
-        handler = Handler(role=role)
+        role = SystemRole.get("ShellGPT")
+        handler = Handler(role=role, markdown=False)
         assert handler.color == color
         os.environ["DEFAULT_COLOR"] = "red"
-        handler = Handler(role=role)
+        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.stdout
+        assert "Berlin" in result.output
 
     def test_shell_stdin_with_prompt(self):
         dict_arguments = {
@@ -384,16 +417,20 @@ class TestShellGpt(TestCase):
         }
         stdin = "What is in current folder\n"
         result = runner.invoke(app, self.get_arguments(**dict_arguments), input=stdin)
-        assert result.stdout == "ls | sort\n"
+        assert "ls" in result.output
+        assert "sort" in result.output
 
     def test_role(self):
-        test_role = Path(cfg.get("ROLE_STORAGE_PATH")) / "test_json.json"
+        test_role = Path(cfg.get("ROLE_STORAGE_PATH")) / "json_generator.json"
         test_role.unlink(missing_ok=True)
         dict_arguments = {
             "prompt": "test",
-            "--create-role": "test_json",
+            "--create-role": "json_generator",
         }
-        input = "You are a JSON generator, return only JSON as response.\n" "json\n"
+        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
 
@@ -403,24 +440,24 @@ class TestShellGpt(TestCase):
         }
         result = runner.invoke(app, self.get_arguments(**dict_arguments))
         assert result.exit_code == 0
-        assert "test_json" in result.stdout
+        assert "json_generator" in result.output
 
         dict_arguments = {
             "prompt": "test",
-            "--show-role": "test_json",
+            "--show-role": "json_generator",
         }
         result = runner.invoke(app, self.get_arguments(**dict_arguments))
         assert result.exit_code == 0
-        assert "You are a JSON generator" in result.stdout
+        assert "You are json_generator" in result.output
 
         # Test with command line argument prompt.
         dict_arguments = {
             "prompt": "random username, password, email",
-            "--role": "test_json",
+            "--role": "json_generator",
         }
         result = runner.invoke(app, self.get_arguments(**dict_arguments))
         assert result.exit_code == 0
-        generated_json = json.loads(result.stdout)
+        generated_json = json.loads(result.output)
         assert "username" in generated_json
         assert "password" in generated_json
         assert "email" in generated_json
@@ -428,12 +465,12 @@ class TestShellGpt(TestCase):
         # Test with stdin prompt.
         dict_arguments = {
             "prompt": "",
-            "--role": "test_json",
+            "--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.stdout)
+        generated_json = json.loads(result.output)
         assert "username" in generated_json
         assert "password" in generated_json
         assert "email" in generated_json
@@ -446,6 +483,16 @@ class TestShellGpt(TestCase):
         }
         result = runner.invoke(app, self.get_arguments(**dict_arguments), input="d\n")
         assert result.exit_code == 0
-        # Cant really test it since stdin in disable for --shell flag.
+        # Can't really test it since stdin in disable for --shell flag.
         # for word in ("prints", "hello", "console"):
-        #     assert word in result.stdout
+        #     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.

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")

tests/test_code.py

@@ -0,0 +1,144 @@
+from pathlib import Path
+from unittest.mock import patch
+
+from sgpt.config import cfg
+from sgpt.role import DefaultRoles, SystemRole
+
+from .utils import app, cmd_args, comp_args, mock_comp, runner
+
+role = SystemRole.get(DefaultRoles.CODE.value)
+
+
+@patch("sgpt.handlers.handler.completion")
+def test_code_generation(completion):
+    completion.return_value = mock_comp("print('Hello World')")
+
+    args = {"prompt": "hello world python", "--code": True}
+    result = runner.invoke(app, cmd_args(**args))
+
+    completion.assert_called_once_with(**comp_args(role, args["prompt"]))
+    assert result.exit_code == 0
+    assert "print('Hello World')" in result.output
+
+
+@patch("sgpt.printer.TextPrinter.live_print")
+@patch("sgpt.printer.MarkdownPrinter.live_print")
+@patch("sgpt.handlers.handler.completion")
+def test_code_generation_no_markdown(completion, markdown_printer, text_printer):
+    completion.return_value = mock_comp("print('Hello World')")
+
+    args = {"prompt": "make a commit using git", "--code": True, "--md": True}
+    result = runner.invoke(app, cmd_args(**args))
+
+    assert result.exit_code == 0
+    # Should ignore --md for --code option and output code without markdown.
+    markdown_printer.assert_not_called()
+    text_printer.assert_called()
+
+
+@patch("sgpt.handlers.handler.completion")
+def test_code_generation_stdin(completion):
+    completion.return_value = mock_comp("# Hello\nprint('Hello')")
+
+    args = {"prompt": "make comments for code", "--code": True}
+    stdin = "print('Hello')"
+    result = runner.invoke(app, cmd_args(**args), input=stdin)
+
+    expected_prompt = f"{stdin}\n\n{args['prompt']}"
+    completion.assert_called_once_with(**comp_args(role, expected_prompt))
+    assert result.exit_code == 0
+    assert "# Hello" in result.output
+    assert "print('Hello')" in result.output
+
+
+@patch("sgpt.handlers.handler.completion")
+def test_code_chat(completion):
+    completion.side_effect = [
+        mock_comp("print('hello')"),
+        mock_comp("print('hello')\nprint('world')"),
+    ]
+    chat_name = "_test"
+    chat_path = Path(cfg.get("CHAT_CACHE_PATH")) / chat_name
+    chat_path.unlink(missing_ok=True)
+
+    args = {"prompt": "print hello", "--code": True, "--chat": chat_name}
+    result = runner.invoke(app, cmd_args(**args))
+    assert result.exit_code == 0
+    assert "print('hello')" in result.output
+    assert chat_path.exists()
+
+    args["prompt"] = "also print world"
+    result = runner.invoke(app, cmd_args(**args))
+    assert result.exit_code == 0
+    assert "print('hello')" in result.output
+    assert "print('world')" in result.output
+
+    expected_messages = [
+        {"role": "system", "content": role.role},
+        {"role": "user", "content": "print hello"},
+        {"role": "assistant", "content": "print('hello')"},
+        {"role": "user", "content": "also print world"},
+        {"role": "assistant", "content": "print('hello')\nprint('world')"},
+    ]
+    expected_args = comp_args(role, "", messages=expected_messages)
+    completion.assert_called_with(**expected_args)
+    assert completion.call_count == 2
+
+    args["--shell"] = True
+    result = runner.invoke(app, cmd_args(**args))
+    assert result.exit_code == 2
+    assert "Error" in result.output
+    chat_path.unlink()
+    # TODO: Code chat can be recalled without --code option.
+
+
+@patch("sgpt.handlers.handler.completion")
+def test_code_repl(completion):
+    completion.side_effect = [
+        mock_comp("print('hello')"),
+        mock_comp("print('hello')\nprint('world')"),
+    ]
+    chat_name = "_test"
+    chat_path = Path(cfg.get("CHAT_CACHE_PATH")) / chat_name
+    chat_path.unlink(missing_ok=True)
+
+    args = {"--repl": chat_name, "--code": True}
+    inputs = ["__sgpt__eof__", "print hello", "also print world", "exit()"]
+    result = runner.invoke(app, cmd_args(**args), input="\n".join(inputs))
+
+    expected_messages = [
+        {"role": "system", "content": role.role},
+        {"role": "user", "content": "print hello"},
+        {"role": "assistant", "content": "print('hello')"},
+        {"role": "user", "content": "also print world"},
+        {"role": "assistant", "content": "print('hello')\nprint('world')"},
+    ]
+    expected_args = comp_args(role, "", messages=expected_messages)
+    completion.assert_called_with(**expected_args)
+    assert completion.call_count == 2
+
+    assert result.exit_code == 0
+    assert ">>> print hello" in result.output
+    assert "print('hello')" in result.output
+    assert ">>> also print world" in result.output
+    assert "print('world')" in result.output
+
+
+@patch("sgpt.handlers.handler.completion")
+def test_code_and_shell(completion):
+    args = {"--code": True, "--shell": True}
+    result = runner.invoke(app, cmd_args(**args))
+
+    completion.assert_not_called()
+    assert result.exit_code == 2
+    assert "Error" in result.output
+
+
+@patch("sgpt.handlers.handler.completion")
+def test_code_and_describe_shell(completion):
+    args = {"--code": True, "--describe-shell": True}
+    result = runner.invoke(app, cmd_args(**args))
+
+    completion.assert_not_called()
+    assert result.exit_code == 2
+    assert "Error" in result.output

tests/test_default.py

@@ -0,0 +1,250 @@
+from pathlib import Path
+from unittest.mock import patch
+
+import typer
+from typer.testing import CliRunner
+
+from sgpt import config, main
+from sgpt.__version__ import __version__
+from sgpt.role import DefaultRoles, SystemRole
+
+from .utils import app, cmd_args, comp_args, mock_comp, runner
+
+role = SystemRole.get(DefaultRoles.DEFAULT.value)
+cfg = config.cfg
+
+
+@patch("sgpt.handlers.handler.completion")
+def test_default(completion):
+    completion.return_value = mock_comp("Prague")
+
+    args = {"prompt": "capital of the Czech Republic?"}
+    result = runner.invoke(app, cmd_args(**args))
+
+    completion.assert_called_once_with(**comp_args(role, **args))
+    assert result.exit_code == 0
+    assert "Prague" in result.output
+
+
+@patch("sgpt.handlers.handler.completion")
+def test_default_stdin(completion):
+    completion.return_value = mock_comp("Prague")
+
+    stdin = "capital of the Czech Republic?"
+    result = runner.invoke(app, cmd_args(), input=stdin)
+
+    completion.assert_called_once_with(**comp_args(role, stdin))
+    assert result.exit_code == 0
+    assert "Prague" in result.output
+
+
+@patch("rich.console.Console.print")
+@patch("sgpt.handlers.handler.completion")
+def test_show_chat_use_markdown(completion, console_print):
+    completion.return_value = mock_comp("ok")
+    chat_name = "_test"
+    chat_path = Path(cfg.get("CHAT_CACHE_PATH")) / chat_name
+    chat_path.unlink(missing_ok=True)
+
+    args = {"prompt": "my number is 2", "--chat": chat_name}
+    result = runner.invoke(app, cmd_args(**args))
+    assert result.exit_code == 0
+    assert chat_path.exists()
+
+    result = runner.invoke(app, ["--show-chat", chat_name])
+    assert result.exit_code == 0
+    console_print.assert_called()
+
+
+@patch("rich.console.Console.print")
+@patch("sgpt.handlers.handler.completion")
+def test_show_chat_no_use_markdown(completion, console_print):
+    completion.return_value = mock_comp("ok")
+    chat_name = "_test"
+    chat_path = Path(cfg.get("CHAT_CACHE_PATH")) / chat_name
+    chat_path.unlink(missing_ok=True)
+
+    # Flag '--code' doesn't use markdown
+    args = {"prompt": "my number is 2", "--chat": chat_name, "--code": True}
+    result = runner.invoke(app, cmd_args(**args))
+    assert result.exit_code == 0
+    assert chat_path.exists()
+
+    result = runner.invoke(app, ["--show-chat", chat_name, "--no-md"])
+    assert result.exit_code == 0
+    console_print.assert_not_called()
+
+
+@patch("sgpt.handlers.handler.completion")
+def test_default_chat(completion):
+    completion.side_effect = [mock_comp("ok"), mock_comp("4")]
+    chat_name = "_test"
+    chat_path = Path(cfg.get("CHAT_CACHE_PATH")) / chat_name
+    chat_path.unlink(missing_ok=True)
+
+    args = {"prompt": "my number is 2", "--chat": chat_name}
+    result = runner.invoke(app, cmd_args(**args))
+    assert result.exit_code == 0
+    assert "ok" in result.output
+    assert chat_path.exists()
+
+    args["prompt"] = "my number + 2?"
+    result = runner.invoke(app, cmd_args(**args))
+    assert result.exit_code == 0
+    assert "4" in result.output
+
+    expected_messages = [
+        {"role": "system", "content": role.role},
+        {"role": "user", "content": "my number is 2"},
+        {"role": "assistant", "content": "ok"},
+        {"role": "user", "content": "my number + 2?"},
+        {"role": "assistant", "content": "4"},
+    ]
+    expected_args = comp_args(role, "", messages=expected_messages)
+    completion.assert_called_with(**expected_args)
+    assert completion.call_count == 2
+
+    result = runner.invoke(app, ["--list-chats"])
+    assert result.exit_code == 0
+    assert "_test" in result.output
+
+    result = runner.invoke(app, ["--show-chat", chat_name])
+    assert result.exit_code == 0
+    assert "my number is 2" in result.output
+    assert "ok" in result.output
+    assert "my number + 2?" in result.output
+    assert "4" in result.output
+
+    args["--shell"] = True
+    result = runner.invoke(app, cmd_args(**args))
+    assert result.exit_code == 2
+    assert "Error" in result.output
+
+    args["--code"] = True
+    result = runner.invoke(app, cmd_args(**args))
+    assert result.exit_code == 2
+    assert "Error" in result.output
+    chat_path.unlink()
+
+
+@patch("sgpt.handlers.handler.completion")
+def test_default_repl(completion):
+    completion.side_effect = [mock_comp("ok"), mock_comp("8")]
+    chat_name = "_test"
+    chat_path = Path(cfg.get("CHAT_CACHE_PATH")) / chat_name
+    chat_path.unlink(missing_ok=True)
+
+    args = {"--repl": chat_name}
+    inputs = ["__sgpt__eof__", "my number is 6", "my number + 2?", "exit()"]
+    result = runner.invoke(app, cmd_args(**args), input="\n".join(inputs))
+
+    expected_messages = [
+        {"role": "system", "content": role.role},
+        {"role": "user", "content": "my number is 6"},
+        {"role": "assistant", "content": "ok"},
+        {"role": "user", "content": "my number + 2?"},
+        {"role": "assistant", "content": "8"},
+    ]
+    expected_args = comp_args(role, "", messages=expected_messages)
+    completion.assert_called_with(**expected_args)
+    assert completion.call_count == 2
+
+    assert result.exit_code == 0
+    assert ">>> my number is 6" in result.output
+    assert "ok" in result.output
+    assert ">>> my number + 2?" in result.output
+    assert "8" in result.output
+
+
+@patch("sgpt.handlers.handler.completion")
+def test_default_repl_stdin(completion):
+    completion.side_effect = [mock_comp("ok init"), mock_comp("ok another")]
+    chat_name = "_test"
+    chat_path = Path(cfg.get("CHAT_CACHE_PATH")) / chat_name
+    chat_path.unlink(missing_ok=True)
+
+    my_runner = CliRunner()
+    my_app = typer.Typer()
+    my_app.command()(main)
+
+    args = {"--repl": chat_name}
+    inputs = ["this is stdin", "__sgpt__eof__", "prompt", "another", "exit()"]
+    result = my_runner.invoke(my_app, cmd_args(**args), input="\n".join(inputs))
+
+    expected_messages = [
+        {"role": "system", "content": role.role},
+        {"role": "user", "content": "this is stdin\n\n\n\nprompt"},
+        {"role": "assistant", "content": "ok init"},
+        {"role": "user", "content": "another"},
+        {"role": "assistant", "content": "ok another"},
+    ]
+    expected_args = comp_args(role, "", messages=expected_messages)
+    completion.assert_called_with(**expected_args)
+    assert completion.call_count == 2
+
+    assert result.exit_code == 0
+    assert "this is stdin" in result.output
+    assert ">>> prompt" in result.output
+    assert "ok init" in result.output
+    assert ">>> another" in result.output
+    assert "ok another" in result.output
+
+
+@patch("sgpt.handlers.handler.completion")
+def test_llm_options(completion):
+    completion.return_value = mock_comp("Berlin")
+
+    args = {
+        "prompt": "capital of the Germany?",
+        "--model": "gpt-4-test",
+        "--temperature": 0.5,
+        "--top-p": 0.5,
+        "--no-functions": True,
+    }
+    result = runner.invoke(app, cmd_args(**args))
+
+    expected_args = comp_args(
+        role=role,
+        prompt=args["prompt"],
+        model=args["--model"],
+        temperature=args["--temperature"],
+        top_p=args["--top-p"],
+    )
+    completion.assert_called_once_with(**expected_args)
+    assert result.exit_code == 0
+    assert "Berlin" in result.output
+
+
+@patch("sgpt.handlers.handler.completion")
+def test_version(completion):
+    args = {"--version": True}
+    result = runner.invoke(app, cmd_args(**args))
+
+    completion.assert_not_called()
+    assert __version__ in result.output
+
+
+@patch("sgpt.printer.TextPrinter.live_print")
+@patch("sgpt.printer.MarkdownPrinter.live_print")
+@patch("sgpt.handlers.handler.completion")
+def test_markdown(completion, markdown_printer, text_printer):
+    completion.return_value = mock_comp("pong")
+
+    args = {"prompt": "ping", "--md": True}
+    result = runner.invoke(app, cmd_args(**args))
+    assert result.exit_code == 0
+    markdown_printer.assert_called()
+    text_printer.assert_not_called()
+
+
+@patch("sgpt.printer.TextPrinter.live_print")
+@patch("sgpt.printer.MarkdownPrinter.live_print")
+@patch("sgpt.handlers.handler.completion")
+def test_no_markdown(completion, markdown_printer, text_printer):
+    completion.return_value = mock_comp("pong")
+
+    args = {"prompt": "ping", "--no-md": True}
+    result = runner.invoke(app, cmd_args(**args))
+    assert result.exit_code == 0
+    markdown_printer.assert_not_called()
+    text_printer.assert_called()

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)

tests/test_shell.py

@@ -0,0 +1,189 @@
+import os
+from pathlib import Path
+from unittest.mock import patch
+
+from sgpt.config import cfg
+from sgpt.role import DefaultRoles, SystemRole
+
+from .utils import app, cmd_args, comp_args, mock_comp, runner
+
+
+@patch("sgpt.handlers.handler.completion")
+def test_shell(completion):
+    role = SystemRole.get(DefaultRoles.SHELL.value)
+    completion.return_value = mock_comp("git commit -m test")
+
+    args = {"prompt": "make a commit using git", "--shell": True}
+    result = runner.invoke(app, cmd_args(**args))
+
+    completion.assert_called_once_with(**comp_args(role, args["prompt"]))
+    assert "git commit" in result.output
+    assert "[E]xecute, [M]odify, [D]escribe, [A]bort:" in result.output
+
+
+@patch("sgpt.printer.TextPrinter.live_print")
+@patch("sgpt.printer.MarkdownPrinter.live_print")
+@patch("sgpt.handlers.handler.completion")
+def test_shell_no_markdown(completion, markdown_printer, text_printer):
+    completion.return_value = mock_comp("git commit -m test")
+
+    args = {"prompt": "make a commit using git", "--shell": True, "--md": True}
+    runner.invoke(app, cmd_args(**args))
+
+    # Should ignore --md for --shell option and output text without markdown.
+    markdown_printer.assert_not_called()
+    text_printer.assert_called()
+
+
+@patch("sgpt.handlers.handler.completion")
+def test_shell_stdin(completion):
+    completion.return_value = mock_comp("ls -l | sort")
+    role = SystemRole.get(DefaultRoles.SHELL.value)
+
+    args = {"prompt": "Sort by name", "--shell": True}
+    stdin = "What is in current folder"
+    result = runner.invoke(app, cmd_args(**args), input=stdin)
+
+    expected_prompt = f"{stdin}\n\n{args['prompt']}"
+    completion.assert_called_once_with(**comp_args(role, expected_prompt))
+    assert "ls -l | sort" in result.output
+    assert "[E]xecute, [M]odify, [D]escribe, [A]bort:" in result.output
+
+
+@patch("sgpt.handlers.handler.completion")
+def test_describe_shell(completion):
+    completion.return_value = mock_comp("lists the contents of a folder")
+    role = SystemRole.get(DefaultRoles.DESCRIBE_SHELL.value)
+
+    args = {"prompt": "ls", "--describe-shell": True}
+    result = runner.invoke(app, cmd_args(**args))
+
+    completion.assert_called_once_with(**comp_args(role, args["prompt"]))
+    assert result.exit_code == 0
+    assert "lists" in result.output
+
+
+@patch("sgpt.handlers.handler.completion")
+def test_describe_shell_stdin(completion):
+    completion.return_value = mock_comp("lists the contents of a folder")
+    role = SystemRole.get(DefaultRoles.DESCRIBE_SHELL.value)
+
+    args = {"--describe-shell": True}
+    stdin = "What is in current folder"
+    result = runner.invoke(app, cmd_args(**args), input=stdin)
+
+    expected_prompt = f"{stdin}"
+    completion.assert_called_once_with(**comp_args(role, expected_prompt))
+    assert result.exit_code == 0
+    assert "lists" in result.output
+
+
+@patch("os.system")
+@patch("sgpt.handlers.handler.completion")
+def test_shell_run_description(completion, system):
+    completion.side_effect = [mock_comp("echo hello"), mock_comp("prints hello")]
+    args = {"prompt": "echo hello", "--shell": True}
+    inputs = "__sgpt__eof__\nd\ne\n"
+    result = runner.invoke(app, cmd_args(**args), input=inputs)
+    shell = os.environ.get("SHELL", "/bin/sh")
+    system.assert_called_once_with(f"{shell} -c 'echo hello'")
+    assert result.exit_code == 0
+    assert "echo hello" in result.output
+    assert "prints hello" in result.output
+
+
+@patch("sgpt.handlers.handler.completion")
+def test_shell_chat(completion):
+    completion.side_effect = [mock_comp("ls"), mock_comp("ls | sort")]
+    role = SystemRole.get(DefaultRoles.SHELL.value)
+    chat_name = "_test"
+    chat_path = Path(cfg.get("CHAT_CACHE_PATH")) / chat_name
+    chat_path.unlink(missing_ok=True)
+
+    args = {"prompt": "list folder", "--shell": True, "--chat": chat_name}
+    result = runner.invoke(app, cmd_args(**args))
+    assert "ls" in result.output
+    assert chat_path.exists()
+
+    args["prompt"] = "sort by name"
+    result = runner.invoke(app, cmd_args(**args))
+    assert "ls | sort" in result.output
+
+    expected_messages = [
+        {"role": "system", "content": role.role},
+        {"role": "user", "content": "list folder"},
+        {"role": "assistant", "content": "ls"},
+        {"role": "user", "content": "sort by name"},
+        {"role": "assistant", "content": "ls | sort"},
+    ]
+    expected_args = comp_args(role, "", messages=expected_messages)
+    completion.assert_called_with(**expected_args)
+    assert completion.call_count == 2
+
+    args["--code"] = True
+    result = runner.invoke(app, cmd_args(**args))
+    assert result.exit_code == 2
+    assert "Error" in result.output
+    chat_path.unlink()
+    # TODO: Shell chat can be recalled without --shell option.
+
+
+@patch("os.system")
+@patch("sgpt.handlers.handler.completion")
+def test_shell_repl(completion, mock_system):
+    completion.side_effect = [mock_comp("ls"), mock_comp("ls | sort")]
+    role = SystemRole.get(DefaultRoles.SHELL.value)
+    chat_name = "_test"
+    chat_path = Path(cfg.get("CHAT_CACHE_PATH")) / chat_name
+    chat_path.unlink(missing_ok=True)
+
+    args = {"--repl": chat_name, "--shell": True}
+    inputs = ["__sgpt__eof__", "list folder", "sort by name", "e", "exit()"]
+    result = runner.invoke(app, cmd_args(**args), input="\n".join(inputs))
+    shell = os.environ.get("SHELL", "/bin/sh")
+    mock_system.assert_called_once_with(f"{shell} -c 'ls | sort'")
+
+    expected_messages = [
+        {"role": "system", "content": role.role},
+        {"role": "user", "content": "list folder"},
+        {"role": "assistant", "content": "ls"},
+        {"role": "user", "content": "sort by name"},
+        {"role": "assistant", "content": "ls | sort"},
+    ]
+    expected_args = comp_args(role, "", messages=expected_messages)
+    completion.assert_called_with(**expected_args)
+    assert completion.call_count == 2
+
+    assert result.exit_code == 0
+    assert ">>> list folder" in result.output
+    assert "ls" in result.output
+    assert ">>> sort by name" in result.output
+    assert "ls | sort" in result.output
+
+
+@patch("sgpt.handlers.handler.completion")
+def test_shell_and_describe_shell(completion):
+    args = {"prompt": "ls", "--describe-shell": True, "--shell": True}
+    result = runner.invoke(app, cmd_args(**args))
+
+    completion.assert_not_called()
+    assert result.exit_code == 2
+    assert "Error" in result.output
+
+
+@patch("sgpt.handlers.handler.completion")
+def test_shell_no_interaction(completion):
+    completion.return_value = mock_comp("git commit -m test")
+    role = SystemRole.get(DefaultRoles.SHELL.value)
+
+    args = {
+        "prompt": "make a commit using git",
+        "--shell": True,
+        "--no-interaction": True,
+    }
+    result = runner.invoke(app, cmd_args(**args))
+
+    completion.assert_called_once_with(**comp_args(role, args["prompt"]))
+    assert result.exit_code == 0
+    assert "git commit" in result.output
+    assert "[E]xecute" not in result.output

tests/test_unit.py

@@ -1,73 +0,0 @@
-import os
-import unittest
-
-import requests
-import requests_mock
-
-from sgpt.client import OpenAIClient
-
-
-class TestMain(unittest.TestCase):
-    API_HOST = os.getenv("OPENAI_HOST", "https://api.openai.com")
-    API_URL = f"{API_HOST}/v1/chat/completions"
-    # TODO: Fix tests.
-
-    def setUp(self):
-        self.api_key = os.environ["OPENAI_API_KEY"] = "test key"
-        self.prompt = "What is the capital of France?"
-        self.shell = False
-        self.execute = False
-        self.code = False
-        self.animation = True
-        self.spinner = True
-        self.temperature = 1.0
-        self.top_p = 1.0
-        self.response_text = "Paris"
-        self.model = "gpt-3.5-turbo"
-        self.client = OpenAIClient(self.API_HOST, self.api_key)
-
-    @requests_mock.Mocker()
-    def test_openai_request(self, mock):
-        # TODO: Fix tests.
-        mocked_json = {"choices": [{"message": {"content": self.response_text}}]}
-        mock.post(self.API_URL, json=mocked_json, status_code=200)
-        result = yield from self.client.get_completion(
-            messages=[{"role": "user", "content": self.prompt}],
-            model=self.model,
-            temperature=self.temperature,
-            top_probability=self.top_p,
-            caching=False,
-        )
-        # TODO: Fix tests with generators.
-        self.assertEqual(result, self.response_text)
-        expected_json = {
-            "messages": [{"role": "user", "content": self.prompt}],
-            "model": "gpt-3.5-turbo",
-            "temperature": self.temperature,
-            "top_p": self.top_p,
-        }
-        expected_headers = {
-            "Content-Type": "application/json",
-            "Authorization": f"Bearer {self.api_key}",
-        }
-        request = mock.request_history[0]
-        self.assertEqual(request.json(), expected_json)
-        for key, value in expected_headers.items():
-            self.assertEqual(request.headers[key], value)
-
-    @requests_mock.Mocker()
-    def test_openai_request_fail(self, mock):
-        # TODO: Fix tests.
-        mock.post(self.API_URL, status_code=400)
-        with self.assertRaises(requests.exceptions.HTTPError):
-            yield from self.client.get_completion(
-                messages=[{"role": "user", "content": self.prompt}],
-                model=self.model,
-                temperature=self.temperature,
-                top_probability=self.top_p,
-                caching=False,
-            )
-
-
-if __name__ == "__main__":
-    unittest.main()

tests/utils.py

@@ -0,0 +1,59 @@
+from datetime import datetime
+
+import typer
+from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
+from openai.types.chat.chat_completion_chunk import Choice as StreamChoice
+from openai.types.chat.chat_completion_chunk import ChoiceDelta
+from typer.testing import CliRunner
+
+from sgpt import main
+from sgpt.config import cfg
+
+runner = CliRunner()
+app = typer.Typer()
+app.command()(main)
+
+
+def mock_comp(tokens_string):
+    return [
+        ChatCompletionChunk(
+            id="foo",
+            model=cfg.get("DEFAULT_MODEL"),
+            object="chat.completion.chunk",
+            choices=[
+                StreamChoice(
+                    index=0,
+                    finish_reason=None,
+                    delta=ChoiceDelta(content=token, role="assistant"),
+                ),
+            ],
+            created=int(datetime.now().timestamp()),
+        )
+        for token in tokens_string
+    ]
+
+
+def cmd_args(prompt="", **kwargs):
+    arguments = [prompt]
+    for key, value in kwargs.items():
+        arguments.append(key)
+        if isinstance(value, bool):
+            continue
+        arguments.append(value)
+    arguments.append("--no-cache")
+    arguments.append("--no-functions")
+    return arguments
+
+
+def comp_args(role, prompt, **kwargs):
+    return {
+        "messages": [
+            {"role": "system", "content": role.role},
+            {"role": "user", "content": prompt},
+        ],
+        "model": cfg.get("DEFAULT_MODEL"),
+        "temperature": 0.0,
+        "top_p": 1.0,
+        "stream": True,
+        **kwargs,
+    }