Skip to content

Allow for session continuation #57

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Allow for session continuation #57

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
142 changes: 81 additions & 61 deletions codeaide/logic/chat_handler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,10 +18,11 @@
from codeaide.utils.cost_tracker import CostTracker
from codeaide.utils.file_handler import FileHandler
from codeaide.utils.terminal_manager import TerminalManager
from codeaide.utils.general_utils import generate_session_id
from codeaide.utils.logging_config import get_logger, setup_logger
from codeaide.utils.general_utils import get_project_root
from codeaide.utils.logging_config import get_logger
from PyQt5.QtCore import QObject, pyqtSignal
from codeaide.utils.environment_manager import EnvironmentManager
from codeaide.utils.session_manager import SessionManager


class ChatHandler(QObject):
Expand All @@ -31,31 +32,21 @@ class ChatHandler(QObject):
) # Signal to update chat with (role, message)
show_code_signal = pyqtSignal(str, str) # Signal to show code with (code, version)
traceback_occurred = pyqtSignal(str)
session_updated = pyqtSignal()

def __init__(self):
super().__init__()
"""
Initialize the ChatHandler class.
base_dir = get_project_root()
self.file_handler = FileHandler(base_dir=base_dir)
self.session_manager = SessionManager(base_dir, self.file_handler)

Args:
None
self.session_id = None
self.session_dir = None

Returns:
None
"""
self.session_id = generate_session_id()
self.cost_tracker = CostTracker()
self.file_handler = FileHandler(session_id=self.session_id)
self.session_dir = (
self.file_handler.session_dir
) # Store the specific session directory
self.logger = get_logger()
self.conversation_history = self.file_handler.load_chat_history()
self.environment_manager = EnvironmentManager(self.session_id)
self.terminal_manager = TerminalManager(
environment_manager=self.environment_manager,
traceback_callback=self.emit_traceback_signal,
)
self.environment_manager = None
self.terminal_manager = None
self.latest_version = "0.0"
self.api_client = None
self.api_key_set = False
Expand All @@ -66,17 +57,26 @@ def __init__(self):
self.max_tokens = AI_PROVIDERS[self.current_provider]["models"][
self.current_model
]["max_tokens"]
self.env_manager = EnvironmentManager(self.session_id)

self.api_key_valid, self.api_key_message = self.check_api_key()
self.logger.info(f"New session started with ID: {self.session_id}")
self.logger.info(f"Session directory: {self.session_dir}")
self.chat_window = None

def start_application(self):
from codeaide.ui.chat_window import (
ChatWindow,
) # Import here to avoid circular imports
from codeaide.ui.chat_window import ChatWindow

# Create initial session
self.session_id = self.session_manager.create_new_session()
self.session_dir = self.file_handler.session_dir

self.conversation_history = self.file_handler.load_chat_history()
self.environment_manager = EnvironmentManager(self.session_id)
self.terminal_manager = TerminalManager(
environment_manager=self.environment_manager,
traceback_callback=self.emit_traceback_signal,
)

self.logger.info(f"New session started with ID: {self.session_id}")
self.logger.info(f"Session directory: {self.session_dir}")

self.chat_window = ChatWindow(self)
self.connect_signals()
Expand Down Expand Up @@ -333,8 +333,13 @@ def process_ai_response(self, response):
code_version,
version_description,
requirements,
session_summary,
) = parsed_response

# Update the session summary
if session_summary:
self.update_session_summary(session_summary)

