quantalogic 0.35.0__py3-none-any.whl → 0.40.0__py3-none-any.whl

quantalogic/agent.py

@@ -1,5 +1,6 @@
 """Enhanced QuantaLogic agent implementing the ReAct framework."""
 
+import asyncio
 from collections.abc import Callable
 from datetime import datetime
 from typing import Any
@@ -28,6 +29,9 @@ MAX_RESPONSE_LENGTH = 1024 * 32
 DEFAULT_MAX_INPUT_TOKENS = 128 * 1024
 DEFAULT_MAX_OUTPUT_TOKENS = 4096
 
+# Maximum recursion depth for variable interpolation
+MAX_INTERPOLATION_DEPTH = 10
+
 
 class AgentConfig(BaseModel):
     """Configuration settings for the Agent."""
@@ -50,7 +54,12 @@ class ObserveResponseResult(BaseModel):
 
 
 class Agent(BaseModel):
-    """Enhanced QuantaLogic agent implementing ReAct framework."""
+    """Enhanced QuantaLogic agent implementing ReAct framework.
+
+    Supports both synchronous and asynchronous operations for task solving.
+    Use `solve_task` for synchronous contexts (e.g., CLI tools) and `async_solve_task`
+    for asynchronous contexts (e.g., web servers).
+    """
 
     model_config = ConfigDict(arbitrary_types_allowed=True, validate_assignment=True, extra="forbid")
 
@@ -87,13 +96,29 @@ class Agent(BaseModel):
         get_environment: Callable[[], str] = get_environment,
         compact_every_n_iterations: int | None = None,
         max_tokens_working_memory: int | None = None,
+        event_emitter: EventEmitter | None = None,
     ):
-        """Initialize the agent with model, memory, tools, and configurations."""
+        """Initialize the agent with model, memory, tools, and configurations.
+        
+        Args:
+            model_name: Name of the model to use
+            memory: AgentMemory instance for storing conversation history
+            variable_store: VariableMemory instance for storing variables
+            tools: List of Tool instances 
+            ask_for_user_validation: Function to ask for user validation
+            task_to_solve: Initial task to solve
+            specific_expertise: Description of the agent's expertise
+            get_environment: Function to get environment details
+            compact_every_n_iterations: How often to compact memory
+            max_tokens_working_memory: Maximum token count for working memory
+            event_emitter: EventEmitter instance for event handling
+        """
         try:
             logger.debug("Initializing agent...")
 
-            # Create event emitter
-            event_emitter = EventEmitter()
+            # Create or use provided event emitter
+            if event_emitter is None:
+                event_emitter = EventEmitter()
 
             # Add TaskCompleteTool to the tools list if not already present
             if not any(isinstance(t, TaskCompleteTool) for t in tools):
@@ -155,173 +180,345 @@ class Agent(BaseModel):
 
     @model_name.setter
     def model_name(self, value: str) -> None:
-        """Set the model name."""
+        """Set the model name and update the model instance."""
         self._model_name = value
         # Update the model instance with the new name
         self.model = GenerativeModel(model=value, event_emitter=self.event_emitter)
 
-    def clear_memory(self):
+    def clear_memory(self) -> None:
         """Clear the memory and reset the session."""
         self._reset_session(clear_memory=True)
 
     def solve_task(
         self, task: str, max_iterations: int = 30, streaming: bool = False, clear_memory: bool = True
     ) -> str:
-        """Solve the given task using the ReAct framework.
+        """Solve the given task using the ReAct framework (synchronous version).
+
+        Ideal for synchronous applications. For asynchronous contexts, use `async_solve_task`.
 
         Args:
-            task (str): The task description.
-            max_iterations (int, optional): Maximum number of iterations to attempt solving the task.
-                Defaults to 30 to prevent infinite loops and ensure timely task completion.
-            streaming (bool, optional): Whether to use streaming mode for generating responses.
-            clear_memory (bool, optional): Whether to clear the memory before solving the task.
+            task: The task description
+            max_iterations: Maximum number of iterations
+            streaming: Whether to use streaming mode
+            clear_memory: Whether to clear memory before solving
 
         Returns:
