astraagent 2.25.6 → 2.26.0

package/astra/core/agent.py

@@ -27,7 +27,7 @@ class AstraAgent:
     An elite AI agent more powerful than ChatGPT, Claude, or Gemini.
     """
     
-    def __init__(self, config: AgentConfig = None):
+    def __init__(self, config: AgentConfig = None):
         self.config = config or AgentConfig()
         self.state = AgentState()
         self.logger = self._setup_logging()
@@ -45,9 +45,9 @@ class AstraAgent:
         self.tools = create_default_registry()
         self.llm: Optional[LLMProvider] = None
         
-        # Message history
-        self.messages: List[Message] = []
-        self._format_retry_count: int = 0
+        # Message history
+        self.messages: List[Message] = []
+        self._format_retry_count: int = 0
         
         # Check for API key upfront
         self._validate_config()
@@ -105,9 +105,9 @@ class AstraAgent:
                 api_base=self.config.llm.api_base
             )
     
-    def _build_system_prompt(self, mode: str = "default", goal: Optional[str] = None) -> str:
-        """Build the enhanced system prompt with context."""
-        tool_names = ", ".join(self.tools.list_enabled())
+    def _build_system_prompt(self, mode: str = "default", goal: Optional[str] = None) -> str:
+        """Build the enhanced system prompt with context."""
+        tool_names = ", ".join(self.tools.list_enabled())
         
         # Build base prompt
         prompt = build_system_prompt(
@@ -116,39 +116,44 @@ class AstraAgent:
             mode=mode
         )
         
-        # Add memory context if available
-        try:
-            user_facts = self.memory.get_user_facts()
-            if user_facts:
-                prompt += "\n\n=== KNOWN ABOUT USER ===\n"
-                for fact in user_facts[:5]:
-                    prompt += f"- {fact.content}\n"
-
-            # Add compact goal-aware memory context
-            if goal:
-                goal_context = self.memory.get_goal_context(goal, max_items=12)
-                if goal_context:
-                    prompt += "\n\n=== MEMORY CONTEXT FOR CURRENT GOAL ===\n"
-                    prompt += goal_context[:3000]
-        except (AttributeError, ValueError, KeyError) as e:
-            self.logger.warning(f"Failed to retrieve user facts: {e}")
-        except Exception as e:
-            self.logger.error(f"Unexpected error retrieving user facts: {e}", exc_info=True)
-        
-        return prompt
-
-    def _is_unstructured_final(self, raw_content: str, parsed: Dict[str, Any]) -> bool:
-        """Detect plain-text final answers that violate required JSON format."""
-        if "final" not in parsed:
-            return False
-        content = (raw_content or "").strip()
-        if not content:
-            return False
-        if "```json" in content:
-            return False
-        if content.startswith("{") and content.endswith("}"):
-            return False
-        return True
+        # Add memory context if available
+        try:
+            user_facts = self.memory.get_user_facts()
+            if user_facts:
+                prompt += "\n\n=== KNOWN ABOUT USER ===\n"
+                for fact in user_facts[:5]:
+                    prompt += f"- {fact.content}\n"
+
+            # Add compact goal-aware memory context
+            if goal:
+                goal_context = self.memory.get_goal_context(goal, max_items=12)
+                if goal_context:
+                    prompt += "\n\n=== MEMORY CONTEXT FOR CURRENT GOAL ===\n"
+                    prompt += goal_context[:3000]
+        except (AttributeError, ValueError, KeyError) as e:
+            self.logger.warning(f"Failed to retrieve user facts: {e}")
+        except Exception as e:
+            self.logger.error(f"Unexpected error retrieving user facts: {e}", exc_info=True)
+        
+        return prompt
+
+    def _is_unstructured_final(self, raw_content: str, parsed: Dict[str, Any]) -> bool:
+        """Detect plain-text final answers that violate required JSON format."""
+        if "final" not in parsed:
+            return False
+        content = (raw_content or "").strip()
+        if not content:
+            return False
+        if "```json" in content:
+            return False
+        if content.startswith("{") and content.endswith("}"):
+            return False
+        
+        # Don't force JSON for very short responses (likely greetings or simple confirmations)
+        if len(content) < 100 and not any(kw in content.lower() for kw in ["tool", "action", "exec", "run"]):
+            return False
+            
+        return True
     
     def _parse_response(self, content: str) -> Optional[Dict[str, Any]]:
         """Parse response from LLM. Handles both JSON and plain text."""
@@ -244,30 +249,48 @@ class AstraAgent:
         
         return result
     
+    def _manage_history(self):
+        """Manage conversation history to prevent token bloat while keeping system instructions."""
+        if len(self.messages) <= 20:
+            return
+        
+        # Keep system prompt, initial goal, and last context
+        preserved = []
+        if self.messages and self.messages[0].role == "system":
+            preserved.append(self.messages[0])
+            
+        # Keep most recent messages (sliding window)
+        recent = self.messages[-12:]
+        
+        # Merge - ensuring we don't duplicate the system prompt if it was in 'recent'
+        self.messages = preserved + [m for m in recent if m.role != "system"]
+        self.logger.debug(f"Managed history: reduced to {len(self.messages)} messages")
+
     async def _think(self, goal: str) -> LLMResponse:
         """Generate next action from LLM."""
         self._init_llm()
+        self._manage_history()
         
-        # Update messages if needed
-        if not self.messages:
-            self.messages.append(Message(role="system", content=self._build_system_prompt(goal=goal)))
-            self.messages.append(Message(role="user", content=f"Goal: {goal}"))
-
-            # Add memory context as explicit runtime instruction message
-            try:
-                memory_context = self.memory.get_goal_context(goal, max_items=10)
-                if memory_context:
-                    self.messages.append(
-                        Message(
-                            role="user",
-                            content=(
-                                "Memory context (must be considered while planning):\n"
-                                f"{memory_context[:2500]}"
-                            )
-                        )
-                    )
-            except Exception as e:
-                self.logger.warning(f"Failed to attach memory context: {e}")
+        # Update messages if needed
+        if not self.messages:
+            self.messages.append(Message(role="system", content=self._build_system_prompt(mode=self.config.prompt_mode, goal=goal)))
+            self.messages.append(Message(role="user", content=f"Goal: {goal}"))
+
+            # Add memory context as explicit runtime instruction message
+            try:
+                memory_context = self.memory.get_goal_context(goal, max_items=10)
+                if memory_context:
+                    self.messages.append(
+                        Message(
+                            role="user",
+                            content=(
+                                "Memory context (must be considered while planning):\n"
+                                f"{memory_context[:2500]}"
+                            )
+                        )
+                    )
+            except Exception as e:
+                self.logger.warning(f"Failed to attach memory context: {e}")
         
         # Add context about recent actions
         if self.state.action_history:
@@ -282,12 +305,25 @@ class AstraAgent:
             if len(self.messages) <= 2:
                 self.messages.append(Message(role="assistant", content=context))
         
-        response = await self.llm.generate(
-            self.messages,
-            tools=self.tools.get_all_schemas(),
-            temperature=self.config.llm.temperature,
-            max_tokens=self.config.llm.max_tokens
-        )
+        # Retry logic for rate limits
+        retry_delay = 5
+        for attempt in range(3):
+            try:
+                response = await self.llm.generate(
+                    self.messages,
+                    tools=self.tools.get_all_schemas(),
+                    temperature=self.config.llm.temperature,
+                    max_tokens=self.config.llm.max_tokens
+                )
+                break
+            except RuntimeError as e:
+                err = str(e)
+                if ("429" in err or "rate limit" in err.lower()) and attempt < 2:
+                    self.logger.warning(f"Rate limit hit. Retrying in {retry_delay}s...")
+                    await asyncio.sleep(retry_delay)
+                    retry_delay *= 2
+                    continue
+                raise e
         
         self.state.total_tokens_used += response.tokens_used
         return response
@@ -302,6 +338,12 @@ class AstraAgent:
             self.logger.warning("Max iterations reached")
             return True
         
+        # Check token usage safety
+        if self.state.total_tokens_used > self.config.safety.max_session_tokens:
+            self.logger.warning(f"Token limit reached ({self.state.total_tokens_used}). Stopping for safety.")
+            self.state.last_error = "Token usage threshold exceeded."
+            return True
+        
         # Get LLM response
         try:
             response = await self._think(self.state.current_goal)
@@ -316,11 +358,16 @@ class AstraAgent:
             return False
         
         # Handle tool calls from LLM
-        if response.tool_calls:
-            if response.content:
-                self.memory.remember_conversation("assistant", response.content, metadata={"kind": "tool_call"})
-            for tc in response.tool_calls:
-                action = Action(
+        if response.tool_calls:
+            self.messages.append(Message(
+                role="assistant",
+                content=response.content or "",
+                tool_calls=response.tool_calls
+            ))
+            if response.content:
+                self.memory.remember_conversation("assistant", response.content, metadata={"kind": "tool_call"})
+            for tc in response.tool_calls:
+                action = Action(
                     tool=tc["name"],
                     args=tc["arguments"],
                     thought=response.content
@@ -353,35 +400,35 @@ class AstraAgent:
             self.messages.append(Message(role="assistant", content=response.content))
             return False
         
-        # Check for final response
-        if "final" in parsed:
-            if self._is_unstructured_final(response.content, parsed) and self._format_retry_count < 2:
-                self._format_retry_count += 1
-                self.messages.append(Message(
-                    role="user",
-                    content=(
-                        "FORMAT ERROR: Reply ONLY with valid JSON in the required schema. "
-                        "Do not output plain text."
-                    )
-                ))
-                return False
-
-            self._format_retry_count = 0
-            self.logger.debug(f"Task complete: {parsed['final']}")
-            self.memory.remember_conversation("assistant", str(parsed["final"]))
-            return True
+        # Check for final response
+        if "final" in parsed:
+            if self._is_unstructured_final(response.content, parsed) and self._format_retry_count < 2:
+                self._format_retry_count += 1
+                self.messages.append(Message(
+                    role="user",
+                    content=(
+                        "FORMAT ERROR: Reply ONLY with valid JSON in the required schema. "
+                        "Do not output plain text."
+                    )
+                ))
+                return False
+
+            self._format_retry_count = 0
+            self.logger.debug(f"Task complete: {parsed['final']}")
+            self.memory.remember_conversation("assistant", str(parsed["final"]))
+            return True
         
         # Check for action
-        if "action" in parsed:
-            self._format_retry_count = 0
-            action = Action(
-                tool=parsed["action"],
-                args=parsed.get("args", {}),
-                thought=parsed.get("thought", "")
-            )
-            if action.thought:
-                self.memory.remember_conversation("assistant", action.thought, metadata={"kind": "reasoning"})
-            action.mark_executing()
+        if "action" in parsed:
+            self._format_retry_count = 0
+            action = Action(
+                tool=parsed["action"],
+                args=parsed.get("args", {}),
+                thought=parsed.get("thought", "")
+            )
+            if action.thought:
+                self.memory.remember_conversation("assistant", action.thought, metadata={"kind": "reasoning"})
+            action.mark_executing()
             
             result = await self._execute_tool(action.tool, action.args)
             
@@ -401,11 +448,11 @@ class AstraAgent:
             
             # Add action result to messages
             self.messages.append(Message(role="assistant", content=json.dumps(parsed)))
-            self.messages.append(Message(
-                role="user",
-                content=f"Tool result: {result.output or result.error}"
-            ))
-            return False
+            self.messages.append(Message(
+                role="user",
+                content=f"Tool result: {result.output or result.error}"
+            ))
+            return False
         
         # No clear action, continue
         self.messages.append(Message(role="assistant", content=response.content))
@@ -420,20 +467,20 @@ class AstraAgent:
             self.logger.debug(f"Configuration error: {e}")
             return f"Failed to start: {e}"
         
-        self.logger.debug(f"Starting AstraAgent with goal: {goal}")
-        self.state.set_goal(goal)
-        self.state.is_running = True
-        self.messages = []  # Reset messages
-        self._format_retry_count = 0
+        self.logger.debug(f"Starting AstraAgent with goal: {goal}")
+        self.state.set_goal(goal)
+        self.state.is_running = True
+        self.messages = []  # Reset messages
+        self._format_retry_count = 0
         
         # Remember the goal
-        self.memory.remember(
-            content=f"User goal: {goal}",
-            memory_type="conversation",
-            importance=0.8,
-            tags=["goal", "user_request"]
-        )
-        self.memory.remember_conversation("user", goal, metadata={"kind": "goal"})
+        self.memory.remember(
+            content=f"User goal: {goal}",
+            memory_type="conversation",
+            importance=0.8,
+            tags=["goal", "user_request"]
+        )
+        self.memory.remember_conversation("user", goal, metadata={"kind": "goal"})
         
         try:
             while not await self.step():

package/astra/core/config.py

@@ -53,21 +53,35 @@ class LLMConfig:
     
     def __post_init__(self):
         # Try to get API key from environment if not provided
-        if not self.api_key:
-            self.api_key = os.getenv("LOCAL_API_KEY")
         if not self.api_base:
             self.api_base = os.getenv("LOCAL_API_BASE") or "http://localhost:8000/api/v1"
+
+        if not self.api_key:
+            # Check provider specific keys first
+            if self.provider == "openai":
+                self.api_key = os.getenv("OPENAI_API_KEY")
+            elif self.provider == "gemini":
+                self.api_key = os.getenv("GEMINI_API_KEY")
+            elif self.provider == "anthropic":
+                self.api_key = os.getenv("ANTHROPIC_API_KEY")
+            elif self.provider == "openrouter":
+                self.api_key = os.getenv("OPENROUTER_API_KEY")
+            elif self.provider == "groq":
+                self.api_key = os.getenv("GROQ_API_KEY")
+            else:
+                self.api_key = os.getenv("LOCAL_API_KEY")
     
     def validate(self) -> None:
         """Validate configuration is complete."""
+        # Local provider often doesn't need a key (e.g. Ollama)
+        if self.provider == "local":
+            return
+
         if not self.api_key or not self.api_key.strip():
             raise ValueError(
-                "LOCAL_API_KEY is not set. Please:\n"
-                "  1. Copy .env.template to .env\n"
-                "  2. Set LOCAL_API_KEY=your-key in .env\n"
-                "  Or set environment variable:\n"
-                f"     Windows: set LOCAL_API_KEY=your-key\n"
-                f"     Linux/Mac: export LOCAL_API_KEY=your-key"
+                f"API key for provider '{self.provider}' is not set.\n"
+                "Please run 'astra settings' locally to configure it.\n"
+                f"Or set the environment variable (e.g., {self.provider.upper()}_API_KEY)."
             )
 
 
@@ -124,6 +138,7 @@ class SafetyConfig:
     allowed_domains: List[str] = field(default_factory=list)  # Empty = all allowed
     max_file_size_mb: int = 100
     enable_content_filter: bool = True
+    max_session_tokens: int = 50000  # Prevent runaway costs
 
 
 @dataclass
@@ -144,11 +159,12 @@ class AgentConfig:
     """Master configuration for AstraAgent."""
     # Identity
     name: str = "AstraAgent"
-    version: str = "1.0.0"
+    version: str = "2.26.0"
     description: str = "Autonomous AI Agent"
     
     # Execution
     mode: ExecutionMode = ExecutionMode.AUTONOMOUS
+    prompt_mode: str = "default"  # "default", "engineer", "research", "code", etc.
     max_iterations: int = 50
     reflection_enabled: bool = True
     planning_enabled: bool = True

package/astra/core/memory.py

@@ -335,86 +335,86 @@ class UnifiedMemorySystem:
             metadata={"action": action, "result": result, "success": success}
         )
     
-    def learn(self, knowledge: str, source: str = "inference") -> MemoryItem:
-        """Add to knowledge base."""
-        return self.remember(
-            content=knowledge,
-            memory_type="knowledge",
-            importance=0.7,
-            tags=["knowledge", source],
-            source=source
-        )
-
-    # Backward-compatibility and convenience API
-    def store(
-        self,
-        content: str,
-        memory_type: str = "long_term",
-        importance: float = 0.5,
-        tags: List[str] = None
-    ) -> MemoryItem:
-        """
-        Store memory with legacy memory types.
-
-        Legacy memory types:
-        - short_term -> conversation
-        - long_term -> knowledge
-        - semantic -> fact
-        """
-        type_map = {
-            "short_term": "conversation",
-            "long_term": "knowledge",
-            "semantic": "fact",
-            "conversation": "conversation",
-            "knowledge": "knowledge",
-            "fact": "fact",
-            "action": "action",
-            "observation": "observation",
-        }
-        mapped_type = type_map.get(memory_type, "knowledge")
-        return self.remember(
-            content=content,
-            memory_type=mapped_type,
-            importance=importance,
-            tags=tags or [],
-            source="store_api"
-        )
-
-    def get_recent(self, n: int = 20) -> List[MemoryItem]:
-        """Return recent memories across all stores."""
-        with self._lock:
-            combined = (
-                list(self.conversations.values())
-                + list(self.user_facts.values())
-                + list(self.knowledge.values())
-                + list(self.actions.values())
-                + list(self.observations.values())
-            )
-            combined.sort(key=lambda x: x.created_at, reverse=True)
-            return combined[:n]
-
-    def get_goal_context(self, goal: str, max_items: int = 15) -> str:
-        """Get compact memory context relevant to a specific goal."""
-        if not goal:
-            return self.get_context_summary(max_items=max_items)
-
-        lines: List[str] = []
-
-        # Most important user facts first
-        facts = self.get_user_facts()[:5]
-        if facts:
-            lines.append("Known facts:")
-            for item in facts:
-                lines.append(f"- {item.content[:180]}")
-
-        # Relevant memories by search
-        relevant = self.search(query=goal, limit=max_items)
-        if relevant:
-            lines.append("Relevant memory:")
-            for item in relevant:
-                lines.append(f"- [{item.memory_type}] {item.content[:200]}")
-
-        return "\n".join(lines).strip()
+    def learn(self, knowledge: str, source: str = "inference") -> MemoryItem:
+        """Add to knowledge base."""
+        return self.remember(
+            content=knowledge,
+            memory_type="knowledge",
+            importance=0.7,
+            tags=["knowledge", source],
+            source=source
+        )
+
+    # Backward-compatibility and convenience API
+    def store(
+        self,
+        content: str,
+        memory_type: str = "long_term",
+        importance: float = 0.5,
+        tags: List[str] = None
+    ) -> MemoryItem:
+        """
+        Store memory with legacy memory types.
+
+        Legacy memory types:
+        - short_term -> conversation
+        - long_term -> knowledge
+        - semantic -> fact
+        """
+        type_map = {
+            "short_term": "conversation",
+            "long_term": "knowledge",
+            "semantic": "fact",
+            "conversation": "conversation",
+            "knowledge": "knowledge",
+            "fact": "fact",
+            "action": "action",
+            "observation": "observation",
+        }
+        mapped_type = type_map.get(memory_type, "knowledge")
+        return self.remember(
+            content=content,
+            memory_type=mapped_type,
+            importance=importance,
+            tags=tags or [],
+            source="store_api"
+        )
+
+    def get_recent(self, n: int = 20) -> List[MemoryItem]:
+        """Return recent memories across all stores."""
+        with self._lock:
+            combined = (
+                list(self.conversations.values())
+                + list(self.user_facts.values())
+                + list(self.knowledge.values())
+                + list(self.actions.values())
+                + list(self.observations.values())
+            )
+            combined.sort(key=lambda x: x.created_at, reverse=True)
+            return combined[:n]
+
+    def get_goal_context(self, goal: str, max_items: int = 15) -> str:
+        """Get compact memory context relevant to a specific goal."""
+        if not goal:
+            return self.get_context_summary(max_items=max_items)
+
+        lines: List[str] = []
+
+        # Most important user facts first
+        facts = self.get_user_facts()[:5]
+        if facts:
+            lines.append("Known facts:")
+            for item in facts:
+                lines.append(f"- {item.content[:180]}")
+
+        # Relevant memories by search
+        relevant = self.search(query=goal, limit=max_items)
+        if relevant:
+            lines.append("Relevant memory:")
+            for item in relevant:
+                lines.append(f"- [{item.memory_type}] {item.content[:200]}")
+
+        return "\n".join(lines).strip()
     
     # ===========================================
     # RECALL & SEARCH
@@ -498,11 +498,16 @@ class UnifiedMemorySystem:
                         continue
                     
                     results.append(item)
-                    item.access()
             
             # Sort by relevance
             results.sort(key=lambda x: x.relevance_score, reverse=True)
-            return results[:limit]
+            top_results = results[:limit]
+            
+            # Update access tracking for returned items
+            for item in top_results:
+                item.access()
+                
+            return top_results
     
     def get_user_facts(self, category: str = None) -> List[MemoryItem]:
         """Get all user facts, optionally filtered by category."""