Minor handler improvements

Author: TheR1D

sgpt/app.py

@@ -243,11 +243,11 @@ def main(
     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),
+            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,
-        ).lower()
+        )
 
         if option in ("e", "y"):
             # "y" option is for keeping compatibility with old version.

sgpt/config.py

@@ -6,6 +6,8 @@
 
 import typer
 
+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"
@@ -85,7 +87,7 @@ 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 typer.Abort(f"Missing config key: {key}")
+            raise UsageError(f"Missing config key: {key}")
         return value
 
 

sgpt/handlers/handler.py

@@ -137,19 +137,16 @@ def get_completion(
                 if tool_calls:
                     for tool_call in tool_calls:
                         if use_litellm:
-                            if tool_call.get("id"):
-                                tool_call_id = tool_call.get("id")
-                            if tool_call.get("function", {}).get("name"):
-                                name = tool_call["function"]["name"]
-                            if tool_call.get("function", {}).get("arguments"):
-                                arguments += tool_call["function"]["arguments"]
+                            # 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:
-                            if tool_call.id:
-                                tool_call_id = tool_call.id
-                            if tool_call.function.name:
-                                name = tool_call.function.name
-                            if tool_call.function.arguments:
-                                arguments += tool_call.function.arguments
+                            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

sgpt/llm_functions/common/execute_shell.py

@@ -13,7 +13,7 @@ class Function(BaseModel):
         ...,
         example="ls -la",
         description="Shell command to execute.",
-    ) # type: ignore
+    )  # type: ignore
 
     @classmethod
     def execute(cls, shell_command: str) -> str:

sgpt/utils.py

@@ -5,7 +5,7 @@
 from typing import Any, Callable
 
 import typer
-from click import BadParameter
+from click import BadParameter, UsageError
 
 from sgpt.__version__ import __version__
 from sgpt.integration import bash_integration, zsh_integration
@@ -82,7 +82,7 @@ def install_shell_integration(*_args: Any) -> None:
         with open(os.path.expanduser("~/.bashrc"), "a", encoding="utf-8") as file:
             file.write(bash_integration)
     else:
-        raise typer.Abort("ShellGPT integrations only available for ZSH and Bash.")
+        raise UsageError("ShellGPT integrations only available for ZSH and Bash.")
 
     typer.echo("Done! Restart your shell to apply changes.")