-            str: The final response after task completion.
+            The final response after task completion
         """
         logger.debug(f"Solving task... {task}")
-        self._reset_session(task_to_solve=task, max_iterations=max_iterations, clear_memory=clear_memory)
+        try:
+            loop = asyncio.get_event_loop()
+        except RuntimeError:
+            # Create a new event loop if one doesn't exist
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
 
-        # Generate task summary
-        self.task_to_solve_summary = self._generate_task_summary(task)
+        return loop.run_until_complete(self.async_solve_task(task, max_iterations, streaming, clear_memory))
+
+    async def async_solve_task(
+        self, task: str, max_iterations: int = 30, streaming: bool = False, clear_memory: bool = True
+    ) -> str:
+        """Solve the given task using the ReAct framework (asynchronous version).
+
+        Ideal for asynchronous applications. For synchronous contexts, use `solve_task`.
+
+        Args:
+            task: The task description
+            max_iterations: Maximum number of iterations
+            streaming: Whether to use streaming mode
+            clear_memory: Whether to clear memory before solving
+
+        Returns:
+            The final response after task completion
+        """
+        logger.debug(f"Solving task asynchronously... {task}")
+        self._reset_session(task_to_solve=task, max_iterations=max_iterations, clear_memory=clear_memory)
+        self.task_to_solve_summary = await self._async_generate_task_summary(task)
 
-        # Add system prompt to memory
-        # Check if system prompt is already in memory
-        # if not add it
-        # The system message is always the first message in memory
         if not self.memory.memory or self.memory.memory[0].role != "system":
             self.memory.add(Message(role="system", content=self.config.system_prompt))
 
-        self._emit_event(
-            "session_start",
-            {"system_prompt": self.config.system_prompt, "content": task},
-        )
+        self._emit_event("session_start", {"system_prompt": self.config.system_prompt, "content": task})
 
         self.max_output_tokens = self.model.get_model_max_output_tokens() or DEFAULT_MAX_OUTPUT_TOKENS
         self.max_input_tokens = self.model.get_model_max_input_tokens() or DEFAULT_MAX_INPUT_TOKENS
 
         done = False
         current_prompt = self._prepare_prompt_task(task)
-
         self.current_iteration = 1
-
-        # Emit event: Task Solve Start
-        self._emit_event(
-            "task_solve_start",
-            {"initial_prompt": current_prompt, "task": task},
-        )
-
-        answer: str = ""
+        answer = ""
 
         while not done:
             try:
-                self._update_total_tokens(message_history=self.memory.memory, prompt=current_prompt)
-
-                # Emit event: Task Think Start after updating total tokens
+                self._update_total_tokens(self.memory.memory, current_prompt)
                 self._emit_event("task_think_start", {"prompt": current_prompt})
-
-                self._compact_memory_if_needed(current_prompt)
+                await self._async_compact_memory_if_needed(current_prompt)
 
                 if streaming:
-                    # For streaming, collect the response chunks
                     content = ""
-                    for chunk in self.model.generate_with_history(
+                    async_stream = await self.model.async_generate_with_history(
                         messages_history=self.memory.memory,
                         prompt=current_prompt,
                         streaming=True,
-                    ):
+                    )
+                    async for chunk in async_stream:
                         content += chunk
-
-                    # Create a response object similar to non-streaming mode
                     result = ResponseStats(
                         response=content,
-                        usage=TokenUsage(
-                            prompt_tokens=0,  # We don't have token counts in streaming mode
-                            completion_tokens=0,
-                            total_tokens=0,
-                        ),
+                        usage=TokenUsage(prompt_tokens=0, completion_tokens=0, total_tokens=0),
                         model=self.model.model,
                         finish_reason="stop",
                     )
                 else:
-                    result = self.model.generate_with_history(
-                        messages_history=self.memory.memory, prompt=current_prompt, streaming=False,
-                        stop_words=["thinking"]
+                    result = await self.model.async_generate_with_history(
+                        messages_history=self.memory.memory,
+                        prompt=current_prompt,
+                        streaming=False,
+                        stop_words=["thinking"],
                     )
 
                 content = result.response
-                if not streaming:  # Only update tokens for non-streaming mode
+                if not streaming:
                     token_usage = result.usage
                     self.total_tokens = token_usage.total_tokens
 
-                # Emit event: Task Think End
-                self._emit_event(
-                    "task_think_end",
-                    {
-                        "response": content,
-                    },
-                )
-
-                # Process the assistant's response
-                result = self._observe_response(content, iteration=self.current_iteration)
-
+                self._emit_event("task_think_end", {"response": content})
+                result = await self._async_observe_response(content, iteration=self.current_iteration)
                 current_prompt = result.next_prompt
 
                 if result.executed_tool == "task_complete":
-                    self._emit_event(
-                        "task_complete",
-                        {
-                            "response": result.answer,
-                        },
-                    )
-                    answer = result.answer
+                    self._emit_event("task_complete", {"response": result.answer})
+                    answer = result.answer or ""  # Ensure answer is never None
                     done = True
 
                 self._update_session_memory(current_prompt, content)
-
                 self.current_iteration += 1
                 if self.current_iteration >= self.max_iterations:
                     done = True
                     self._emit_event("error_max_iterations_reached")
 
             except Exception as e:
-                logger.error(f"Error during task solving: {str(e)}")
-                # Optionally, decide to continue or break based on exception type
+                logger.error(f"Error during async task solving: {str(e)}")
                 answer = f"Error: {str(e)}"
                 done = True
 
-        # Emit event: Task Solve End
         self._emit_event("task_solve_end")
+        return answer
 
-        logger.debug(f"Task solved: {answer}")
+    def _observe_response(self, content: str, iteration: int = 1) -> ObserveResponseResult:
+        """Analyze the assistant's response and determine next steps (synchronous wrapper).
+        
+        Args:
+            content: The response content to analyze
+            iteration: Current iteration number
+            
+        Returns:
+            ObserveResponseResult with next steps information
+        """
+        try:
+            loop = asyncio.get_event_loop()
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
 
-        return answer
+        return loop.run_until_complete(self._async_observe_response(content, iteration))
 
-    def _reset_session(self, task_to_solve: str = "", max_iterations: int = 30, clear_memory: bool = True):
-        """Reset the agent's session."""
-        logger.debug("Resetting session...")
-        self.task_to_solve = task_to_solve
-        if clear_memory:
-            logger.debug("Clearing memory...")
-            self.memory.reset()
-            self.variable_store.reset()
-            self.total_tokens = 0
-        self.current_iteration = 0
-        self.max_output_tokens = self.model.get_model_max_output_tokens() or DEFAULT_MAX_OUTPUT_TOKENS
-        self.max_input_tokens = self.model.get_model_max_input_tokens() or DEFAULT_MAX_INPUT_TOKENS
-        self.max_iterations = max_iterations
+    async def _async_observe_response(self, content: str, iteration: int = 1) -> ObserveResponseResult:
+        """Analyze the assistant's response and determine next steps (asynchronous).
+        
+        Args:
+            content: The response content to analyze
+            iteration: Current iteration number
+            
+        Returns:
+            ObserveResponseResult with next steps information
+        """
+        try:
+            parsed_content = self._parse_tool_usage(content)
+            if not parsed_content:
+                return self._handle_no_tool_usage()
 
