diff --git a/README.md b/README.md index 2217e38e..7c06d23c 100644 --- a/README.md +++ b/README.md @@ -7,9 +7,12 @@ https://github.com/TheR1D/shell_gpt/assets/16740832/9197283c-db6a-4b46-bfea-3eb7 ```shell pip install shell-gpt ``` +By default, ShellGPT uses OpenAI's API and GPT-4 model. You'll need an API key, you can generate one [here](https://beta.openai.com/account/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. -You'll need an OpenAI API key, you can generate one [here](https://beta.openai.com/account/api-keys). -You will be prompted for your key which will then be stored in `~/.config/shell_gpt/.sgptrc`. +> [!TIP] +> Alternatively, you can use locally hosted open source models which are available for free. To use local models, you will need to run your own LLM backend server such as [Ollama](https://github.com/ollama/ollama). To set up ShellGPT with Ollama, please follow this comprehensive [guide](https://github.com/TheR1D/shell_gpt/wiki/Ollama). +> +> **❗️Note that ShellGPT is not optimized for local models and may not work as expected.** ## Usage **ShellGPT** is designed to quickly analyse and retrieve information. It's useful for straightforward requests ranging from technical configurations to general knowledge. @@ -24,7 +27,7 @@ git diff | sgpt "Generate git commit message, for my changes" # -> Added main feature details into README.md ``` -You can analyze logs from various sources by passing them using stdin, along with a prompt. 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 my_app | sgpt "check logs, find errors, provide possible solutions" ``` @@ -40,7 +43,7 @@ You can also use all kind of redirection operators to pass input: sgpt "summarise" < document.txt # -> The document discusses the impact... sgpt << EOF -What is the best way to lear Golang. +What is the best way to lear Golang? Provide simple hello world example. EOF # -> The best way to learn Golang... @@ -444,9 +447,6 @@ Possible options for `CODE_THEME`: https://pygments.org/styles/ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────╯ ``` -## LocalAI -By default, ShellGPT leverages OpenAI's large language models. However, it also provides the flexibility to use locally hosted models, which can be a cost-effective alternative. To use local models, you will need to run your own API server. You can accomplish this by using [LocalAI](https://github.com/go-skynet/LocalAI), a self-hosted, OpenAI-compatible API. Setting up LocalAI allows you to run language models on your own hardware, potentially without the need for an internet connection, depending on your usage. To set up your LocalAI, please follow this comprehensive [guide](https://github.com/TheR1D/shell_gpt/wiki/LocalAI). Remember that the performance of your local models may depend on the specifications of your hardware and the specific language model you choose to deploy. - ## Docker Run the container using the `OPENAI_API_KEY` environment variable, and a docker volume to store cache: ```shell diff --git a/pyproject.toml b/pyproject.toml index 2642605c..82726a9f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,8 +4,8 @@ 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" @@ -28,24 +28,15 @@ classifiers = [ "Programming Language :: Python :: 3.11", ] dependencies = [ - "requests >= 2.28.2, < 3.0.0", + "litellm >= 1.20.1, < 2.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", - "openai >= 1.6.1, < 2.0.0", "instructor >= 0.4.5, < 1.0.0", 'pyreadline3 >= 3.4.1, < 4.0.0; sys_platform == "win32"', ] -[project.scripts] -sgpt = "sgpt:cli" - -[project.urls] -homepage = "https://github.com/ther1d/shell_gpt" -repository = "https://github.com/ther1d/shell_gpt" -documentation = "https://github.com/TheR1D/shell_gpt/blob/main/README.md" - [project.optional-dependencies] test = [ "pytest >= 7.2.2, < 8.0.0", @@ -61,6 +52,14 @@ dev = [ "pre-commit >= 3.1.1, < 4.0.0", ] +[project.scripts] +sgpt = "sgpt:cli" + +[project.urls] +homepage = "https://github.com/ther1d/shell_gpt" +repository = "https://github.com/ther1d/shell_gpt" +documentation = "https://github.com/TheR1D/shell_gpt/blob/main/README.md" + [tool.hatch.version] path = "sgpt/__version__.py" diff --git a/sgpt/__version__.py b/sgpt/__version__.py index c68196d1..67bc602a 100644 --- a/sgpt/__version__.py +++ b/sgpt/__version__.py @@ -1 +1 @@ -__version__ = "1.2.0" +__version__ = "1.3.0" diff --git a/sgpt/handlers/handler.py b/sgpt/handlers/handler.py index 7d3b0d25..8cecc56d 100644 --- a/sgpt/handlers/handler.py +++ b/sgpt/handlers/handler.py @@ -2,7 +2,7 @@ from pathlib import Path from typing import Any, Dict, Generator, List, Optional -from openai import OpenAI +import litellm # type: ignore from ..cache import Cache from ..config import cfg @@ -10,16 +10,13 @@ from ..printer import MarkdownPrinter, Printer, TextPrinter from ..role import DefaultRoles, SystemRole +litellm.suppress_debug_info = True + class Handler: cache = Cache(int(cfg.get("CACHE_LENGTH")), Path(cfg.get("CACHE_PATH"))) def __init__(self, role: SystemRole) -> None: - self.client = OpenAI( - base_url=cfg.get("OPENAI_BASE_URL"), - api_key=cfg.get("OPENAI_API_KEY"), - timeout=int(cfg.get("REQUEST_TIMEOUT")), - ) self.role = role @property @@ -73,28 +70,30 @@ def get_completion( if is_shell_role or is_code_role or is_dsc_shell_role: functions = None - for chunk in self.client.chat.completions.create( + for chunk in litellm.completion( model=model, temperature=temperature, top_p=top_p, - messages=messages, # type: ignore - functions=functions, # type: ignore + messages=messages, + functions=functions, stream=True, + api_key=cfg.get("OPENAI_API_KEY"), ): - delta = chunk.choices[0].delta # type: ignore - if delta.function_call: - if delta.function_call.name: - name = delta.function_call.name - if delta.function_call.arguments: - arguments += delta.function_call.arguments - if chunk.choices[0].finish_reason == "function_call": # type: ignore + delta = chunk.choices[0].delta + function_call = delta.get("function_call") + if function_call: + if function_call.name: + name = function_call.name + if function_call.arguments: + arguments += function_call.arguments + if chunk.choices[0].finish_reason == "function_call": yield from self.handle_function_call(messages, name, arguments) yield from self.get_completion( model, temperature, top_p, messages, functions, caching=False ) return - yield delta.content or "" + yield delta.get("content") or "" def handle( self, diff --git a/tests/test_code.py b/tests/test_code.py index 5df2c421..1b49f8da 100644 --- a/tests/test_code.py +++ b/tests/test_code.py @@ -4,14 +4,14 @@ from sgpt.config import cfg from sgpt.role import DefaultRoles, SystemRole -from .utils import app, cmd_args, comp_args, comp_chunks, runner +from .utils import app, cmd_args, comp_args, mock_comp, runner role = SystemRole.get(DefaultRoles.CODE.value) -@patch("openai.resources.chat.Completions.create") +@patch("litellm.completion") def test_code_generation(mock): - mock.return_value = comp_chunks("print('Hello World')") + mock.return_value = mock_comp("print('Hello World')") args = {"prompt": "hello world python", "--code": True} result = runner.invoke(app, cmd_args(**args)) @@ -21,9 +21,9 @@ def test_code_generation(mock): assert "print('Hello World')" in result.stdout -@patch("openai.resources.chat.Completions.create") +@patch("litellm.completion") def test_code_generation_stdin(completion): - completion.return_value = comp_chunks("# Hello\nprint('Hello')") + completion.return_value = mock_comp("# Hello\nprint('Hello')") args = {"prompt": "make comments for code", "--code": True} stdin = "print('Hello')" @@ -36,11 +36,11 @@ def test_code_generation_stdin(completion): assert "print('Hello')" in result.stdout -@patch("openai.resources.chat.Completions.create") +@patch("litellm.completion") def test_code_chat(completion): completion.side_effect = [ - comp_chunks("print('hello')"), - comp_chunks("print('hello')\nprint('world')"), + mock_comp("print('hello')"), + mock_comp("print('hello')\nprint('world')"), ] chat_name = "_test" chat_path = Path(cfg.get("CHAT_CACHE_PATH")) / chat_name @@ -77,11 +77,11 @@ def test_code_chat(completion): # TODO: Code chat can be recalled without --code option. -@patch("openai.resources.chat.Completions.create") +@patch("litellm.completion") def test_code_repl(completion): completion.side_effect = [ - comp_chunks("print('hello')"), - comp_chunks("print('hello')\nprint('world')"), + mock_comp("print('hello')"), + mock_comp("print('hello')\nprint('world')"), ] chat_name = "_test" chat_path = Path(cfg.get("CHAT_CACHE_PATH")) / chat_name @@ -109,7 +109,7 @@ def test_code_repl(completion): assert "print('world')" in result.stdout -@patch("openai.resources.chat.Completions.create") +@patch("litellm.completion") def test_code_and_shell(completion): args = {"--code": True, "--shell": True} result = runner.invoke(app, cmd_args(**args)) @@ -119,7 +119,7 @@ def test_code_and_shell(completion): assert "Error" in result.stdout -@patch("openai.resources.chat.Completions.create") +@patch("litellm.completion") def test_code_and_describe_shell(completion): args = {"--code": True, "--describe-shell": True} result = runner.invoke(app, cmd_args(**args)) diff --git a/tests/test_default.py b/tests/test_default.py index 418fb470..52de465e 100644 --- a/tests/test_default.py +++ b/tests/test_default.py @@ -8,15 +8,15 @@ from sgpt.__version__ import __version__ from sgpt.role import DefaultRoles, SystemRole -from .utils import app, cmd_args, comp_args, comp_chunks, runner +from .utils import app, cmd_args, comp_args, mock_comp, runner role = SystemRole.get(DefaultRoles.DEFAULT.value) cfg = config.cfg -@patch("openai.resources.chat.Completions.create") +@patch("litellm.completion") def test_default(completion): - completion.return_value = comp_chunks("Prague") + completion.return_value = mock_comp("Prague") args = {"prompt": "capital of the Czech Republic?"} result = runner.invoke(app, cmd_args(**args)) @@ -26,9 +26,9 @@ def test_default(completion): assert "Prague" in result.stdout -@patch("openai.resources.chat.Completions.create") +@patch("litellm.completion") def test_default_stdin(completion): - completion.return_value = comp_chunks("Prague") + completion.return_value = mock_comp("Prague") stdin = "capital of the Czech Republic?" result = runner.invoke(app, cmd_args(), input=stdin) @@ -38,9 +38,9 @@ def test_default_stdin(completion): assert "Prague" in result.stdout -@patch("openai.resources.chat.Completions.create") +@patch("litellm.completion") def test_default_chat(completion): - completion.side_effect = [comp_chunks("ok"), comp_chunks("4")] + 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) @@ -90,9 +90,9 @@ def test_default_chat(completion): chat_path.unlink() -@patch("openai.resources.chat.Completions.create") +@patch("litellm.completion") def test_default_repl(completion): - completion.side_effect = [comp_chunks("ok"), comp_chunks("8")] + 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) @@ -119,9 +119,9 @@ def test_default_repl(completion): assert "8" in result.stdout -@patch("openai.resources.chat.Completions.create") +@patch("litellm.completion") def test_default_repl_stdin(completion): - completion.side_effect = [comp_chunks("ok init"), comp_chunks("ok another")] + 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) @@ -153,9 +153,9 @@ def test_default_repl_stdin(completion): assert "ok another" in result.stdout -@patch("openai.resources.chat.Completions.create") +@patch("litellm.completion") def test_llm_options(completion): - completion.return_value = comp_chunks("Berlin") + completion.return_value = mock_comp("Berlin") args = { "prompt": "capital of the Germany?", @@ -179,7 +179,7 @@ def test_llm_options(completion): assert "Berlin" in result.stdout -@patch("openai.resources.chat.Completions.create") +@patch("litellm.completion") def test_version(completion): args = {"--version": True} result = runner.invoke(app, cmd_args(**args)) diff --git a/tests/test_roles.py b/tests/test_roles.py index d184bfcd..dcf1b94a 100644 --- a/tests/test_roles.py +++ b/tests/test_roles.py @@ -5,12 +5,12 @@ from sgpt.config import cfg from sgpt.role import SystemRole -from .utils import app, cmd_args, comp_args, comp_chunks, runner +from .utils import app, cmd_args, comp_args, mock_comp, runner -@patch("openai.resources.chat.Completions.create") +@patch("litellm.completion") def test_role(completion): - completion.return_value = comp_chunks('{"foo": "bar"}') + 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"} @@ -44,6 +44,7 @@ def test_role(completion): 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) diff --git a/tests/test_shell.py b/tests/test_shell.py index 346d4b42..a008625a 100644 --- a/tests/test_shell.py +++ b/tests/test_shell.py @@ -5,13 +5,13 @@ from sgpt.config import cfg from sgpt.role import DefaultRoles, SystemRole -from .utils import app, cmd_args, comp_args, comp_chunks, runner +from .utils import app, cmd_args, comp_args, mock_comp, runner -@patch("openai.resources.chat.Completions.create") +@patch("litellm.completion") def test_shell(completion): role = SystemRole.get(DefaultRoles.SHELL.value) - completion.return_value = comp_chunks("git commit -m test") + 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)) @@ -22,9 +22,9 @@ def test_shell(completion): assert "[E]xecute, [D]escribe, [A]bort:" in result.stdout -@patch("openai.resources.chat.Completions.create") +@patch("litellm.completion") def test_shell_stdin(completion): - completion.return_value = comp_chunks("ls -l | sort") + completion.return_value = mock_comp("ls -l | sort") role = SystemRole.get(DefaultRoles.SHELL.value) args = {"prompt": "Sort by name", "--shell": True} @@ -38,9 +38,9 @@ def test_shell_stdin(completion): assert "[E]xecute, [D]escribe, [A]bort:" in result.stdout -@patch("openai.resources.chat.Completions.create") +@patch("litellm.completion") def test_describe_shell(completion): - completion.return_value = comp_chunks("lists the contents of a folder") + completion.return_value = mock_comp("lists the contents of a folder") role = SystemRole.get(DefaultRoles.DESCRIBE_SHELL.value) args = {"prompt": "ls", "--describe-shell": True} @@ -51,9 +51,9 @@ def test_describe_shell(completion): assert "lists" in result.stdout -@patch("openai.resources.chat.Completions.create") +@patch("litellm.completion") def test_describe_shell_stdin(completion): - completion.return_value = comp_chunks("lists the contents of a folder") + completion.return_value = mock_comp("lists the contents of a folder") role = SystemRole.get(DefaultRoles.DESCRIBE_SHELL.value) args = {"--describe-shell": True} @@ -67,9 +67,9 @@ def test_describe_shell_stdin(completion): @patch("os.system") -@patch("openai.resources.chat.Completions.create") +@patch("litellm.completion") def test_shell_run_description(completion, system): - completion.side_effect = [comp_chunks("echo hello"), comp_chunks("prints hello")] + 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) @@ -80,9 +80,9 @@ def test_shell_run_description(completion, system): assert "prints hello" in result.stdout -@patch("openai.resources.chat.Completions.create") +@patch("litellm.completion") def test_shell_chat(completion): - completion.side_effect = [comp_chunks("ls"), comp_chunks("ls | sort")] + 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 @@ -119,9 +119,9 @@ def test_shell_chat(completion): @patch("os.system") -@patch("openai.resources.chat.Completions.create") +@patch("litellm.completion") def test_shell_repl(completion, mock_system): - completion.side_effect = [comp_chunks("ls"), comp_chunks("ls | sort")] + 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 @@ -151,7 +151,7 @@ def test_shell_repl(completion, mock_system): assert "ls | sort" in result.stdout -@patch("openai.resources.chat.Completions.create") +@patch("litellm.completion") def test_shell_and_describe_shell(completion): args = {"prompt": "ls", "--describe-shell": True, "--shell": True} result = runner.invoke(app, cmd_args(**args)) @@ -161,9 +161,9 @@ def test_shell_and_describe_shell(completion): assert "Error" in result.stdout -@patch("openai.resources.chat.Completions.create") +@patch("litellm.completion") def test_shell_no_interaction(completion): - completion.return_value = comp_chunks("git commit -m test") + completion.return_value = mock_comp("git commit -m test") role = SystemRole.get(DefaultRoles.SHELL.value) args = { diff --git a/tests/utils.py b/tests/utils.py index 13a7ce2f..ed0998c1 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -1,9 +1,7 @@ -import datetime +from unittest.mock import ANY 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 litellm import completion as completion from typer.testing import CliRunner from sgpt import main @@ -14,23 +12,9 @@ app.command()(main) -def comp_chunks(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.datetime.now().timestamp()), - ) - for token in tokens_string - ] +def mock_comp(tokens_string): + model = cfg.get("DEFAULT_MODEL") + return completion(model=model, mock_response=tokens_string, stream=True) def cmd_args(prompt="", **kwargs): @@ -56,5 +40,6 @@ def comp_args(role, prompt, **kwargs): "top_p": 1.0, "functions": None, "stream": True, + "api_key": ANY, **kwargs, }