shell-gpt-plus 1.1.0__py3-none-any.whl

sgpt/__init__.py

@@ -0,0 +1,5 @@
+from .init_config import deploy_user_config
+deploy_user_config()
+
+from .app import main as main
+from .app import entry_point as cli  # noqa: F401

sgpt/__main__.py

@@ -0,0 +1,3 @@
+from .app import entry_point
+
+entry_point()

sgpt/__version__.py

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

sgpt/app.py

@@ -0,0 +1,276 @@
+import os
+
+# To allow users to use arrow keys in the REPL.
+import readline  # noqa: F401
+import sys
+
+import typer
+from click import UsageError
+from click.types import Choice
+from prompt_toolkit import PromptSession
+
+from sgpt.config import cfg
+from sgpt.function import get_openai_schemas
+from sgpt.handlers.chat_handler import ChatHandler
+from sgpt.handlers.default_handler import DefaultHandler
+from sgpt.handlers.repl_handler import ReplHandler
+from sgpt.llm_functions.init_functions import install_functions as inst_funcs
+from sgpt.role import DefaultRoles, SystemRole
+from sgpt.utils import (
+    get_edited_prompt,
+    get_sgpt_version,
+    install_shell_integration,
+    run_command,
+)
+
+
+def main(
+    prompt: str = typer.Argument(
+        "",
+        show_default=False,
+        help="The prompt to generate completions for.",
+    ),
+    model: str = typer.Option(
+        cfg.get("DEFAULT_MODEL"),
+        help="Large language model to use.",
+    ),
+    temperature: float = typer.Option(
+        0.0,
+        min=0.0,
+        max=2.0,
+        help="Randomness of generated output.",
+    ),
+    top_p: float = typer.Option(
+        1.0,
+        min=0.0,
+        max=1.0,
+        help="Limits highest probable tokens (words).",
+    ),
+    md: bool = typer.Option(
+        cfg.get("PRETTIFY_MARKDOWN") == "true",
+        help="Prettify markdown output.",
+    ),
+    shell: bool = typer.Option(
+        False,
+        "--shell",
+        "-s",
+        help="Generate and execute shell commands.",
+        rich_help_panel="Assistance Options",
+    ),
+    interaction: bool = typer.Option(
+        cfg.get("SHELL_INTERACTION") == "true",
+        help="Interactive mode for --shell option.",
+        rich_help_panel="Assistance Options",
+    ),
+    describe_shell: bool = typer.Option(
+        False,
+        "--describe-shell",
+        "-d",
+        help="Describe a shell command.",
+        rich_help_panel="Assistance Options",
+    ),
+    code: bool = typer.Option(
+        False,
+        "--code",
+        "-c",
+        help="Generate only code.",
+        rich_help_panel="Assistance Options",
+    ),
+    functions: bool = typer.Option(
+        cfg.get("OPENAI_USE_FUNCTIONS") == "true",
+        help="Allow function calls.",
+        rich_help_panel="Assistance Options",
+    ),
+    editor: bool = typer.Option(
+        False,
+        help="Open $EDITOR to provide a prompt.",
+    ),
+    cache: bool = typer.Option(
+        True,
+        help="Cache completion results.",
+    ),
+    version: bool = typer.Option(
+        False,
+        "--version",
+        help="Show version.",
+        callback=get_sgpt_version,
+    ),
+    chat: str = typer.Option(
+        None,
+        help="Follow conversation with id, " 'use "temp" for quick session.',
+        rich_help_panel="Chat Options",
+    ),
+    repl: str = typer.Option(
+        None,
+        help="Start a REPL (Read-eval-print loop) session.",
+        rich_help_panel="Chat Options",
+    ),
+    show_chat: str = typer.Option(
+        None,
+        help="Show all messages from provided chat id.",
+        rich_help_panel="Chat Options",
+    ),
+    list_chats: bool = typer.Option(
+        False,
+        "--list-chats",
+        "-lc",
+        help="List all existing chat ids.",
+        callback=ChatHandler.list_ids,
+        rich_help_panel="Chat Options",
+    ),
+    role: str = typer.Option(
+        None,
+        help="System role for GPT model.",
+        rich_help_panel="Role Options",
+    ),
+    create_role: str = typer.Option(
+        None,
+        help="Create role.",
+        callback=SystemRole.create,
+        rich_help_panel="Role Options",
+    ),
+    show_role: str = typer.Option(
+        None,
+        help="Show role.",
+        callback=SystemRole.show,
+        rich_help_panel="Role Options",
+    ),
+    list_roles: bool = typer.Option(
+        False,
+        "--list-roles",
+        "-lr",
+        help="List roles.",
+        callback=SystemRole.list,
+        rich_help_panel="Role Options",
+    ),
+    install_integration: bool = typer.Option(
+        False,
+        help="Install shell integration (ZSH and Bash only)",
+        callback=install_shell_integration,
+        hidden=True,  # Hiding since should be used only once.
+    ),
+    install_functions: bool = typer.Option(
+        False,
+        help="Install default functions.",
+        callback=inst_funcs,
+        hidden=True,  # Hiding since should be used only once.
+    ),
+) -> None:
+    stdin_passed = not sys.stdin.isatty()
+
+    if stdin_passed:
+        stdin = ""
+        # TODO: This is very hacky.
+        # In some cases, we need to pass stdin along with inputs.
+        # When we want part of stdin to be used as a init prompt,
+        # but rest of the stdin to be used as a inputs. For example:
+        # echo "hello\n__sgpt__eof__\nThis is input" | sgpt --repl temp
+        # In this case, "hello" will be used as a init prompt, and
+        # "This is input" will be used as "interactive" input to the REPL.
+        # This is useful to test REPL with some initial context.
+        for line in sys.stdin:
+            if "__sgpt__eof__" in line:
+                break
+            stdin += line
+        prompt = f"{stdin}\n\n{prompt}" if prompt else stdin
+        try:
+            # Switch to stdin for interactive input.
+            if os.name == "posix":
+                sys.stdin = open("/dev/tty", "r")
+            elif os.name == "nt":
+                sys.stdin = open("CON", "r")
+        except OSError:
+            # Non-interactive shell.
+            pass
+
+    if show_chat:
+        ChatHandler.show_messages(show_chat, md)
+
+    if sum((shell, describe_shell, code)) > 1:
+        raise UsageError(
+            "Only one of --shell, --describe-shell, and --code options can be used at a time."
+        )
+
+    if chat and repl:
+        raise UsageError("--chat and --repl options cannot be used together.")
+
+    if editor and stdin_passed:
+        raise UsageError("--editor option cannot be used with stdin input.")
+
+    if editor:
+        prompt = get_edited_prompt()
+
+    role_class = (
+        DefaultRoles.check_get(shell, describe_shell, code)
+        if not role
+        else SystemRole.get(role)
+    )
+
+    function_schemas = (get_openai_schemas() or None) if functions else None
+
+    if repl:
+        # Will be in infinite loop here until user exits with Ctrl+C.
+        ReplHandler(repl, role_class, md).handle(
+            init_prompt=prompt,
+            model=model,
+            temperature=temperature,
+            top_p=top_p,
+            caching=cache,
+            functions=function_schemas,
+        )
+
+    if chat:
+        full_completion = ChatHandler(chat, role_class, md).handle(
+            prompt=prompt,
+            model=model,
+            temperature=temperature,
+            top_p=top_p,
+            caching=cache,
+            functions=function_schemas,
+        )
+    else:
+        full_completion = DefaultHandler(role_class, md).handle(
+            prompt=prompt,
+            model=model,
+            temperature=temperature,
+            top_p=top_p,
+            caching=cache,
+            functions=function_schemas,
+        )
+
+    session: PromptSession[str] = PromptSession()
+
+    while shell and interaction:
+        option = typer.prompt(
+            text="[E]xecute, [M]odify, [D]escribe, [A]bort",
+            type=Choice(("e", "m", "d", "a", "y"), case_sensitive=False),
+            default="e" if cfg.get("DEFAULT_EXECUTE_SHELL_CMD") == "true" else "a",
+            show_choices=False,
+            show_default=False,
+        )
+
+        if option in ("e", "y"):
+            # "y" option is for keeping compatibility with old version.
+            run_command(full_completion)
+        elif option == "m":
+            full_completion = session.prompt("", default=full_completion)
+            continue
+        elif option == "d":
+            DefaultHandler(DefaultRoles.DESCRIBE_SHELL.get_role(), md).handle(
+                full_completion,
+                model=model,
+                temperature=temperature,
+                top_p=top_p,
+                caching=cache,
+                functions=function_schemas,
+            )
+            continue
+        break
+
+
+def entry_point() -> None:
+    typer.run(main)
+
+
+if __name__ == "__main__":
+    entry_point()

