Ollama integration 🦙 (#463)

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

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"
 

sgpt/__version__.py

@@ -1 +1 @@
-__version__ = "1.2.0"
+__version__ = "1.3.0"

sgpt/handlers/handler.py

@@ -2,7 +2,7 @@ import json
 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 ..function import get_function
 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 @@ class Handler:
         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,

tests/test_code.py

@@ -4,14 +4,14 @@ from unittest.mock import patch
 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))

tests/test_default.py

@@ -8,15 +8,15 @@ from sgpt import config, main
 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))

tests/test_roles.py

@@ -5,12 +5,12 @@ from unittest.mock import patch
 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)

tests/test_shell.py

@@ -5,13 +5,13 @@ from unittest.mock import patch
 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 = {

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 = typer.Typer()
 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,
     }