chat-console 0.2.8__py3-none-any.whl → 0.2.9__py3-none-any.whl

app/__init__.py

@@ -3,4 +3,4 @@ Chat CLI
 A command-line interface for chatting with various LLM providers like ChatGPT and Claude.
 """
 
-__version__ = "0.2.8"
+__version__ = "0.2.9"

app/api/ollama.py

@@ -22,6 +22,9 @@ class OllamaClient(BaseModelClient):
         # Track active stream session
         self._active_stream_session = None
         
+        # Track model loading state
+        self._model_loading = False
+        
         # Path to the cached models file
         self.models_cache_path = Path(__file__).parent.parent / "data" / "ollama-models.json"
         
@@ -191,6 +194,10 @@ class OllamaClient(BaseModelClient):
                                 raise aiohttp.ClientError("Model not ready")
                     except (aiohttp.ClientError, asyncio.TimeoutError) as e:
                         logger.info(f"Model cold start detected: {str(e)}")
+                        # Set model loading flag
+                        self._model_loading = True
+                        logger.info("Setting model_loading state to True")
+                        
                         # Model might need loading, try pulling it
                         async with session.post(
                             f"{self.base_url}/api/pull",
@@ -199,8 +206,10 @@ class OllamaClient(BaseModelClient):
                         ) as pull_response:
                             if pull_response.status != 200:
                                 logger.error("Failed to pull model")
+                                self._model_loading = False  # Reset flag on failure
                                 raise Exception("Failed to pull model")
                             logger.info("Model pulled successfully")
+                            self._model_loading = False  # Reset flag after successful pull
                 
                 # Now proceed with actual generation
                 session = aiohttp.ClientSession()
@@ -208,7 +217,7 @@ class OllamaClient(BaseModelClient):
                 
                 try:
                     logger.debug(f"Sending streaming request to {self.base_url}/api/generate")
-                    async with session.post(
+                    response = await session.post(
                         f"{self.base_url}/api/generate",
                         json={
                             "model": model,
@@ -217,19 +226,36 @@ class OllamaClient(BaseModelClient):
                             "stream": True
                         },
                         timeout=60  # Longer timeout for actual generation
-                    ) as response:
-                        response.raise_for_status()
-                        async for line in response.content:
-                            if line:
-                                chunk = line.decode().strip()
-                                try:
-                                    data = json.loads(chunk)
-                                    if "response" in data:
-                                        yield data["response"]
-                                except json.JSONDecodeError:
-                                    continue
-                        logger.info("Streaming completed successfully")
-                        return
+                    )
+                    response.raise_for_status()
+                    
+                    # Process the response stream
+                    while True:
+                        if not self._active_stream_session:
+                            logger.info("Stream session was closed externally")
+                            break
+                            
+                        try:
+                            line = await asyncio.wait_for(response.content.readline(), timeout=0.5)
+                            if not line:  # End of stream
+                                break
+                                
+                            chunk = line.decode().strip()
+                            try:
+                                data = json.loads(chunk)
+                                if "response" in data:
+                                    yield data["response"]
+                            except json.JSONDecodeError:
+                                continue
+                        except asyncio.TimeoutError:
+                            # This allows checking for cancellation regularly
+                            continue
+                        except asyncio.CancelledError:
+                            logger.info("Stream processing was cancelled")
+                            raise
+                            
+                    logger.info("Streaming completed successfully")
+                    return
                 finally:
                     self._active_stream_session = None  # Clear reference when done
                     await session.close()  # Ensure session is closed
@@ -260,6 +286,12 @@ class OllamaClient(BaseModelClient):
             logger.info("Cancelling active stream session")
             await self._active_stream_session.close()
             self._active_stream_session = None
+            self._model_loading = False
+            logger.info("Stream session closed successfully")
+            
+    def is_loading_model(self) -> bool:
+        """Check if Ollama is currently loading a model"""
+        return self._model_loading
             
     async def get_model_details(self, model_id: str) -> Dict[str, Any]:
         """Get detailed information about a specific Ollama model"""

app/main.py

@@ -477,30 +477,45 @@ class SimpleChatApp(App): # Keep SimpleChatApp class definition
             log("Attempting to cancel generation task")
             if self.current_generation_task and not self.current_generation_task.done():
                 log("Cancelling active generation task.")
-                # Get the client for the current model
+                
+                # Get the client for the current model first and cancel the connection
                 try:
                     model = self.selected_model
                     client = BaseModelClient.get_client_for_model(model)
-                    # Call the client's cancel method if it's an Ollama client
+                    
+                    # Call the client's cancel method if it's supported
                     if hasattr(client, 'cancel_stream'):
                         log("Calling client.cancel_stream() to terminate API session")
-                        await client.cancel_stream()
+                        try:
+                            # This will close the HTTP connection to Ollama server
+                            await client.cancel_stream()
+                            log("Client stream cancelled successfully")
+                        except Exception as e:
+                            log.error(f"Error in client.cancel_stream(): {str(e)}")
+                except Exception as e:
+                    log.error(f"Error setting up client cancellation: {str(e)}")
+                
+                # Now cancel the asyncio task - this should raise CancelledError in the task
+                try:
+                    log("Cancelling asyncio task")
+                    self.current_generation_task.cancel()
+                    # Give a moment for cancellation to propagate
+                    await asyncio.sleep(0.1)
+                    log(f"Task cancelled. Task done: {self.current_generation_task.done()}")
                 except Exception as e:
-                    log.error(f"Error cancelling client stream: {str(e)}")
+                    log.error(f"Error cancelling task: {str(e)}")
                 
-                # Now cancel the asyncio 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
+                # Notify user that we're stopping
+                self.notify("Stopping generation...", severity="warning", timeout=2)
             else:
-                # This case might happen if is_generating is True, but no active task found to cancel. Resetting flag.")
-                self.is_generating = False # Reset flag manually if task is missing
+                # This happens if is_generating is True, but no active task found to cancel
+                log("No active generation task found, but is_generating=True. Resetting state.")
+                self.is_generating = False
                 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.
+            # Optionally add other escape behaviors here if needed
 
     def update_app_info(self) -> None:
         """Update the displayed app information."""

app/ui/chat_interface.py

@@ -204,6 +204,11 @@ class ChatInterface(Container):
         display: none;
         padding: 0 1;
     }
+    
+    #loading-indicator.model-loading {
+        background: $warning;
+        color: $text;
+    }
     """
     
     class MessageSent(Message):