-    def _update_total_tokens(self, message_history: list[Message], prompt: str) -> None:
-        self.total_tokens = self.model.token_counter_with_history(message_history, prompt)
+            for tool_name, tool_input in parsed_content.items():
+                tool = self.tools.get(tool_name)
+                if not tool:
+                    return self._handle_tool_not_found(tool_name)
 
-    def _compact_memory_if_needed(self, current_prompt: str = ""):
-        """Compacts the memory if it exceeds the maximum occupancy or token limit."""
+                arguments_with_values = self._parse_tool_arguments(tool, tool_input)
+                is_repeated_call = self._is_repeated_tool_call(tool_name, arguments_with_values)
+
+                if is_repeated_call:
+                    executed_tool, response = self._handle_repeated_tool_call(tool_name, arguments_with_values)
+                else:
+                    executed_tool, response = await self._async_execute_tool(tool_name, tool, arguments_with_values)
+
+                if not executed_tool:
+                    return self._handle_tool_execution_failure(response)
+
+                variable_name = self.variable_store.add(response)
+                new_prompt = self._format_observation_response(response, executed_tool, variable_name, iteration)
+
+                return ObserveResponseResult(
+                    next_prompt=new_prompt,
+                    executed_tool=executed_tool,
+                    answer=response if executed_tool == "task_complete" else None,
+                )
+        except Exception as e:
+            return self._handle_error(e)
+
+    def _execute_tool(self, tool_name: str, tool: Tool, arguments_with_values: dict) -> tuple[str, Any]:
+        """Execute a tool with validation if required (synchronous wrapper).
+        
+        Args:
+            tool_name: Name of the tool to execute
+            tool: Tool instance
+            arguments_with_values: Tool arguments
+            
+        Returns:
+            Tuple of (executed_tool_name, response)
+        """
+        try:
+            loop = asyncio.get_event_loop()
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+
+        return loop.run_until_complete(self._async_execute_tool(tool_name, tool, arguments_with_values))
+
+    async def _async_execute_tool(self, tool_name: str, tool: Tool, arguments_with_values: dict) -> tuple[str, Any]:
+        """Execute a tool with validation if required (asynchronous).
+        
+        Args:
+            tool_name: Name of the tool to execute
+            tool: Tool instance
+            arguments_with_values: Tool arguments
+            
+        Returns:
+            Tuple of (executed_tool_name, response)
+        """
+        if tool.need_validation:
+            question_validation = (
+                "Do you permit the execution of this tool?\n"
+                f"Tool: {tool_name}\nArguments:\n"
+                "<arguments>\n"
+                + "\n".join([f"    <{key}>{value}</{key}>" for key, value in arguments_with_values.items()])
+                + "\n</arguments>\nYes or No"
+            )
+            permission_granted = self.ask_for_user_validation(question_validation)
+            if not permission_granted:
+                return "", f"Error: execution of tool '{tool_name}' was denied by the user."
+
+        self._emit_event("tool_execution_start", {"tool_name": tool_name, "arguments": arguments_with_values})
+
+        try:
+            arguments_with_values_interpolated = {
+                key: await self._async_interpolate_variables(value) for key, value in arguments_with_values.items()
+            }
+            if tool.need_variables:
+                arguments_with_values_interpolated["variables"] = self.variable_store
+            if tool.need_caller_context_memory:
+                arguments_with_values_interpolated["caller_context_memory"] = self.memory.memory
+
+            converted_args = self.tools.validate_and_convert_arguments(tool_name, arguments_with_values_interpolated)
+            injectable_properties = tool.get_injectable_properties_in_execution()
+            for key, value in injectable_properties.items():
+                converted_args[key] = value
+
+            if hasattr(tool, "async_execute") and callable(tool.async_execute):
+                response = await tool.async_execute(**converted_args)
+            else:
+                # Fall back to synchronous execution if async is not available
+                response = tool.execute(**converted_args)
+            executed_tool = tool.name
+        except Exception as e:
+            response = f"Error executing tool: {tool_name}: {str(e)}\n"
+            executed_tool = ""
+
+        self._emit_event(
+            "tool_execution_end", {"tool_name": tool_name, "arguments": arguments_with_values, "response": response}
+        )
+        return executed_tool, response
+
+    async def _async_interpolate_variables(self, text: str, depth: int = 0) -> str:
+        """Interpolate variables using $var$ syntax in the given text with recursion protection.
+        
+        Args:
+            text: Text containing variable references
+            depth: Current recursion depth
+            
+        Returns:
+            Text with variables interpolated
+        """
+        if not isinstance(text, str):
+            return str(text)
+            
+        if depth > MAX_INTERPOLATION_DEPTH:
+            logger.warning(f"Max interpolation depth ({MAX_INTERPOLATION_DEPTH}) reached, stopping recursion")
+            return text
+            
+        try:
+            import re
+            
+            # Process each variable in the store
+            for var in self.variable_store.keys():
+                # Properly escape the variable name for regex using re.escape
+                # but handle $ characters separately since they're part of our syntax
+                escaped_var = re.escape(var).replace('\\$', '$')
+                pattern = f"\\${escaped_var}\\$"
+                
+                # Get variable value as string
+                replacement = str(self.variable_store[var])
+                
+                # Replace all occurrences
+                text = re.sub(pattern, lambda m: replacement, text)
+                
+            # Check if there are still variables to interpolate (for nested variables)
+            if '$' in text and depth < MAX_INTERPOLATION_DEPTH:
+                return await self._async_interpolate_variables(text, depth + 1)
+                
+            return text
+        except Exception as e:
+            logger.error(f"Error in _async_interpolate_variables: {str(e)}")
+            return text
+
+    def _interpolate_variables(self, text: str) -> str:
+        """Interpolate variables using $var$ syntax in the given text (synchronous wrapper).
+        
+        Args:
+            text: Text containing variable references
+            
+        Returns:
+            Text with variables interpolated
+        """
+        try:
+            loop = asyncio.get_event_loop()
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+
+        return loop.run_until_complete(self._async_interpolate_variables(text))
+
+    def _compact_memory_if_needed(self, current_prompt: str = "") -> None:
+        """Compacts the memory if it exceeds the maximum occupancy (synchronous wrapper).
+        
+        Args:
+            current_prompt: Current prompt to calculate token usage
+        """
+        try:
+            loop = asyncio.get_event_loop()
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+
+        return loop.run_until_complete(self._async_compact_memory_if_needed(current_prompt))
+
+    async def _async_compact_memory_if_needed(self, current_prompt: str = "") -> None:
+        """Compacts the memory if it exceeds the maximum occupancy or token limit.
+        
+        Args:
+            current_prompt: Current prompt to calculate token usage
+        """
         ratio_occupied = self._calculate_context_occupancy()
 
