Merge pull request #1 from AnswerDotAI/tmux-multi-pane

add option to include all panes from the session
AnswerDotAI · Dec 3, 2024 · fdb710a · fdb710a
2 parents d8a5155 + 1c38d2a
commit fdb710a
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 7 deletions.
diff --git a/nbs/00_core.ipynb b/nbs/00_core.ipynb
@@ -196,13 +196,23 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "bb809e08",
+   "id": "0d70591e",
    "metadata": {},
    "outputs": [],
    "source": [
     "#| exports\n",
-    "def get_history(n):\n",
-    "    try: return subprocess.check_output(['tmux', 'capture-pane', '-p', '-S', f'-{n}'], text=True)\n",
+    "def get_history(n, all_panes=False):\n",
+    "    try:\n",
+    "        if not all_panes: return subprocess.check_output(['tmux', 'capture-pane', '-p', '-S', f'-{n}'], text=True)\n",
+    "        current = subprocess.check_output(['tmux', 'display-message', '-p', '#{pane_id}'], text=True).strip()\n",
+    "        panes = [p for p in subprocess.check_output(['tmux', 'list-panes', '-F', '#{pane_id}'], text=True).splitlines() if p != current]        \n",
+    "        outputs = []\n",
+    "        for pane in panes:\n",
+    "            out = subprocess.check_output(['tmux', 'capture-pane', '-p', '-t', pane, '-S', f'-{n}'], text=True)\n",
+    "            outputs.append(f\"=== Pane {pane} ===\\n{out.strip()}\")            \n",
+    "        out = subprocess.check_output(['tmux', 'capture-pane', '-p', '-t', current, '-S', f'-{n}'], text=True)\n",
+    "        outputs.append(f\"=== Current pane ===\\n{out.strip()}\")\n",
+    "        return '\\n\\n'.join(outputs)\n",
     "    except subprocess.CalledProcessError: return None"
    ]
   },
@@ -249,6 +259,7 @@
     "    query: Param('The query to send to the LLM', str, nargs='+'),\n",
     "    action: bool = False, # Run ShellSage in action mode\n",
     "    NH: bool = False, # Don't include terminal history\n",
+    "    AP: bool = False, # Whether to include all panes in session\n",
     "    n: int = 200, # Number of history lines\n",
     "    code_theme: str = 'monokai', # The code theme to use when rendering ShellSage's responses\n",
     "    code_lexer: str = 'python', # The lexer to use for inline code markdown blocks\n",
@@ -258,7 +269,7 @@
     "    ctxt = ''\n",
     "    # Get tmux history if requested and available\n",
     "    if not NH:\n",
-    "        history = get_history(n)\n",
+    "        history = get_history(n,AP)\n",
     "        if history: ctxt += f'<terminal_history>\\n{history}\\n</terminal_history>'\n",
     "\n",
     "    # Read from stdin if available\n",

diff --git a/shell_sage/core.py b/shell_sage/core.py
@@ -107,8 +107,18 @@
 ssa = chat.toolloop
 
 # %% ../nbs/00_core.ipynb 11
-def get_history(n):
-    try: return subprocess.check_output(['tmux', 'capture-pane', '-p', '-S', f'-{n}'], text=True)
+def get_history(n, all_panes=False):
+    try:
+        if not all_panes: return subprocess.check_output(['tmux', 'capture-pane', '-p', '-S', f'-{n}'], text=True)
+        current = subprocess.check_output(['tmux', 'display-message', '-p', '#{pane_id}'], text=True).strip()
+        panes = [p for p in subprocess.check_output(['tmux', 'list-panes', '-F', '#{pane_id}'], text=True).splitlines() if p != current]        
+        outputs = []
+        for pane in panes:
+            out = subprocess.check_output(['tmux', 'capture-pane', '-p', '-t', pane, '-S', f'-{n}'], text=True)
+            outputs.append(f"=== Pane {pane} ===\n{out.strip()}")            
+        out = subprocess.check_output(['tmux', 'capture-pane', '-p', '-t', current, '-S', f'-{n}'], text=True)
+        outputs.append(f"=== Current pane ===\n{out.strip()}")
+        return '\n\n'.join(outputs)
     except subprocess.CalledProcessError: return None
 
 # %% ../nbs/00_core.ipynb 12
@@ -132,6 +142,7 @@ def main(
     query: Param('The query to send to the LLM', str, nargs='+'),
     action: bool = False, # Run ShellSage in action mode
     NH: bool = False, # Don't include terminal history
+    AP: bool = False, # Whether to include all panes in session
     n: int = 200, # Number of history lines
     code_theme: str = 'monokai', # The code theme to use when rendering ShellSage's responses
     code_lexer: str = 'python', # The lexer to use for inline code markdown blocks
@@ -141,7 +152,7 @@ def main(
     ctxt = ''
     # Get tmux history if requested and available
     if not NH:
-        history = get_history(n)
+        history = get_history(n,AP)
         if history: ctxt += f'<terminal_history>\n{history}\n</terminal_history>'
 
     # Read from stdin if available