@@ -238,7 +243,7 @@ class ChatInterface(Container):
                 yield MessageDisplay(message, highlight_code=CONFIG["highlight_code"])
         with Container(id="input-area"):
             yield Container(
-                Label("Generating response...", id="loading-text"),
+                Label("▪▪▪ Generating response...", id="loading-text", markup=True),
                 id="loading-indicator"
             )
             with Container(id="controls"):
@@ -328,16 +333,30 @@ class ChatInterface(Container):
         if input_widget.has_focus:
             input_widget.focus()
         
-    def start_loading(self) -> None:
-        """Show loading indicator"""
+    def start_loading(self, model_loading: bool = False) -> None:
+        """Show loading indicator
+        
+        Args:
+            model_loading: If True, indicates Ollama is loading a model
+        """
         self.is_loading = True
         loading = self.query_one("#loading-indicator")
+        loading_text = self.query_one("#loading-text")
+        
+        if model_loading:
+            loading.add_class("model-loading")
+            loading_text.update("⚙️ Loading Ollama model...")
+        else:
+            loading.remove_class("model-loading")
+            loading_text.update("▪▪▪ Generating response...")
+            
         loading.display = True
         
     def stop_loading(self) -> None:
         """Hide loading indicator"""
         self.is_loading = False
         loading = self.query_one("#loading-indicator")