-        # Compact memory if any of these conditions are met:
-        # 1. Memory occupancy exceeds MAX_OCCUPANCY, or
-        # 2. Current iteration is a multiple of compact_every_n_iterations, or
-        # 3. Working memory exceeds max_tokens_working_memory (if set)
         should_compact_by_occupancy = ratio_occupied >= MAX_OCCUPANCY
         should_compact_by_iteration = (
             self.compact_every_n_iterations is not None
@@ -329,7 +526,8 @@ class Agent(BaseModel):
             and self.current_iteration % self.compact_every_n_iterations == 0
         )
         should_compact_by_token_limit = (
-            self.max_tokens_working_memory is not None and self.total_tokens > self.max_tokens_working_memory
+            self.max_tokens_working_memory is not None 
+            and self.total_tokens > self.max_tokens_working_memory
         )
 
         if should_compact_by_occupancy or should_compact_by_iteration or should_compact_by_token_limit:
@@ -341,83 +539,162 @@ class Agent(BaseModel):
                     f"Memory compaction triggered: Iteration {self.current_iteration} is a multiple of {self.compact_every_n_iterations}"
                 )
 
+            if should_compact_by_token_limit:
+                logger.debug(
+                    f"Memory compaction triggered: Token count {self.total_tokens} exceeds limit {self.max_tokens_working_memory}"
+                )
+
             self._emit_event("memory_full")
-            self.memory.compact()
+            await self._async_compact_memory()
             self.total_tokens = self.model.token_counter_with_history(self.memory.memory, current_prompt)
             self._emit_event("memory_compacted")
 
-    def _emit_event(self, event_type: str, data: dict[str, Any] | None = None) -> None:
-        """
-        Emit an event with system context and optional additional data.
+    async def _async_compact_memory(self) -> None:
+        """Compact memory asynchronously."""
+        self.memory.compact()
 
-        Why: Provides a standardized way to track and log system events
-        with consistent contextual information.
+    async def _async_compact_memory_with_summary(self) -> str:
+        """Generate a summary and compact memory asynchronously.
+        
+        Returns:
+            Generated summary text
         """
-        # Use empty dict as default to avoid mutable default argument
-        event_data = {
-            "iteration": self.current_iteration,
-            "total_tokens": self.total_tokens,
-            "context_occupancy": self._calculate_context_occupancy(),
-            "max_input_tokens": self.max_input_tokens,
-            "max_output_tokens": self.max_output_tokens,
-        }
+        prompt_summary = (
+            "Summarize the conversation concisely:\n"
+            "format in markdown:\n"
+            "<thinking>\n"
+            " - 1. **Completed Steps**: Briefly describe the steps.\n"
+            " - 2. **Variables Used**: List the variables.\n"
+            " - 3. **Progress Analysis**: Assess progress.\n"
+            "</thinking>\n"
+            "Keep the summary clear and actionable.\n"
+        )
 
-        # Merge additional data if provided
-        if data:
-            event_data.update(data)
+        memory_copy = self.memory.memory.copy()
 
-        self.event_emitter.emit(event_type, event_data)
+        if len(memory_copy) < 3:
+            logger.warning("Not enough messages to compact memory with summary")
+            return "Memory compaction skipped: not enough messages"
 
-    def _observe_response(self, content: str, iteration: int = 1) -> ObserveResponseResult:
-        """Analyze the assistant's response and determine next steps.
+        user_message = memory_copy.pop()
+        assistant_message = memory_copy.pop()
+        summary = await self.model.async_generate_with_history(messages_history=memory_copy, prompt=prompt_summary)
+        
+        # Remove last system message if present
+        if memory_copy and memory_copy[-1].role == "system":
+            memory_copy.pop()
+            
+        memory_copy.append(Message(role="user", content=summary.response))
+        memory_copy.append(assistant_message)
+        memory_copy.append(user_message)
+        self.memory.memory = memory_copy
+        return summary.response
 
+    def _generate_task_summary(self, content: str) -> str:
+        """Generate a concise task-focused summary (synchronous wrapper).
+        
         Args:
-            content (str): The assistant's response content.
-            iteration (int, optional): The current iteration number of task solving.
-                Helps track the progress and prevent infinite loops. Defaults to 1.
-
+            content: The content to summarize
+            
         Returns:
-            ObserveResponseResult: A result indicating if the task is done and the next prompt.
+            Generated task summary
         """
         try:
-            parsed_content = self._parse_tool_usage(content)
-            if not parsed_content:
-                return self._handle_no_tool_usage()
+            loop = asyncio.get_event_loop()
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
 
-            for tool_name, tool_input in parsed_content.items():
-                tool = self.tools.get(tool_name)
-                if not tool:
-                    return self._handle_tool_not_found(tool_name)
+        return loop.run_until_complete(self._async_generate_task_summary(content))
 
