astraagent 2.25.6 → 2.26.0

package/astra/chat.py

@@ -1,6 +1,7 @@
 """Chat Mode - Enhanced with keyboard shortcuts, file access, and settings."""
 
 import asyncio
+import sys
 import json
 import os
 from pathlib import Path
@@ -15,11 +16,50 @@ from astra.tasks import TaskExecutor, format_task_result
 from astra.updater import UpdateManager, format_update_info
 
 
+class ThinkingIndicator:
+    """Animated thinking indicator."""
+    def __init__(self, message: str = "Thinking"):
+        self.message = message
+        self.running = False
+        self.task = None
+
+    async def _animate(self):
+        chars = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]
+        idx = 0
+        while self.running:
+            sys.stdout.write(f"\r🤖 {chars[idx]} {self.message}...")
+            sys.stdout.flush()
+            idx = (idx + 1) % len(chars)
+            await asyncio.sleep(0.1)
+
+    async def __aenter__(self):
+        self.running = True
+        self.task = asyncio.create_task(self._animate())
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        self.running = False
+        if self.task:
+            self.task.cancel()
+            try:
+                await self.task
+            except asyncio.CancelledError:
+                pass
+        # Clear line completely
+        sys.stdout.write("\r" + " " * 40 + "\r") 
+        sys.stdout.flush()
+
+
 class ChatHistory:
     """Manages chat history and logging."""
     
-    def __init__(self, logs_dir: str = "./logs/chats"):
-        self.logs_dir = Path(logs_dir)
+    def __init__(self, logs_dir: str = None):
+        if logs_dir:
+            self.logs_dir = Path(logs_dir)
+        else:
+            # Use user's home directory for logs
+            self.logs_dir = Path.home() / ".astra" / "logs" / "chats"
+            
         self.logs_dir.mkdir(parents=True, exist_ok=True)
         self.messages: List[dict] = []
         self.provider_name: str = ""
@@ -192,8 +232,7 @@ TIPS:
   • Search works across entire device
   • Tasks handle file operations safely
   • Updates check npm and PyPI simultaneously
-""")
-    print("="*70 + "\n")
+
   • Use Ctrl+C to interrupt response
   • Use Tab for auto-completion (if available)
   • Multi-line input: Use Shift+Enter
@@ -396,7 +435,11 @@ async def chat_with_provider(provider: LLMProvider, history: ChatHistory, provid
     print(f"Provider: {history.provider_name} | Model: {history.model_name}")
     print(f"Type '/help' for commands or just start typing\n")
     
-    conversation = []
+    from astra.prompts import SYSTEM_PROMPT
+    
+    conversation = [
+        {"role": "system", "content": SYSTEM_PROMPT.format(workspace="./", tools="Internal Tools")}
+    ]
     should_close = False
     
     while True:
@@ -699,22 +742,40 @@ async def chat_with_provider(provider: LLMProvider, history: ChatHistory, provid
         history.add_message("user", user_input)
         
         # Get AI response
-        try:
-            print("🤖 ", end="", flush=True)
-            messages = [Message(role=msg["role"], content=msg["content"]) for msg in conversation]
-            response = await provider.generate(messages, temperature=0.7)
-            
-            ai_response = response.content
-            print(ai_response, end="\n\n", flush=True)
-            
-            # Add to history
-            conversation.append({"role": "assistant", "content": ai_response})
-            history.add_message("assistant", ai_response)
-            
-        except RuntimeError as e:
-            print(f"\n❌ Error: {str(e)}\n")
-        except Exception as e:
-            print(f"\n❌ Unexpected error: {str(e)}\n")
+        messages = [Message(role=msg["role"], content=msg["content"]) for msg in conversation]
+        max_retries = 3
+        retry_delay = 2
+
+        for attempt in range(max_retries):
+            try:
+                # Use Thinking Animation
+                async with ThinkingIndicator("Thinking"):
+                    response = await provider.generate(messages, temperature=0.7)
+                
+                # Print response (Animation clears line automatically)
+                ai_response = response.content
+                print(f"🤖 {ai_response}\n")
+                
+                # Add to history
+                conversation.append({"role": "assistant", "content": ai_response})
+                history.add_message("assistant", ai_response)
+                break # Success
+
+            except RuntimeError as e:
+                err_msg = str(e)
+                if "429" in err_msg or "rate limit" in err_msg.lower():
+                    if attempt < max_retries - 1:
+                        print(f"\r⚠️  Rate limit hit. Retrying in {retry_delay}s... (Attempt {attempt+1}/{max_retries})")
+                        await asyncio.sleep(retry_delay)
+                        retry_delay *= 2
+                        continue
+                
+                print(f"\n❌ Error: {err_msg}\n")
+                break
+                
+            except Exception as e:
+                print(f"\n❌ Unexpected error: {str(e)}\n")
+                break
     
     return should_close