+        loading.remove_class("model-loading")
         loading.display = False
         
     def clear_messages(self) -> None:

app/utils.py

@@ -86,12 +86,71 @@ 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):
+        # Update UI with model loading state if it's an Ollama client
+        if hasattr(client, 'is_loading_model'):
+            # Send signal to update UI for model loading if needed
+            try:
+                # The client might be in model loading state even before generating
+                model_loading = client.is_loading_model()
+                logger.info(f"Initial model loading state: {model_loading}")
+                
+                # Get the chat interface and update loading indicator
+                if hasattr(app, 'query_one'):
+                    loading = app.query_one("#loading-indicator")
+                    if model_loading:
+                        loading.add_class("model-loading")
+                        app.query_one("#loading-text").update("Loading Ollama model...")
+                    else:
+                        loading.remove_class("model-loading")
+            except Exception as e:
+                logger.error(f"Error setting initial loading state: {str(e)}")
+        
+        stream_generator = client.generate_stream(messages, model, style)
+        
+        # Check if we just entered model loading state
+        if hasattr(client, 'is_loading_model') and client.is_loading_model():
+            logger.info("Model loading started during generation")
+            try:
+                if hasattr(app, 'query_one'):
+                    loading = app.query_one("#loading-indicator")
+                    loading.add_class("model-loading")
+                    app.query_one("#loading-text").update("Loading Ollama model...")
+            except Exception as e:
+                logger.error(f"Error updating UI for model loading: {str(e)}")
+        
+        # Use asyncio.shield to ensure we can properly interrupt the stream processing
+        async for chunk in stream_generator:
+            # Check for cancellation frequently
+            if asyncio.current_task().cancelled():
+                logger.info("Task cancellation detected during chunk processing")
+                # Close the client stream if possible
+                if hasattr(client, 'cancel_stream'):
+                    await client.cancel_stream()
+                raise asyncio.CancelledError()
+                
+            # Check if model loading state changed
+            if hasattr(client, 'is_loading_model'):
+                model_loading = client.is_loading_model()
+                try:
+                    if hasattr(app, 'query_one'):
+                        loading = app.query_one("#loading-indicator")
+                        loading_text = app.query_one("#loading-text")
+                        
+                        if model_loading and not loading.has_class("model-loading"):
+                            # Model loading started
+                            logger.info("Model loading started during streaming")
+                            loading.add_class("model-loading")
+                            loading_text.update("⚙️ Loading Ollama model...")
+                        elif not model_loading and loading.has_class("model-loading"):
+                            # Model loading finished
+                            logger.info("Model loading finished during streaming")
+                            loading.remove_class("model-loading")
+                            loading_text.update("▪▪▪ Generating response...")
+                except Exception as e:
+                    logger.error(f"Error updating loading state during streaming: {str(e)}")
+                
             if chunk:  # Only process non-empty chunks
                 buffer.append(chunk)
                 current_time = time.time()
@@ -100,7 +159,7 @@ 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
-                    # No need to check app.is_generating here, rely on CancelledError
+                    # Send content to UI
                     await callback(full_response)
                     buffer = []
                     last_update = current_time
