diff --git a/nbs/00_core.ipynb b/nbs/00_core.ipynb index 032b885..a039dc0 100644 --- a/nbs/00_core.ipynb +++ b/nbs/00_core.ipynb @@ -266,7 +266,15 @@ "source": [ "#| export\n", "def _aliases(shell):\n", - " return co([shell, '-ic', 'alias'], text=True).strip()" + " try:\n", + " if sys.platform == 'win32':\n", + " try:\n", + " return co(['powershell', '-Command', 'Get-Alias'], text=True).strip()\n", + " except:\n", + " return \"Windows aliases not available\"\n", + " return co([shell, '-ic', 'alias'], text=True).strip()\n", + " except:\n", + " return \"Aliases not available\"\n" ] }, { @@ -289,13 +297,22 @@ "source": [ "#| export\n", "def _sys_info():\n", - " sys = co(['uname', '-a'], text=True).strip()\n", - " ssys = f'{sys}'\n", - " shell = co('echo $SHELL', shell=True, text=True).strip()\n", - " sshell = f'{shell}'\n", - " aliases = _aliases(shell)\n", - " saliases = f'\\n{aliases}\\n'\n", - " return f'\\n{ssys}\\n{sshell}\\n{saliases}\\n'" + " try:\n", + " if sys.platform == 'win32':\n", + " sys_info = co(['systeminfo'], text=True).strip()\n", + " ssys = f'Windows: {sys_info}'\n", + " shell = os.environ.get('SHELL') or os.environ.get('ComSpec', 'cmd.exe')\n", + " else:\n", + " sys_info = co(['uname', '-a'], text=True).strip()\n", + " ssys = f'{sys_info}'\n", + " shell = co('echo $SHELL', shell=True, text=True).strip()\n", + " \n", + " sshell = f'{shell}'\n", + " aliases = _aliases(shell)\n", + " saliases = f'\\n{aliases}\\n'\n", + " return f'\\n{ssys}\\n{sshell}\\n{saliases}\\n'\n", + " except Exception as e:\n", + " return f'\\nSystem info not available: {str(e)}\\n'" ] }, { @@ -316,6 +333,22 @@ "## Tmux" ] }, + { + "cell_type": "code", + "execution_count": null, + "id": "f6e43e12", + "metadata": {}, + "outputs": [], + "source": [ + "#| export\n", + "def is_tmux_available():\n", + " try:\n", + " co(['tmux', 'has-session'], stderr=subprocess.DEVNULL)\n", + " return True\n", + " except (subprocess.CalledProcessError, FileNotFoundError):\n", + " return False" + ] + }, { "cell_type": "code", "execution_count": null, @@ -326,9 +359,13 @@ "#| export\n", "def get_pane(n, pid=None):\n", " \"Get output from a tmux pane\"\n", - " cmd = ['tmux', 'capture-pane', '-p', '-S', f'-{n}']\n", - " if pid: cmd += ['-t', pid]\n", - " return co(cmd, text=True)" + " if not is_tmux_available(): return None\n", + " try:\n", + " cmd = ['tmux', 'capture-pane', '-p', '-S', f'-{n}']\n", + " if pid: cmd += ['-t', pid]\n", + " return co(cmd, text=True)\n", + " except subprocess.CalledProcessError:\n", + " return None" ] }, { @@ -351,9 +388,13 @@ "source": [ "#| export\n", "def get_panes(n):\n", - " cid = co(['tmux', 'display-message', '-p', '#{pane_id}'], text=True).strip()\n", - " pids = [p for p in co(['tmux', 'list-panes', '-F', '#{pane_id}'], text=True).splitlines()] \n", - " return '\\n'.join(f\"{get_pane(n, p)}\" for p in pids) " + " if not is_tmux_available(): return None\n", + " try:\n", + " cid = co(['tmux', 'display-message', '-p', '#{pane_id}'], text=True).strip()\n", + " pids = [p for p in co(['tmux', 'list-panes', '-F', '#{pane_id}'], text=True).splitlines()] \n", + " return '\\n'.join(f\"{get_pane(n, p)}\" for p in pids)\n", + " except subprocess.CalledProcessError:\n", + " return None" ] }, { @@ -397,8 +438,12 @@ "source": [ "#| export\n", "def tmux_history_lim():\n", - " lim = co(['tmux', 'display-message', '-p', '#{history-limit}'], text=True).strip()\n", - " return int(lim) if lim.isdigit() else 3000\n" + " if not is_tmux_available(): return 2000\n", + " try:\n", + " lim = co(['tmux', 'display-message', '-p', '#{history-limit}'], text=True).strip()\n", + " return int(lim) if lim.isdigit() else 3000\n", + " except (subprocess.CalledProcessError, ValueError):\n", + " return 2000" ] }, { @@ -422,6 +467,23 @@ "tmux_history_lim()" ] }, + { + "cell_type": "code", + "execution_count": null, + "id": "a0648292", + "metadata": {}, + "outputs": [], + "source": [ + "#| export\n", + "def get_powershell_history(n):\n", + " \"Get history from PowerShell\"\n", + " try:\n", + " cmd = ['powershell', '-Command', f'Get-History -Count {n} | Format-List CommandLine,Status']\n", + " return co(cmd, text=True).strip()\n", + " except (subprocess.CalledProcessError, FileNotFoundError):\n", + " return None" + ] + }, { "cell_type": "code", "execution_count": null, @@ -431,11 +493,16 @@ "source": [ "#| export\n", "def get_history(n, pid='current'):\n", - " try:\n", - " if pid=='current': return get_pane(n)\n", - " if pid=='all': return get_panes(n)\n", - " return get_pane(n, pid)\n", - " except subprocess.CalledProcessError: return None" + " if is_tmux_available():\n", + " try:\n", + " if pid=='current': return get_pane(n)\n", + " if pid=='all': return get_panes(n)\n", + " return get_pane(n, pid)\n", + " except subprocess.CalledProcessError:\n", + " return None\n", + " elif sys.platform == 'win32':\n", + " return get_powershell_history(n)\n", + " return None" ] }, { @@ -813,7 +880,7 @@ "\n", " if mode not in ['default', 'command', 'agent', 'sassy']:\n", " raise Exception(f\"{mode} is not valid. Must be one of the following: ['default', 'command', 'agent', 'sassy']\")\n", - " if mode == 'command' and os.environ.get('TMUX') is None:\n", + " if mode == 'command' and not is_tmux_available():\n", " raise Exception('Must be in a tmux session to use command mode.')\n", "\n", " if verbosity > 0: print(f\"{datetime.now()} | Starting ShellSage request with options {opts}\")\n", @@ -823,11 +890,11 @@ " query = ' '.join(query)\n", " ctxt = '' if skip_system else _sys_info()\n", "\n", - " # Get tmux history if in a tmux session\n", - " if os.environ.get('TMUX'):\n", - " if verbosity > 0: print(f\"{datetime.now()} | Adding TMUX history to prompt\")\n", + " # Get history from tmux or PowerShell\n", + " if is_tmux_available() or sys.platform == 'win32':\n", + " if verbosity > 0: print(f\"{datetime.now()} | Adding terminal history to prompt\")\n", " if opts.history_lines is None or opts.history_lines < 0:\n", - " opts.history_lines = tmux_history_lim()\n", + " opts.history_lines = tmux_history_lim() if is_tmux_available() else 50\n", " history = get_history(opts.history_lines, pid)\n", " if history: ctxt += f'\\n{history}\\n'\n", "\n", diff --git a/shell_sage/_modidx.py b/shell_sage/_modidx.py index fa0b488..0eb05d7 100644 --- a/shell_sage/_modidx.py +++ b/shell_sage/_modidx.py @@ -15,8 +15,10 @@ 'shell_sage.core.get_opts': ('core.html#get_opts', 'shell_sage/core.py'), 'shell_sage.core.get_pane': ('core.html#get_pane', 'shell_sage/core.py'), 'shell_sage.core.get_panes': ('core.html#get_panes', 'shell_sage/core.py'), + 'shell_sage.core.get_powershell_history': ('core.html#get_powershell_history', 'shell_sage/core.py'), 'shell_sage.core.get_res': ('core.html#get_res', 'shell_sage/core.py'), 'shell_sage.core.get_sage': ('core.html#get_sage', 'shell_sage/core.py'), + 'shell_sage.core.is_tmux_available': ('core.html#is_tmux_available', 'shell_sage/core.py'), 'shell_sage.core.main': ('core.html#main', 'shell_sage/core.py'), 'shell_sage.core.mk_db': ('core.html#mk_db', 'shell_sage/core.py'), 'shell_sage.core.tmux_history_lim': ('core.html#tmux_history_lim', 'shell_sage/core.py'), diff --git a/shell_sage/core.py b/shell_sage/core.py index 2aa70bf..3b767ff 100644 --- a/shell_sage/core.py +++ b/shell_sage/core.py @@ -1,9 +1,9 @@ # AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_core.ipynb. # %% auto 0 -__all__ = ['print', 'sp', 'csp', 'asp', 'ssp', 'default_cfg', 'chats', 'clis', 'sps', 'conts', 'p', 'log_path', 'get_pane', - 'get_panes', 'tmux_history_lim', 'get_history', 'get_opts', 'get_sage', 'trace', 'get_res', 'Log', 'mk_db', - 'main'] +__all__ = ['print', 'sp', 'csp', 'asp', 'ssp', 'default_cfg', 'chats', 'clis', 'sps', 'conts', 'p', 'log_path', + 'is_tmux_available', 'get_pane', 'get_panes', 'tmux_history_lim', 'get_powershell_history', 'get_history', + 'get_opts', 'get_sage', 'trace', 'get_res', 'Log', 'mk_db', 'main'] # %% ../nbs/00_core.ipynb 3 from anthropic.types import ToolUseBlock @@ -174,46 +174,97 @@ # %% ../nbs/00_core.ipynb 11 def _aliases(shell): - return co([shell, '-ic', 'alias'], text=True).strip() + try: + if sys.platform == 'win32': + try: + return co(['powershell', '-Command', 'Get-Alias'], text=True).strip() + except: + return "Windows aliases not available" + return co([shell, '-ic', 'alias'], text=True).strip() + except: + return "Aliases not available" + # %% ../nbs/00_core.ipynb 13 def _sys_info(): - sys = co(['uname', '-a'], text=True).strip() - ssys = f'{sys}' - shell = co('echo $SHELL', shell=True, text=True).strip() - sshell = f'{shell}' - aliases = _aliases(shell) - saliases = f'\n{aliases}\n' - return f'\n{ssys}\n{sshell}\n{saliases}\n' + try: + if sys.platform == 'win32': + sys_info = co(['systeminfo'], text=True).strip() + ssys = f'Windows: {sys_info}' + shell = os.environ.get('SHELL') or os.environ.get('ComSpec', 'cmd.exe') + else: + sys_info = co(['uname', '-a'], text=True).strip() + ssys = f'{sys_info}' + shell = co('echo $SHELL', shell=True, text=True).strip() + + sshell = f'{shell}' + aliases = _aliases(shell) + saliases = f'\n{aliases}\n' + return f'\n{ssys}\n{sshell}\n{saliases}\n' + except Exception as e: + return f'\nSystem info not available: {str(e)}\n' # %% ../nbs/00_core.ipynb 16 +def is_tmux_available(): + try: + co(['tmux', 'has-session'], stderr=subprocess.DEVNULL) + return True + except (subprocess.CalledProcessError, FileNotFoundError): + return False + +# %% ../nbs/00_core.ipynb 17 def get_pane(n, pid=None): "Get output from a tmux pane" - cmd = ['tmux', 'capture-pane', '-p', '-S', f'-{n}'] - if pid: cmd += ['-t', pid] - return co(cmd, text=True) + if not is_tmux_available(): return None + try: + cmd = ['tmux', 'capture-pane', '-p', '-S', f'-{n}'] + if pid: cmd += ['-t', pid] + return co(cmd, text=True) + except subprocess.CalledProcessError: + return None -# %% ../nbs/00_core.ipynb 18 +# %% ../nbs/00_core.ipynb 19 def get_panes(n): - cid = co(['tmux', 'display-message', '-p', '#{pane_id}'], text=True).strip() - pids = [p for p in co(['tmux', 'list-panes', '-F', '#{pane_id}'], text=True).splitlines()] - return '\n'.join(f"{get_pane(n, p)}" for p in pids) + if not is_tmux_available(): return None + try: + cid = co(['tmux', 'display-message', '-p', '#{pane_id}'], text=True).strip() + pids = [p for p in co(['tmux', 'list-panes', '-F', '#{pane_id}'], text=True).splitlines()] + return '\n'.join(f"{get_pane(n, p)}" for p in pids) + except subprocess.CalledProcessError: + return None -# %% ../nbs/00_core.ipynb 21 +# %% ../nbs/00_core.ipynb 22 def tmux_history_lim(): - lim = co(['tmux', 'display-message', '-p', '#{history-limit}'], text=True).strip() - return int(lim) if lim.isdigit() else 3000 - - -# %% ../nbs/00_core.ipynb 23 -def get_history(n, pid='current'): + if not is_tmux_available(): return 2000 + try: + lim = co(['tmux', 'display-message', '-p', '#{history-limit}'], text=True).strip() + return int(lim) if lim.isdigit() else 3000 + except (subprocess.CalledProcessError, ValueError): + return 2000 + +# %% ../nbs/00_core.ipynb 24 +def get_powershell_history(n): + "Get history from PowerShell" try: - if pid=='current': return get_pane(n) - if pid=='all': return get_panes(n) - return get_pane(n, pid) - except subprocess.CalledProcessError: return None + cmd = ['powershell', '-Command', f'Get-History -Count {n} | Format-List CommandLine,Status'] + return co(cmd, text=True).strip() + except (subprocess.CalledProcessError, FileNotFoundError): + return None # %% ../nbs/00_core.ipynb 25 +def get_history(n, pid='current'): + if is_tmux_available(): + try: + if pid=='current': return get_pane(n) + if pid=='all': return get_panes(n) + return get_pane(n, pid) + except subprocess.CalledProcessError: + return None + elif sys.platform == 'win32': + return get_powershell_history(n) + return None + +# %% ../nbs/00_core.ipynb 27 default_cfg = asdict(ShellSageConfig()) def get_opts(**opts): cfg = get_cfg() @@ -221,7 +272,7 @@ def get_opts(**opts): if v is None: opts[k] = cfg.get(k, default_cfg.get(k)) return AttrDict(opts) -# %% ../nbs/00_core.ipynb 27 +# %% ../nbs/00_core.ipynb 29 chats = {'anthropic': cla.Chat, 'openai': cos.Chat} clis = {'anthropic': cla.Client, 'openai': cos.Client} sps = {'default': sp, 'command': csp, 'sassy': ssp, 'agent': asp} @@ -237,7 +288,7 @@ def get_sage(provider, model, base_url=None, api_key=None, mode='default'): else: cli = clis[provider](model) return partial(cli, sp=sps[mode]) -# %% ../nbs/00_core.ipynb 31 +# %% ../nbs/00_core.ipynb 33 def trace(msgs): for m in msgs: if isinstance(m.content, str): continue @@ -248,7 +299,7 @@ def trace(msgs): tool_use = cla.find_block(m, ToolUseBlock) if tool_use: print(f'Tool use: {tool_use.name}\nTool input: {tool_use.input}') -# %% ../nbs/00_core.ipynb 33 +# %% ../nbs/00_core.ipynb 35 conts = {'anthropic': cla.contents, 'openai': cos.contents} p = r'```(?:bash\n|\n)?([^`]+)```' def get_res(sage, q, provider, mode='default', verbosity=0): @@ -259,7 +310,7 @@ def get_res(sage, q, provider, mode='default', verbosity=0): return conts[provider](sage.toolloop(q, trace_func=trace if verbosity else None)) else: return conts[provider](sage(q)) -# %% ../nbs/00_core.ipynb 38 +# %% ../nbs/00_core.ipynb 40 class Log: id:int; timestamp:str; query:str; response:str; model:str; mode:str log_path = Path("~/.shell_sage/logs/").expanduser() @@ -269,7 +320,7 @@ def mk_db(): db.logs = db.create(Log) return db -# %% ../nbs/00_core.ipynb 41 +# %% ../nbs/00_core.ipynb 43 @call_parse def main( query: Param('The query to send to the LLM', str, nargs='+'), @@ -293,7 +344,7 @@ def main( if mode not in ['default', 'command', 'agent', 'sassy']: raise Exception(f"{mode} is not valid. Must be one of the following: ['default', 'command', 'agent', 'sassy']") - if mode == 'command' and os.environ.get('TMUX') is None: + if mode == 'command' and not is_tmux_available(): raise Exception('Must be in a tmux session to use command mode.') if verbosity > 0: print(f"{datetime.now()} | Starting ShellSage request with options {opts}") @@ -303,11 +354,11 @@ def main( query = ' '.join(query) ctxt = '' if skip_system else _sys_info() - # Get tmux history if in a tmux session - if os.environ.get('TMUX'): - if verbosity > 0: print(f"{datetime.now()} | Adding TMUX history to prompt") + # Get history from tmux or PowerShell + if is_tmux_available() or sys.platform == 'win32': + if verbosity > 0: print(f"{datetime.now()} | Adding terminal history to prompt") if opts.history_lines is None or opts.history_lines < 0: - opts.history_lines = tmux_history_lim() + opts.history_lines = tmux_history_lim() if is_tmux_available() else 50 history = get_history(opts.history_lines, pid) if history: ctxt += f'\n{history}\n'