-                arguments_with_values = self._parse_tool_arguments(tool, tool_input)
-                is_repeated_call = self._is_repeated_tool_call(tool_name, arguments_with_values)
+    async def _async_generate_task_summary(self, content: str) -> str:
+        """Generate a concise task-focused summary using the generative model.
 
-                if is_repeated_call:
-                    executed_tool, response = self._handle_repeated_tool_call(tool_name, arguments_with_values)
-                else:
-                    executed_tool, response = self._execute_tool(tool_name, tool, arguments_with_values)
+        Args:
+            content: The content to summarize
 
-                if not executed_tool:
-                    return self._handle_tool_execution_failure(response)
+        Returns:
+            Generated task summary
+        """
+        try:
+            if len(content) < 1024 * 4:
+                return content
+                
+            prompt = (
+                "Create a task summary that captures ONLY: \n"
+                "1. Primary objective/purpose\n"
+                "2. Core actions/requirements\n"
+                "3. Desired end-state/outcome\n\n"
+                "Guidelines:\n"
+                "- Use imperative voice\n"
+                f"Input Task Description:\n{content}\n\n"
+            )
+            result = await self.model.async_generate(prompt=prompt)
+            logger.debug(f"Generated summary: {result.response}")
+            return result.response.strip() + "\n🚨 The FULL task is in <task> tag in the previous messages.\n"
+        except Exception as e:
+            logger.error(f"Error generating summary: {str(e)}")
+            return f"Summary generation failed: {str(e)}"
 
-                variable_name = self.variable_store.add(response)
-                new_prompt = self._format_observation_response(response, executed_tool, variable_name, iteration)
+    def _reset_session(self, task_to_solve: str = "", max_iterations: int = 30, clear_memory: bool = True) -> None:
+        """Reset the agent's session.
+        
+        Args:
+            task_to_solve: New task to solve
+            max_iterations: Maximum number of iterations
+            clear_memory: Whether to clear memory
+        """
+        logger.debug("Resetting session...")
+        self.task_to_solve = task_to_solve
+        if clear_memory:
+            logger.debug("Clearing memory...")
+            self.memory.reset()
+            self.variable_store.reset()
+            self.total_tokens = 0
+        self.current_iteration = 0
+        self.max_output_tokens = self.model.get_model_max_output_tokens() or DEFAULT_MAX_OUTPUT_TOKENS
+        self.max_input_tokens = self.model.get_model_max_input_tokens() or DEFAULT_MAX_INPUT_TOKENS
+        self.max_iterations = max_iterations
 
-                return ObserveResponseResult(
-                    next_prompt=new_prompt,
-                    executed_tool=executed_tool,
-                    answer=response if executed_tool == "task_complete" else None,
-                )
+    def _update_total_tokens(self, message_history: list[Message], prompt: str) -> None:
+        """Update the total tokens count based on message history and prompt.
+        
+        Args:
+            message_history: List of messages
+            prompt: Current prompt
+        """
+        self.total_tokens = self.model.token_counter_with_history(message_history, prompt)
 
-        except Exception as e:
-            return self._handle_error(e)
+    def _emit_event(self, event_type: str, data: dict[str, Any] | None = None) -> None:
+        """Emit an event with system context and optional additional data.
+        
+        Args:
+            event_type: Type of event
+            data: Additional event data
+        """
+        event_data = {
+            "iteration": self.current_iteration,
+            "total_tokens": self.total_tokens,
+            "context_occupancy": self._calculate_context_occupancy(),
+            "max_input_tokens": self.max_input_tokens,
+            "max_output_tokens": self.max_output_tokens,
+        }
+        if data:
+            event_data.update(data)
+        self.event_emitter.emit(event_type, event_data)
 
     def _parse_tool_usage(self, content: str) -> dict:
-        """Extract tool usage from the response content."""
+        """Extract tool usage from the response content.
+        
+        Args:
+            content: Response content
+            
+        Returns:
+            Dictionary mapping tool names to inputs
+        """
         if not content or not isinstance(content, str):
             return {}
 
-        # Extract action 
         xml_parser = ToleranceXMLParser()
         action = xml_parser.extract_elements(text=content, element_names=["action"])
 
@@ -426,17 +703,31 @@ class Agent(BaseModel):
         if action:
             return xml_parser.extract_elements(text=action["action"], element_names=tool_names)
         else:
-            # Fallback to extracting tool usage directly
             return xml_parser.extract_elements(text=content, element_names=tool_names)
 
-
-    def _parse_tool_arguments(self, tool, tool_input: str) -> dict:
-        """Parse the tool arguments from the tool input."""
+    def _parse_tool_arguments(self, tool: Tool, tool_input: str) -> dict:
+        """Parse the tool arguments from the tool input.
+        
+        Args:
+            tool: Tool instance
+            tool_input: Raw tool input text
+            
+        Returns:
+            Dictionary of parsed arguments
+        """
         tool_parser = ToolParser(tool=tool)
         return tool_parser.parse(tool_input)
 
     def _is_repeated_tool_call(self, tool_name: str, arguments_with_values: dict) -> bool:
-        """Check if the tool call is repeated."""
+        """Check if the tool call is repeated.
+        
+        Args:
+            tool_name: Name of the tool
+            arguments_with_values: Tool arguments
+            
+        Returns:
+            True if call is repeated, False otherwise
+        """
         current_call = {
             "tool_name": tool_name,
             "arguments": arguments_with_values,
@@ -455,16 +746,27 @@ class Agent(BaseModel):
             current_call["count"] = 1
 
         self.last_tool_call = current_call
-        return is_repeated_call and repeat_count >= 2
+        return is_repeated_call and current_call.get("count", 0) >= 2
 
     def _handle_no_tool_usage(self) -> ObserveResponseResult:
-        """Handle the case where no tool usage is found in the response."""
+        """Handle the case where no tool usage is found in the response.
+        
+        Returns:
+            ObserveResponseResult with error message
+        """
         return ObserveResponseResult(
             next_prompt="Error: No tool usage found in response.", executed_tool=None, answer=None
         )
 
     def _handle_tool_not_found(self, tool_name: str) -> ObserveResponseResult:
-        """Handle the case where the tool is not found."""
+        """Handle the case where the tool is not found.
+        
+        Args:
+            tool_name: Name of the tool
+            
+        Returns:
+            ObserveResponseResult with error message
+        """
         logger.warning(f"Tool '{tool_name}' not found in tool manager.")
         return ObserveResponseResult(
             next_prompt=f"Error: Tool '{tool_name}' not found in tool manager.",
@@ -472,8 +774,16 @@ class Agent(BaseModel):
             answer=None,
         )
 
-    def _handle_repeated_tool_call(self, tool_name: str, arguments_with_values: dict) -> (str, str):
-        """Handle the case where a tool call is repeated."""
+    def _handle_repeated_tool_call(self, tool_name: str, arguments_with_values: dict) -> tuple[str, str]:
+        """Handle the case where a tool call is repeated.
+        
+        Args:
+            tool_name: Name of the tool
+            arguments_with_values: Tool arguments
+            
+        Returns:
+            Tuple of (executed_tool_name, error_message)
+        """
         repeat_count = self.last_tool_call.get("count", 0)
         error_message = (
             "Error: Detected repeated identical tool call pattern.\n"
@@ -489,7 +799,14 @@ class Agent(BaseModel):
         return tool_name, error_message
 
     def _handle_tool_execution_failure(self, response: str) -> ObserveResponseResult:
-        """Handle the case where tool execution fails."""
+        """Handle the case where tool execution fails.
+        
+        Args:
+            response: Error response
+            
+        Returns:
+            ObserveResponseResult with error message
+        """
         return ObserveResponseResult(
             next_prompt=response,
             executed_tool="",
@@ -497,7 +814,14 @@ class Agent(BaseModel):
         )
 
     def _handle_error(self, error: Exception) -> ObserveResponseResult:
-        """Handle any exceptions that occur during response observation."""
+        """Handle any exceptions that occur during response observation.
+        
+        Args:
+            error: Exception that occurred
+            
+        Returns:
+            ObserveResponseResult with error message
+        """
         logger.error(f"Error in _observe_response: {str(error)}")
         return ObserveResponseResult(
             next_prompt=f"An error occurred while processing the response: {str(error)}",
@@ -506,9 +830,19 @@ class Agent(BaseModel):
         )
 
     def _format_observation_response(
-        self, response: str, last_exectured_tool: str, variable_name: str, iteration: int
+        self, response: str, last_executed_tool: str, variable_name: str, iteration: int
     ) -> str:
-        """Format the observation response with the given response, variable name, and iteration."""
+        """Format the observation response with the given response, variable name, and iteration.
+        
+        Args:
+            response: Tool execution response
+            last_executed_tool: Name of last executed tool
+            variable_name: Name of variable storing response
+            iteration: Current iteration number
+            
+        Returns:
+            Formatted observation response
+        """
         response_display = response
         if len(response) > MAX_RESPONSE_LENGTH:
             response_display = response[:MAX_RESPONSE_LENGTH]
@@ -516,8 +850,7 @@ class Agent(BaseModel):
                 f"... content was truncated full content available by interpolation in variable {variable_name}"
             )
 
-        # Format the response message
-        formatted_response = formatted_response = (
+        formatted_response = (
             "# Analysis and Next Action Decision Point\n\n"
             f"📊 Progress: Iteration {iteration}/{self.max_iterations}\n\n"
             "## Global Task summary:\n"
@@ -535,7 +868,7 @@ class Agent(BaseModel):
             "1. Your analysis of the progression resulting from the execution of the tool in <thinking> tags, don't include <context_analysis/>\n"
             "2. Your tool execution plan in <tool_name> tags\n\n"
             "## Last executed action result\n"
-            f"Last executed tool {last_exectured_tool} Execution Result:\n"
+            f"Last executed tool {last_executed_tool} Execution Result:\n"
             f"\n<{variable_name}>\n{response_display}\n</{variable_name}>\n"
             "## Response Format\n"
             "```xml\n"
@@ -561,132 +894,14 @@ class Agent(BaseModel):
 
         return formatted_response
 
-    def _execute_tool(self, tool_name: str, tool, arguments_with_values: dict) -> tuple[str, Any]:
-        """Execute a tool with validation if required.
-
-        Args:
-            tool_name: Name of the tool to execute
-            tool: Tool instance
-            arguments_with_values: Dictionary of argument names and values
-
-        Returns:
-            tuple containing:
-                - executed_tool name (str)
-                - tool execution response (Any)
-
-        Note:
-            Predefined variable properties take precedence over dynamically provided values.
-            This ensures consistent behavior when tools have both predefined properties
-            and runtime-provided arguments. The precedence order is:
-            1. Predefined properties from tool.get_injectable_properties_in_execution()
-            2. Runtime-provided arguments from arguments_with_values
-        """
-        # Handle tool validation if required
-        if tool.need_validation:
-            logger.debug(f"Tool '{tool_name}' requires validation.")
-            self._emit_event(
-                "tool_execute_validation_start",
-                {"tool_name": tool_name, "arguments": arguments_with_values},
-            )
-
-            question_validation: str = (
-                "Do you permit the execution of this tool?\n"
-                f"Tool: {tool_name}\n"
-                "Arguments:\n"
-                "<arguments>\n"
-                + "\n".join([f"    <{key}>{value}</{key}>" for key, value in arguments_with_values.items()])
-                + "\n</arguments>\n"
-                "Yes or No"
-            )
-            permission_granted = self.ask_for_user_validation(question_validation)
-
-            self._emit_event(
-                "tool_execute_validation_end",
-                {"tool_name": tool_name, "arguments": arguments_with_values},
-            )
-
-            if not permission_granted:
-                logger.debug(f"Execution of tool '{tool_name}' was denied by the user.")
-                return "", f"Error: execution of tool '{tool_name}' was denied by the user."
-
-        # Emit event: Tool Execution Start
-        self._emit_event(
-            "tool_execution_start",
-            {"tool_name": tool_name, "arguments": arguments_with_values},
-        )
-
-        try:
-            # Execute the tool synchronously
-            arguments_with_values_interpolated = {
-                key: self._interpolate_variables(value) for key, value in arguments_with_values.items()
-            }
-
-            arguments_with_values_interpolated = arguments_with_values_interpolated
-
-            # test if tool need variables in context
-            if tool.need_variables:
-                # Inject variables into the tool if needed
-                arguments_with_values_interpolated["variables"] = self.variable_store
-            if tool.need_caller_context_memory:
-                # Inject caller context into the tool if needed
-                arguments_with_values_interpolated["caller_context_memory"] = self.memory.memory
-
-            try:
-                # Convert arguments to proper types
-                converted_args = self.tools.validate_and_convert_arguments(
-                    tool_name, arguments_with_values_interpolated
-                )
-            except ValueError as e:
-                return "", f"Argument Error: {str(e)}"
-
-            # Add injectable variables
-            injectable_properties = tool.get_injectable_properties_in_execution()
-            for key, value in injectable_properties.items():
-                converted_args[key] = value
-
-            # Call tool execute with named arguments
-            response = tool.execute(**converted_args)
-            executed_tool = tool.name
-        except Exception as e:
-            response = f"Error executing tool: {tool_name}: {str(e)}\n"
-            executed_tool = ""
-
-        # Emit event: Tool Execution End
-        self._emit_event(
-            "tool_execution_end",
-            {
-                "tool_name": tool_name,
-                "arguments": arguments_with_values,
-                "response": response,
-            },
-        )
-
-        return executed_tool, response
-
-    def _interpolate_variables(self, text: str) -> str:
-        """Interpolate variables using $var$ syntax in the given text."""
-        try:
-            import re
-
-            for var in self.variable_store.keys():
-                # Create safe pattern without double-escaping backslashes
-                safe_var = re.sub(r"([\\\.\^\$\*\+\?\{\}\[\]\|\(\)])", r"\\\1", var)
-                pattern = rf"\${safe_var}\$"
-                replacement = self.variable_store[var]
-                text = re.sub(pattern, replacement, text)
-            return text
-        except Exception as e:
-            logger.error(f"Error in _interpolate_variables: {str(e)}")
-            return text
-
     def _prepare_prompt_task(self, task: str) -> str:
         """Prepare the initial prompt for the task.
 
         Args:
-            task (str): The task description.
+            task: The task description
 
         Returns:
-            str: The formatted task prompt.
+            The formatted task prompt
         """
         prompt_task: str = (
             "## Your task to solve:\n"
@@ -702,7 +917,11 @@ class Agent(BaseModel):
         return prompt_task
 
     def _get_tools_names_prompt(self) -> str:
-        """Construct a detailed prompt that lists the available tools for task execution."""
+        """Construct a detailed prompt that lists the available tools for task execution.
+        
+        Returns:
+            Formatted tools prompt
+        """
         prompt_use_tools: str = (
             "To accomplish this task, you have access to these tools:\n"
             "\n"
@@ -712,15 +931,22 @@ class Agent(BaseModel):
             "1. Select ONE tool per message\n"
             "2. You will receive the tool's output in the next user response\n"
             "3. Choose the most appropriate tool for each step\n"
-            "4. Use task_complete tool to confirm task completion\n"
+            "4. If it's not asked to write on files, don't use write_file tool\n"
+            "5. If files are written, then use tool to display the prepared download link\n"
+            "6. Give the final full answer using all the variables\n"
+            "7. Use task_complete tool to confirm task completion with the full content of the final answer\n"
         )
         return prompt_use_tools
 
     def _get_variable_prompt(self) -> str:
-        """Construct a prompt that explains how to use variables."""
+        """Construct a prompt that explains how to use variables.
+        
+        Returns:
+            Formatted variables prompt
+        """
         prompt_use_variables: str = (
             "To use a variable interpolation, use the format $variable_name$ in function arguments.\n"
-            "Example: <write_file><file_path>/path/to/file.txt</file_path><content>$var1$</write_file>\n"
+            "Example: <write_file><file_path>/path/to/file.txt</file_path><content>$var1$</content></write_file>\n"
             "\n"
             "Available variables:\n"
             "\n"
@@ -731,96 +957,110 @@ class Agent(BaseModel):
         return prompt_use_variables
 
     def _calculate_context_occupancy(self) -> float:
-        """Calculate the number of tokens in percentages for prompt and completion."""
+        """Calculate the number of tokens in percentages for prompt and completion.
+        
+        Returns:
+            Percentage of context window occupied
+        """
         total_tokens = self.total_tokens
-        # Calculate token usage of prompt
         max_tokens = self.model.get_model_max_input_tokens()
 
-        # Handle None value and prevent division by zero
         if max_tokens is None or max_tokens <= 0:
             logger.warning(f"Invalid max tokens value: {max_tokens}. Using default of {DEFAULT_MAX_INPUT_TOKENS}.")
             max_tokens = DEFAULT_MAX_INPUT_TOKENS
 
         return round((total_tokens / max_tokens) * 100, 2)
 
-    def _compact_memory_with_summary(self) -> str:
-        prompt_summary = (
-            "Summarize the conversation concisely:\n"
-            "format in markdown:\n"
-            "<thinking>\n"
-            " - 1. **Completed Steps**: Briefly describe the steps.\n"
-            " - 2. **Variables Used**: List the variables.\n"
-            " - 3. **Progress Analysis**: Assess progress.\n"
-            "</thinking>\n"
-            "Keep the summary clear and actionable.\n"
-        )
-
-        # Get all message system, except the last assistant / user message
-        memory_copy = self.memory.memory.copy()
-
-        # Remove the last assistant / user message
-        user_message = memory_copy.pop()
-        assistant_message = memory_copy.pop()
-        summary = self.model.generate_with_history(messages_history=memory_copy, prompt=prompt_summary)
-        # Remove user message
-        memory_copy.pop()
-        # Replace by summary
-        memory_copy.append(Message(role="user", content=summary.response))
-        memory_copy.append(assistant_message)
-        memory_copy.append(user_message)
-        self.memory.memory = memory_copy
-        return summary.response
-
-    def _generate_task_summary(self, content: str) -> str:
-        """Generate a concise task-focused summary using the generative model.
-
-        Args:
-            content (str): The content to summarize
-
-        Returns:
-            str: Generated task summary
-        """
-        try:
-            if len(content) < 1024*4:
-                return content
-            prompt = (
-                "Create a task summary that captures ONLY: \n"
-                "1. Primary objective/purpose\n"
-                "2. Core actions/requirements\n"
-                "3. Desired end-state/outcome\n\n"
-                "Guidelines:\n"
-                "- Use imperative voice\n"
-                f"Input Task Description:\n{content}\n\n"
-            )
-            result = self.model.generate(prompt=prompt)
-            logger.debug(f"Generated summary: {result.response}")
-            return result.response.strip() + "\n🚨 The FULL task is in <task> tag in the previous messages.\n"
-        except Exception as e:
-            logger.error(f"Error generating summary: {str(e)}")
-            return f"Summary generation failed: {str(e)}"
-
     def _update_session_memory(self, user_content: str, assistant_content: str) -> None:
-        """
-        Log session messages to memory and emit events.
+        """Log session messages to memory and emit events.
 
         Args:
-            user_content (str): The user's content.
-            assistant_content (str): The assistant's content.
+            user_content: The user's content
+            assistant_content: The assistant's content
         """
         self.memory.add(Message(role="user", content=user_content))
-        self._emit_event(
-            "session_add_message",
-            {"role": "user", "content": user_content},
-        )
+        self._emit_event("session_add_message", {"role": "user", "content": user_content})
 
         self.memory.add(Message(role="assistant", content=assistant_content))
-
-        self._emit_event(
-            "session_add_message",
-            {"role": "assistant", "content": assistant_content},
-        )
+        self._emit_event("session_add_message", {"role": "assistant", "content": assistant_content})
 
     def update_model(self, new_model_name: str) -> None:
-        """Update the model name and recreate the model instance."""
+        """Update the model name and recreate the model instance.
+        
+        Args:
+            new_model_name: New model name to use
+        """
         self.model_name = new_model_name
         self.model = GenerativeModel(model=new_model_name, event_emitter=self.event_emitter)
+
+    def add_tool(self, tool: Tool) -> None:
+        """Add a new tool to the agent's tool manager.
+        
+        Args:
+            tool: The tool instance to add
+            
+        Raises:
+            ValueError: If a tool with the same name already exists
+        """
+        if tool.name in self.tools.tool_names():
+            raise ValueError(f"Tool with name '{tool.name}' already exists")
+            
+        self.tools.add(tool)
+        # Update tools markdown in config
+        self.config = AgentConfig(
+            environment_details=self.config.environment_details,
+            tools_markdown=self.tools.to_markdown(),
+            system_prompt=self.config.system_prompt,
+        )
+        logger.debug(f"Added tool: {tool.name}")
+        
+    def remove_tool(self, tool_name: str) -> None:
+        """Remove a tool from the agent's tool manager.
+        
+        Args:
+            tool_name: Name of the tool to remove
+            
+        Raises:
+            ValueError: If tool doesn't exist or is TaskCompleteTool
+        """
+        if tool_name not in self.tools.tool_names():
+            raise ValueError(f"Tool '{tool_name}' does not exist")
+            
+        tool = self.tools.get(tool_name)
+        if isinstance(tool, TaskCompleteTool):
+            raise ValueError("Cannot remove TaskCompleteTool as it is required")
+            
+        self.tools.remove(tool_name)
+        # Update tools markdown in config
+        self.config = AgentConfig(
+            environment_details=self.config.environment_details,
+            tools_markdown=self.tools.to_markdown(),
+            system_prompt=self.config.system_prompt,
+        )
+        logger.debug(f"Removed tool: {tool_name}")
+        
+    def set_tools(self, tools: list[Tool]) -> None:
+        """Set/replace all tools for the agent.
+        
+        Args:
+            tools: List of tool instances to set
+            
+        Note:
+            TaskCompleteTool will be automatically added if not present
+        """
+        # Ensure TaskCompleteTool is present
+        if not any(isinstance(t, TaskCompleteTool) for t in tools):
+            tools.append(TaskCompleteTool())
+            
+        # Create new tool manager and add tools
+        tool_manager = ToolManager()
+        tool_manager.add_list(tools)
+        self.tools = tool_manager
+        
+        # Update config with new tools markdown
+        self.config = AgentConfig(
+            environment_details=self.config.environment_details,
+            tools_markdown=self.tools.to_markdown(),
+            system_prompt=self.config.system_prompt,
+        )
+        logger.debug(f"Set {len(tools)} tools")