@@ -114,23 +173,25 @@ async def generate_streaming_response(app: 'SimpleChatApp', messages: List[Dict]
             full_response += new_content
             await callback(full_response)
 
-        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)}")
+        logger.info(f"Streaming response completed successfully. 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)
+        logger.info(f"Streaming response task cancelled. Partial response length: {len(full_response)}")
+        # Ensure the client stream is closed
+        if hasattr(client, 'cancel_stream'):
+            await client.cancel_stream()
+        # Return whatever was collected so far
+        return full_response
+        
     except Exception as e:
-        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
+        logger.error(f"Error during streaming response: {str(e)}")
+        # Close the client stream if possible
+        if hasattr(client, 'cancel_stream'):
+            await client.cancel_stream()
+        # Re-raise the exception for the caller to handle
+        raise
 
 def ensure_ollama_running() -> bool:
     """

{chat_console-0.2.8.dist-info → chat_console-0.2.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: chat-console
-Version: 0.2.8
+Version: 0.2.9
 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.8.dist-info → chat_console-0.2.9.dist-info}/RECORD

@@ -1,24 +1,24 @@
-app/__init__.py,sha256=LLUTUQ0_sgXPrC2d7wZMLj-d74v6qfpS2myh6lOjoaY,130
+app/__init__.py,sha256=g2BzewDN5X96Dl5Zzw8uag1TBEdPIU1ceTm7u-BJrjM,130
 app/config.py,sha256=sKNp6Za4ZfW-CZBOvEv0TncAS77AnKi86hTM51C4KQ4,5227
 app/database.py,sha256=nt8CVuDpy6zw8mOYqDcfUmNw611t7Ln7pz22M0b6-MI,9967
-app/main.py,sha256=QTDXkVweTX8fYb4LGte71RiEBv81l7jwbZwzWBQpUoc,48157
+app/main.py,sha256=k726xRBcuPgbUsUg4s-REhtaljccjDLNzA_C-fPkQk4,48866
 app/models.py,sha256=4-y9Lytay2exWPFi0FDlVeRL3K2-I7E-jBqNzTfokqY,2644
-app/utils.py,sha256=6XSIJBcJPOXPIHnvKvRwnttdRnN9BSlodcKVj57RLeM,8861
+app/utils.py,sha256=IyINMrM6oGXtN5HRPuKoFEyfKg0fR4FVXIi_0e2KxI0,11798
 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=NgfETreb7EdFIux9fvkDfIBj77wcJvic77ObUV95TlI,49866
+app/api/ollama.py,sha256=FTIlgZmvpZd6K4HL2nUD19-p9Xb1TA859LfnCgewpcU,51354
 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
+app/ui/chat_interface.py,sha256=R8tdy72TcT7veemUzcJOjbPY32WizBdNHgfmq69EFfA,14275
 app/ui/chat_list.py,sha256=WQTYVNSSXlx_gQal3YqILZZKL9UiTjmNMIDX2I9pAMM,11205
 app/ui/model_browser.py,sha256=5h3gVsuGIUrXjYVF-QclZFhYtX2kH14LvT22Ufm9etg,49453
 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.8.dist-info/licenses/LICENSE,sha256=srHZ3fvcAuZY1LHxE7P6XWju2njRCHyK6h_ftEbzxSE,1057
-chat_console-0.2.8.dist-info/METADATA,sha256=rQ7BLE3Ne3YsLxrL8SkIfqNpUpxvUwMm2Rhrb5uDirY,2921
-chat_console-0.2.8.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-chat_console-0.2.8.dist-info/entry_points.txt,sha256=kkVdEc22U9PAi2AeruoKklfkng_a_aHAP6VRVwrAD7c,67
-chat_console-0.2.8.dist-info/top_level.txt,sha256=io9g7LCbfmTG1SFKgEOGXmCFB9uMP2H5lerm0HiHWQE,4
-chat_console-0.2.8.dist-info/RECORD,,
+chat_console-0.2.9.dist-info/licenses/LICENSE,sha256=srHZ3fvcAuZY1LHxE7P6XWju2njRCHyK6h_ftEbzxSE,1057
+chat_console-0.2.9.dist-info/METADATA,sha256=zTSJePqMsi0n6fEz8s4gtLwHe_726-ijfTjPwH_Mumw,2921
+chat_console-0.2.9.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+chat_console-0.2.9.dist-info/entry_points.txt,sha256=kkVdEc22U9PAi2AeruoKklfkng_a_aHAP6VRVwrAD7c,67
+chat_console-0.2.9.dist-info/top_level.txt,sha256=io9g7LCbfmTG1SFKgEOGXmCFB9uMP2H5lerm0HiHWQE,4
+chat_console-0.2.9.dist-info/RECORD,,