chat-console 0.2.0__py3-none-any.whl → 0.2.2.dev1__py3-none-any.whl

app/api/ollama.py

@@ -15,6 +15,9 @@ class OllamaClient(BaseModelClient):
         self.base_url = OLLAMA_BASE_URL.rstrip('/')
         logger.info(f"Initializing Ollama client with base URL: {self.base_url}")
         
+        # Track active stream session
+        self._active_stream_session = None
+        
         # Try to start Ollama if not running
         if not ensure_ollama_running():
             raise Exception(f"Failed to start Ollama server. Please ensure Ollama is installed and try again.")
@@ -158,6 +161,7 @@ class OllamaClient(BaseModelClient):
         prompt = self._prepare_messages(messages, style)
         retries = 2
         last_error = None
+        self._active_stream_session = None  # Track the active session
         
         while retries >= 0:
             try:
@@ -192,7 +196,10 @@ class OllamaClient(BaseModelClient):
                             logger.info("Model pulled successfully")
                 
                 # Now proceed with actual generation
-                async with aiohttp.ClientSession() as session:
+                session = aiohttp.ClientSession()
+                self._active_stream_session = session  # Store reference to active session
+                
+                try:
                     logger.debug(f"Sending streaming request to {self.base_url}/api/generate")
                     async with session.post(
                         f"{self.base_url}/api/generate",
@@ -216,6 +223,9 @@ class OllamaClient(BaseModelClient):
                                     continue
                         logger.info("Streaming completed successfully")
                         return
+                finally:
+                    self._active_stream_session = None  # Clear reference when done
+                    await session.close()  # Ensure session is closed
                         
             except aiohttp.ClientConnectorError:
                 last_error = "Could not connect to Ollama server. Make sure Ollama is running and accessible at " + self.base_url
@@ -223,6 +233,9 @@ class OllamaClient(BaseModelClient):
                 last_error = f"Ollama API error: {e.status} - {e.message}"
             except aiohttp.ClientTimeout:
                 last_error = "Request to Ollama server timed out"
+            except asyncio.CancelledError:
+                logger.info("Streaming cancelled by client")
+                raise  # Propagate cancellation
             except Exception as e:
                 last_error = f"Error streaming completion: {str(e)}"
             
@@ -233,3 +246,10 @@ class OllamaClient(BaseModelClient):
                 await asyncio.sleep(1)
                 
         raise Exception(last_error)
+        
+    async def cancel_stream(self) -> None:
+        """Cancel any active streaming request"""
+        if self._active_stream_session:
+            logger.info("Cancelling active stream session")
+            await self._active_stream_session.close()
+            self._active_stream_session = None

app/main.py

@@ -291,6 +291,7 @@ class SimpleChatApp(App): # Keep SimpleChatApp class definition
 
     current_conversation = reactive(None) # Keep SimpleChatApp reactive var
     is_generating = reactive(False) # Keep SimpleChatApp reactive var
+    current_generation_task: Optional[asyncio.Task] = None # Add task reference
 
     def __init__(self, initial_text: Optional[str] = None): # Keep SimpleChatApp __init__
         super().__init__() # Keep SimpleChatApp __init__
@@ -441,26 +442,33 @@ class SimpleChatApp(App): # Keep SimpleChatApp class definition
         await self.create_new_conversation() # Keep SimpleChatApp action_new_conversation
         log("action_new_conversation finished") # Added log
 
-    def action_escape(self) -> None: # Modify SimpleChatApp action_escape
+    def action_escape(self) -> None:
         """Handle escape key globally."""
-        log("action_escape triggered") # Added log
+        log("action_escape triggered")
         settings_panel = self.query_one("#settings-panel")
-        log(f"Settings panel visible: {settings_panel.has_class('visible')}") # Added log
+        log(f"Settings panel visible: {settings_panel.has_class('visible')}")
+
         if settings_panel.has_class("visible"):
-            log("Hiding settings panel") # Added log
-            # If settings panel is visible, hide it
+            log("Hiding settings panel")
             settings_panel.remove_class("visible")
-            self.query_one("#message-input").focus() # Focus input after closing settings
+            self.query_one("#message-input").focus()
         elif self.is_generating:
-            log("Stopping generation") # Added log
-            # Otherwise, stop generation if running
-            self.is_generating = False # Keep SimpleChatApp action_escape
-            self.notify("Generation stopped", severity="warning") # Keep SimpleChatApp action_escape
-            loading = self.query_one("#loading-indicator") # Keep SimpleChatApp action_escape
-            loading.add_class("hidden") # Keep SimpleChatApp action_escape
-        else: # Optional: Add other escape behavior for the main screen if desired # Keep SimpleChatApp action_escape comment
-            log("Escape pressed, but settings not visible and not generating.") # Added log
-            # pass # Keep SimpleChatApp action_escape comment
+            log("Attempting to cancel generation task")
+            if self.current_generation_task and not self.current_generation_task.done():
+                log("Cancelling active generation task.")
+                self.current_generation_task.cancel()
+                # The finally block in generate_response will handle is_generating = False and UI updates
+                self.notify("Stopping generation...", severity="warning", timeout=2) # Notify user immediately
+            else:
+                # This case might happen if is_generating is True but the task is already done or None
+                log("is_generating is True, but no active task found to cancel. Resetting flag.")
+                self.is_generating = False # Reset flag manually if task is missing
+                loading = self.query_one("#loading-indicator")
+                loading.add_class("hidden")
+        else:
+            log("Escape pressed, but settings not visible and not actively generating.")
+            # Optionally add other escape behaviors here if needed for the main screen
+            # e.g., clear input, deselect item, etc.
 
     def update_app_info(self) -> None:
         """Update the displayed app information."""
@@ -651,10 +659,10 @@ class SimpleChatApp(App): # Keep SimpleChatApp class definition
                         log.error(f"Error updating UI: {str(e)}") # Use log instead of logger
 
             # Generate the response with timeout and cleanup # Keep SimpleChatApp generate_response
-            generation_task = None # Keep SimpleChatApp generate_response
+            self.current_generation_task = None # Clear previous task reference
             try: # Keep SimpleChatApp generate_response
                 # Create a task for the response generation # Keep SimpleChatApp generate_response
-                generation_task = asyncio.create_task( # Keep SimpleChatApp generate_response
+                self.current_generation_task = asyncio.create_task( # Keep SimpleChatApp generate_response
                     generate_streaming_response( # Keep SimpleChatApp generate_response
                         self, # Pass the app instance
                         api_messages, # Keep SimpleChatApp generate_response
@@ -664,13 +672,16 @@ class SimpleChatApp(App): # Keep SimpleChatApp class definition
                         update_ui # Keep SimpleChatApp generate_response
                     ) # Keep SimpleChatApp generate_response
                 ) # Keep SimpleChatApp generate_response
+                generation_task = self.generation_task  # Keep local reference for later checks
 
                 # Wait for response with timeout # Keep SimpleChatApp generate_response
-                full_response = await asyncio.wait_for(generation_task, timeout=60)  # Longer timeout # Keep SimpleChatApp generate_response
+                log.info(f"Waiting for generation task {self.current_generation_task} with timeout...") # Add log
+                full_response = await asyncio.wait_for(self.current_generation_task, timeout=60)  # Longer timeout # Keep SimpleChatApp generate_response
+                log.info(f"Generation task {self.current_generation_task} completed. Full response length: {len(full_response) if full_response else 0}") # Add log
 
-                # Save to database only if we got a complete response # Keep SimpleChatApp generate_response
-                if self.is_generating and full_response: # Keep SimpleChatApp generate_response
-                    log("Generation finished, saving full response to DB") # Added log
+                # Save to database only if we got a complete response and weren't cancelled
+                if self.is_generating and full_response: # Check is_generating flag here
+                    log("Generation finished normally, saving full response to DB") # Added log
                     self.db.add_message( # Keep SimpleChatApp generate_response
                         self.current_conversation.id, # Keep SimpleChatApp generate_response
                         "assistant", # Keep SimpleChatApp generate_response
@@ -679,11 +690,24 @@ class SimpleChatApp(App): # Keep SimpleChatApp class definition
                     # Force a final refresh # Keep SimpleChatApp generate_response
                     self.refresh(layout=True) # Keep SimpleChatApp generate_response
                     await asyncio.sleep(0.1)  # Wait for UI to update # Keep SimpleChatApp generate_response
-                elif not full_response:
+                elif not full_response and self.is_generating: # Only log if not cancelled
                     log("Generation finished but full_response is empty/None") # Added log
+                else:
+                    # This case handles cancellation where full_response might be partial or None
+                    log("Generation was cancelled or finished without a full response.")
+
+            except asyncio.CancelledError: # Handle cancellation explicitly
+                log.warning("Generation task was cancelled.")
+                self.notify("Generation stopped by user.", severity="warning")
+                # Remove the potentially incomplete message from UI state
+                if self.messages and self.messages[-1].role == "assistant":
+                    self.messages.pop()
+                await self.update_messages_ui() # Update UI to remove partial message
 
             except asyncio.TimeoutError: # Keep SimpleChatApp generate_response
-                log.error("Response generation timed out") # Use log instead of logger
+                log.error(f"Response generation timed out waiting for task {self.current_generation_task}") # Use log instead of logger
+                # Log state at timeout
+                log.error(f"Timeout state: is_generating={self.is_generating}, task_done={self.current_generation_task.done() if self.current_generation_task else 'N/A'}")
                 error_msg = "Response generation timed out. The model may be busy or unresponsive. Please try again." # Keep SimpleChatApp generate_response
                 self.notify(error_msg, severity="error") # Keep SimpleChatApp generate_response
 
@@ -695,31 +719,36 @@ class SimpleChatApp(App): # Keep SimpleChatApp class definition
                 await self.update_messages_ui() # Keep SimpleChatApp generate_response
 
             finally: # Keep SimpleChatApp generate_response
-                # Ensure task is properly cancelled and cleaned up # Keep SimpleChatApp generate_response
-                if generation_task: # Keep SimpleChatApp generate_response
-                    if not generation_task.done(): # Keep SimpleChatApp generate_response
-                        log("Cancelling generation task") # Added log
-                        generation_task.cancel() # Keep SimpleChatApp generate_response
-                        try: # Keep SimpleChatApp generate_response
-                            await generation_task # Keep SimpleChatApp generate_response
-                        except (asyncio.CancelledError, Exception) as e: # Keep SimpleChatApp generate_response
-                            log.error(f"Error cleaning up generation task: {str(e)}") # Use log instead of logger
-
+                # Ensure flag is reset and task reference is cleared
+                log(f"Setting is_generating to False in finally block") # Added log
+                self.is_generating = False # Keep SimpleChatApp generate_response
+                self.current_generation_task = None # Clear task reference
+                loading = self.query_one("#loading-indicator") # Keep SimpleChatApp generate_response
+                loading.add_class("hidden") # Keep SimpleChatApp generate_response
                 # Force a final UI refresh # Keep SimpleChatApp generate_response
                 self.refresh(layout=True) # Keep SimpleChatApp generate_response
 
         except Exception as e: # Keep SimpleChatApp generate_response
-            log.error(f"Exception during generate_response: {str(e)}") # Added log
+            # Catch any other unexpected errors during generation setup/handling
+            log.error(f"Unexpected exception during generate_response: {str(e)}") # Added log
             self.notify(f"Error generating response: {str(e)}", severity="error") # Keep SimpleChatApp generate_response
-            # Add error message # Keep SimpleChatApp generate_response
-            error_msg = f"Error generating response: {str(e)}" # Keep SimpleChatApp generate_response
+            # Add error message to UI # Keep SimpleChatApp generate_response
+            error_msg = f"Error: {str(e)}" # Keep SimpleChatApp generate_response
             self.messages.append(Message(role="assistant", content=error_msg)) # Keep SimpleChatApp generate_response
             await self.update_messages_ui() # Keep SimpleChatApp generate_response
-        finally: # Keep SimpleChatApp generate_response
-            log(f"Setting is_generating to False in finally block") # Added log
-            self.is_generating = False # Keep SimpleChatApp generate_response
+            # The finally block below will handle resetting is_generating and hiding loading
+
+        finally: # Keep SimpleChatApp generate_response - This finally block now primarily handles cleanup
+            log(f"Ensuring is_generating is False and task is cleared in outer finally block") # Added log
+            self.is_generating = False # Ensure flag is always reset
+            self.current_generation_task = None # Ensure task ref is cleared
             loading = self.query_one("#loading-indicator") # Keep SimpleChatApp generate_response
-            loading.add_class("hidden") # Keep SimpleChatApp generate_response
+            loading.add_class("hidden") # Ensure loading indicator is hidden
+            # Re-focus input after generation attempt (success, failure, or cancel)
+            try:
+                self.query_one("#message-input").focus()
+            except Exception:
+                pass # Ignore if input not found
 
     def on_model_selector_model_selected(self, event: ModelSelector.ModelSelected) -> None: # Keep SimpleChatApp on_model_selector_model_selected
         """Handle model selection""" # Keep SimpleChatApp on_model_selector_model_selected docstring

app/utils.py

@@ -86,14 +86,12 @@ async def generate_streaming_response(app: 'SimpleChatApp', messages: List[Dict]
     buffer = []
     last_update = time.time()
     update_interval = 0.1  # Update UI every 100ms
+    generation_task = None
     
     try:
+        # The cancellation is now handled by cancelling the asyncio Task in main.py
+        # which will raise CancelledError here, interrupting the loop.
         async for chunk in client.generate_stream(messages, model, style):
-            # Check if generation was cancelled by the app (e.g., via escape key)
-            if not app.is_generating:
-                logger.info("Generation cancelled by app flag.")
-                break # Exit the loop immediately
-
             if chunk:  # Only process non-empty chunks
                 buffer.append(chunk)
                 current_time = time.time()
@@ -102,34 +100,37 @@ async def generate_streaming_response(app: 'SimpleChatApp', messages: List[Dict]
                 if current_time - last_update >= update_interval or len(''.join(buffer)) > 100:
                     new_content = ''.join(buffer)
                     full_response += new_content
-                    # Check again before calling callback, in case it was cancelled during chunk processing
-                    if not app.is_generating:
-                        logger.info("Generation cancelled before UI update.")
-                        break
+                    # No need to check app.is_generating here, rely on CancelledError
                     await callback(full_response)
                     buffer = []
                     last_update = current_time
                     
                     # Small delay to let UI catch up
                     await asyncio.sleep(0.05)
-        
-        # Send any remaining content if generation wasn't cancelled
-        if buffer and app.is_generating:
+
+        # Send any remaining content if the loop finished normally
+        if buffer:
             new_content = ''.join(buffer)
             full_response += new_content
             await callback(full_response)
-        
-        if app.is_generating:
-            logger.info("Streaming response completed normally.")
-        else:
-             logger.info("Streaming response loop exited due to cancellation.")
-             
+
+        logger.info("Streaming response loop finished normally.") # Clarify log message
+        # Add log before returning
+        logger.info(f"generate_streaming_response returning normally. Full response length: {len(full_response)}")
         return full_response
+    except asyncio.CancelledError:
+        # This is expected when the user cancels via Escape
+        logger.info("Streaming response task cancelled.") # Clarify log message
+        # Add log before returning
+        logger.info(f"generate_streaming_response returning after cancellation. Partial response length: {len(full_response)}")
+        # Do not re-raise CancelledError, let the caller handle it
+        return full_response # Return whatever was collected so far (might be partial)
     except Exception as e:
-        logger.error(f"Error in streaming response: {str(e)}")
-        # Ensure the app knows generation stopped on error
-        app.is_generating = False
-        raise
+        logger.error(f"Error during streaming response: {str(e)}") # Clarify log message
+        # Ensure the app knows generation stopped on error ONLY if it wasn't cancelled
+        if not isinstance(e, asyncio.CancelledError):
+             app.is_generating = False # Reset flag on other errors
+        raise # Re-raise other exceptions
 
 def ensure_ollama_running() -> bool:
     """

{chat_console-0.2.0.dist-info → chat_console-0.2.2.dev1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: chat-console
-Version: 0.2.0
+Version: 0.2.2.dev1
 Summary: A command-line interface for chatting with LLMs, storing chats and (future) rag interactions
 Home-page: https://github.com/wazacraftrfid/chat-console
 Author: Johnathan Greenaway

{chat_console-0.2.0.dist-info → chat_console-0.2.2.dev1.dist-info}/RECORD

@@ -1,13 +1,13 @@
 app/__init__.py,sha256=OeqboIrx_Kjea0CY9Be8nLI-No1YWQfqbWIp-4lMOOI,131
 app/config.py,sha256=sKNp6Za4ZfW-CZBOvEv0TncAS77AnKi86hTM51C4KQ4,5227
 app/database.py,sha256=nt8CVuDpy6zw8mOYqDcfUmNw611t7Ln7pz22M0b6-MI,9967
-app/main.py,sha256=PMx9DcD2MAVyqeYanqu6mqr43RG4oXA9jFIdBwoQFQY,48174
+app/main.py,sha256=PCc1clX9sdJh8cpzcgUlnrRyt4o8k5KCJ9UwTQ8Thas,50133
 app/models.py,sha256=4-y9Lytay2exWPFi0FDlVeRL3K2-I7E-jBqNzTfokqY,2644
-app/utils.py,sha256=AgmLrmyikt1Y7KturNmZVK2eC6de-RfRadVbp3HmUAg,8434
+app/utils.py,sha256=6XSIJBcJPOXPIHnvKvRwnttdRnN9BSlodcKVj57RLeM,8861
 app/api/__init__.py,sha256=A8UL84ldYlv8l7O-yKzraVFcfww86SgWfpl4p7R03-w,62
 app/api/anthropic.py,sha256=x5PmBXEKe_ow2NWk8XdqSPR0hLOdCc_ypY5QAySeA78,4234
 app/api/base.py,sha256=-6RSxSpqe-OMwkaq1wVWbu3pVkte-ZYy8rmdvt-Qh48,3953
-app/api/ollama.py,sha256=zFZ3g2sYncvMgcvx92jTCLkigIaDvTuhILcLiCrwisc,11640
+app/api/ollama.py,sha256=Wm7IJPZhp4pygwEYIA4zk0mRsDaU87LkOmYUXPt7MyM,12551
 app/api/openai.py,sha256=1fYgFXXL6yj_7lQ893Yj28RYG4M8d6gt_q1gzhhjcig,3641
 app/ui/__init__.py,sha256=RndfbQ1Tv47qdSiuQzvWP96lPS547SDaGE-BgOtiP_w,55
 app/ui/chat_interface.py,sha256=VwmVvltxS9l18DI9U7kL43t8kSPPNsrkkrrUSoGu16Q,13623
@@ -15,9 +15,9 @@ app/ui/chat_list.py,sha256=WQTYVNSSXlx_gQal3YqILZZKL9UiTjmNMIDX2I9pAMM,11205
 app/ui/model_selector.py,sha256=Aj1irAs9DQMn8wfcPsFZGxWmx0JTzHjSe7pVdDMwqTQ,13182
 app/ui/search.py,sha256=b-m14kG3ovqW1-i0qDQ8KnAqFJbi5b1FLM9dOnbTyIs,9763
 app/ui/styles.py,sha256=04AhPuLrOd2yenfRySFRestPeuTPeMLzhmMB67NdGvw,5615
-chat_console-0.2.0.dist-info/licenses/LICENSE,sha256=srHZ3fvcAuZY1LHxE7P6XWju2njRCHyK6h_ftEbzxSE,1057
-chat_console-0.2.0.dist-info/METADATA,sha256=DPnX1Q1e7iTfNyyQRhJezaUXuHtq1E9mqb7k7gDe0GI,2921
-chat_console-0.2.0.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-chat_console-0.2.0.dist-info/entry_points.txt,sha256=kkVdEc22U9PAi2AeruoKklfkng_a_aHAP6VRVwrAD7c,67
-chat_console-0.2.0.dist-info/top_level.txt,sha256=io9g7LCbfmTG1SFKgEOGXmCFB9uMP2H5lerm0HiHWQE,4
-chat_console-0.2.0.dist-info/RECORD,,
+chat_console-0.2.2.dev1.dist-info/licenses/LICENSE,sha256=srHZ3fvcAuZY1LHxE7P6XWju2njRCHyK6h_ftEbzxSE,1057
+chat_console-0.2.2.dev1.dist-info/METADATA,sha256=HcWEG6ahXupJhz9POYMoGI90alNTHFLL0mfEPQ4xecg,2926
+chat_console-0.2.2.dev1.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+chat_console-0.2.2.dev1.dist-info/entry_points.txt,sha256=kkVdEc22U9PAi2AeruoKklfkng_a_aHAP6VRVwrAD7c,67
+chat_console-0.2.2.dev1.dist-info/top_level.txt,sha256=io9g7LCbfmTG1SFKgEOGXmCFB9uMP2H5lerm0HiHWQE,4
+chat_console-0.2.2.dev1.dist-info/RECORD,,