praisonaiagents 0.0.96__tar.gz → 0.0.97__tar.gz

{praisonaiagents-0.0.96 → praisonaiagents-0.0.97}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: praisonaiagents
-Version: 0.0.96
+Version: 0.0.97
 Summary: Praison AI agents for completing complex tasks with Self Reflection Agents
 Author: Mervin Praison
 Requires-Python: >=3.10
@@ -19,6 +19,9 @@ Requires-Dist: mem0ai>=0.1.0; extra == "knowledge"
 Requires-Dist: chromadb>=1.0.0; extra == "knowledge"
 Requires-Dist: markitdown[all]>=0.1.0; extra == "knowledge"
 Requires-Dist: chonkie>=1.0.2; extra == "knowledge"
+Provides-Extra: graph
+Requires-Dist: mem0ai[graph]>=0.1.0; extra == "graph"
+Requires-Dist: chromadb>=1.0.0; extra == "graph"
 Provides-Extra: llm
 Requires-Dist: litellm>=1.50.0; extra == "llm"
 Requires-Dist: pydantic>=2.4.2; extra == "llm"
@@ -28,6 +31,7 @@ Requires-Dist: uvicorn>=0.34.0; extra == "api"
 Provides-Extra: all
 Requires-Dist: praisonaiagents[memory]; extra == "all"
 Requires-Dist: praisonaiagents[knowledge]; extra == "all"
+Requires-Dist: praisonaiagents[graph]; extra == "all"
 Requires-Dist: praisonaiagents[llm]; extra == "all"
 Requires-Dist: praisonaiagents[mcp]; extra == "all"
 Requires-Dist: praisonaiagents[api]; extra == "all"

{praisonaiagents-0.0.96 → praisonaiagents-0.0.97}/praisonaiagents/__init__.py

@@ -11,6 +11,7 @@ from .agents.autoagents import AutoAgents
 from .knowledge.knowledge import Knowledge
 from .knowledge.chunking import Chunking
 from .mcp.mcp import MCP