if code and self.compare_versions(code_version, self.latest_version) <= 0:
raise ValueError(
f"New version {code_version} is not higher than the latest version {self.latest_version}"
Expand Down Expand Up @@ -567,55 +572,70 @@ def get_latest_version(self):
def set_latest_version(self, version):
self.latest_version = version

def start_new_session(self, chat_window):
self.logger.info("Starting new session")

# Log the previous session path correctly
self.logger.info(f"Previous session path: {self.session_dir}")

# Generate new session ID
new_session_id = generate_session_id()

# Create new FileHandler with new session ID
new_file_handler = FileHandler(session_id=new_session_id)

# Copy existing log to new session and set up new logger
self.file_handler.copy_log_to_new_session(new_session_id)
setup_logger(new_file_handler.session_dir)

# Update instance variables
def start_new_session(self, chat_window, based_on=None, initial_session=False):
new_session_id = self.session_manager.create_new_session(based_on=based_on)
self.session_id = new_session_id
self.file_handler = new_file_handler
self.session_dir = new_file_handler.session_dir # Update the session directory
self.file_handler.set_session_id(new_session_id)
self.session_dir = self.file_handler.session_dir

# Clear conversation history
self.conversation_history = []

# Clear chat display in UI
chat_window.clear_chat_display()

# Close code pop-up if it exists
chat_window.close_code_popup()

# Add system message about previous session
system_message = f"A new session has been started. The previous chat will not be visible to the agent. Previous session data saved in: {self.session_dir}"
chat_window.add_to_chat("System", system_message)
chat_window.add_to_chat("AI", INITIAL_MESSAGE)
if not initial_session:
# Inform the user that a new session is starting if this is done after an existing session was underway
system_message = f"A new session has been started. The previous chat will not be visible to the agent. Previous session data saved in: {self.session_dir}"
chat_window.add_to_chat("System", system_message)
chat_window.add_to_chat("AI", INITIAL_MESSAGE)

self.logger.info(f"New session started with ID: {self.session_id}")
self.logger.info(f"New session directory: {self.session_dir}")

# New method to load a previous session
return new_session_id

def load_previous_session(self, session_id, chat_window):
if session_id is None:
self.logger.error("Attempted to load a session with None id")
raise ValueError("Invalid session id: None")

self.logger.info(f"Loading previous session: {session_id}")
self.session_id = session_id
self.file_handler = FileHandler(session_id=session_id)
new_session_id = self.session_manager.load_previous_session(session_id)
self.session_id = new_session_id
self.file_handler.set_session_id(new_session_id)
self.session_dir = self.file_handler.session_dir

# Load chat contents
chat_window.load_chat_contents()
# Load chat history from the original session
self.conversation_history = self.file_handler.load_chat_history(session_id)

chat_window.clear_chat_display()
# Load chat contents from the original session
chat_contents = self.file_handler.load_chat_contents(session_id)
for content in chat_contents:
if isinstance(content, dict) and "role" in content and "content" in content:
chat_window.add_to_chat(content["role"], content["content"])
elif (
isinstance(content, dict)
and "sender" in content
and "message" in content
):
chat_window.add_to_chat(content["sender"], content["message"])
else:
self.logger.warning(f"Unexpected chat content format: {content}")

system_message = f"A new session has been created based on session {session_id}. Previous session data copied to: {self.session_dir}"
chat_window.add_to_chat("System", system_message)

self.logger.info(f"Loaded previous session with ID: {session_id}")
self.logger.info(f"Created new session with ID: {new_session_id}")

return new_session_id

def update_session_summary(self, summary):
self.session_manager.update_session_summary(summary)
self.session_updated.emit()

self.logger.info(f"Loaded previous session with ID: {self.session_id}")
def get_all_sessions(self):
return self.session_manager.get_all_sessions()

def emit_traceback_signal(self, traceback_text):
self.logger.info(
Expand All @@ -638,4 +658,4 @@ def send_traceback_to_agent(self, traceback_text):
self.chat_window.on_submit()

def cleanup(self):
self.env_manager.cleanup()
self.environment_manager.cleanup()
113 changes: 107 additions & 6 deletions codeaide/ui/chat_window.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -167,6 +167,7 @@ def __init__(self, chat_handler):
self.timer.timeout.connect(lambda: None)

self.logger.info("Chat window initialized")
self.chat_handler.session_updated.connect(self.update_session_dropdown)

def setup_ui(self):
central_widget = QWidget(self)
Expand All @@ -175,33 +176,51 @@ def setup_ui(self):
main_layout.setSpacing(5)
main_layout.setContentsMargins(8, 8, 8, 8)

# Create a widget for the dropdowns
# Create a widget for the dropdowns (keep this as is)
dropdown_widget = QWidget()
dropdown_layout = QHBoxLayout(dropdown_widget)
dropdown_layout.setContentsMargins(0, 0, 0, 0)
dropdown_layout.setSpacing(5)

# Provider dropdown
# Provider dropdown (keep this as is)
self.provider_dropdown = QComboBox()
self.provider_dropdown.addItems(AI_PROVIDERS.keys())
self.provider_dropdown.setCurrentText(DEFAULT_PROVIDER)
self.provider_dropdown.currentTextChanged.connect(self.update_model_dropdown)
dropdown_layout.addWidget(QLabel("Provider:"))
dropdown_layout.addWidget(self.provider_dropdown)

# Model dropdown
# Model dropdown (keep this as is)
self.model_dropdown = QComboBox()
self.update_model_dropdown(DEFAULT_PROVIDER, add_message_to_chat=False)
self.model_dropdown.currentTextChanged.connect(self.update_chat_handler)
dropdown_layout.addWidget(QLabel("Model:"))
dropdown_layout.addWidget(self.model_dropdown)

# Add stretch to push everything to the left
# Add stretch to push everything to the left (keep this as is)
dropdown_layout.addStretch(1)

# Add the dropdown widget to the main layout
# Add the dropdown widget to the main layout (keep this as is)
main_layout.addWidget(dropdown_widget)

# Create a new widget for the session selector
session_widget = QWidget()
session_layout = QHBoxLayout(session_widget)
session_layout.setContentsMargins(0, 0, 0, 0)
session_layout.setSpacing(5)

# Session selector dropdown
self.session_dropdown = QComboBox()
session_layout.addWidget(QLabel("Session Selector:"))
session_layout.addWidget(self.session_dropdown)
session_layout.addStretch(1) # Add stretch to keep alignment consistent

# Add the session widget to the main layout
main_layout.addWidget(session_widget)

# Now that self.session_dropdown is initialized, we can update it
self.update_session_dropdown()

# Chat display
self.chat_display = QTextEdit(self)
self.chat_display.setReadOnly(True)
Expand Down Expand Up @@ -532,7 +551,7 @@ def on_new_session_clicked(self):

if reply == QMessageBox.Yes:
self.logger.info("User confirmed starting a new session")
self.chat_handler.start_new_session(self)
self.chat_handler.start_new_session(self, initial_session=False)
else:
self.logger.info("User cancelled starting a new session")

Expand Down Expand Up @@ -750,3 +769,85 @@ def scroll_to_bottom(self):
# Scroll to the bottom
scrollbar = self.input_text.verticalScrollBar()
scrollbar.setValue(scrollbar.maximum())

def update_session_dropdown(self):
self.session_dropdown.blockSignals(True)
self.session_dropdown.clear()
sessions = self.chat_handler.get_all_sessions()
current_session_id = self.chat_handler.session_id

# Always add the current session at the top
current_session = next(
(s for s in sessions if s["id"] == current_session_id), None
)
if current_session:
self.session_dropdown.addItem(
f"{current_session_id} (Current) - {current_session['summary']}",
current_session_id,
)
else:
# If the current session is not in the list (e.g., it's new and empty), add it manually
current_summary = (
self.chat_handler.session_manager.get_session_summary(
current_session_id
)
or "New session"
)
self.session_dropdown.addItem(
f"{current_session_id} (Current) - {current_summary}",
current_session_id,
)

# Add other non-empty sessions
for session in sessions:
if session["id"] != current_session_id:
self.session_dropdown.addItem(
f"{session['id']} - {session['summary']}", session["id"]
)

self.session_dropdown.setCurrentIndex(0)
self.session_dropdown.blockSignals(False)

# Reconnect the signal
self.session_dropdown.currentIndexChanged.connect(self.on_session_selected)

def on_session_selected(self, index):
if index == 0: # Current session
return

selected_session_id = self.session_dropdown.itemData(index)
if selected_session_id is None:
self.logger.error(f"Invalid session selected at index {index}")
self.session_dropdown.setCurrentIndex(0)
return

reply = QMessageBox.question(
self,
"Load Previous Session",
f"This will create a new session based on session {selected_session_id}. Are you sure you'd like to proceed?",
QMessageBox.Yes | QMessageBox.No,
QMessageBox.No,
)

if reply == QMessageBox.Yes:
self.logger.info(
f"User confirmed loading previous session: {selected_session_id}"
)
try:
self.update_session_dropdown()
# Disconnect and reconnect the signal to prevent multiple calls
self.session_dropdown.currentIndexChanged.disconnect()
self.session_dropdown.setCurrentIndex(0)
self.session_dropdown.currentIndexChanged.connect(
self.on_session_selected
)
except Exception as e:
self.logger.error(f"Error loading previous session: {str(e)}")
QMessageBox.critical(
self, "Error", f"Failed to load previous session: {str(e)}"
)
self.session_dropdown.setCurrentIndex(0)
else:
self.logger.info("User cancelled loading previous session")
# Reset the dropdown to the current session
self.session_dropdown.setCurrentIndex(0)
11 changes: 10 additions & 1 deletion codeaide/utils/api_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -165,12 +165,21 @@ def parse_response(response, provider):
version_description = outer_json.get("version_description")
requirements = outer_json.get("requirements", [])
questions = outer_json.get("questions", [])
session_summary = outer_json.get("session_summary", "")

# Clean the code if it exists
if code:
code = clean_code(code)

return text, questions, code, code_version, version_description, requirements
return (
text,
questions,
code,
code_version,
version_description,
requirements,
session_summary,
)


def clean_code(code):
Expand Down
Loading
Loading