Add SQLite database support for conversation history

- Introduce SQLite database to store conversation history
- Implement functions to save and load conversations from the database
- Replace previous file-based history storage with database-backed solution
- Add new command to list and resume previous conversations
- Update chat logic to generate topics for saved conversations
- Add database initialization and management functions
- Modify command line interface to support resuming conversations
This commit is contained in:
Hayden Johnson 2025-05-17 10:18:50 -07:00
parent 892fdc317e
commit 780b755837

View file

@ -3,8 +3,8 @@
import argparse import argparse
import os import os
import sys import sys
import datetime
import sqlite3 import sqlite3
import warnings
import json import json
from ollama import Client from ollama import Client
@ -19,6 +19,10 @@ from prompt_toolkit.key_binding import KeyBindings
from prompt_toolkit import PromptSession from prompt_toolkit import PromptSession
default_assistant = "qwen3:14b"
default_topic_llm = "qwen3:1.7b"
warnings.filterwarnings("ignore")
class AIAssistant: class AIAssistant:
def __init__(self, server="http://localhost:11434", model="qwen3:14b"): def __init__(self, server="http://localhost:11434", model="qwen3:14b"):
@ -76,14 +80,32 @@ class AIAssistant:
else: else:
self.history = [self.system_prompt()] self.history = [self.system_prompt()]
def _stringify_history(self):
s = ""
for item in self.history:
s += item["role"] + ":\n" + item["content"]
return s
def _strip_thinking_tags(self, text: str):
start_idx = text.find("<think>")
if (start_idx < 0):
return text
end_idx = text.find("</think>") + len("</think>")
stripped_text = text[end_idx:]
stripped_text = stripped_text.strip()
return stripped_text
def save_history(self): def save_history(self):
"""Save the current conversation to the database with a generated topic.""" """Save the current conversation to the database with a generated topic."""
# Only save if this is the first user message # Only save if this is the first user message
if len(self.history) == 3: if len(self.history) == 3:
# Generate a topic using the AI # Generate a topic using the AI
system_prompt = self.system_prompt() system_prompt = self.system_prompt()
user_prompt = "Generate a concise, descriptive topic for this conversation based on the following content:\n" user_prompt = "/no_think Generate a concise, 5 word descriptive topic for this conversation based on the following content. Do not use markdown, just plaintext. KEEP IT TO 5 WORDS OR LESS.:\n\n"
topic = self.client.chat(model=self.model, messages=[system_prompt, {"role": "user", "content": user_prompt}], stream=False)['message']['content'].strip() user_prompt += self._stringify_history()
topic = self.client.chat(model=default_topic_llm, messages=[system_prompt, {"role": "user", "content": user_prompt}], stream=False)['message']['content'].strip()
topic = self._strip_thinking_tags(topic)
self._save_to_db(topic) self._save_to_db(topic)
else: else:
# For subsequent messages, we can update the topic in the future # For subsequent messages, we can update the topic in the future
@ -93,6 +115,7 @@ class AIAssistant:
"""Load a conversation from the database by ID. If no ID, start a new one.""" """Load a conversation from the database by ID. If no ID, start a new one."""
if conversation_id: if conversation_id:
self._load_from_db(conversation_id) self._load_from_db(conversation_id)
print(self._stringify_history())
else: else:
self.history = [self.system_prompt()] self.history = [self.system_prompt()]
@ -223,15 +246,25 @@ class InputHandler:
except: except:
return return
def follow_up(self, args, query):
if type(args.follow_up) is str:
second_input = args.follow_up
else:
second_input = self.improved_input()
query += f'\n\nUser Question:\n{second_input}'
return query
def handle_piped_input(self, args): def handle_piped_input(self, args):
all_input = sys.stdin.read() all_input = sys.stdin.read()
query = f'Use the following context to answer the question. There will be no follow up questions from the user so make sure your answer is complete:\n{all_input}\n' query = f'Use the following context to answer the question. There will be no follow up questions from the user so make sure your answer is complete:\n{all_input}\n'
if args.copy: if args.copy:
query += 'Answer the question using a codeblock for any code or shell scripts\n' query += 'Answer the question using a codeblock for any code or shell scripts\n'
if args.follow_up: if args.follow_up:
second_input = self.improved_input() query = self.follow_up(args, query)
query += f'\n{second_input}'
result = self.assistant.chat(query, stream=False) result = self.assistant.chat(query, stream=False)
result = self.assistant._strip_thinking_tags(result)
print(result)
blocks = self.extract_code_block(result) blocks = self.extract_code_block(result)
if args.copy and len(blocks): if args.copy and len(blocks):
self.copy_string_to_clipboard(blocks[0]) self.copy_string_to_clipboard(blocks[0])
@ -256,6 +289,14 @@ Description:\n
self.assistant.history = [self.assistant.system_prompt()] self.assistant.history = [self.assistant.system_prompt()]
self.arg_shell(args) self.arg_shell(args)
exit() exit()
if args.follow_up:
query = "Provide a complete answer to the user's question, there will be no follow up questions from the user.\n\n"
query = self.follow_up(args, query)
result = self.assistant.chat(query, stream=False)
result = self.assistant._strip_thinking_tags(result)
print(result)
exit()
print("\033[91massistant\033[0m: Type your message (press Ctrl+D to send):") print("\033[91massistant\033[0m: Type your message (press Ctrl+D to send):")
while True: while True:
@ -313,7 +354,8 @@ class CommandParser:
'/clear': self.handle_clear, '/clear': self.handle_clear,
'/clipboard': None, '/clipboard': None,
'/exit': self.handle_exit, '/exit': self.handle_exit,
'/copy': self.handle_copy '/copy': self.handle_copy,
'/list': self.handle_list
} }
def parse_commands(self, text): def parse_commands(self, text):
@ -368,6 +410,36 @@ class CommandParser:
self.input_handler.copy_string_to_clipboard(block) self.input_handler.copy_string_to_clipboard(block)
return True return True
def handle_list(self):
assistant = self.assistant
db_path = assistant.db_path
# Connect to the database and fetch saved conversations
conn = sqlite3.connect(db_path)
cursor = conn.cursor()
cursor.execute("SELECT id, topic FROM conversations")
rows = cursor.fetchall()
conn.close()
if not rows:
print("No saved conversations.")
return
# Display saved conversations
print("\nSaved conversations:")
for row in rows:
print(f"ID: {row[0]}, Topic: {row[1]}")
# Prompt user to select an ID
choice = input("\nEnter conversation ID to load (or 'q' to quit): ").strip()
if choice.lower() == 'q':
return
try:
conv_id = int(choice)
assistant.load_history(conversation_id=conv_id)
print(f"\nLoaded conversation with ID {conv_id}.")
except ValueError:
print("Invalid ID. Please enter a number.")
except Exception as e:
print(f"Error loading conversation: {e}")
def handle_exit(self): def handle_exit(self):
sys.exit(0) sys.exit(0)