From b1a7c63fd1d4a5854c7046b8d78cf706c07aac5d Mon Sep 17 00:00:00 2001 From: Lachlan Date: Tue, 1 Oct 2024 22:08:18 +1300 Subject: [PATCH 1/2] basic sqlite and datasette working --- .gitignore | 4 +- aider/coders/base_coder.py | 7 +- aider/commands.py | 60 ++++++++- aider/io.py | 267 ++++++++++++++++++++++++++++++++++++- 4 files changed, 332 insertions(+), 6 deletions(-) diff --git a/.gitignore b/.gitignore index 7767cef8f28..01e7c0450dc 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,5 @@ +*.db +env/ .DS_Store .vscode/ aider.code-workspace @@ -11,4 +13,4 @@ _site .jekyll-cache/ .jekyll-metadata aider/__version__.py -.venv/ \ No newline at end of file +.venv/ diff --git a/aider/coders/base_coder.py b/aider/coders/base_coder.py index 2b90dbb5a6c..20f7c088ec1 100755 --- a/aider/coders/base_coder.py +++ b/aider/coders/base_coder.py @@ -703,6 +703,7 @@ def get_images_message(self): def run_stream(self, user_message): self.io.user_input(user_message) + self.io.append_chat_history(user_message, role="user") self.init_before_message() yield from self.send_message(user_message) @@ -1432,11 +1433,13 @@ def send(self, messages, model=None, functions=None): if self.partial_response_content: self.io.ai_output(self.partial_response_content) + self.io.append_chat_history(self.partial_response_content, role="assistant") elif self.partial_response_function_call: - # TODO: push this into subclasses args = self.parse_partial_args() if args: - self.io.ai_output(json.dumps(args, indent=4)) + content = json.dumps(args, indent=4) + self.io.ai_output(content) + self.io.append_chat_history(content, role="assistant") self.calculate_and_show_tokens_and_cost(messages, completion) diff --git a/aider/commands.py b/aider/commands.py index e177e77ca39..285e7c54a45 100644 --- a/aider/commands.py +++ b/aider/commands.py @@ -11,7 +11,7 @@ from prompt_toolkit.completion import Completion, PathCompleter from prompt_toolkit.document import Document -from aider import models, prompts, voice +from aider import models, prompts, utils, voice from aider.format_settings import format_settings from aider.help import Help, install_help_extra from aider.llm import litellm @@ -22,7 +22,6 @@ from .dump import dump # noqa: F401 - class SwitchCoder(Exception): def __init__(self, **kwargs): self.kwargs = kwargs @@ -1246,6 +1245,63 @@ def cmd_report(self, args): report_github_issue(issue_text, title=title, confirm=False) + def cmd_datasette(self, args): + "Launch or stop Datasette with the current chat history" + if not utils.check_pip_install_extra( + self.io, + "datasette", + "You need to install Datasette and its query assistant", + ["datasette", "datasette-query-assistant"], + ): + return + + if not self.io.use_sqlite: + self.io.tool_error( + "SQLite integration is not enabled. Please enable it to use Datasette." + ) + return + + db_path = self.io.chat_history_file.with_suffix(".db") + if not db_path.exists(): + self.io.tool_error("Chat history database not found.") + return + + datasette_process = getattr(self.coder, "datasette_process", None) + if datasette_process and datasette_process.poll() is None: + datasette_process.terminate() + datasette_process.wait() + self.coder.datasette_process = None + self.io.tool_output("Stopped Datasette.") + return + + try: + datasette_process = subprocess.Popen( + ["datasette", str(db_path)], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL + ) + self.coder.datasette_process = datasette_process + self.io.tool_output("Launched Datasette. Access it at http://localhost:8001/") + self.io.tool_output( + "For help, visit: https://docs.datasette.io/en/1.0a14/getting_started.html" + ) + except Exception as e: + self.io.tool_error(f"Failed to launch Datasette: {str(e)}") + + def cmd_search(self, args): + "Search the chat history using full-text search" + if not args.strip(): + self.io.tool_error("Please provide a search query.") + return + + results = self.io.search_chat_history(args.strip()) + if not results: + self.io.tool_output("No results found.") + return + + self.io.tool_output(f"Search results for '{args.strip()}':") + for timestamp, role, content in results: + self.io.tool_output(f"\n[{timestamp}] {role.upper()}:") + self.io.tool_output(content.strip()) + def expand_subdir(file_path): if file_path.is_file(): diff --git a/aider/io.py b/aider/io.py index 81829a83bdf..ffc1bf36cd4 100644 --- a/aider/io.py +++ b/aider/io.py @@ -1,5 +1,9 @@ import base64 import os +import re +import sqlite3 +import subprocess +import uuid from collections import defaultdict from dataclasses import dataclass from datetime import datetime @@ -25,6 +29,13 @@ from .dump import dump # noqa: F401 from .utils import is_image_file +try: + import datasette + + DATASETTE_AVAILABLE = True +except ImportError: + DATASETTE_AVAILABLE = False + @dataclass class ConfirmGroup: @@ -192,6 +203,7 @@ def __init__( dry_run=False, llm_history_file=None, editingmode=EditingMode.EMACS, + use_sqlite=True, ): self.never_prompts = set() self.editingmode = editingmode @@ -229,9 +241,50 @@ def __init__( self.encoding = encoding self.dry_run = dry_run + self.use_sqlite = use_sqlite + + self.conn = None + self.cursor = None current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S") self.append_chat_history(f"\n# aider chat started at {current_time}\n\n") + self.console = Console(force_terminal=False, no_color=True) # non-pretty + self.prompt_session = None + if self.pretty: + # Initialize PromptSession + session_kwargs = { + "input": self.input, + "output": self.output, + "lexer": PygmentsLexer(MarkdownLexer), + "editing_mode": self.editingmode, + "cursor": ModalCursorShapeConfig(), + } + if self.input_history_file is not None: + session_kwargs["history"] = FileHistory(self.input_history_file) + try: + self.prompt_session = PromptSession(**session_kwargs) + self.console = Console() # pretty console + except Exception as err: + self.console = Console(force_terminal=False, no_color=True) + self.tool_error(f"Can't initialize prompt toolkit: {err}") # non-pretty + else: + self.console = Console(force_terminal=False, no_color=True) # non-pretty + + self.init_sqlite() + self.start_conversation(os.getcwd(), self.get_git_commit_hash()) + + def get_git_commit_hash(self): + try: + result = subprocess.run( + ["git", "rev-parse", "HEAD"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + check=True, + ) + return result.stdout.strip() + except subprocess.CalledProcessError: + return None self.prompt_session = None if self.pretty: @@ -254,6 +307,131 @@ def __init__( else: self.console = Console(force_terminal=False, no_color=True) # non-pretty + def init_sqlite(self): + if not self.chat_history_file: + return + db_path = self.chat_history_file.with_suffix(".db") + self.conn = sqlite3.connect(db_path) + self.cursor = self.conn.cursor() + self.cursor.execute(""" + CREATE TABLE IF NOT EXISTS conversations ( + id TEXT PRIMARY KEY, + start_time TEXT, + end_time TEXT, + working_directory TEXT, + git_commit_hash TEXT + ) + """) + self.cursor.execute(""" + CREATE TABLE IF NOT EXISTS chat_history ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + conversation_id TEXT, + timestamp TEXT, + role TEXT, + content TEXT, + tokens_sent INTEGER, + tokens_received INTEGER, + cost REAL, + files_in_context TEXT, + FOREIGN KEY (conversation_id) REFERENCES conversations(id) + ) + """) + self.cursor.execute(""" + CREATE VIRTUAL TABLE IF NOT EXISTS chat_history_fts USING fts5( + timestamp, role, content, content='chat_history', content_rowid='id' + ) + """) + self.cursor.execute(""" + CREATE TRIGGER IF NOT EXISTS chat_history_ai AFTER INSERT ON chat_history BEGIN + INSERT INTO chat_history_fts(rowid, timestamp, role, content) + VALUES (new.id, new.timestamp, new.role, new.content); + END; + """) + self.cursor.execute(""" + CREATE TRIGGER IF NOT EXISTS chat_history_ad AFTER DELETE ON chat_history BEGIN + INSERT INTO chat_history_fts(chat_history_fts, rowid, timestamp, role, content) + VALUES('delete', old.id, old.timestamp, old.role, old.content); + END; + """) + self.cursor.execute(""" + CREATE TRIGGER IF NOT EXISTS chat_history_au AFTER UPDATE ON chat_history BEGIN + INSERT INTO chat_history_fts(chat_history_fts, rowid, timestamp, role, content) + VALUES('delete', old.id, old.timestamp, old.role, old.content); + INSERT INTO chat_history_fts(rowid, timestamp, role, content) + VALUES (new.id, new.timestamp, new.role, new.content); + END; + """) + self.conn.commit() + self.tool_output(f"SQLite database with FTS initialized at {db_path}") + self.current_conversation_id = None + try: + result = subprocess.run( + ["git", "rev-parse", "HEAD"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + check=True, + ) + return result.stdout.strip() + except subprocess.CalledProcessError: + return None + if not self.chat_history_file: + return + db_path = self.chat_history_file.with_suffix(".db") + self.conn = sqlite3.connect(db_path) + self.cursor = self.conn.cursor() + self.cursor.execute(""" + CREATE TABLE IF NOT EXISTS conversations ( + id TEXT PRIMARY KEY, + start_time TEXT, + end_time TEXT, + working_directory TEXT, + git_commit_hash TEXT + ) + """) + self.cursor.execute(""" + CREATE TABLE IF NOT EXISTS chat_history ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + conversation_id TEXT, + timestamp TEXT, + role TEXT, + content TEXT, + tokens_sent INTEGER, + tokens_received INTEGER, + cost REAL, + files_in_context TEXT, + FOREIGN KEY (conversation_id) REFERENCES conversations(id) + ) + """) + self.cursor.execute(""" + CREATE VIRTUAL TABLE IF NOT EXISTS chat_history_fts USING fts5( + timestamp, role, content, content='chat_history', content_rowid='id' + ) + """) + self.cursor.execute(""" + CREATE TRIGGER IF NOT EXISTS chat_history_ai AFTER INSERT ON chat_history BEGIN + INSERT INTO chat_history_fts(rowid, timestamp, role, content) + VALUES (new.id, new.timestamp, new.role, new.content); + END; + """) + self.cursor.execute(""" + CREATE TRIGGER IF NOT EXISTS chat_history_ad AFTER DELETE ON chat_history BEGIN + INSERT INTO chat_history_fts(chat_history_fts, rowid, timestamp, role, content) + VALUES('delete', old.id, old.timestamp, old.role, old.content); + END; + """) + self.cursor.execute(""" + CREATE TRIGGER IF NOT EXISTS chat_history_au AFTER UPDATE ON chat_history BEGIN + INSERT INTO chat_history_fts(chat_history_fts, rowid, timestamp, role, content) + VALUES('delete', old.id, old.timestamp, old.role, old.content); + INSERT INTO chat_history_fts(rowid, timestamp, role, content) + VALUES (new.id, new.timestamp, new.role, new.content); + END; + """) + self.conn.commit() + self.tool_output(f"SQLite database with FTS initialized at {db_path}") + self.current_conversation_id = None + def _get_style(self): style_dict = {} if not self.pretty: @@ -425,6 +603,7 @@ def _(event): print() self.user_input(inp) + self.append_chat_history(inp, role="user") return inp def add_to_input_history(self, inp): @@ -674,7 +853,40 @@ def assistant_output(self, message, pretty=None): def print(self, message=""): print(message) - def append_chat_history(self, text, linebreak=False, blockquote=False, strip=True): + def start_conversation(self, working_directory, git_commit_hash): + self.current_conversation_id = str(uuid.uuid4()) + start_time = datetime.now().isoformat() + self.cursor.execute( + ( + "INSERT INTO conversations (id, start_time, working_directory, git_commit_hash)" + " VALUES (?, ?, ?, ?)" + ), + (self.current_conversation_id, start_time, working_directory, git_commit_hash), + ) + self.conn.commit() + + def end_conversation(self): + if self.current_conversation_id: + end_time = datetime.now().isoformat() + self.cursor.execute( + "UPDATE conversations SET end_time = ? WHERE id = ?", + (end_time, self.current_conversation_id), + ) + self.conn.commit() + self.current_conversation_id = None + + def append_chat_history( + self, + text, + linebreak=False, + blockquote=False, + strip=True, + role=None, + tokens_sent=0, + tokens_received=0, + cost=0.0, + files_in_context=None, + ): if blockquote: if strip: text = text.strip() @@ -685,6 +897,7 @@ def append_chat_history(self, text, linebreak=False, blockquote=False, strip=Tru text = text + " \n" if not text.endswith("\n"): text += "\n" + if self.chat_history_file is not None: try: with self.chat_history_file.open("a", encoding=self.encoding, errors="ignore") as f: @@ -695,3 +908,55 @@ def append_chat_history(self, text, linebreak=False, blockquote=False, strip=Tru " Permission denied." ) self.chat_history_file = None # Disable further attempts to write + + if self.use_sqlite and role: + db_path = self.chat_history_file.with_suffix(".db") + try: + with sqlite3.connect(db_path) as conn: + cursor = conn.cursor() + timestamp = datetime.now().isoformat() + files_in_context_str = ",".join(files_in_context) if files_in_context else None + cursor.execute( + """ + INSERT INTO chat_history + (conversation_id, timestamp, role, content, tokens_sent, tokens_received, cost, files_in_context) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + self.current_conversation_id, + timestamp, + role, + text, + tokens_sent, + tokens_received, + cost, + files_in_context_str, + ), + ) + conn.commit() + except sqlite3.Error as e: + self.tool_error(f"SQLite error: {e}") + + def search_chat_history(self, query): + if not self.use_sqlite: + self.tool_error("SQLite integration is not enabled. Unable to search chat history.") + return [] + + db_path = self.chat_history_file.with_suffix(".db") + try: + with sqlite3.connect(db_path) as conn: + cursor = conn.cursor() + cursor.execute( + """ + SELECT timestamp, role, content + FROM chat_history_fts + WHERE chat_history_fts MATCH ? + ORDER BY rank + """, + (query,), + ) + results = cursor.fetchall() + return results + except sqlite3.Error as e: + self.tool_error(f"SQLite error during search: {e}") + return [] From b67bf9c6b66f50c4bc0bdb23284b32210fd4fe02 Mon Sep 17 00:00:00 2001 From: Lachlan Date: Tue, 1 Oct 2024 22:08:41 +1300 Subject: [PATCH 2/2] basic sqlite and datasette working --- aider/commands.py | 1 + aider/io.py | 9 --------- 2 files changed, 1 insertion(+), 9 deletions(-) diff --git a/aider/commands.py b/aider/commands.py index 285e7c54a45..17e32ff4d93 100644 --- a/aider/commands.py +++ b/aider/commands.py @@ -22,6 +22,7 @@ from .dump import dump # noqa: F401 + class SwitchCoder(Exception): def __init__(self, **kwargs): self.kwargs = kwargs diff --git a/aider/io.py b/aider/io.py index ffc1bf36cd4..07f37380693 100644 --- a/aider/io.py +++ b/aider/io.py @@ -1,6 +1,5 @@ import base64 import os -import re import sqlite3 import subprocess import uuid @@ -29,14 +28,6 @@ from .dump import dump # noqa: F401 from .utils import is_image_file -try: - import datasette - - DATASETTE_AVAILABLE = True -except ImportError: - DATASETTE_AVAILABLE = False - - @dataclass class ConfirmGroup: preference: str = None