+from .session import Session
 from .main import (
     TaskOutput,
     ReflectionOutput,
@@ -40,6 +41,7 @@ __all__ = [
     'TaskOutput',
     'ReflectionOutput',
     'AutoAgents',
+    'Session',
     'display_interaction',
     'display_self_reflection',
     'display_instruction',

{praisonaiagents-0.0.96 → praisonaiagents-0.0.97}/praisonaiagents/agent/agent.py

@@ -16,7 +16,8 @@ from ..main import (
     display_self_reflection,
     ReflectionOutput,
     client,
-    adisplay_instruction
+    adisplay_instruction,
+    approval_callback
 )
 import inspect
 import uuid
@@ -570,6 +571,35 @@ Your Goal: {self.goal}
         """
         logging.debug(f"{self.name} executing tool {function_name} with arguments: {arguments}")
 
+        # Check if approval is required for this tool
+        from ..approval import is_approval_required, console_approval_callback, get_risk_level, mark_approved, ApprovalDecision
+        if is_approval_required(function_name):
+            risk_level = get_risk_level(function_name)
+            logging.info(f"Tool {function_name} requires approval (risk level: {risk_level})")
+            
+            # Use global approval callback or default console callback
+            callback = approval_callback or console_approval_callback
+            
+            try:
+                decision = callback(function_name, arguments, risk_level)
+                if not decision.approved:
+                    error_msg = f"Tool execution denied: {decision.reason}"
+                    logging.warning(error_msg)
+                    return {"error": error_msg, "approval_denied": True}
+                
+                # Mark as approved in context to prevent double approval in decorator
+                mark_approved(function_name)
+                
+                # Use modified arguments if provided
+                if decision.modified_args:
+                    arguments = decision.modified_args
+                    logging.info(f"Using modified arguments: {arguments}")
+                    
+            except Exception as e:
+                error_msg = f"Error during approval process: {str(e)}"
+                logging.error(error_msg)
+                return {"error": error_msg, "approval_error": True}
+
         # Special handling for MCP tools
         # Check if tools is an MCP instance with the requested function name
         from ..mcp.mcp import MCP
@@ -982,43 +1012,7 @@ Your Goal: {self.goal}
                     if not response:
                         return None
 
-                    tool_calls = getattr(response.choices[0].message, 'tool_calls', None)
                     response_text = response.choices[0].message.content.strip()
-                    if tool_calls: ## TODO: Most likely this tool call is already called in _chat_completion, so maybe we can remove this.
-                        messages.append({
-                            "role": "assistant",
-                            "content": response_text,
-                            "tool_calls": tool_calls
-                        })
-                        
-                        for tool_call in tool_calls:
-                            function_name = tool_call.function.name
-                            arguments = json.loads(tool_call.function.arguments)
-
-                            if self.verbose:
-                                display_tool_call(f"Agent {self.name} is calling function '{function_name}' with arguments: {arguments}", console=self.console)
-
-                            tool_result = self.execute_tool(function_name, arguments)
-
-                            if tool_result:
-                                if self.verbose:
-                                    display_tool_call(f"Function '{function_name}' returned: {tool_result}", console=self.console)
-                                messages.append({
-                                    "role": "tool",
-                                    "tool_call_id": tool_call.id,
-                                    "content": json.dumps(tool_result)
-                                })
-                            else:
-                                messages.append({
-                                    "role": "tool",
-                                    "tool_call_id": tool_call.id,
-                                    "content": "Function returned an empty output"
-                                })
-                            
-                        response = self._chat_completion(messages, temperature=temperature, stream=stream)
-                        if not response:
-                            return None
-                        response_text = response.choices[0].message.content.strip()
 
                     # Handle output_json or output_pydantic if specified
                     if output_json or output_pydantic:
@@ -1418,6 +1412,21 @@ Your Goal: {self.goal}
         """Async version of execute_tool"""
         try:
             logging.info(f"Executing async tool: {function_name} with arguments: {arguments}")
+            
+            # Check if approval is required for this tool
+            from ..approval import is_approval_required, request_approval
+            if is_approval_required(function_name):
+                decision = await request_approval(function_name, arguments)
+                if not decision.approved:
+                    error_msg = f"Tool execution denied: {decision.reason}"
+                    logging.warning(error_msg)
+                    return {"error": error_msg, "approval_denied": True}
+                
+                # Use modified arguments if provided
+                if decision.modified_args:
+                    arguments = decision.modified_args
+                    logging.info(f"Using modified arguments: {arguments}")
+            
             # Try to find the function in the agent's tools list first
             func = None
             for tool in self.tools:

{praisonaiagents-0.0.96 → praisonaiagents-0.0.97}/praisonaiagents/agents/agents.py

@@ -63,7 +63,6 @@ def process_task_context(context_item, verbose=0, user_id=None):
     """
     Process a single context item for task execution.
     This helper function avoids code duplication between async and sync execution methods.
-    
     Args:
         context_item: The context item to process (can be string, list, task object, or dict)
         verbose: Verbosity level for logging
@@ -203,7 +202,6 @@ class PraisonAIAgents:
                 mem_cfg = memory_config
                 if not mem_cfg:
                     mem_cfg = next((t.config.get('memory_config') for t in tasks if hasattr(t, 'config') and t.config), None)
-                    
                 # Set default memory config if none provided
                 if not mem_cfg:
                     mem_cfg = {
@@ -215,7 +213,6 @@ class PraisonAIAgents:
                         },
                         "rag_db_path": "./.praison/chroma_db"
                     }
-                
                 # Add embedder config if provided
                 if embedder:
                     if isinstance(embedder, dict):
@@ -231,17 +228,14 @@ class PraisonAIAgents:
                     self.shared_memory = Memory(config=mem_cfg, verbose=verbose)
                     if verbose >= 5:
                         logger.info("Initialized shared memory for PraisonAIAgents")
-                    
                     # Distribute memory to tasks
                     for task in tasks:
                         if not task.memory:
                             task.memory = self.shared_memory
                             if verbose >= 5:
                                 logger.info(f"Assigned shared memory to task {task.id}")
-                            
             except Exception as e:
                 logger.error(f"Failed to initialize shared memory: {e}")
-            
         # Update tasks with shared memory
         if self.shared_memory:
             for task in tasks:
@@ -898,6 +892,105 @@ Context:
     def clear_state(self) -> None:
         """Clear all state values"""
         self._state.clear()
+    
+    # Convenience methods for enhanced state management
+    def has_state(self, key: str) -> bool:
+        """Check if a state key exists"""
+        return key in self._state
+    
+    def get_all_state(self) -> Dict[str, Any]:
+        """Get a copy of the entire state dictionary"""
+        return self._state.copy()
+    
+    def delete_state(self, key: str) -> bool:
+        """Delete a state key if it exists. Returns True if deleted, False if key didn't exist."""
+        if key in self._state:
+            del self._state[key]
+            return True
+        return False
+    
+    def increment_state(self, key: str, amount: float = 1, default: float = 0) -> float:
+        """Increment a numeric state value. Creates the key with default if it doesn't exist."""
+        current = self._state.get(key, default)
+        if not isinstance(current, (int, float)):
+            raise TypeError(f"Cannot increment non-numeric value at key '{key}': {type(current).__name__}")
+        new_value = current + amount
+        self._state[key] = new_value
+        return new_value
+    
+    def append_to_state(self, key: str, value: Any, max_length: Optional[int] = None) -> List[Any]:
+        """Append a value to a list state. Creates the list if it doesn't exist.
+        
+        Args:
+            key: State key
+            value: Value to append
+            max_length: Optional maximum length for the list
+            
+        Returns:
+            The updated list
+            
+        Raises:
+            TypeError: If the existing value is not a list and convert_to_list=False
+        """
+        if key not in self._state:
+            self._state[key] = []
+        elif not isinstance(self._state[key], list):
+            # Be explicit about type conversion for better user experience
+            current_value = self._state[key]
+            self._state[key] = [current_value]
+        
+        self._state[key].append(value)
+        
+        # Trim list if max_length is specified
+        if max_length and len(self._state[key]) > max_length:
+            self._state[key] = self._state[key][-max_length:]
+        
+        return self._state[key]
+    
+    def save_session_state(self, session_id: str, include_memory: bool = True) -> None:
+        """Save current state to memory for session persistence"""
+        if self.shared_memory and include_memory:
+            state_data = {
+                "session_id": session_id,
+                "user_id": self.user_id,
+                "run_id": self.run_id,
+                "state": self._state,
+                "agents": [agent.name for agent in self.agents],
+                "process": self.process
+            }
+            self.shared_memory.store_short_term(
+                text=f"Session state for {session_id}",
+                metadata={
+                    "type": "session_state",
+                    "session_id": session_id,
+                    "user_id": self.user_id,
+                    "state_data": state_data
+                }
+            )
+    
+    def restore_session_state(self, session_id: str) -> bool:
+        """Restore state from memory for session persistence. Returns True if restored."""
+        if not self.shared_memory:
+            return False
+        
+        # Use metadata-based search for better SQLite compatibility
+        results = self.shared_memory.search_short_term(
+            query=f"type:session_state",
+            limit=10  # Get more results to filter by session_id
+        )
+        
+        # Filter results by session_id in metadata
+        for result in results:
+            metadata = result.get("metadata", {})
+            if (metadata.get("type") == "session_state" and 
+                metadata.get("session_id") == session_id):
+                state_data = metadata.get("state_data", {})
+                if "state" in state_data:
+                    # Merge with existing state instead of replacing
+                    self._state.update(state_data["state"])
+                    return True
+        
+        return False
         
     def launch(self, path: str = '/agents', port: int = 8000, host: str = '0.0.0.0', debug: bool = False, protocol: str = "http"):
         """

praisonaiagents-0.0.97/praisonaiagents/approval.py

@@ -0,0 +1,263 @@
+"""
+Human Approval Framework for PraisonAI Agents
+
+This module provides a minimal human-in-the-loop approval system for dangerous tool operations.
+It extends the existing callback system to require human approval before executing high-risk tools.
+"""
+
+import logging
+import asyncio
+from typing import Dict, Set, Optional, Callable, Any, Literal
+from functools import wraps
+from contextvars import ContextVar
+from rich.console import Console
+from rich.panel import Panel
+from rich.text import Text
+from rich.prompt import Confirm
+
+# Global registries for approval requirements
+APPROVAL_REQUIRED_TOOLS: Set[str] = set()
+TOOL_RISK_LEVELS: Dict[str, str] = {}
+
+# Risk levels
+RiskLevel = Literal["critical", "high", "medium", "low"]
+
+# Global approval callback
+approval_callback: Optional[Callable] = None
+
+# Context variable to track if we're in an approved execution context
+_approved_context: ContextVar[Set[str]] = ContextVar('approved_context', default=set())
+
+class ApprovalDecision:
+    """Result of an approval request"""
+    def __init__(self, approved: bool, modified_args: Optional[Dict[str, Any]] = None, reason: str = ""):
+        self.approved = approved
+        self.modified_args = modified_args or {}
+        self.reason = reason
+
+def set_approval_callback(callback_fn: Callable):
+    """Set a custom approval callback function.
+    
+    The callback should accept (function_name, arguments, risk_level) and return ApprovalDecision.
+    """
+    global approval_callback
+    approval_callback = callback_fn
+
+def mark_approved(tool_name: str):
+    """Mark a tool as approved in the current context."""
+    approved = _approved_context.get(set())
+    approved.add(tool_name)
+    _approved_context.set(approved)
+
+def is_already_approved(tool_name: str) -> bool:
+    """Check if a tool is already approved in the current context."""
+    approved = _approved_context.get(set())
+    return tool_name in approved
+
+def clear_approval_context():
+    """Clear the approval context."""
+    _approved_context.set(set())
+
+def require_approval(risk_level: RiskLevel = "high"):
+    """Decorator to mark a tool as requiring human approval.
+    
+    Args:
+        risk_level: The risk level of the tool ("critical", "high", "medium", "low")
+    """
+    def decorator(func):
+        tool_name = getattr(func, '__name__', str(func))
+        APPROVAL_REQUIRED_TOOLS.add(tool_name)
+        TOOL_RISK_LEVELS[tool_name] = risk_level
+        
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            # Skip approval if already approved in current context
+            if is_already_approved(tool_name):
+                return func(*args, **kwargs)
+            
+            # Request approval before executing the function
+            try:
+                # Try to check if we're in an async context
+                try:
+                    asyncio.get_running_loop()
+                    # We're in an async context, but this is a sync function
+                    # Fall back to sync approval to avoid loop conflicts
+                    raise RuntimeError("Use sync fallback in async context")
+                except RuntimeError:
+                    # Either no running loop or we want sync fallback
+                    # Use asyncio.run for clean async execution
+                    decision = asyncio.run(request_approval(tool_name, kwargs))
+            except Exception as e:
+                # Fallback to sync approval if async fails
+                logging.warning(f"Async approval failed, using sync fallback: {e}")
+                callback = approval_callback or console_approval_callback
+                decision = callback(tool_name, kwargs, risk_level)
+            
+            if not decision.approved:
+                raise PermissionError(f"Execution of {tool_name} denied: {decision.reason}")
+            
+            # Mark as approved and merge modified args
+            mark_approved(tool_name)
+            kwargs.update(decision.modified_args)
+            return func(*args, **kwargs)
+        
+        @wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            # Skip approval if already approved in current context
+            if is_already_approved(tool_name):
+                return await func(*args, **kwargs)
+            
+            # Request approval before executing the function
+            decision = await request_approval(tool_name, kwargs)
+            if not decision.approved:
+                raise PermissionError(f"Execution of {tool_name} denied: {decision.reason}")
+            
+            # Mark as approved and merge modified args
+            mark_approved(tool_name)
+            kwargs.update(decision.modified_args)
+            return await func(*args, **kwargs)
+        
+        # Return the appropriate wrapper based on function type
+        if asyncio.iscoroutinefunction(func):
+            return async_wrapper
+        else:
+            return wrapper
+    
+    return decorator
+
+def console_approval_callback(function_name: str, arguments: Dict[str, Any], risk_level: str) -> ApprovalDecision:
+    """Default console-based approval callback.
+    
+    Displays tool information and prompts user for approval via console.
+    """
+    console = Console()
+    
+    # Create risk level styling
+    risk_colors = {
+        "critical": "bold red",
+        "high": "red", 
+        "medium": "yellow",
+        "low": "blue"
+    }
+    risk_color = risk_colors.get(risk_level, "white")
+    
+    # Display tool information
+    tool_info = f"[bold]Function:[/] {function_name}\n"
+    tool_info += f"[bold]Risk Level:[/] [{risk_color}]{risk_level.upper()}[/{risk_color}]\n"
+    tool_info += f"[bold]Arguments:[/]\n"
+    
+    for key, value in arguments.items():
+        # Truncate long values for display
+        str_value = str(value)
+        if len(str_value) > 100:
+            str_value = str_value[:97] + "..."
+        tool_info += f"  {key}: {str_value}\n"
+    
+    console.print(Panel(
+        tool_info.strip(),
+        title="🔒 Tool Approval Required",
+        border_style=risk_color,
+        title_align="left"
+    ))
+    
+    # Get user approval
+    try:
+        approved = Confirm.ask(
+            f"[{risk_color}]Do you want to execute this {risk_level} risk tool?[/{risk_color}]",
+            default=False
+        )
+        
+        if approved:
+            console.print("[green]✅ Tool execution approved[/green]")
+            return ApprovalDecision(approved=True, reason="User approved")
+        else:
+            console.print("[red]❌ Tool execution denied[/red]")
+            return ApprovalDecision(approved=False, reason="User denied")
+            
+    except KeyboardInterrupt:
+        console.print("\n[red]❌ Tool execution cancelled by user[/red]")
+        return ApprovalDecision(approved=False, reason="User cancelled")
+    except Exception as e:
+        console.print(f"[red]Error during approval: {e}[/red]")
+        return ApprovalDecision(approved=False, reason=f"Approval error: {e}")
+
+async def request_approval(function_name: str, arguments: Dict[str, Any]) -> ApprovalDecision:
+    """Request approval for a tool execution.
+    
+    Args:
+        function_name: Name of the function to execute
+        arguments: Arguments to pass to the function
+        
+    Returns:
+        ApprovalDecision with approval status and any modifications
+    """
+    # Check if approval is required
+    if function_name not in APPROVAL_REQUIRED_TOOLS:
+        return ApprovalDecision(approved=True, reason="No approval required")
+    
+    risk_level = TOOL_RISK_LEVELS.get(function_name, "medium")
+    
+    # Use custom callback if set, otherwise use console callback
+    callback = approval_callback or console_approval_callback
+    
+    try:
+        # Handle async callbacks
+        if asyncio.iscoroutinefunction(callback):
+            decision = await callback(function_name, arguments, risk_level)
+        else:
+            # Run sync callback in executor to avoid blocking
+            loop = asyncio.get_event_loop()
+            decision = await loop.run_in_executor(None, callback, function_name, arguments, risk_level)
+        
+        return decision
+        
+    except Exception as e:
+        logging.error(f"Error in approval callback: {e}")
+        return ApprovalDecision(approved=False, reason=f"Approval callback error: {e}")
+
+# Default dangerous tools - can be configured at runtime
+DEFAULT_DANGEROUS_TOOLS = {
+    # Critical risk tools
+    "execute_command": "critical",
+    "kill_process": "critical", 
+    "execute_code": "critical",
+    
+    # High risk tools
+    "write_file": "high",
+    "delete_file": "high",
+    "move_file": "high",
+    "copy_file": "high",
+    "execute_query": "high",
+    
+    # Medium risk tools
+    "evaluate": "medium",
+    "crawl": "medium",
+    "scrape_page": "medium",
+}
+
+def configure_default_approvals():
+    """Configure default dangerous tools to require approval."""
+    for tool_name, risk_level in DEFAULT_DANGEROUS_TOOLS.items():
+        APPROVAL_REQUIRED_TOOLS.add(tool_name)
+        TOOL_RISK_LEVELS[tool_name] = risk_level
+
+def add_approval_requirement(tool_name: str, risk_level: RiskLevel = "high"):
+    """Dynamically add approval requirement for a tool."""
+    APPROVAL_REQUIRED_TOOLS.add(tool_name)
+    TOOL_RISK_LEVELS[tool_name] = risk_level
+
+def remove_approval_requirement(tool_name: str):
+    """Remove approval requirement for a tool."""
+    APPROVAL_REQUIRED_TOOLS.discard(tool_name)
+    TOOL_RISK_LEVELS.pop(tool_name, None)
+
+def is_approval_required(tool_name: str) -> bool:
+    """Check if a tool requires approval."""
+    return tool_name in APPROVAL_REQUIRED_TOOLS
+
+def get_risk_level(tool_name: str) -> Optional[str]:
+    """Get the risk level of a tool."""
+    return TOOL_RISK_LEVELS.get(tool_name)
+
+# Initialize with defaults
+configure_default_approvals()

{praisonaiagents-0.0.96 → praisonaiagents-0.0.97}/praisonaiagents/knowledge/knowledge.py

@@ -17,9 +17,13 @@ class CustomMemory:
         }).from_config(config)
 
     @staticmethod
-    def _add_to_vector_store(self, messages, metadata, filters, infer):
+    def _add_to_vector_store(self, messages, metadata=None, filters=None, infer=None):
         # Custom implementation that doesn't use LLM
-        parsed_messages = "\n".join([msg["content"] for msg in messages])
+        # Handle different message formats for backward compatibility
+        if isinstance(messages, list):
+            parsed_messages = "\n".join([msg.get("content", str(msg)) if isinstance(msg, dict) else str(msg) for msg in messages])
+        else:
+            parsed_messages = str(messages)
         
         # Create a simple fact without using LLM
         new_retrieved_facts = [parsed_messages]
@@ -34,7 +38,7 @@ class CustomMemory:
         memory_id = self._create_memory(
             data=parsed_messages,
             existing_embeddings=new_message_embeddings,
-            metadata=metadata
+            metadata=metadata or {}
         )
         
         return [{
@@ -137,6 +141,10 @@ class Knowledge:
             # Merge reranker config if provided
             if "reranker" in self._config:
                 base_config["reranker"].update(self._config["reranker"])
+            
+            # Merge graph_store config if provided (for graph memory support)
+            if "graph_store" in self._config:
+                base_config["graph_store"] = self._config["graph_store"]
         return base_config
 
     @cached_property
@@ -184,7 +192,22 @@ class Knowledge:
                 if not content:
                     return []
                 
-            result = self.memory.add(content, user_id=user_id, agent_id=agent_id, run_id=run_id, metadata=metadata)
+            # Try new API format first, fall back to old format for backward compatibility
+            try:
+                # Convert content to messages format for mem0 API compatibility
+                if isinstance(content, str):
+                    messages = [{"role": "user", "content": content}]
+                else:
+                    messages = content if isinstance(content, list) else [{"role": "user", "content": str(content)}]
+                
+                result = self.memory.add(messages=messages, user_id=user_id, agent_id=agent_id, run_id=run_id, metadata=metadata)
+            except TypeError as e:
+                # Fallback to old API format if messages parameter is not supported
+                if "unexpected keyword argument" in str(e) or "positional argument" in str(e):
+                    self._log(f"Falling back to legacy API format due to: {e}")
+                    result = self.memory.add(content, user_id=user_id, agent_id=agent_id, run_id=run_id, metadata=metadata)
+                else:
+                    raise
             self._log(f"Store operation result: {result}")
             return result
         except Exception as e:

{praisonaiagents-0.0.96 → praisonaiagents-0.0.97}/praisonaiagents/llm/llm.py

@@ -1512,6 +1512,10 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
         if self.stop_phrases:
             params["stop"] = self.stop_phrases
         
+        # Add extra settings for provider-specific parameters (e.g., num_ctx for Ollama)
+        if self.extra_settings:
+            params.update(self.extra_settings)
+        
         # Override with any provided parameters
         params.update(override_params)
         

{praisonaiagents-0.0.96 → praisonaiagents-0.0.97}/praisonaiagents/main.py

@@ -43,12 +43,18 @@ error_logs = []
 sync_display_callbacks = {}
 async_display_callbacks = {}
 
+# Global approval callback registry
+approval_callback = None
+
 # At the top of the file, add display_callbacks to __all__
 __all__ = [
     'error_logs',
     'register_display_callback',
+    'register_approval_callback',
     'sync_display_callbacks',
     'async_display_callbacks',
+    'execute_callback',
+    'approval_callback',
     # ... other exports
 ]
 
@@ -65,6 +71,15 @@ def register_display_callback(display_type: str, callback_fn, is_async: bool = F
     else:
         sync_display_callbacks[display_type] = callback_fn
 
+def register_approval_callback(callback_fn):
+    """Register a global approval callback function for dangerous tool operations.
+    
+    Args:
+        callback_fn: Function that takes (function_name, arguments, risk_level) and returns ApprovalDecision
+    """
+    global approval_callback
+    approval_callback = callback_fn
+
 async def execute_callback(display_type: str, **kwargs):
     """Execute both sync and async callbacks for a given display type.