sgpt/cache.py

@@ -0,0 +1,61 @@
+import json
+from hashlib import md5
+from pathlib import Path
+from typing import Any, Callable, Generator, no_type_check
+
+
+class Cache:
+    """
+    Decorator class that adds caching functionality to a function.
+    """
+
+    def __init__(self, length: int, cache_path: Path) -> None:
+        """
+        Initialize the Cache decorator.
+
+        :param length: Integer, maximum number of cache files to keep.
+        """
+        self.length = length
+        self.cache_path = cache_path
+        self.cache_path.mkdir(parents=True, exist_ok=True)
+
+    def __call__(self, func: Callable[..., Any]) -> Callable[..., Any]:
+        """
+        The Cache decorator.
+
+        :param func: The function to cache.
+        :return: Wrapped function with caching.
+        """
+
+        def wrapper(*args: Any, **kwargs: Any) -> Generator[str, None, None]:
+            key = md5(json.dumps((args[1:], kwargs)).encode("utf-8")).hexdigest()
+            file = self.cache_path / key
+            if kwargs.pop("caching") and file.exists():
+                yield file.read_text()
+                return
+            result = ""
+            for i in func(*args, **kwargs):
+                result += i
+                yield i
+            if "@FunctionCall" not in result:
+                file.write_text(result, encoding="utf-8")
+            self._delete_oldest_files(self.length)  # type: ignore
+
+        return wrapper
+
+    @no_type_check
+    def _delete_oldest_files(self, max_files: int) -> None:
+        """
+        Class method to delete the oldest cached files in the CACHE_DIR folder.
+
+        :param max_files: Integer, the maximum number of files to keep in the CACHE_DIR folder.
+        """
+        # Get all files in the folder.
+        files = self.cache_path.glob("*")
+        # Sort files by last modification time in ascending order.
+        files = sorted(files, key=lambda f: f.stat().st_mtime)
+        # Delete the oldest files if the number of files exceeds the limit.
+        if len(files) > max_files:
+            num_files_to_delete = len(files) - max_files
+            for i in range(num_files_to_delete):
+                files[i].unlink()

