Merge pull request #76 from AnswerDotAI/live

Fix duplicated output during tool use by pausing Live display

Author: ncoop57

nbs/00_core.ipynb

@@ -34,8 +34,8 @@
    "outputs": [],
    "source": [
     "#| export\n",
+    "from contextlib import contextmanager\n",
     "from datetime import datetime\n",
-    "from itertools import accumulate\n",
     "from fastcore.script import *\n",
     "from fastcore.tools import *\n",
     "from fastcore.utils import *\n",
@@ -89,7 +89,10 @@
    "source": [
     "#| export\n",
     "console = Console()\n",
-    "print = console.print"
+    "print = console.print\n",
+    "_live = None\n",
+    "_res = \"\"\n",
+    "_md = None"
    ]
   },
   {
@@ -508,30 +511,50 @@
     "opts = get_opts(model=None, log=None, api_base=None, api_key=''); opts"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "ae529a86",
+   "metadata": {},
+   "source": [
+    "Rich's `Live` display and `input()` can't coexist — `Live` manages the terminal cursor, and `input()` needs it too. When `Live.stop()` is called, it prints its current renderable as static text, which causes duplication since `get_res` accumulates the full response. The fix: clear Live's renderable before stopping, flush the accumulated response as static markdown, and reset the buffer. This way streaming resumes fresh after the tool interaction with no overlap."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "5b3ae69b",
+   "id": "7eba8e55",
    "metadata": {},
    "outputs": [],
    "source": [
     "#| export\n",
     "_always_allow = set()  # Session-level tracking of auto-approved tools\n",
     "\n",
+    "@contextmanager\n",
+    "def _pause_live():\n",
+    "    global _res\n",
+    "    _live.update('', refresh=True)\n",
+    "    if _res: print(_md(_res))\n",
+    "    _res = ''\n",
+    "    _live.stop()\n",
+    "    try: yield\n",
+    "    finally: _live.start()\n",
+    "\n",
     "def with_permission(action_desc):\n",
     "    def decorator(func):\n",
     "        @wraps(func)\n",
     "        def wrapper(*args, **kwargs):\n",
+    "            global _res\n",
     "            if IN_NOTEBOOK or func.__name__ in _always_allow: return func(*args, **kwargs)\n",
     "            limit = 50\n",
     "            details_dict = {\n",
     "                \"args\": [str(arg)[:limit] + (\"...\" if len(str(arg)) > limit else \"\") for arg in args],\n",
     "                \"kwargs\": {k: str(v)[:limit] + (\"...\" if len(str(v)) > limit else \"\") for k, v in kwargs.items()}\n",
     "            }\n",
-    "            print(f\"About to {action_desc} with the following arguments:\")\n",
-    "            print(details_dict if args else kwargs)\n",
-    "            res = input(\"Execute this? (y/n/a=always/suggestion): \").lower().strip()\n",
-    "\n",
+    "            with _pause_live():\n",
+    "                print(f\"About to {action_desc} with the following arguments:\", str(details_dict) if args else str(kwargs))\n",
+    "                res = input(\"Execute this? (y/n/a=always/suggestion): \").lower().strip()\n",
+    "                print()\n",
+    "            \n",
     "            if res == 'a': _always_allow.add(func.__name__)\n",
     "            if res in ('y','a'): return func(*args, **kwargs)\n",
     "            elif res == 'n': return \"[Command cancelled by user]\"\n",
@@ -666,10 +689,14 @@
    "source": [
     "#| export\n",
     "def get_res(sage, q, opts):\n",
-    "    from litellm.types.utils import ModelResponseStream # lazy load\n",
-    "    # need to use stream=True to get search citations\n",
-    "    gen = sage(q, max_steps=10, stream=True, api_base=opts.api_base, api_key=opts.api_key, think=opts.think) \n",
-    "    yield from accumulate(o.choices[0].delta.content or \"\" for o in gen if isinstance(o, ModelResponseStream))"
+    "    global _res\n",
+    "    from litellm.types.utils import ModelResponseStream\n",
+    "    _res = \"\"\n",
+    "    gen = sage(q, max_steps=10, stream=True, api_base=opts.api_base, api_key=opts.api_key, think=opts.think)\n",
+    "    for o in gen:\n",
+    "        if isinstance(o, ModelResponseStream):\n",
+    "            _res += o.choices[0].delta.content or \"\"\n",
+    "            yield _res"
    ]
   },
   {
@@ -864,10 +891,12 @@
     "    res=\"\"\n",
     "    try:\n",
     "        with Live(Spinner(\"dots\", text=\"Connecting...\"), auto_refresh=False) as live:\n",
+    "            global _live, _md\n",
+    "            _live = live\n",
     "            if mode not in ['default', 'sassy']:\n",
     "                raise Exception(f\"{mode} is not valid. Must be one of the following: ['default', 'sassy']\")\n",
     "            \n",
-    "            md = partial(Markdown, code_theme=opts.code_theme, inline_code_lexer=opts.code_lexer,\n",
+    "            _md = partial(Markdown, code_theme=opts.code_theme, inline_code_lexer=opts.code_lexer,\n",
     "                         inline_code_theme=opts.code_theme)\n",
     "            query = ' '.join(query)\n",
     "            ctxt = '' if skip_system else _sys_info()\n",
@@ -889,7 +918,7 @@
     "            query = f'{ctxt}\\n<query>\\n{query}\\n</query>'\n",
     "\n",
     "            sage = get_sage(opts.model, mode, search=opts.search, use_safecmd=opts.safecmd)\n",
-    "            for res in get_res(sage, query, opts): live.update(md(res), refresh=True)\n",
+    "            for res in get_res(sage, query, opts): live.update(_md(res), refresh=True)\n",
     "            \n",
     "        # Handle logging if the log flag is set\n",
     "        if opts.log:\n",

shell_sage/_modidx.py

@@ -13,6 +13,7 @@
                                                                                  'shell_sage/core.py'),
                                  'shell_sage.core.Log': ('core.html#log', 'shell_sage/core.py'),
                                  'shell_sage.core._aliases': ('core.html#_aliases', 'shell_sage/core.py'),
+                                 'shell_sage.core._pause_live': ('core.html#_pause_live', 'shell_sage/core.py'),
                                  'shell_sage.core._sys_info': ('core.html#_sys_info', 'shell_sage/core.py'),
                                  'shell_sage.core.extract': ('core.html#extract', 'shell_sage/core.py'),
                                  'shell_sage.core.extract_cf': ('core.html#extract_cf', 'shell_sage/core.py'),

shell_sage/core.py

@@ -6,8 +6,8 @@
            'main', 'extract_cf', 'extract']
 
 # %% ../nbs/00_core.ipynb #d7c5634a
+from contextlib import contextmanager
 from datetime import datetime
-from itertools import accumulate
 from fastcore.script import *
 from fastcore.tools import *
 from fastcore.utils import *
@@ -34,6 +34,9 @@ def __rich_console__(self:CodeBlock, console, options):
 # %% ../nbs/00_core.ipynb #9d52ca34
 console = Console()
 print = console.print
+_live = None
+_res = ""
+_md = None
 
 # %% ../nbs/00_core.ipynb #977bd215
 def Chat(*arg, **kw):
@@ -179,23 +182,35 @@ def get_opts(**opts):
         if v is None: opts[k] = cfg.get(k, default_cfg.get(k))
     return AttrDict(opts)
 
-# %% ../nbs/00_core.ipynb #5b3ae69b
+# %% ../nbs/00_core.ipynb #7eba8e55
 _always_allow = set()  # Session-level tracking of auto-approved tools
 
+@contextmanager
+def _pause_live():
+    global _res
+    _live.update('', refresh=True)
+    if _res: print(_md(_res))
+    _res = ''
+    _live.stop()
+    try: yield
+    finally: _live.start()
+
 def with_permission(action_desc):
     def decorator(func):
         @wraps(func)
         def wrapper(*args, **kwargs):
+            global _res
             if IN_NOTEBOOK or func.__name__ in _always_allow: return func(*args, **kwargs)
             limit = 50
             details_dict = {
                 "args": [str(arg)[:limit] + ("..." if len(str(arg)) > limit else "") for arg in args],
                 "kwargs": {k: str(v)[:limit] + ("..." if len(str(v)) > limit else "") for k, v in kwargs.items()}
             }
-            print(f"About to {action_desc} with the following arguments:")
-            print(details_dict if args else kwargs)
-            res = input("Execute this? (y/n/a=always/suggestion): ").lower().strip()
-
+            with _pause_live():
+                print(f"About to {action_desc} with the following arguments:", str(details_dict) if args else str(kwargs))
+                res = input("Execute this? (y/n/a=always/suggestion): ").lower().strip()
+                print()
+            
             if res == 'a': _always_allow.add(func.__name__)
             if res in ('y','a'): return func(*args, **kwargs)
             elif res == 'n': return "[Command cancelled by user]"
@@ -219,10 +234,14 @@ def get_sage(model, mode='default', search=False, use_safecmd=False):
 
 # %% ../nbs/00_core.ipynb #68be9484
 def get_res(sage, q, opts):
-    from litellm.types.utils import ModelResponseStream # lazy load
-    # need to use stream=True to get search citations
-    gen = sage(q, max_steps=10, stream=True, api_base=opts.api_base, api_key=opts.api_key, think=opts.think) 
-    yield from accumulate(o.choices[0].delta.content or "" for o in gen if isinstance(o, ModelResponseStream))
+    global _res
+    from litellm.types.utils import ModelResponseStream
+    _res = ""
+    gen = sage(q, max_steps=10, stream=True, api_base=opts.api_base, api_key=opts.api_key, think=opts.think)
+    for o in gen:
+        if isinstance(o, ModelResponseStream):
+            _res += o.choices[0].delta.content or ""
+            yield _res
 
 # %% ../nbs/00_core.ipynb #4e6e4d92
 class Log: id:int; timestamp:str; query:str; response:str; model:str; mode:str
@@ -260,10 +279,12 @@ def main(
     res=""
     try:
         with Live(Spinner("dots", text="Connecting..."), auto_refresh=False) as live:
+            global _live, _md
+            _live = live
             if mode not in ['default', 'sassy']:
                 raise Exception(f"{mode} is not valid. Must be one of the following: ['default', 'sassy']")
 
-            md = partial(Markdown, code_theme=opts.code_theme, inline_code_lexer=opts.code_lexer,
+            _md = partial(Markdown, code_theme=opts.code_theme, inline_code_lexer=opts.code_lexer,
                          inline_code_theme=opts.code_theme)
             query = ' '.join(query)
             ctxt = '' if skip_system else _sys_info()
@@ -285,7 +306,7 @@ def main(
             query = f'{ctxt}\n<query>\n{query}\n</query>'
 
             sage = get_sage(opts.model, mode, search=opts.search, use_safecmd=opts.safecmd)
-            for res in get_res(sage, query, opts): live.update(md(res), refresh=True)
+            for res in get_res(sage, query, opts): live.update(_md(res), refresh=True)
 
         # Handle logging if the log flag is set
         if opts.log: