diff --git a/README.md b/README.md index 7908555..819445b 100644 --- a/README.md +++ b/README.md @@ -1,7 +1,7 @@ # ShellGPT A command-line productivity tool powered by AI large language models (LLM). This command-line tool offers streamlined generation of **shell commands, code snippets, documentation**, eliminating the need for external resources (like Google search). Supports Linux, macOS, Windows and compatible with all major Shells like PowerShell, CMD, Bash, Zsh, etc. -https://github.com/TheR1D/shell_gpt/assets/16740832/9197283c-db6a-4b46-bfea-3eb776dd9093 +https://github.com/TheR1D/shell_gpt/assets/16740832/721ddb19-97e7-428f-a0ee-107d027ddd59 ## Installation ```shell @@ -290,28 +290,7 @@ The snippet of code you've provided is written in Python. It prompts the user... 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 the following syntax: -```python -# execute_shell_command.py -import subprocess -from pydantic import Field -from instructor import OpenAISchema - - -class Function(OpenAISchema): - """ - Executes a shell command and returns the output (result). - """ - shell_command: str = Field(..., example="ls -la", descriptions="Shell command to execute.") - - class Config: - title = "execute_shell_command" - - @classmethod - def execute(cls, shell_command: str) -> str: - result = subprocess.run(shell_command.split(), capture_output=True, text=True) - return f"Exit code: {result.returncode}, Output:\n{result.stdout}" -``` +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 diff --git a/pyproject.toml b/pyproject.toml index 6e33435..a51a6a3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,12 +24,10 @@ classifiers = [ "Programming Language :: Python :: 3.13", ] dependencies = [ - "openai >= 1.34.0, < 2.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", - "instructor >= 1.0.0, < 2.0.0", 'pyreadline3 >= 3.4.1, < 4.0.0; sys_platform == "win32"', "prompt_toolkit >= 3.0.51", ] diff --git a/sgpt/__version__.py b/sgpt/__version__.py index 56dadec..5b60188 100644 --- a/sgpt/__version__.py +++ b/sgpt/__version__.py @@ -1 +1 @@ -__version__ = "1.4.5" +__version__ = "1.5.0" diff --git a/sgpt/app.py b/sgpt/app.py index 7201a6c..e4b4d63 100644 --- a/sgpt/app.py +++ b/sgpt/app.py @@ -5,7 +5,7 @@ import readline # noqa: F401 import sys import typer -from click import BadArgumentUsage +from click import UsageError from click.types import Choice from prompt_toolkit import PromptSession @@ -187,15 +187,15 @@ def main( 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() @@ -248,6 +248,7 @@ def main( show_choices=False, show_default=False, ) + if option in ("e", "y"): # "y" option is for keeping compatibility with old version. run_command(full_completion) diff --git a/sgpt/function.py b/sgpt/function.py index d1156ae..0e33196 100644 --- a/sgpt/function.py +++ b/sgpt/function.py @@ -1,9 +1,10 @@ import importlib.util import sys -from abc import ABCMeta from pathlib import Path from typing import Any, Callable, Dict, List +from pydantic import BaseModel + from .config import cfg @@ -11,8 +12,8 @@ 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["name"] + self._openai_schema = module.Function.openai_schema() + self._name = self._openai_schema["function"]["name"] @property def name(self) -> str: @@ -34,13 +35,17 @@ class Function: sys.modules[module_name] = module spec.loader.exec_module(module) # type: ignore - if not isinstance(module.Function, ABCMeta): + 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 a 'execute' static method" + 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 @@ -59,15 +64,4 @@ def get_function(name: str) -> Callable[..., Any]: def get_openai_schemas() -> List[Dict[str, Any]]: - transformed_schemas = [] - for function in functions: - schema = { - "type": "function", - "function": { - "name": function.openai_schema["name"], - "description": function.openai_schema.get("description", ""), - "parameters": function.openai_schema.get("parameters", {}), - }, - } - transformed_schemas.append(schema) - return transformed_schemas + return [function.openai_schema for function in functions] diff --git a/sgpt/handlers/chat_handler.py b/sgpt/handlers/chat_handler.py index 6fcfc72..30930d7 100644 --- a/sgpt/handlers/chat_handler.py +++ b/sgpt/handlers/chat_handler.py @@ -3,7 +3,7 @@ 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 @@ -151,15 +151,13 @@ class ChatHandler(Handler): if self.initiated: 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}"' - ) + 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.' ) diff --git a/sgpt/handlers/handler.py b/sgpt/handlers/handler.py index 2f4df95..0f5f1a2 100644 --- a/sgpt/handlers/handler.py +++ b/sgpt/handlers/handler.py @@ -59,14 +59,22 @@ class Handler: 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": "", - "function_call": {"name": name, "arguments": arguments}, + "content": None, + "tool_calls": [ + { + "id": tool_call_id, + "type": "function", + "function": {"name": name, "arguments": arguments}, + } + ], } ) @@ -80,7 +88,11 @@ class Handler: result = get_function(name)(**dict_args) if cfg.get("SHOW_FUNCTIONS_OUTPUT") == "true": yield f"```text\n{result}\n```\n" - messages.append({"role": "function", "content": result, "name": name}) + + # Add tool response message + messages.append( + {"role": "tool", "content": result, "tool_call_id": tool_call_id} + ) @cache def get_completion( @@ -91,7 +103,7 @@ class Handler: messages: List[Dict[str, Any]], functions: Optional[List[Dict[str, str]]], ) -> Generator[str, None, None]: - name = arguments = "" + 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 @@ -124,12 +136,21 @@ class Handler: ) if tool_calls: for tool_call in tool_calls: - if tool_call.function.name: - name = tool_call.function.name - if tool_call.function.arguments: - arguments += tool_call.function.arguments + 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, name, arguments) + yield from self.handle_function_call( + messages, tool_call_id, name, arguments + ) yield from self.get_completion( model=model, temperature=temperature, diff --git a/sgpt/llm_functions/common/execute_shell.py b/sgpt/llm_functions/common/execute_shell.py index 55c626c..3f5d290 100644 --- a/sgpt/llm_functions/common/execute_shell.py +++ b/sgpt/llm_functions/common/execute_shell.py @@ -1,10 +1,10 @@ import subprocess +from typing import Any, Dict -from instructor import OpenAISchema -from pydantic import Field +from pydantic import BaseModel, Field -class Function(OpenAISchema): +class Function(BaseModel): """ Executes a shell command and returns the output (result). """ @@ -12,11 +12,8 @@ class Function(OpenAISchema): shell_command: str = Field( ..., example="ls -la", - descriptions="Shell command to execute.", - ) - - class Config: - title = "execute_shell_command" + description="Shell command to execute.", + ) # type: ignore @classmethod def execute(cls, shell_command: str) -> str: @@ -26,3 +23,20 @@ class Function(OpenAISchema): output, _ = process.communicate() exit_code = process.returncode return f"Exit code: {exit_code}, Output:\n{output.decode()}" + + @classmethod + def openai_schema(cls) -> Dict[str, Any]: + """Generate OpenAI function schema from Pydantic model.""" + schema = cls.model_json_schema() + return { + "type": "function", + "function": { + "name": "execute_shell_command", + "description": cls.__doc__.strip() if cls.__doc__ else "", + "parameters": { + "type": "object", + "properties": schema.get("properties", {}), + "required": schema.get("required", []), + }, + }, + } diff --git a/sgpt/llm_functions/mac/apple_script.py b/sgpt/llm_functions/mac/apple_script.py index 6d07b39..2277a63 100644 --- a/sgpt/llm_functions/mac/apple_script.py +++ b/sgpt/llm_functions/mac/apple_script.py @@ -1,23 +1,20 @@ import subprocess +from typing import Any, Dict -from instructor import OpenAISchema -from pydantic import Field +from pydantic import BaseModel, Field -class Function(OpenAISchema): +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', - descriptions="Apple Script to execute.", - ) - - class Config: - title = "execute_apple_script" + description="Apple Script to execute.", + ) # type: ignore @classmethod def execute(cls, apple_script): @@ -31,3 +28,20 @@ class Function(OpenAISchema): return f"Output: {output}" except Exception as e: return f"Error: {e}" + + @classmethod + def openai_schema(cls) -> Dict[str, Any]: + """Generate OpenAI function schema from Pydantic model.""" + schema = cls.model_json_schema() + return { + "type": "function", + "function": { + "name": "execute_apple_script", + "description": cls.__doc__.strip() if cls.__doc__ else "", + "parameters": { + "type": "object", + "properties": schema.get("properties", {}), + "required": schema.get("required", []), + }, + }, + } diff --git a/sgpt/role.py b/sgpt/role.py index 2ae34e1..31371e4 100644 --- a/sgpt/role.py +++ b/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 @@ -76,7 +76,7 @@ class SystemRole: 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