sgpt/config.py

@@ -0,0 +1,92 @@
+import os
+from getpass import getpass
+from pathlib import Path
+from tempfile import gettempdir
+from typing import Any
+
+from click import UsageError
+
+CONFIG_FOLDER = os.path.expanduser("~/.config")
+SHELL_GPT_CONFIG_FOLDER = Path(CONFIG_FOLDER) / "shell_gpt"
+SHELL_GPT_CONFIG_PATH = SHELL_GPT_CONFIG_FOLDER / ".sgptrc"
+ROLE_STORAGE_PATH = SHELL_GPT_CONFIG_FOLDER / "roles"
+FUNCTIONS_PATH = SHELL_GPT_CONFIG_FOLDER / "functions"
+CHAT_CACHE_PATH = Path(gettempdir()) / "chat_cache"
+CACHE_PATH = Path(gettempdir()) / "cache"
+
+# TODO: Refactor ENV variables with SGPT_ prefix.
+DEFAULT_CONFIG = {
+    # TODO: Refactor it to CHAT_STORAGE_PATH.
+    "CHAT_CACHE_PATH": os.getenv("CHAT_CACHE_PATH", str(CHAT_CACHE_PATH)),
+    "CACHE_PATH": os.getenv("CACHE_PATH", str(CACHE_PATH)),
+    "CHAT_CACHE_LENGTH": int(os.getenv("CHAT_CACHE_LENGTH", "100")),
+    "CACHE_LENGTH": int(os.getenv("CHAT_CACHE_LENGTH", "100")),
+    "REQUEST_TIMEOUT": int(os.getenv("REQUEST_TIMEOUT", "60")),
+    "DEFAULT_MODEL": os.getenv("DEFAULT_MODEL", "deepseek-v4-flash"),
+    "DEFAULT_COLOR": os.getenv("DEFAULT_COLOR", "magenta"),
+    "ROLE_STORAGE_PATH": os.getenv("ROLE_STORAGE_PATH", str(ROLE_STORAGE_PATH)),
+    "DEFAULT_EXECUTE_SHELL_CMD": os.getenv("DEFAULT_EXECUTE_SHELL_CMD", "false"),
+    "DISABLE_STREAMING": os.getenv("DISABLE_STREAMING", "false"),
+    "CODE_THEME": os.getenv("CODE_THEME", "dracula"),
+    "OPENAI_FUNCTIONS_PATH": os.getenv("OPENAI_FUNCTIONS_PATH", str(FUNCTIONS_PATH)),
+    "OPENAI_USE_FUNCTIONS": os.getenv("OPENAI_USE_FUNCTIONS", "true"),
+    "SHOW_FUNCTIONS_OUTPUT": os.getenv("SHOW_FUNCTIONS_OUTPUT", "false"),
+    "API_BASE_URL": os.getenv("API_BASE_URL", "https://api.deepseek.com"),
+    "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.
+}
+
+
+class Config(dict):  # type: ignore
+    def __init__(self, config_path: Path, **defaults: Any):
+        self.config_path = config_path
+
+        if self._exists:
+            self._read()
+            has_new_config = False
+            for key, value in defaults.items():
+                if key not in self:
+                    has_new_config = True
+                    self[key] = value
+            if has_new_config:
+                self._write()
+        else:
+            config_path.parent.mkdir(parents=True, exist_ok=True)
+            # Don't write API key to config file if it is in the environment.
+            if not defaults.get("OPENAI_API_KEY") and not os.getenv("OPENAI_API_KEY"):
+                __api_key = getpass(prompt="Please enter your DeepSeek API key: ")
+                defaults["OPENAI_API_KEY"] = __api_key
+            super().__init__(**defaults)
+            self._write()
+
+    @property
+    def _exists(self) -> bool:
+        return self.config_path.exists()
+
+    def _write(self) -> None:
+        with open(self.config_path, "w", encoding="utf-8") as file:
+            string_config = ""
+            for key, value in self.items():
+                string_config += f"{key}={value}\n"
+            file.write(string_config)
+
+    def _read(self) -> None:
+        with open(self.config_path, "r", encoding="utf-8") as file:
+            for line in file:
+                if line.strip() and not line.startswith("#"):
+                    key, value = line.strip().split("=", 1)
+                    self[key] = value
+
+    def get(self, key: str) -> str:  # type: ignore
+        # Prioritize environment variables over config file.
+        value = os.getenv(key) or super().get(key)
+        if not value:
+            raise UsageError(f"Missing config key: {key}")
+        return value
+
+
+cfg = Config(SHELL_GPT_CONFIG_PATH, **DEFAULT_CONFIG)

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/functions/execute_shell.py

