11# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_core.ipynb.
22
33# %% auto 0
4- __all__ = ['print' , 'sp' , 'csp' , 'ssp' , 'default_cfg' , 'clis' , 'sps' , 'conts' , 'p' , 'get_pane ' , 'get_panes ' , 'tmux_history_lim ' ,
5- 'get_history' , 'get_opts' , 'get_sage' , 'get_res' , 'main' ]
4+ __all__ = ['print' , 'sp' , 'csp' , 'ssp' , 'default_cfg' , 'clis' , 'sps' , 'conts' , 'p' , 'log_path ' , 'get_pane ' , 'get_panes ' ,
5+ 'tmux_history_lim' , ' get_history''get_opts' , 'get_sage' , 'get_res' , 'Log' , 'mk_db ''main' ]
66
77# %% ../nbs/00_core.ipynb 3
88from datetime import datetime
1616from . import __version__
1717from .config import *
1818from subprocess import check_output as co
19+ from fastlite import database
1920
2021import os ,re ,subprocess ,sys
2122import claudette as cla , cosette as cos
@@ -200,26 +201,37 @@ def get_res(sage, q, provider, is_command=False):
200201 else : return conts [provider ](sage (q ))
201202
202203# %% ../nbs/00_core.ipynb 34
204+ class Log : id :int ; timestamp :str ; query :str ; response :str ; model :str ; mode :str
205+
206+ log_path = Path ("~/.shell_sage/logs/" ).expanduser ()
207+ def mk_db ():
208+ log_path .mkdir (parents = True , exist_ok = True )
209+ db = database (log_path / "logs.db" )
210+ db .logs = db .create (Log )
211+ return db
212+
213+ # %% ../nbs/00_core.ipynb 37
203214@call_parse
204215def main (
205216 query : Param ('The query to send to the LLM' , str , nargs = '+' ),
206217 v : Param ("Print version" , action = 'version' ) = '%(prog)s ' + __version__ ,
207- pid : str = 'current' , # `current`, `all` or tmux pane_id (e.g. %0) for context
208- skip_system : bool = False , # Whether to skip system information in the AI's context
209- history_lines : int = None , # Number of history lines. Defaults to tmux scrollback history length
210- s : bool = False , # Enable sassy mode
211- c : bool = False , # Enable command mode
212- provider : str = None , # The LLM Provider
213- model : str = None , # The LLM model that will be invoked on the LLM provider
218+ pid : str = 'current' , # `current`, `all` or tmux pane_id (e.g. %0) for context
219+ skip_system : bool = False , # Whether to skip system information in the AI's context
220+ history_lines : int = None , # Number of history lines. Defaults to tmux scrollback history length
221+ s : bool = False , # Enable sassy mode
222+ c : bool = False , # Enable command mode
223+ log : bool = False , # Enable logging
224+ provider : str = None , # The LLM Provider
225+ model : str = None , # The LLM model that will be invoked on the LLM provider
214226 base_url : str = None ,
215227 api_key : str = None ,
216- code_theme : str = None , # The code theme to use when rendering ShellSage's responses
217- code_lexer : str = None , # The lexer to use for inline code markdown blocks
218- verbosity : int = 0 # Level of verbosity (0 or 1)
228+ code_theme : str = None , # The code theme to use when rendering ShellSage's responses
229+ code_lexer : str = None , # The lexer to use for inline code markdown blocks
230+ verbosity : int = 0 # Level of verbosity (0 or 1)
219231):
220232 opts = get_opts (history_lines = history_lines , provider = provider , model = model ,
221233 base_url = base_url , api_key = api_key , code_theme = code_theme ,
222- code_lexer = code_lexer )
234+ code_lexer = code_lexer , log = log )
223235
224236 mode = 'default'
225237 if s : mode = 'sassy'
@@ -228,32 +240,42 @@ def main(
228240 raise Exception ('Must be in a tmux session to use command mode.' )
229241 mode = 'command'
230242
231- if verbosity > 0 :
232- print (f"{ datetime .now ()} { opts } )
233- md = partial (Markdown , code_theme = opts .code_theme , inline_code_lexer = opts .code_lexer , inline_code_theme = opts .code_theme )
243+ if verbosity > 0 : print (f"{ datetime .now ()} { opts } )
244+
245+ md = partial (Markdown , code_theme = opts .code_theme , inline_code_lexer = opts .code_lexer ,
246+ inline_code_theme = opts .code_theme )
234247 query = ' ' .join (query )
235248 ctxt = '' if skip_system else _sys_info ()
236249
237250 # Get tmux history if in a tmux session
238251 if os .environ .get ('TMUX' ):
239- if verbosity > 0 : print (f"{ datetime .now ()} )
252+ if verbosity > 0 : print (f"{ datetime .now ()} )
240253 if opts .history_lines is None or opts .history_lines < 0 :
241254 opts .history_lines = tmux_history_lim ()
242- history = get_history (opts .history_lines ,pid )
255+ history = get_history (opts .history_lines , pid )
243256 if history : ctxt += f'<terminal_history>\n { history } \n </terminal_history>'
244257
245258 # Read from stdin if available
246- if not sys .stdin .isatty ():
247- if verbosity > 0 : print (f"{ datetime .now ()} )
259+ if not sys .stdin .isatty ():
260+ if verbosity > 0 : print (f"{ datetime .now ()} )
248261 ctxt += f'\n <context>\n { sys .stdin .read ()}
249262
250- if verbosity > 0 : print (f"{ datetime .now ()} )
263+ if verbosity > 0 : print (f"{ datetime .now ()} )
264+
251265 query = f'{ ctxt } \n <query>\n { query } \n </query>'
252266 query = [mk_msg (query )] if opts .provider == 'openai' else query
253267
254- if verbosity > 0 : print (f"{ datetime .now ()} )
255- sage = get_sage (opts .provider , opts .model , opts .base_url , opts .api_key , mode )
256- res = get_res (sage , query , opts .provider , is_command = c )
268+ if verbosity > 0 :
269+ print (f"{ datetime .now ()} )
270+
271+ sage = get_sage (opts .provider , opts .model , opts .base_url , opts .api_key , mode )
272+ res = get_res (sage , query , opts .provider , is_command = c )
257273
274+ # Handle logging if the log flag is set
275+ if opts .log :
276+ db = mk_db ()
277+ db .logs .insert (Log (timestamp = datetime .now ().isoformat (), query = query ,
278+ response = res , model = opts .model , mode = mode ))
279+
258280 if c : co (['tmux' , 'send-keys' , res ], text = True )
259281 else : print (md (res ))
0 commit comments