@@ -0,0 +1,86 @@
+import subprocess
+import locale
+import re
+from typing import Any, ClassVar, Dict, List  # 新增 ClassVar, List 导入
+
+from pydantic import BaseModel, Field
+
+
+class Function(BaseModel):
+    """
+    执行一个 Shell 命令并返回输出（结果）。
+    注意：危险命令会被自动拦截，确保系统安全。
+    """
+
+    shell_command: str = Field(
+        ...,
+        example="dir",
+        description="要执行的 shell 命令。",
+    )
+
+    # 关键修改：添加 ClassVar[List[str]] 类型注解
+    DANGEROUS_PATTERNS: ClassVar[List[str]] = [
+        r"\bdel\b",
+        r"\berase\b",
+        r"\brmdir\b",
+        r"\brd\b",
+        r"\bformat\b",
+        r"\bshutdown\b",
+        r"\brestart\b",
+        r"\blogoff\b",
+        r"\breg\s+delete\b",
+        r"\breg\s+add\b.*\/f",
+        r"\bicacls\b",
+        r"\btakeown\b",
+        r"\bnet\s+user\b.*\/delete",
+        r"\bnet\s+localgroup\b.*\/delete",
+        r"\bsc\s+stop\b",
+        r"\bsc\s+delete\b",
+        r"\bwmic\s+process\s+call\s+create\b",
+        r"\bschtasks\b.*\/delete",
+        r"\bpowershell\s+.*-EncodedCommand\b",
+        r"\bcmd\s+/c\s+.*del\b",
+        r"&&\s*del\b",
+        r"\|\s*del\b",
+    ]
+
+    @classmethod
+    def _is_dangerous(cls, command: str) -> bool:
+        """检查命令是否命中危险模式"""
+        command_lower = command.lower()
+        for pattern in cls.DANGEROUS_PATTERNS:
+            if re.search(pattern, command_lower):
+                return True
+        return False
+
+    @classmethod
+    def execute(cls, shell_command: str) -> str:
+        if cls._is_dangerous(shell_command):
+            return (
+                "Error: Potentially dangerous command detected and blocked.\n"
+                "If you believe this is a safe operation, please rephrase it "
+                "or execute it manually with proper authorization."
+            )
+        process = subprocess.Popen(
+            shell_command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT
+        )
+        output, _ = process.communicate()
+        exit_code = process.returncode
+        encoding = locale.getpreferredencoding()
+        return f"Exit code: {exit_code}, Output:\n{output.decode(encoding, errors='replace')}"
+
+    @classmethod
+    def openai_schema(cls) -> Dict[str, Any]:
+        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/functions/fetch_url.py

@@ -0,0 +1,44 @@
+import requests
+from pydantic import BaseModel, Field
+from typing import Dict, Any
+from bs4 import BeautifulSoup
+
+class Function(BaseModel):
+    """
+    抓取指定 URL 的网页文本内容（提取正文）。
+    当需要查看某个链接的具体内容时调用此函数。
+    """
+
+    url: str = Field(..., description="要抓取的网页完整网址，例如 'https://example.com'。")
+
+    @classmethod
+    def execute(cls, url: str) -> str:
+        try:
+            headers = {"User-Agent": "Mozilla/5.0"}
+            resp = requests.get(url, headers=headers, timeout=15)
+            resp.raise_for_status()
+            soup = BeautifulSoup(resp.text, "html.parser")
+            # 移除脚本和样式
+            for tag in soup(["script", "style", "nav", "footer", "header"]):
+                tag.decompose()
+            text = soup.get_text(separator="\n", strip=True)
+            # 限制长度，避免超出 token 限制
+            return text[:3000]
+        except Exception as e:
+            return f"抓取网页出错：{e}"
+
+    @classmethod
+    def openai_schema(cls) -> Dict[str, Any]:
+        schema = cls.model_json_schema()
+        return {
+            "type": "function",
+            "function": {
+                "name": "fetch_url",
+                "description": cls.__doc__.strip() if cls.__doc__ else "",
+                "parameters": {
+                    "type": "object",
+                    "properties": schema.get("properties", {}),
+                    "required": schema.get("required", []),
+                },
+            },
+        }

sgpt/functions/recall.py

@@ -0,0 +1,39 @@
+import json
+from pathlib import Path
+from typing import Any, Dict
+
+from pydantic import BaseModel, Field
+
+MEMORY_FILE = Path.home() / ".shell_gpt_memory.txt"
+
+
+class Function(BaseModel):
+    query: str = Field(
+        ...,
+        description="用于回忆记忆的关键词或问题（所有记忆都会被返回）",
+    )
+
+    @classmethod
+    def execute(cls, query: str) -> str:
+        if not MEMORY_FILE.exists():
+            return "当前没有任何保存的记忆。"
+        content = MEMORY_FILE.read_text(encoding="utf-8")
+        if not content.strip():
+            return "当前没有任何保存的记忆。"
+        return f"已保存的记忆：\n{content.strip()}"
+
+    @classmethod
+    def openai_schema(cls) -> Dict[str, Any]:
+        schema = cls.model_json_schema()
+        return {
+            "type": "function",
+            "function": {
+                "name": "recall",
+                "description": cls.__doc__.strip() if cls.__doc__ else "",
+                "parameters": {
+                    "type": "object",
+                    "properties": schema.get("properties", {}),
+                    "required": schema.get("required", []),
+                },
+            },
+        }