import os
import time
import json
import copy
import logging
import asyncio
from typing import List, Optional, Any, Dict, Union, Literal, TYPE_CHECKING, Callable, Tuple, Generator
from rich.console import Console
from rich.live import Live
from ..llm import (
    get_openai_client,
    ChatCompletionMessage,
    Choice,
    CompletionTokensDetails,
    PromptTokensDetails,
    CompletionUsage,
    ChatCompletion,
    ToolCall,
    process_stream_chunks
)
from ..main import (
    display_error,
    display_tool_call,
    display_instruction,
    display_interaction,
    display_generating,
    display_self_reflection,
    ReflectionOutput,
    adisplay_instruction,
    approval_callback,
    execute_sync_callback
)
import inspect
import uuid

# Global variables for API server
_server_started = {}  # Dict of port -> started boolean
_registered_agents = {}  # Dict of port -> Dict of path -> agent_id
_shared_apps = {}  # Dict of port -> FastAPI app

# Don't import FastAPI dependencies here - use lazy loading instead

if TYPE_CHECKING:
    from ..task.task import Task
    from ..main import TaskOutput
    from ..handoff import Handoff

class Agent:
    @classmethod
    def _configure_logging(cls):
        """Configure logging settings once for all agent instances."""
        # Configure logging to suppress unwanted outputs
        logging.getLogger("litellm").setLevel(logging.WARNING)
        
        # Allow httpx logging when LOGLEVEL=debug, otherwise suppress it
        loglevel = os.environ.get('LOGLEVEL', 'INFO').upper()
        if loglevel == 'DEBUG':
            logging.getLogger("httpx").setLevel(logging.INFO)
            logging.getLogger("httpcore").setLevel(logging.INFO)
        else:
            logging.getLogger("httpx").setLevel(logging.WARNING)
            logging.getLogger("httpcore").setLevel(logging.WARNING)
    
    def _generate_tool_definition(self, function_name):
        """
        Generate a tool definition from a function name by inspecting the function.
        """
        logging.debug(f"Attempting to generate tool definition for: {function_name}")
        
        # First try to get the tool definition if it exists
        tool_def_name = f"{function_name}_definition"
        tool_def = globals().get(tool_def_name)
        logging.debug(f"Looking for {tool_def_name} in globals: {tool_def is not None}")
        
        if not tool_def:
            import __main__
            tool_def = getattr(__main__, tool_def_name, None)
            logging.debug(f"Looking for {tool_def_name} in __main__: {tool_def is not None}")
        
        if tool_def:
            logging.debug(f"Found tool definition: {tool_def}")
            return tool_def

        # Try to find the function in the agent's tools list first
        func = None
        for tool in self.tools:
            if callable(tool) and getattr(tool, '__name__', '') == function_name:
                func = tool
                break
        
        logging.debug(f"Looking for {function_name} in agent tools: {func is not None}")
        
        # If not found in tools, try globals and main
        if not func:
            func = globals().get(function_name)
            logging.debug(f"Looking for {function_name} in globals: {func is not None}")
            
            if not func:
                import __main__
                func = getattr(__main__, function_name, None)
                logging.debug(f"Looking for {function_name} in __main__: {func is not None}")

        if not func or not callable(func):
            logging.debug(f"Function {function_name} not found or not callable")
            return None

        import inspect
        # Langchain tools
        if inspect.isclass(func) and hasattr(func, 'run') and not hasattr(func, '_run'):
            original_func = func
            func = func.run
            function_name = original_func.__name__
        # CrewAI tools
        elif inspect.isclass(func) and hasattr(func, '_run'):
            original_func = func
            func = func._run
            function_name = original_func.__name__

        sig = inspect.signature(func)
        logging.debug(f"Function signature: {sig}")
        
        # Skip self, *args, **kwargs, so they don't get passed in arguments
        parameters_list = []
        for name, param in sig.parameters.items():
            if name == "self":
                continue
            if param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD):
                continue
            parameters_list.append((name, param))

        parameters = {
            "type": "object",
            "properties": {},
            "required": []
        }
        
        # Parse docstring for parameter descriptions
        docstring = inspect.getdoc(func)
        logging.debug(f"Function docstring: {docstring}")
        
        param_descriptions = {}
        if docstring:
            import re
            param_section = re.split(r'\s*Args:\s*', docstring)
            logging.debug(f"Param section split: {param_section}")
            if len(param_section) > 1:
                param_lines = param_section[1].split('\n')
                for line in param_lines:
                    line = line.strip()
                    if line and ':' in line:
                        param_name, param_desc = line.split(':', 1)
                        param_descriptions[param_name.strip()] = param_desc.strip()
        
        logging.debug(f"Parameter descriptions: {param_descriptions}")

        for name, param in parameters_list:
            param_type = "string"  # Default type
            if param.annotation != inspect.Parameter.empty:
                if param.annotation == int:
                    param_type = "integer"
                elif param.annotation == float:
                    param_type = "number"
                elif param.annotation == bool:
                    param_type = "boolean"
                elif param.annotation == list:
                    param_type = "array"
                elif param.annotation == dict:
                    param_type = "object"
            
            param_info = {"type": param_type}
            if name in param_descriptions:
                param_info["description"] = param_descriptions[name]
            
            parameters["properties"][name] = param_info
            if param.default == inspect.Parameter.empty:
                parameters["required"].append(name)
        
        logging.debug(f"Generated parameters: {parameters}")

        # Extract description from docstring
        description = docstring.split('\n')[0] if docstring else f"Function {function_name}"
        
        tool_def = {
            "type": "function",
            "function": {
                "name": function_name,
                "description": description,
                "parameters": parameters
            }
        }
        logging.debug(f"Generated tool definition: {tool_def}")
        return tool_def

    def __init__(
        self,
        name: Optional[str] = None,
        role: Optional[str] = None,
        goal: Optional[str] = None,
        backstory: Optional[str] = None,
        instructions: Optional[str] = None,
        llm: Optional[Union[str, Any]] = None,
        tools: Optional[List[Any]] = None,
        function_calling_llm: Optional[Any] = None,
        max_iter: int = 20,
        max_rpm: Optional[int] = None,
        max_execution_time: Optional[int] = None,
        memory: Optional[Any] = None,
        verbose: bool = True,
        allow_delegation: bool = False,
        step_callback: Optional[Any] = None,
        cache: bool = True,
        system_template: Optional[str] = None,
        prompt_template: Optional[str] = None,
        response_template: Optional[str] = None,
        allow_code_execution: Optional[bool] = False,
        max_retry_limit: int = 2,
        respect_context_window: bool = True,
        code_execution_mode: Literal["safe", "unsafe"] = "safe",
        embedder_config: Optional[Dict[str, Any]] = None,
        knowledge: Optional[List[str]] = None,
        knowledge_config: Optional[Dict[str, Any]] = None,
        use_system_prompt: Optional[bool] = True,
        markdown: bool = True,
        stream: bool = False,
        metrics: bool = False,
        self_reflect: bool = False,
        max_reflect: int = 3,
        min_reflect: int = 1,
        reflect_llm: Optional[str] = None,
        reflect_prompt: Optional[str] = None,
        user_id: Optional[str] = None,
        reasoning_steps: bool = False,
        guardrail: Optional[Union[Callable[['TaskOutput'], Tuple[bool, Any]], str]] = None,
        max_guardrail_retries: int = 3,
        handoffs: Optional[List[Union['Agent', 'Handoff']]] = None,
        base_url: Optional[str] = None,
        api_key: Optional[str] = None,
        web_search: Optional[Union[bool, Dict[str, Any]]] = None,
        web_fetch: Optional[Union[bool, Dict[str, Any]]] = None,
        prompt_caching: Optional[bool] = None,
        claude_memory: Optional[Union[bool, Any]] = None,
        plan_mode: bool = False,
        planning: bool = False,
        planning_tools: Optional[List[Any]] = None,
        planning_reasoning: bool = False
    ):
        """Initialize an Agent instance.

        Args:
            name (Optional[str], optional): Name of the agent used for identification and logging.
                If None, defaults to "Agent". Defaults to None.
            role (Optional[str], optional): Role or job title that defines the agent's expertise
                and behavior patterns. Examples: "Data Analyst", "Content Writer". Defaults to None.
            goal (Optional[str], optional): Primary objective or goal the agent aims to achieve.
                Defines the agent's purpose and success criteria. Defaults to None.
            backstory (Optional[str], optional): Background story or context that shapes the agent's
                personality and decision-making approach. Defaults to None.
            instructions (Optional[str], optional): Direct instructions that override role, goal,
                and backstory when provided. Used for simple, task-specific agents. Defaults to None.
            llm (Optional[Union[str, Any]], optional): Language model configuration. Can be a model
                name string (e.g., "gpt-5-nano", "anthropic/claude-3-sonnet") or a configured LLM object.
                Defaults to environment variable OPENAI_MODEL_NAME or "gpt-5-nano".
            tools (Optional[List[Any]], optional): List of tools, functions, or capabilities
                available to the agent for task execution. Can include callables, tool objects,
                or MCP instances. Defaults to None.
            function_calling_llm (Optional[Any], optional): Dedicated language model for function
                calling operations. If None, uses the main llm parameter. Defaults to None.
            max_iter (int, optional): Maximum number of iterations the agent can perform during
                task execution to prevent infinite loops. Defaults to 20.
            max_rpm (Optional[int], optional): Maximum requests per minute to rate limit API calls
                and prevent quota exhaustion. If None, no rate limiting is applied. Defaults to None.
            max_execution_time (Optional[int], optional): Maximum execution time in seconds for
                agent operations before timeout. If None, no time limit is enforced. Defaults to None.
            memory (Optional[Any], optional): Memory system for storing and retrieving information
                across conversations. Requires memory dependencies to be installed. Defaults to None.
            verbose (bool, optional): Enable detailed logging and status updates during agent
                execution for debugging and monitoring. Defaults to True.
            allow_delegation (bool, optional): Allow the agent to delegate tasks to other agents
                or sub-processes when appropriate. Defaults to False.
            step_callback (Optional[Any], optional): Callback function called after each step
                of agent execution for custom monitoring or intervention. Defaults to None.
            cache (bool, optional): Enable caching of responses and computations to improve
                performance and reduce API costs. Defaults to True.
            system_template (Optional[str], optional): Custom template for system prompts that
                overrides the default system prompt generation. Defaults to None.
            prompt_template (Optional[str], optional): Template for formatting user prompts
                before sending to the language model. Defaults to None.
            response_template (Optional[str], optional): Template for formatting agent responses
                before returning to the user. Defaults to None.
            allow_code_execution (Optional[bool], optional): Enable the agent to execute code
                snippets during task completion. Use with caution for security. Defaults to False.
            max_retry_limit (int, optional): Maximum number of retry attempts for failed operations
                before giving up. Helps handle transient errors. Defaults to 2.
            respect_context_window (bool, optional): Automatically manage context window size
                to prevent token limit errors with large conversations. Defaults to True.
            code_execution_mode (Literal["safe", "unsafe"], optional): Safety mode for code execution.
                "safe" restricts dangerous operations, "unsafe" allows full code execution. Defaults to "safe".
            embedder_config (Optional[Dict[str, Any]], optional): Configuration dictionary for
                text embedding models used in knowledge retrieval and similarity search. Defaults to None.
            knowledge (Optional[List[str]], optional): List of knowledge sources (file paths, URLs,
                or text content) to be processed and made available to the agent. Defaults to None.
            knowledge_config (Optional[Dict[str, Any]], optional): Configuration for knowledge
                processing and retrieval system including chunking and indexing parameters. Defaults to None.
            use_system_prompt (Optional[bool], optional): Whether to include system prompts in
                conversations to establish agent behavior and context. Defaults to True.
            markdown (bool, optional): Enable markdown formatting in agent responses for better
                readability and structure. Defaults to True.
            stream (bool, optional): Enable streaming responses from the language model for real-time
                output when using Agent.start() method. Defaults to False for backward compatibility.
            metrics (bool, optional): Enable automatic token usage tracking and display summary
                when tasks complete. Simplifies token monitoring for cost optimization. Defaults to False.
            self_reflect (bool, optional): Enable self-reflection capabilities where the agent
                evaluates and improves its own responses. Defaults to False.
            max_reflect (int, optional): Maximum number of self-reflection iterations to prevent
                excessive reflection loops. Defaults to 3.
            min_reflect (int, optional): Minimum number of self-reflection iterations required
                before accepting a response as satisfactory. Defaults to 1.
            reflect_llm (Optional[str], optional): Dedicated language model for self-reflection
                operations. If None, uses the main llm parameter. Defaults to None.
            reflect_prompt (Optional[str], optional): Custom prompt template for self-reflection
                that guides the agent's self-evaluation process. Defaults to None.
            user_id (Optional[str], optional): Unique identifier for the user or session to
                enable personalized responses and memory isolation. Defaults to "praison".
            reasoning_steps (bool, optional): Enable step-by-step reasoning output to show the
                agent's thought process during problem solving. Defaults to False.
            guardrail (Optional[Union[Callable[['TaskOutput'], Tuple[bool, Any]], str]], optional):
                Safety mechanism to validate agent outputs. Can be a validation function or
                description string for LLM-based validation. Defaults to None.
            max_guardrail_retries (int, optional): Maximum number of retry attempts when guardrail
                validation fails before giving up. Defaults to 3.
            handoffs (Optional[List[Union['Agent', 'Handoff']]], optional): List of agents or
                handoff configurations that this agent can delegate tasks to. Enables agent-to-agent
                collaboration and task specialization. Defaults to None.
            base_url (Optional[str], optional): Base URL for custom LLM endpoints (e.g., Ollama).
                If provided, automatically creates a custom LLM instance. Defaults to None.
            api_key (Optional[str], optional): API key for LLM provider. If not provided,
                falls back to environment variables. Defaults to None.

        Raises:
            ValueError: If all of name, role, goal, backstory, and instructions are None.
            ImportError: If memory or LLM features are requested but dependencies are not installed.
        """
        # Add check at start if memory is requested
        if memory is not None:
            try:
                from ..memory.memory import Memory
                MEMORY_AVAILABLE = True
            except ImportError:
                raise ImportError(
                    "Memory features requested in Agent but memory dependencies not installed. "
                    "Please install with: pip install \"praisonaiagents[memory]\""
                )

        # Handle backward compatibility for required fields
        if all(x is None for x in [name, role, goal, backstory, instructions]):
            raise ValueError("At least one of name, role, goal, backstory, or instructions must be provided")

        # Configure logging only once at the class level
        if not hasattr(Agent, '_logging_configured'):
            Agent._configure_logging()
            Agent._logging_configured = True

        # If instructions are provided, use them to set role, goal, and backstory
        if instructions:
            self.name = name or "Agent"
            self.role = role or "Assistant"
            self.goal = goal or instructions
            self.backstory = backstory or instructions
            # Set self_reflect to False by default for instruction-based agents
            self.self_reflect = False if self_reflect is None else self_reflect
        else:
            # Use provided values or defaults
            self.name = name or "Agent"
            self.role = role or "Assistant"
            self.goal = goal or "Help the user with their tasks"
            self.backstory = backstory or "I am an AI assistant"
            # Default to True for traditional agents if not specified
            self.self_reflect = True if self_reflect is None else self_reflect
        
        self.instructions = instructions
        # Check for model name in environment variable if not provided
        self._using_custom_llm = False
        # Flag to track if final result has been displayed to prevent duplicates
        self._final_display_shown = False
        
        # Store OpenAI client parameters for lazy initialization
        self._openai_api_key = api_key
        self._openai_base_url = base_url
        self.__openai_client = None

        # If base_url is provided, always create a custom LLM instance
        if base_url:
            try:
                from ..llm.llm import LLM
                # Handle different llm parameter types with base_url
                if isinstance(llm, dict):
                    # Merge base_url and api_key into the dict
                    llm_config = llm.copy()
                    llm_config['base_url'] = base_url
                    if api_key:
                        llm_config['api_key'] = api_key
                    llm_config['metrics'] = metrics
                    self.llm_instance = LLM(**llm_config)
                else:
                    # Create LLM with model string and base_url
                    model_name = llm or os.getenv('OPENAI_MODEL_NAME', 'gpt-5-nano')
                    self.llm_instance = LLM(
                        model=model_name,
                        base_url=base_url,
                        api_key=api_key,
                        metrics=metrics,
                        web_search=web_search,
                        web_fetch=web_fetch,
                        prompt_caching=prompt_caching,
                        claude_memory=claude_memory
                    )
                self._using_custom_llm = True
            except ImportError as e:
                raise ImportError(
                    "LLM features requested but dependencies not installed. "
                    "Please install with: pip install \"praisonaiagents[llm]\""
                ) from e
        # If the user passes a dictionary (for advanced configuration)
        elif isinstance(llm, dict) and "model" in llm:
            try:
                from ..llm.llm import LLM
                # Add api_key if provided and not in dict
                if api_key and 'api_key' not in llm:
                    llm = llm.copy()
                    llm['api_key'] = api_key
                # Add metrics parameter
                llm = llm.copy()
                llm['metrics'] = metrics
                self.llm_instance = LLM(**llm)  # Pass all dict items as kwargs
                self._using_custom_llm = True
            except ImportError as e:
                raise ImportError(
                    "LLM features requested but dependencies not installed. "
                    "Please install with: pip install \"praisonaiagents[llm]\""
                ) from e
        # If the user passes a string with a slash (provider/model)
        elif isinstance(llm, str) and "/" in llm:
            try:
                from ..llm.llm import LLM
                # Pass the entire string so LiteLLM can parse provider/model
                llm_params = {'model': llm}
                if api_key:
                    llm_params['api_key'] = api_key
                llm_params['metrics'] = metrics
                llm_params['web_search'] = web_search
                llm_params['web_fetch'] = web_fetch
                llm_params['prompt_caching'] = prompt_caching
                llm_params['claude_memory'] = claude_memory
                self.llm_instance = LLM(**llm_params)
                self._using_custom_llm = True
                
                # Ensure tools are properly accessible when using custom LLM
                if tools:
                    logging.debug(f"Tools passed to Agent with custom LLM: {tools}")
                    # Store the tools for later use
                    self.tools = tools
            except ImportError as e:
                raise ImportError(
                    "LLM features requested but dependencies not installed. "
                    "Please install with: pip install \"praisonaiagents[llm]\""
                ) from e
        # Otherwise, fall back to OpenAI environment/name
        else:
            self.llm = llm or os.getenv('OPENAI_MODEL_NAME', 'gpt-5-nano')
        # Handle tools parameter - ensure it's always a list
        if callable(tools):
            # If a single function/callable is passed, wrap it in a list
            self.tools = [tools]
        elif isinstance(tools, str):
            # Single tool name string - resolve from registry
            self.tools = self._resolve_tool_names([tools])
        elif isinstance(tools, (list, tuple)):
            # Check if list contains strings (tool names) that need resolution
            if tools and all(isinstance(t, str) for t in tools):
                self.tools = self._resolve_tool_names(tools)
            else:
                self.tools = list(tools)
        else:
            # Handle all falsy values (None, False, 0, "", etc.) by defaulting to empty list
            self.tools = tools or []
        self.function_calling_llm = function_calling_llm
        self.max_iter = max_iter
        self.max_rpm = max_rpm
        self.max_execution_time = max_execution_time
        self._memory_instance = None
        self._init_memory(memory, user_id)
        self.verbose = verbose
        self.allow_delegation = allow_delegation
        self.step_callback = step_callback
        self.cache = cache
        self.system_template = system_template
        self.prompt_template = prompt_template
        self.response_template = response_template
        self.allow_code_execution = allow_code_execution
        self.max_retry_limit = max_retry_limit
        self.respect_context_window = respect_context_window
        self.code_execution_mode = code_execution_mode
        self.embedder_config = embedder_config
        self.knowledge = knowledge
        self.use_system_prompt = use_system_prompt
        # NOTE: chat_history is not thread-safe. If concurrent access is needed,
        # consider using threading.Lock or other synchronization mechanisms
        self.chat_history = []
        self.markdown = markdown
        self.stream = stream
        self.metrics = metrics
        self.max_reflect = max_reflect
        self.min_reflect = min_reflect
        self.reflect_prompt = reflect_prompt
        # Use the same model selection logic for reflect_llm
        self.reflect_llm = reflect_llm or os.getenv('OPENAI_MODEL_NAME', 'gpt-5-nano')
        self._console = None  # Lazy load console when needed
        
        # Initialize system prompt
        self.system_prompt = f"""{self.backstory}\n
Your Role: {self.role}\n
Your Goal: {self.goal}
        """

        # Lazy generate unique ID when needed
        self._agent_id = None

        # Store user_id
        self.user_id = user_id or "praison"
        self.reasoning_steps = reasoning_steps
        self.plan_mode = plan_mode  # Read-only mode for planning
        self.planning = planning  # Enable planning mode
        self.planning_tools = planning_tools  # Tools for planning phase
        self.planning_reasoning = planning_reasoning  # Enable reasoning during planning
        self._planning_agent = None  # Lazy loaded PlanningAgent
        self.web_search = web_search
        self.web_fetch = web_fetch
        self.prompt_caching = prompt_caching
        self.claude_memory = claude_memory
        
        # Initialize rules manager for persistent context (like Cursor/Windsurf)
        self._rules_manager = None
        self._init_rules_manager()
        
        # Handle web_search fallback: inject DuckDuckGo tool for unsupported models
        if web_search and not self._model_supports_web_search():
            from ..tools.duckduckgo_tools import internet_search
            # Check if internet_search is not already in tools
            tool_names = [getattr(t, '__name__', str(t)) for t in self.tools]
            if 'internet_search' not in tool_names and 'duckduckgo' not in tool_names:
                self.tools.append(internet_search)
                logging.info("Model does not support native web search. Added DuckDuckGo fallback tool.")
        
        # Log warning if web_fetch is enabled but model doesn't support it
        if web_fetch and not self._model_supports_web_fetch():
            logging.warning(f"Model '{self.llm}' does not support native web fetch. Web fetch will be ignored.")
        
        # Initialize guardrail settings
        self.guardrail = guardrail
        self.max_guardrail_retries = max_guardrail_retries
        self._guardrail_fn = None
        self._setup_guardrail()
        
        # Cache for system prompts and formatted tools
        # Note: In single-threaded usage (common case), these are safe
        # For multi-threaded usage, consider using threading.Lock
        self._system_prompt_cache = {}
        self._formatted_tools_cache = {}
        # Limit cache size to prevent unbounded growth
        self._max_cache_size = 100

        # Process handoffs and convert them to tools
        self.handoffs = handoffs if handoffs else []
        self._process_handoffs()

        # Check if knowledge parameter has any values
        if not knowledge:
            self.knowledge = None
            self._knowledge_sources = None
            self._knowledge_processed = True  # No knowledge to process
        else:
            # Store knowledge sources for lazy processing
            self._knowledge_sources = knowledge
            self._knowledge_processed = False
            self._knowledge_config = knowledge_config
            self.knowledge = None  # Will be initialized on first use

    @property
    def console(self):
        """Lazily initialize Rich Console only when needed."""
        if self._console is None:
            from rich.console import Console
            self._console = Console()
        return self._console
    
    @property
    def _openai_client(self):
        """Lazily initialize OpenAI client only when needed."""
        if self.__openai_client is None:
            try:
                self.__openai_client = get_openai_client(
                    api_key=self._openai_api_key, 
                    base_url=self._openai_base_url
                )
            except ValueError as e:
                # If we're using a custom LLM, we might not need the OpenAI client
                # Return None and let the calling code handle it
                if self._using_custom_llm:
                    return None
                else:
                    raise e
        return self.__openai_client

    @property
    def agent_id(self):
        """Lazily generate agent ID when first accessed."""
        if self._agent_id is None:
            import uuid
            self._agent_id = str(uuid.uuid4())
        return self._agent_id
    
    def get_available_tools(self) -> List[Any]:
        """
        Get tools available to this agent, filtered by plan_mode if enabled.
        
        In plan_mode, only read-only tools are available to prevent
        modifications during the planning phase.
        
        Returns:
            List of available tools
        """
        if not self.plan_mode:
            return self.tools
            
        # Filter to read-only tools only
        from ..planning import READ_ONLY_TOOLS, RESTRICTED_TOOLS
        
        filtered_tools = []
        for tool in self.tools:
            tool_name = getattr(tool, '__name__', str(tool)).lower()
            
            # Check if tool is in restricted list
            is_restricted = any(
                restricted.lower() in tool_name 
                for restricted in RESTRICTED_TOOLS
            )
            
            if not is_restricted:
                filtered_tools.append(tool)
                
        return filtered_tools
    
    def _model_supports_web_search(self) -> bool:
        """
        Check if the agent's model supports native web search via LiteLLM.
        
        Returns:
            bool: True if the model supports native web search, False otherwise
        """
        from ..llm.model_capabilities import supports_web_search
        
        # Get the model name
        if hasattr(self, 'llm_instance') and self.llm_instance:
            model_name = self.llm_instance.model
        elif hasattr(self, 'llm') and self.llm:
            model_name = self.llm
        else:
            model_name = "gpt-5-nano"
        
        return supports_web_search(model_name)
    
    def _model_supports_web_fetch(self) -> bool:
        """
        Check if the agent's model supports web fetch via LiteLLM.
        
        Web fetch allows the model to retrieve full content from specific URLs.
        Currently only supported by Anthropic Claude models.
        
        Returns:
            bool: True if the model supports web fetch, False otherwise
        """
        from ..llm.model_capabilities import supports_web_fetch
        
        # Get the model name
        if hasattr(self, 'llm_instance') and self.llm_instance:
            model_name = self.llm_instance.model
        elif hasattr(self, 'llm') and self.llm:
            model_name = self.llm
        else:
            model_name = "gpt-5-nano"
        
        return supports_web_fetch(model_name)
    
    def _model_supports_prompt_caching(self) -> bool:
        """
        Check if the agent's model supports prompt caching via LiteLLM.
        
        Prompt caching allows caching parts of prompts to reduce costs and latency.
        Supported by OpenAI, Anthropic, Bedrock, and Deepseek.
        
        Returns:
            bool: True if the model supports prompt caching, False otherwise
        """
        from ..llm.model_capabilities import supports_prompt_caching
        
        # Get the model name
        if hasattr(self, 'llm_instance') and self.llm_instance:
            model_name = self.llm_instance.model
        elif hasattr(self, 'llm') and self.llm:
            model_name = self.llm
        else:
            model_name = "gpt-4o-mini"
        
        return supports_prompt_caching(model_name)
    
    def _init_rules_manager(self):
        """
        Initialize RulesManager for persistent rules/instructions.
        
        Automatically discovers rules from:
        - ~/.praison/rules/ (global)
        - .praison/rules/ (workspace)
        - Subdirectory rules
        """
        try:
            from ..memory.rules_manager import RulesManager
            import os
            
            # Get workspace path (current working directory)
            workspace_path = os.getcwd()
            
            self._rules_manager = RulesManager(
                workspace_path=workspace_path,
                verbose=1 if self.verbose else 0
            )
            
            # Log discovered rules
            stats = self._rules_manager.get_stats()
            if stats["total_rules"] > 0:
                logging.debug(f"RulesManager: Discovered {stats['total_rules']} rules")
        except ImportError:
            logging.debug("RulesManager not available")
            self._rules_manager = None
        except Exception as e:
            logging.debug(f"Could not initialize RulesManager: {e}")
            self._rules_manager = None
    
    def get_rules_context(self, file_path: Optional[str] = None, include_manual: Optional[List[str]] = None) -> str:
        """
        Get rules context for the current conversation.
        
        Args:
            file_path: Optional file path for glob-based rule matching
            include_manual: Optional list of manual rule names to include (via @mention)
            
        Returns:
            Formatted rules context string
        """
        if not self._rules_manager:
            return ""
        
        return self._rules_manager.build_rules_context(
            file_path=file_path,
            include_manual=include_manual
        )
    
    def _init_memory(self, memory, user_id: Optional[str] = None):
        """
        Initialize memory based on the memory parameter.
        
        Args:
            memory: Can be:
                - True: Use FileMemory with default settings
                - False/None: No memory
                - "file": Use FileMemory
                - "sqlite": Use existing Memory class with SQLite
                - dict: Configuration for memory
                - Memory/FileMemory instance: Use directly
            user_id: User identifier for memory isolation
        """
        self.memory = memory
        
        if memory is None or memory is False:
            self._memory_instance = None
            return
        
        # Determine user_id
        mem_user_id = user_id or getattr(self, 'user_id', None) or "default"
        
        if memory is True or memory == "file":
            # Use FileMemory (zero dependencies)
            from ..memory.file_memory import FileMemory
            self._memory_instance = FileMemory(
                user_id=mem_user_id,
                verbose=1 if getattr(self, 'verbose', False) else 0
            )
        elif isinstance(memory, str) and memory in ("sqlite", "chromadb", "mem0", "mongodb"):
            # Use full Memory class with specific provider
            try:
                from ..memory.memory import Memory
                config = {"provider": memory if memory != "sqlite" else "rag"}
                self._memory_instance = Memory(config)
            except ImportError:
                logging.warning(f"Memory provider '{memory}' requires additional dependencies. Falling back to FileMemory.")
                from ..memory.file_memory import FileMemory
                self._memory_instance = FileMemory(user_id=mem_user_id)
        elif isinstance(memory, dict):
            # Configuration dict
            provider = memory.get("provider", "file")
            if provider == "file":
                from ..memory.file_memory import FileMemory
                self._memory_instance = FileMemory(
                    user_id=memory.get("user_id", mem_user_id),
                    config=memory
                )
            else:
                try:
                    from ..memory.memory import Memory
                    self._memory_instance = Memory(memory)
                except ImportError:
                    logging.warning("Full Memory class requires additional dependencies. Falling back to FileMemory.")
                    from ..memory.file_memory import FileMemory
                    self._memory_instance = FileMemory(user_id=mem_user_id)
        else:
            # Assume it's already a memory instance
            self._memory_instance = memory
    
    def get_memory_context(self, query: Optional[str] = None) -> str:
        """
        Get memory context for the current conversation.
        
        Args:
            query: Optional query to focus the context
            
        Returns:
            Formatted memory context string
        """
        if not self._memory_instance:
            return ""
        
        if hasattr(self._memory_instance, 'get_context'):
            return self._memory_instance.get_context(query=query)
        
        return ""
    
    def store_memory(self, content: str, memory_type: str = "short_term", **kwargs):
        """
        Store content in memory.
        
        Args:
            content: Content to store
            memory_type: Type of memory (short_term, long_term, entity, episodic)
            **kwargs: Additional arguments for the memory method
        """
        if not self._memory_instance:
            return
        
        if memory_type == "short_term" and hasattr(self._memory_instance, 'add_short_term'):
            self._memory_instance.add_short_term(content, **kwargs)
        elif memory_type == "long_term" and hasattr(self._memory_instance, 'add_long_term'):
            self._memory_instance.add_long_term(content, **kwargs)
        elif memory_type == "entity" and hasattr(self._memory_instance, 'add_entity'):
            self._memory_instance.add_entity(content, **kwargs)
        elif memory_type == "episodic" and hasattr(self._memory_instance, 'add_episodic'):
            self._memory_instance.add_episodic(content, **kwargs)
    
    def _display_memory_info(self):
        """Display memory information to user in a friendly format."""
        if not self._memory_instance:
            return
        
        # Only display once per chat session
        if hasattr(self, '_memory_displayed') and self._memory_displayed:
            return
        self._memory_displayed = True
        
        stats = self._memory_instance.get_stats()
        short_count = stats.get('short_term_count', 0)
        long_count = stats.get('long_term_count', 0)
        entity_count = stats.get('entity_count', 0)
        storage_path = stats.get('storage_path', '')
        
        total_memories = short_count + long_count + entity_count
        
        if total_memories > 0:
            from rich.panel import Panel
            from rich.text import Text
            
            # Build memory info text
            info_parts = []
            if long_count > 0:
                info_parts.append(f"💾 {long_count} long-term")
            if short_count > 0:
                info_parts.append(f"⚡ {short_count} short-term")
            if entity_count > 0:
                info_parts.append(f"👤 {entity_count} entities")
            
            memory_text = Text()
            memory_text.append("🧠 Memory loaded: ", style="bold cyan")
            memory_text.append(" | ".join(info_parts))
            memory_text.append(f"\n📁 Storage: {storage_path}", style="dim")
            
            self.console.print(Panel(
                memory_text,
                title="[bold]Agent Memory[/bold]",
                border_style="cyan",
                expand=False
            ))
    
    @property
    def llm_model(self):
        """Unified property to get the LLM model regardless of configuration type.
        
        Returns:
            The LLM model/instance being used by this agent.
            - For standard models: returns the model string (e.g., "gpt-5-nano")
            - For custom LLM instances: returns the LLM instance object
            - For provider models: returns the LLM instance object
        """
        if hasattr(self, 'llm_instance') and self.llm_instance:
            return self.llm_instance
        elif hasattr(self, 'llm') and self.llm:
            return self.llm
        else:
            # Default fallback
            return "gpt-5-nano"

    def _ensure_knowledge_processed(self):
        """Ensure knowledge is initialized and processed when first accessed."""
        if not self._knowledge_processed and self._knowledge_sources:
            # Initialize Knowledge with provided or default config
            from praisonaiagents.knowledge import Knowledge
            self.knowledge = Knowledge(self._knowledge_config or None)
            
            # Process all knowledge sources
            for source in self._knowledge_sources:
                self._process_knowledge(source)
            
            self._knowledge_processed = True
    
    def _process_knowledge(self, knowledge_item):
        """Process and store knowledge from a file path, URL, or string."""
        try:
            if os.path.exists(knowledge_item):
                # It's a file path
                self.knowledge.add(knowledge_item, user_id=self.user_id, agent_id=self.agent_id)
            elif knowledge_item.startswith("http://") or knowledge_item.startswith("https://"):
                # It's a URL
                pass
            else:
                # It's a string content
                self.knowledge.store(knowledge_item, user_id=self.user_id, agent_id=self.agent_id)
        except Exception as e:
            logging.error(f"Error processing knowledge item: {knowledge_item}, error: {e}")

    def _setup_guardrail(self):
        """Setup the guardrail function based on the provided guardrail parameter."""
        if self.guardrail is None:
            self._guardrail_fn = None
            return
            
        if callable(self.guardrail):
            # Validate function signature
            sig = inspect.signature(self.guardrail)
            positional_args = [
                param for param in sig.parameters.values()
                if param.default is inspect.Parameter.empty
            ]
            if len(positional_args) != 1:
                raise ValueError("Agent guardrail function must accept exactly one parameter (TaskOutput)")
            
            # Check return annotation if present
            from typing import get_args, get_origin
            return_annotation = sig.return_annotation
            if return_annotation != inspect.Signature.empty:
                return_annotation_args = get_args(return_annotation)
                if not (
                    get_origin(return_annotation) is tuple
                    and len(return_annotation_args) == 2
                    and return_annotation_args[0] is bool
                    and (
                        return_annotation_args[1] is Any
                        or return_annotation_args[1] is str
                        or str(return_annotation_args[1]).endswith('TaskOutput')
                        or str(return_annotation_args[1]).startswith('typing.Union')
                    )
                ):
                    raise ValueError(
                        "If return type is annotated, it must be Tuple[bool, Any] or Tuple[bool, Union[str, TaskOutput]]"
                    )
            
            self._guardrail_fn = self.guardrail
        elif isinstance(self.guardrail, str):
            # Create LLM-based guardrail
            from ..guardrails import LLMGuardrail
            llm = getattr(self, 'llm', None) or getattr(self, 'llm_instance', None)
            self._guardrail_fn = LLMGuardrail(description=self.guardrail, llm=llm)
        else:
            raise ValueError("Agent guardrail must be either a callable or a string description")

    def _process_handoffs(self):
        """Process handoffs and convert them to tools that can be used by the agent."""
        if not self.handoffs:
            return
            
        # Import here to avoid circular imports
        from .handoff import Handoff
        
        for handoff_item in self.handoffs:
            try:
                if isinstance(handoff_item, Handoff):
                    # Convert Handoff object to a tool function
                    tool_func = handoff_item.to_tool_function(self)
                    self.tools.append(tool_func)
                elif hasattr(handoff_item, 'name') and hasattr(handoff_item, 'chat'):
                    # Direct agent reference - create a simple handoff
                    from .handoff import handoff
                    handoff_obj = handoff(handoff_item)
                    tool_func = handoff_obj.to_tool_function(self)
                    self.tools.append(tool_func)
                else:
                    logging.warning(
                        f"Invalid handoff item type: {type(handoff_item)}. "
                        "Expected Agent or Handoff instance."
                    )
            except Exception as e:
                logging.error(f"Failed to process handoff item {handoff_item}: {e}")

    def _process_guardrail(self, task_output):
        """Process the guardrail validation for a task output.
        
        Args:
            task_output: The task output to validate
            
        Returns:
            GuardrailResult: The result of the guardrail validation
        """
        from ..guardrails import GuardrailResult
        
        if not self._guardrail_fn:
            return GuardrailResult(success=True, result=task_output)
        
        try:
            # Call the guardrail function
            result = self._guardrail_fn(task_output)
            
            # Convert the result to a GuardrailResult
            return GuardrailResult.from_tuple(result)
            
        except Exception as e:
            logging.error(f"Agent {self.name}: Error in guardrail validation: {e}")
            # On error, return failure
            return GuardrailResult(
                success=False,
                result=None,
                error=f"Agent guardrail validation error: {str(e)}"
            )

    def _apply_guardrail_with_retry(self, response_text, prompt, temperature=1.0, tools=None, task_name=None, task_description=None, task_id=None):
        """Apply guardrail validation with retry logic.
        
        Args:
            response_text: The response to validate
            prompt: Original prompt for regeneration if needed
            temperature: Temperature for regeneration
            tools: Tools for regeneration
            
        Returns:
            str: The validated response text or None if validation fails after retries
        """
        if not self._guardrail_fn:
            return response_text
            
        from ..main import TaskOutput
        
        retry_count = 0
        current_response = response_text
        
        while retry_count <= self.max_guardrail_retries:
            # Create TaskOutput object
            task_output = TaskOutput(
                description="Agent response output",
                raw=current_response,
                agent=self.name
            )
            
            # Process guardrail
            guardrail_result = self._process_guardrail(task_output)
            
            if guardrail_result.success:
                logging.info(f"Agent {self.name}: Guardrail validation passed")
                # Return the potentially modified result
                if guardrail_result.result and hasattr(guardrail_result.result, 'raw'):
                    return guardrail_result.result.raw
                elif guardrail_result.result:
                    return str(guardrail_result.result)
                else:
                    return current_response
            
            # Guardrail failed
            if retry_count >= self.max_guardrail_retries:
                raise Exception(
                    f"Agent {self.name} response failed guardrail validation after {self.max_guardrail_retries} retries. "
                    f"Last error: {guardrail_result.error}"
                )
            
            retry_count += 1
            logging.warning(f"Agent {self.name}: Guardrail validation failed (retry {retry_count}/{self.max_guardrail_retries}): {guardrail_result.error}")
            
            # Regenerate response for retry
            try:
                retry_prompt = f"{prompt}\n\nNote: Previous response failed validation due to: {guardrail_result.error}. Please provide an improved response."
                response = self._chat_completion([{"role": "user", "content": retry_prompt}], temperature, tools, task_name=task_name, task_description=task_description, task_id=task_id)
                if response and response.choices:
                    current_response = response.choices[0].message.content.strip()
                else:
                    raise Exception("Failed to generate retry response")
            except Exception as e:
                logging.error(f"Agent {self.name}: Error during guardrail retry: {e}")
                # If we can't regenerate, fail the guardrail
                raise Exception(
                    f"Agent {self.name} guardrail retry failed: {e}"
                )
        
        return current_response
    
    def _get_tools_cache_key(self, tools):
        """Generate a cache key for tools list."""
        if tools is None:
            return "none"
        if not tools:
            return "empty"
        # Create a simple hash based on tool names
        tool_names = []
        for tool in tools:
            if callable(tool) and hasattr(tool, '__name__'):
                tool_names.append(tool.__name__)
            elif isinstance(tool, dict) and 'function' in tool and 'name' in tool['function']:
                tool_names.append(tool['function']['name'])
            elif isinstance(tool, str):
                tool_names.append(tool)
        return "|".join(sorted(tool_names))
    
    def _build_system_prompt(self, tools=None):
        """Build the system prompt with tool information.
        
        Args:
            tools: Optional list of tools to use (defaults to self.tools)
            
        Returns:
            str: The system prompt or None if use_system_prompt is False
        """
        if not self.use_system_prompt:
            return None
        
        # Check cache first (skip cache if memory is enabled since context is dynamic)
        if not self._memory_instance:
            tools_key = self._get_tools_cache_key(tools)
            cache_key = f"{self.role}:{self.goal}:{tools_key}"
            
            if cache_key in self._system_prompt_cache:
                return self._system_prompt_cache[cache_key]
        else:
            cache_key = None  # Don't cache when memory is enabled
            
        system_prompt = f"""{self.backstory}\n
Your Role: {self.role}\n
Your Goal: {self.goal}"""
        
        # Add rules context if rules manager is enabled
        if self._rules_manager:
            rules_context = self.get_rules_context()
            if rules_context:
                system_prompt += f"\n\n## Rules (Guidelines you must follow)\n{rules_context}"
        
        # Add memory context if memory is enabled
        if self._memory_instance:
            memory_context = self.get_memory_context()
            if memory_context:
                system_prompt += f"\n\n## Memory (Information you remember about the user)\n{memory_context}"
                # Display memory info to user if verbose
                if self.verbose:
                    self._display_memory_info()
        
        # Add tool usage instructions if tools are available
        # Use provided tools or fall back to self.tools
        tools_to_use = tools if tools is not None else self.tools
        if tools_to_use:
            tool_names = []
            for tool in tools_to_use:
                try:
                    if callable(tool) and hasattr(tool, '__name__'):
                        tool_names.append(tool.__name__)
                    elif isinstance(tool, dict) and isinstance(tool.get('function'), dict) and 'name' in tool['function']:
                        tool_names.append(tool['function']['name'])
                    elif isinstance(tool, str):
                        tool_names.append(tool)
                    elif hasattr(tool, "to_openai_tool"):
                        # Handle MCP tools
                        openai_tools = tool.to_openai_tool()
                        if isinstance(openai_tools, list):
                            for t in openai_tools:
                                if isinstance(t, dict) and 'function' in t and 'name' in t['function']:
                                    tool_names.append(t['function']['name'])
                        elif isinstance(openai_tools, dict) and 'function' in openai_tools:
                            tool_names.append(openai_tools['function']['name'])
                except (AttributeError, KeyError, TypeError) as e:
                    logging.warning(f"Could not extract tool name from {tool}: {e}")
                    continue
            
            if tool_names:
                system_prompt += f"\n\nYou have access to the following tools: {', '.join(tool_names)}. Use these tools when appropriate to help complete your tasks. Always use tools when they can help provide accurate information or perform actions."
        
        # Cache the generated system prompt (only if cache_key is set, i.e., memory not enabled)
        # Simple cache size limit to prevent unbounded growth
        if cache_key and len(self._system_prompt_cache) < self._max_cache_size:
            self._system_prompt_cache[cache_key] = system_prompt
        return system_prompt

    def _build_messages(self, prompt, temperature=1.0, output_json=None, output_pydantic=None, tools=None):
        """Build messages list for chat completion.
        
        Args:
            prompt: The user prompt (str or list)
            temperature: Temperature for the chat
            output_json: Optional Pydantic model for JSON output
            output_pydantic: Optional Pydantic model for JSON output (alias)
            tools: Optional list of tools to use (defaults to self.tools)
            
        Returns:
            tuple: (messages list, original prompt)
        """
        # Build system prompt using the helper method
        system_prompt = self._build_system_prompt(tools)
        
        # Use openai_client's build_messages method if available
        if self._openai_client is not None:
            messages, original_prompt = self._openai_client.build_messages(
                prompt=prompt,
                system_prompt=system_prompt,
                chat_history=self.chat_history,
                output_json=output_json,
                output_pydantic=output_pydantic
            )
        else:
            # Fallback implementation for when OpenAI client is not available
            messages = []
            
            # Add system message if provided
            if system_prompt:
                messages.append({"role": "system", "content": system_prompt})
            
            # Add chat history
            messages.extend(self.chat_history)
            
            # Add user prompt
            if isinstance(prompt, list):
                messages.extend(prompt)
                original_prompt = prompt
            else:
                messages.append({"role": "user", "content": str(prompt)})
                original_prompt = str(prompt)
            
            # Add JSON format instruction if needed
            if output_json or output_pydantic:
                model = output_pydantic or output_json
                json_instruction = f"\nPlease respond with valid JSON matching this schema: {model.model_json_schema()}"
                messages[-1]["content"] += json_instruction
        
        return messages, original_prompt

    def _format_tools_for_completion(self, tools=None):
        """Format tools for OpenAI completion API.
        
        Supports:
        - Pre-formatted OpenAI tools (dicts with type='function')
        - Lists of pre-formatted tools
        - Callable functions
        - String function names
        - Objects with to_openai_tool() method
        
        Args:
            tools: List of tools in various formats or None to use self.tools
            
        Returns:
            List of formatted tools or empty list
        """
        if tools is None:
            tools = self.tools
        
        if not tools:
            return []
        
        # Check cache first
        tools_key = self._get_tools_cache_key(tools)
        if tools_key in self._formatted_tools_cache:
            return self._formatted_tools_cache[tools_key]
            
        formatted_tools = []
        for tool in tools:
            # Handle pre-formatted OpenAI tools
            if isinstance(tool, dict) and tool.get('type') == 'function':
                # Validate nested dictionary structure before accessing
                if 'function' in tool and isinstance(tool['function'], dict) and 'name' in tool['function']:
                    formatted_tools.append(tool)
                else:
                    logging.warning(f"Skipping malformed OpenAI tool: missing function or name")
            # Handle lists of tools
            elif isinstance(tool, list):
                for subtool in tool:
                    if isinstance(subtool, dict) and subtool.get('type') == 'function':
                        # Validate nested dictionary structure before accessing
                        if 'function' in subtool and isinstance(subtool['function'], dict) and 'name' in subtool['function']:
                            formatted_tools.append(subtool)
                        else:
                            logging.warning(f"Skipping malformed OpenAI tool in list: missing function or name")
            # Handle string tool names
            elif isinstance(tool, str):
                tool_def = self._generate_tool_definition(tool)
                if tool_def:
                    formatted_tools.append(tool_def)
                else:
                    logging.warning(f"Could not generate definition for tool: {tool}")
            # Handle objects with to_openai_tool method (MCP tools)
            elif hasattr(tool, "to_openai_tool"):
                openai_tools = tool.to_openai_tool()
                # MCP tools can return either a single tool or a list of tools
                if isinstance(openai_tools, list):
                    formatted_tools.extend(openai_tools)
                elif openai_tools is not None:
                    formatted_tools.append(openai_tools)
            # Handle callable functions
            elif callable(tool):
                tool_def = self._generate_tool_definition(tool.__name__)
                if tool_def:
                    formatted_tools.append(tool_def)
            else:
                logging.warning(f"Tool {tool} not recognized")
        
        # Validate JSON serialization before returning
        if formatted_tools:
            try:
                json.dumps(formatted_tools)  # Validate serialization
            except (TypeError, ValueError) as e:
                logging.error(f"Tools are not JSON serializable: {e}")
                return []
        
        # Cache the formatted tools
        # Simple cache size limit to prevent unbounded growth
        if len(self._formatted_tools_cache) < self._max_cache_size:
            self._formatted_tools_cache[tools_key] = formatted_tools
        return formatted_tools

    def generate_task(self) -> 'Task':
        """Generate a Task object from the agent's instructions"""
        from ..task.task import Task
        
        description = self.instructions if self.instructions else f"Execute task as {self.role} with goal: {self.goal}"
        expected_output = "Complete the assigned task successfully"
        
        return Task(
            name=self.name,
            description=description,
            expected_output=expected_output,
            agent=self,
            tools=self.tools
        )

    def _resolve_tool_names(self, tool_names):
        """Resolve tool names to actual tool instances from registry.
        
        Args:
            tool_names: List of tool name strings
            
        Returns:
            List of resolved tool instances
        """
        resolved = []
        try:
            from ..tools.registry import get_registry
            registry = get_registry()
            
            for name in tool_names:
                tool = registry.get(name)
                if tool is not None:
                    resolved.append(tool)
                else:
                    logging.warning(f"Tool '{name}' not found in registry")
        except ImportError:
            logging.warning("Tool registry not available, cannot resolve tool names")
        
        return resolved

    def _cast_arguments(self, func, arguments):
        """Cast arguments to their expected types based on function signature."""
        if not callable(func) or not arguments:
            return arguments
        
        try:
            sig = inspect.signature(func)
            casted_args = {}
            
            for param_name, arg_value in arguments.items():
                if param_name in sig.parameters:
                    param = sig.parameters[param_name]
                    if param.annotation != inspect.Parameter.empty:
                        # Handle common type conversions
                        if param.annotation == int and isinstance(arg_value, (str, float)):
                            try:
                                casted_args[param_name] = int(float(arg_value))
                            except (ValueError, TypeError):
                                casted_args[param_name] = arg_value
                        elif param.annotation == float and isinstance(arg_value, (str, int)):
                            try:
                                casted_args[param_name] = float(arg_value)
                            except (ValueError, TypeError):
                                casted_args[param_name] = arg_value
                        elif param.annotation == bool and isinstance(arg_value, str):
                            casted_args[param_name] = arg_value.lower() in ('true', '1', 'yes', 'on')
                        else:
                            casted_args[param_name] = arg_value
                    else:
                        casted_args[param_name] = arg_value
                else:
                    casted_args[param_name] = arg_value
            
            return casted_args
        except Exception as e:
            logging.debug(f"Type casting failed for {getattr(func, '__name__', 'unknown function')}: {e}")
            return arguments

    def execute_tool(self, function_name, arguments):
        """
        Execute a tool dynamically based on the function name and arguments.
        """
        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
        if isinstance(self.tools, MCP):
            logging.debug(f"Looking for MCP tool {function_name}")
            
            # Handle SSE MCP client
            if hasattr(self.tools, 'is_sse') and self.tools.is_sse:
                if hasattr(self.tools, 'sse_client'):
                    for tool in self.tools.sse_client.tools:
                        if tool.name == function_name:
                            logging.debug(f"Found matching SSE MCP tool: {function_name}")
                            return tool(**arguments)
            # Handle stdio MCP client
            elif hasattr(self.tools, 'runner'):
                # Check if any of the MCP tools match the function name
                for mcp_tool in self.tools.runner.tools:
                    if hasattr(mcp_tool, 'name') and mcp_tool.name == function_name:
                        logging.debug(f"Found matching MCP tool: {function_name}")
                        return self.tools.runner.call_tool(function_name, arguments)

        # Try to find the function in the agent's tools list first
        func = None
        for tool in self.tools if isinstance(self.tools, (list, tuple)) else []:
            # Check for BaseTool instances (plugin system)
            from ..tools.base import BaseTool
            if isinstance(tool, BaseTool) and tool.name == function_name:
                func = tool
                break
            # Check for FunctionTool (decorated functions)
            if hasattr(tool, 'name') and getattr(tool, 'name', None) == function_name:
                func = tool
                break
            if (callable(tool) and getattr(tool, '__name__', '') == function_name) or \
               (inspect.isclass(tool) and tool.__name__ == function_name):
                func = tool
                break
        
        if func is None:
            # Check the global tool registry for plugins
            try:
                from ..tools.registry import get_registry
                registry = get_registry()
                func = registry.get(function_name)
            except ImportError:
                pass
        
        if func is None:
            # If not found in tools or registry, try globals and main
            func = globals().get(function_name)
            if not func:
                import __main__
                func = getattr(__main__, function_name, None)

        if func:
            try:
                # BaseTool instances (plugin system) - call run() method
                from ..tools.base import BaseTool
                if isinstance(func, BaseTool):
                    casted_arguments = self._cast_arguments(func.run, arguments)
                    return func.run(**casted_arguments)
                
                # Langchain: If it's a class with run but not _run, instantiate and call run
                if inspect.isclass(func) and hasattr(func, 'run') and not hasattr(func, '_run'):
                    instance = func()
                    run_params = {k: v for k, v in arguments.items() 
                                  if k in inspect.signature(instance.run).parameters 
                                  and k != 'self'}
                    casted_params = self._cast_arguments(instance.run, run_params)
                    return instance.run(**casted_params)

                # CrewAI: If it's a class with an _run method, instantiate and call _run
                elif inspect.isclass(func) and hasattr(func, '_run'):
                    instance = func()
                    run_params = {k: v for k, v in arguments.items() 
                                  if k in inspect.signature(instance._run).parameters 
                                  and k != 'self'}
                    casted_params = self._cast_arguments(instance._run, run_params)
                    return instance._run(**casted_params)

                # Otherwise treat as regular function
                elif callable(func):
                    casted_arguments = self._cast_arguments(func, arguments)
                    return func(**casted_arguments)
            except Exception as e:
                error_msg = str(e)
                logging.error(f"Error executing tool {function_name}: {error_msg}")
                return {"error": error_msg}
        
        error_msg = f"Tool '{function_name}' is not callable"
        logging.error(error_msg)
        return {"error": error_msg}

    def clear_history(self):
        self.chat_history = []

    def __str__(self):
        return f"Agent(name='{self.name}', role='{self.role}', goal='{self.goal}')"

    def _process_stream_response(self, messages, temperature, start_time, formatted_tools=None, reasoning_steps=False):
        """Process streaming response and return final response"""
        if self._openai_client is None:
            raise ValueError("OpenAI client is not initialized. Please provide OPENAI_API_KEY or use a custom LLM provider.")
        
        return self._openai_client.process_stream_response(
            messages=messages,
            model=self.llm,
            temperature=temperature,
            tools=formatted_tools,
            start_time=start_time,
            console=self.console,
            display_fn=self.display_generating if self.verbose else None,
            reasoning_steps=reasoning_steps
        )

    def _chat_completion(self, messages, temperature=1.0, tools=None, stream=True, reasoning_steps=False, task_name=None, task_description=None, task_id=None):
        start_time = time.time()
        logging.debug(f"{self.name} sending messages to LLM: {messages}")

        # Use the new _format_tools_for_completion helper method
        formatted_tools = self._format_tools_for_completion(tools)

        try:
            # Use the custom LLM instance if available
            if self._using_custom_llm and hasattr(self, 'llm_instance'):
                if stream:
                    # Debug logs for tool info
                    if formatted_tools:
                        logging.debug(f"Passing {len(formatted_tools)} formatted tools to LLM instance: {formatted_tools}")
                    
                    # Use the LLM instance for streaming responses
                    final_response = self.llm_instance.get_response(
                        prompt=messages[1:],  # Skip system message as LLM handles it separately  
                        system_prompt=messages[0]['content'] if messages and messages[0]['role'] == 'system' else None,
                        temperature=temperature,
                        tools=formatted_tools if formatted_tools else None,
                        verbose=self.verbose,
                        markdown=self.markdown,
                        stream=stream,
                        console=self.console,
                        execute_tool_fn=self.execute_tool,
                        agent_name=self.name,
                        agent_role=self.role,
                        agent_tools=[t.__name__ for t in self.tools] if self.tools else None,
                        task_name=task_name,
                        task_description=task_description,
                        task_id=task_id,
                        reasoning_steps=reasoning_steps
                    )
                else:
                    # Non-streaming with custom LLM - don't show streaming-like behavior
                    if False:  # Don't use display_generating when stream=False to avoid streaming-like behavior
                        # This block is disabled to maintain consistency with the OpenAI path fix
                        with Live(
                            display_generating("", start_time),
                            console=self.console,
                            refresh_per_second=4,
                        ) as live:
                            final_response = self.llm_instance.get_response(
                                prompt=messages[1:],
                                system_prompt=messages[0]['content'] if messages and messages[0]['role'] == 'system' else None,
                                temperature=temperature,
                                tools=formatted_tools if formatted_tools else None,
                                verbose=self.verbose,
                                markdown=self.markdown,
                                stream=stream,
                                console=self.console,
                                execute_tool_fn=self.execute_tool,
                                agent_name=self.name,
                                agent_role=self.role,
                                agent_tools=[t.__name__ for t in self.tools] if self.tools else None,
                                task_name=task_name,
                                task_description=task_description,
                                task_id=task_id,
                                reasoning_steps=reasoning_steps
                            )
                    else:
                        final_response = self.llm_instance.get_response(
                            prompt=messages[1:],
                            system_prompt=messages[0]['content'] if messages and messages[0]['role'] == 'system' else None,
                            temperature=temperature,
                            tools=formatted_tools if formatted_tools else None,
                            verbose=self.verbose,
                            markdown=self.markdown,
                            stream=stream,
                            console=self.console,
                            execute_tool_fn=self.execute_tool,
                            agent_name=self.name,
                            agent_role=self.role,
                            agent_tools=[t.__name__ for t in self.tools] if self.tools else None,
                            task_name=task_name,
                            task_description=task_description,
                            task_id=task_id,
                            reasoning_steps=reasoning_steps
                        )
            else:
                # Use the standard OpenAI client approach with tool support
                # Note: openai_client expects tools in various formats and will format them internally
                # But since we already have formatted_tools, we can pass them directly
                if self._openai_client is None:
                    raise ValueError("OpenAI client is not initialized. Please provide OPENAI_API_KEY or use a custom LLM provider.")
                
                final_response = self._openai_client.chat_completion_with_tools(
                    messages=messages,
                    model=self.llm,
                    temperature=temperature,
                    tools=formatted_tools,  # Already formatted for OpenAI
                    execute_tool_fn=self.execute_tool,
                    stream=stream,
                    console=self.console if (self.verbose or stream) else None,
                    display_fn=self.display_generating if self.verbose else None,
                    reasoning_steps=reasoning_steps,
                    verbose=self.verbose,
                    max_iterations=10
                )

            return final_response

        except Exception as e:
            display_error(f"Error in chat completion: {e}")
            return None
    
    def _execute_callback_and_display(self, prompt: str, response: str, generation_time: float, task_name=None, task_description=None, task_id=None):
        """Helper method to execute callbacks and display interaction.
        
        This centralizes the logic for callback execution and display to avoid duplication.
        """
        # Always execute callbacks regardless of verbose setting (only when not using custom LLM)
        if not self._using_custom_llm:
            execute_sync_callback(
                'interaction',
                message=prompt,
                response=response,
                markdown=self.markdown,
                generation_time=generation_time,
                agent_name=self.name,
                agent_role=self.role,
                agent_tools=[t.__name__ for t in self.tools] if self.tools else None,
                task_name=task_name,
                task_description=task_description, 
                task_id=task_id
            )
        # Always display final interaction when verbose is True to ensure consistent formatting
        # This ensures both OpenAI and custom LLM providers (like Gemini) show formatted output
        if self.verbose and not self._final_display_shown:
            display_interaction(prompt, response, markdown=self.markdown, 
                              generation_time=generation_time, console=self.console,
                              agent_name=self.name,
                              agent_role=self.role,
                              agent_tools=[t.__name__ for t in self.tools] if self.tools else None,
                              task_name=None,  # Not available in this context
                              task_description=None,  # Not available in this context
                              task_id=None)  # Not available in this context
            self._final_display_shown = True
    
    def display_generating(self, content: str, start_time: float):
        """Display function for generating animation with agent info."""
        from rich.panel import Panel
        from rich.markdown import Markdown
        elapsed = time.time() - start_time
        
        # Show content if provided (for both streaming and progressive display)
        if content:
            display_content = Markdown(content) if self.markdown else content
            return Panel(
                display_content,
                title=f"[bold]{self.name}[/bold] - Generating... {elapsed:.1f}s",
                border_style="green",
                expand=False
            )
        # else:
        #     # No content yet: show generating message
        #     return Panel(
        #         f"[bold cyan]Generating response...[/bold cyan]",
        #         title=f"[bold]{self.name}[/bold] - {elapsed:.1f}s",
        #         border_style="cyan",
        #         expand=False
        #     )

    def chat(self, prompt, temperature=1.0, tools=None, output_json=None, output_pydantic=None, reasoning_steps=False, stream=None, task_name=None, task_description=None, task_id=None):
        # Reset the final display flag for each new conversation
        self._final_display_shown = False
        
        # Log all parameter values when in debug mode
        if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
            param_info = {
                "prompt": str(prompt)[:100] + "..." if isinstance(prompt, str) and len(str(prompt)) > 100 else str(prompt),
                "temperature": temperature,
                "tools": [t.__name__ if hasattr(t, "__name__") else str(t) for t in tools] if tools else None,
                "output_json": str(output_json.__class__.__name__) if output_json else None,
                "output_pydantic": str(output_pydantic.__class__.__name__) if output_pydantic else None,
                "reasoning_steps": reasoning_steps,
                "agent_name": self.name,
                "agent_role": self.role,
                "agent_goal": self.goal
            }
            logging.debug(f"Agent.chat parameters: {json.dumps(param_info, indent=2, default=str)}")
        
        start_time = time.time()
        reasoning_steps = reasoning_steps or self.reasoning_steps
        # Use agent's stream setting if not explicitly provided
        if stream is None:
            stream = self.stream
        # Search for existing knowledge if any knowledge is provided
        if self._knowledge_sources and not self._knowledge_processed:
            self._ensure_knowledge_processed()
        
        if self.knowledge:
            search_results = self.knowledge.search(prompt, agent_id=self.agent_id)
            if search_results:
                # Check if search_results is a list of dictionaries or strings
                if isinstance(search_results, dict) and 'results' in search_results:
                    # Extract memory content from the results
                    knowledge_content = "\n".join([result['memory'] for result in search_results['results']])
                else:
                    # If search_results is a list of strings, join them directly
                    knowledge_content = "\n".join(search_results)
                
                # Append found knowledge to the prompt
                prompt = f"{prompt}\n\nKnowledge: {knowledge_content}"

        if self._using_custom_llm:
            try:
                # Special handling for MCP tools when using provider/model format
                # Fix: Handle empty tools list properly - use self.tools if tools is None or empty
                if tools is None or (isinstance(tools, list) and len(tools) == 0):
                    tool_param = self.tools
                else:
                    tool_param = tools
                
                # Convert MCP tool objects to OpenAI format if needed
                if tool_param is not None:
                    from ..mcp.mcp import MCP
                    if isinstance(tool_param, MCP) and hasattr(tool_param, 'to_openai_tool'):
                        logging.debug("Converting MCP tool to OpenAI format")
                        openai_tool = tool_param.to_openai_tool()
                        if openai_tool:
                            # Handle both single tool and list of tools
                            if isinstance(openai_tool, list):
                                tool_param = openai_tool
                            else:
                                tool_param = [openai_tool]
                            logging.debug(f"Converted MCP tool: {tool_param}")
                
                # Store chat history length for potential rollback
                chat_history_length = len(self.chat_history)
                
                # Normalize prompt content for consistent chat history storage
                normalized_content = prompt
                if isinstance(prompt, list):
                    # Extract text from multimodal prompts
                    normalized_content = next((item["text"] for item in prompt if item.get("type") == "text"), "")
                
                # Prevent duplicate messages
                if not (self.chat_history and 
                        self.chat_history[-1].get("role") == "user" and 
                        self.chat_history[-1].get("content") == normalized_content):
                    # Add user message to chat history BEFORE LLM call so handoffs can access it
                    self.chat_history.append({"role": "user", "content": normalized_content})
                
                try:
                    # Pass everything to LLM class
                    response_text = self.llm_instance.get_response(
                    prompt=prompt,
                    system_prompt=self._build_system_prompt(tools),
                    chat_history=self.chat_history,
                    temperature=temperature,
                    tools=tool_param,
                    output_json=output_json,
                    output_pydantic=output_pydantic,
                    verbose=self.verbose,
                    markdown=self.markdown,
                    self_reflect=self.self_reflect,
                    max_reflect=self.max_reflect,
                    min_reflect=self.min_reflect,
                    console=self.console,
                    agent_name=self.name,
                    agent_role=self.role,
                    agent_tools=[t.__name__ if hasattr(t, '__name__') else str(t) for t in (tools if tools is not None else self.tools)],
                    task_name=task_name,
                    task_description=task_description,
                    task_id=task_id,
                    execute_tool_fn=self.execute_tool,  # Pass tool execution function
                    reasoning_steps=reasoning_steps,
                    stream=stream  # Pass the stream parameter from chat method
                    )

                    self.chat_history.append({"role": "assistant", "content": response_text})

                    # Log completion time if in debug mode
                    if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
                        total_time = time.time() - start_time
                        logging.debug(f"Agent.chat completed in {total_time:.2f} seconds")

                    # Apply guardrail validation for custom LLM response
                    try:
                        validated_response = self._apply_guardrail_with_retry(response_text, prompt, temperature, tools, task_name, task_description, task_id)
                        # Execute callback and display after validation
                        self._execute_callback_and_display(prompt, validated_response, time.time() - start_time, task_name, task_description, task_id)
                        return validated_response
                    except Exception as e:
                        logging.error(f"Agent {self.name}: Guardrail validation failed for custom LLM: {e}")
                        # Rollback chat history on guardrail failure
                        self.chat_history = self.chat_history[:chat_history_length]
                        return None
                except Exception as e:
                    # Rollback chat history if LLM call fails
                    self.chat_history = self.chat_history[:chat_history_length]
                    display_error(f"Error in LLM chat: {e}")
                    return None
            except Exception as e:
                display_error(f"Error in LLM chat: {e}")
                return None
        else:
            # Use the new _build_messages helper method
            messages, original_prompt = self._build_messages(prompt, temperature, output_json, output_pydantic)
            
            # Store chat history length for potential rollback
            chat_history_length = len(self.chat_history)
            
            # Normalize original_prompt for consistent chat history storage
            normalized_content = original_prompt
            if isinstance(original_prompt, list):
                # Extract text from multimodal prompts
                normalized_content = next((item["text"] for item in original_prompt if item.get("type") == "text"), "")
            
            # Prevent duplicate messages
            if not (self.chat_history and 
                    self.chat_history[-1].get("role") == "user" and 
                    self.chat_history[-1].get("content") == normalized_content):
                # Add user message to chat history BEFORE LLM call so handoffs can access it
                self.chat_history.append({"role": "user", "content": normalized_content})

            reflection_count = 0
            start_time = time.time()
            
            # Wrap entire while loop in try-except for rollback on any failure
            try:
                while True:
                    try:
                        if self.verbose:
                            # Handle both string and list prompts for instruction display
                            display_text = prompt
                            if isinstance(prompt, list):
                                # Extract text content from multimodal prompt
                                display_text = next((item["text"] for item in prompt if item["type"] == "text"), "")
                            
                            if display_text and str(display_text).strip():
                                # Pass agent information to display_instruction
                                agent_tools = [t.__name__ if hasattr(t, '__name__') else str(t) for t in self.tools]
                                display_instruction(
                                    f"Agent {self.name} is processing prompt: {display_text}", 
                                    console=self.console,
                                    agent_name=self.name,
                                    agent_role=self.role,
                                    agent_tools=agent_tools
                                )

                        response = self._chat_completion(messages, temperature=temperature, tools=tools if tools else None, reasoning_steps=reasoning_steps, stream=stream, task_name=task_name, task_description=task_description, task_id=task_id)
                        if not response:
                            # Rollback chat history on response failure
                            self.chat_history = self.chat_history[:chat_history_length]
                            return None

                        response_text = response.choices[0].message.content.strip()

                        # Handle output_json or output_pydantic if specified
                        if output_json or output_pydantic:
                            # Add to chat history and return raw response
                            # User message already added before LLM call via _build_messages
                            self.chat_history.append({"role": "assistant", "content": response_text})
                            # Apply guardrail validation even for JSON output
                            try:
                                validated_response = self._apply_guardrail_with_retry(response_text, original_prompt, temperature, tools, task_name, task_description, task_id)
                                # Execute callback after validation
                                self._execute_callback_and_display(original_prompt, validated_response, time.time() - start_time, task_name, task_description, task_id)
                                return validated_response
                            except Exception as e:
                                logging.error(f"Agent {self.name}: Guardrail validation failed for JSON output: {e}")
                                # Rollback chat history on guardrail failure
                                self.chat_history = self.chat_history[:chat_history_length]
                                return None

                        if not self.self_reflect:
                            # User message already added before LLM call via _build_messages
                            self.chat_history.append({"role": "assistant", "content": response_text})
                            if self.verbose:
                                logging.debug(f"Agent {self.name} final response: {response_text}")
                            # Return only reasoning content if reasoning_steps is True
                            if reasoning_steps and hasattr(response.choices[0].message, 'reasoning_content'):
                                # Apply guardrail to reasoning content
                                try:
                                    validated_reasoning = self._apply_guardrail_with_retry(response.choices[0].message.reasoning_content, original_prompt, temperature, tools, task_name, task_description, task_id)
                                    # Execute callback after validation
                                    self._execute_callback_and_display(original_prompt, validated_reasoning, time.time() - start_time, task_name, task_description, task_id)
                                    return validated_reasoning
                                except Exception as e:
                                    logging.error(f"Agent {self.name}: Guardrail validation failed for reasoning content: {e}")
                                    # Rollback chat history on guardrail failure
                                    self.chat_history = self.chat_history[:chat_history_length]
                                    return None
                            # Apply guardrail to regular response
                            try:
                                validated_response = self._apply_guardrail_with_retry(response_text, original_prompt, temperature, tools, task_name, task_description, task_id)
                                # Execute callback after validation
                                self._execute_callback_and_display(original_prompt, validated_response, time.time() - start_time, task_name, task_description, task_id)
                                return validated_response
                            except Exception as e:
                                logging.error(f"Agent {self.name}: Guardrail validation failed: {e}")
                                # Rollback chat history on guardrail failure
                                self.chat_history = self.chat_history[:chat_history_length]
                                return None

                        reflection_prompt = f"""
Reflect on your previous response: '{response_text}'.
{self.reflect_prompt if self.reflect_prompt else "Identify any flaws, improvements, or actions."}
Provide a "satisfactory" status ('yes' or 'no').
Output MUST be JSON with 'reflection' and 'satisfactory'.
                        """
                        logging.debug(f"{self.name} reflection attempt {reflection_count+1}, sending prompt: {reflection_prompt}")
                        messages.append({"role": "user", "content": reflection_prompt})

                        try:
                            # Check if we're using a custom LLM (like Gemini)
                            if self._using_custom_llm or self._openai_client is None:
                                # For custom LLMs, we need to handle reflection differently
                                # Use non-streaming to get complete JSON response
                                reflection_response = self._chat_completion(messages, temperature=temperature, tools=None, stream=False, reasoning_steps=False, task_name=task_name, task_description=task_description, task_id=task_id)
                                
                                if not reflection_response or not reflection_response.choices:
                                    raise Exception("No response from reflection request")
                                
                                reflection_text = reflection_response.choices[0].message.content.strip()
                                
                                # Clean the JSON output
                                cleaned_json = self.clean_json_output(reflection_text)
                                
                                # Parse the JSON manually
                                reflection_data = json.loads(cleaned_json)
                                
                                # Create a reflection output object manually
                                class CustomReflectionOutput:
                                    def __init__(self, data):
                                        self.reflection = data.get('reflection', '')
                                        self.satisfactory = data.get('satisfactory', 'no').lower()
                                
                                reflection_output = CustomReflectionOutput(reflection_data)
                            else:
                                # Use OpenAI's structured output for OpenAI models
                                reflection_response = self._openai_client.sync_client.beta.chat.completions.parse(
                                    model=self.reflect_llm if self.reflect_llm else self.llm,
                                    messages=messages,
                                    temperature=temperature,
                                    response_format=ReflectionOutput
                                )

                                reflection_output = reflection_response.choices[0].message.parsed

                            if self.verbose:
                                display_self_reflection(f"Agent {self.name} self reflection (using {self.reflect_llm if self.reflect_llm else self.llm}): reflection='{reflection_output.reflection}' satisfactory='{reflection_output.satisfactory}'", console=self.console)

                            messages.append({"role": "assistant", "content": f"Self Reflection: {reflection_output.reflection} Satisfactory?: {reflection_output.satisfactory}"})

                            # Only consider satisfactory after minimum reflections
                            if reflection_output.satisfactory == "yes" and reflection_count >= self.min_reflect - 1:
                                if self.verbose:
                                    display_self_reflection("Agent marked the response as satisfactory after meeting minimum reflections", console=self.console)
                                # User message already added before LLM call via _build_messages
                                self.chat_history.append({"role": "assistant", "content": response_text})
                                # Apply guardrail validation after satisfactory reflection
                                try:
                                    validated_response = self._apply_guardrail_with_retry(response_text, original_prompt, temperature, tools, task_name, task_description, task_id)
                                    # Execute callback after validation
                                    self._execute_callback_and_display(original_prompt, validated_response, time.time() - start_time, task_name, task_description, task_id)
                                    return validated_response
                                except Exception as e:
                                    logging.error(f"Agent {self.name}: Guardrail validation failed after reflection: {e}")
                                    # Rollback chat history on guardrail failure
                                    self.chat_history = self.chat_history[:chat_history_length]
                                    return None

                            # Check if we've hit max reflections
                            if reflection_count >= self.max_reflect - 1:
                                if self.verbose:
                                    display_self_reflection("Maximum reflection count reached, returning current response", console=self.console)
                                # User message already added before LLM call via _build_messages
                                self.chat_history.append({"role": "assistant", "content": response_text})
                                # Apply guardrail validation after max reflections
                                try:
                                    validated_response = self._apply_guardrail_with_retry(response_text, original_prompt, temperature, tools, task_name, task_description, task_id)
                                    # Execute callback after validation
                                    self._execute_callback_and_display(original_prompt, validated_response, time.time() - start_time, task_name, task_description, task_id)
                                    return validated_response
                                except Exception as e:
                                    logging.error(f"Agent {self.name}: Guardrail validation failed after max reflections: {e}")
                                    # Rollback chat history on guardrail failure
                                    self.chat_history = self.chat_history[:chat_history_length]
                                    return None
                            
                            # If not satisfactory and not at max reflections, continue with regeneration
                            logging.debug(f"{self.name} reflection count {reflection_count + 1}, continuing reflection process")
                            messages.append({"role": "user", "content": "Now regenerate your response using the reflection you made"})
                            # For custom LLMs during reflection, always use non-streaming to ensure complete responses
                            use_stream = self.stream if not self._using_custom_llm else False
                            response = self._chat_completion(messages, temperature=temperature, tools=None, stream=use_stream, task_name=task_name, task_description=task_description, task_id=task_id)
                            response_text = response.choices[0].message.content.strip()
                            reflection_count += 1
                            continue  # Continue the loop for more reflections

                        except Exception as e:
                                display_error(f"Error in parsing self-reflection json {e}. Retrying", console=self.console)
                                logging.error("Reflection parsing failed.", exc_info=True)
                                messages.append({"role": "assistant", "content": "Self Reflection failed."})
                                reflection_count += 1
                                continue  # Continue even after error to try again
                    except Exception:
                        # Catch any exception from the inner try block and re-raise to outer handler
                        raise
            except Exception as e:
                # Catch any exceptions that escape the while loop
                display_error(f"Unexpected error in chat: {e}", console=self.console)
                # Rollback chat history
                self.chat_history = self.chat_history[:chat_history_length]
                return None

    def clean_json_output(self, output: str) -> str:
        """Clean and extract JSON from response text."""
        cleaned = output.strip()
        # Remove markdown code blocks if present
        if cleaned.startswith("```json"):
            cleaned = cleaned[len("```json"):].strip()
        if cleaned.startswith("```"):
            cleaned = cleaned[len("```"):].strip()
        if cleaned.endswith("```"):
            cleaned = cleaned[:-3].strip()
        return cleaned  

    async def achat(self, prompt: str, temperature=1.0, tools=None, output_json=None, output_pydantic=None, reasoning_steps=False, task_name=None, task_description=None, task_id=None):
        """Async version of chat method with self-reflection support.""" 
        # Reset the final display flag for each new conversation
        self._final_display_shown = False
        
        # Log all parameter values when in debug mode
        if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
            param_info = {
                "prompt": str(prompt)[:100] + "..." if isinstance(prompt, str) and len(str(prompt)) > 100 else str(prompt),
                "temperature": temperature,
                "tools": [t.__name__ if hasattr(t, "__name__") else str(t) for t in tools] if tools else None,
                "output_json": str(output_json.__class__.__name__) if output_json else None,
                "output_pydantic": str(output_pydantic.__class__.__name__) if output_pydantic else None,
                "reasoning_steps": reasoning_steps,
                "agent_name": self.name,
                "agent_role": self.role,
                "agent_goal": self.goal
            }
            logging.debug(f"Agent.achat parameters: {json.dumps(param_info, indent=2, default=str)}")
        
        start_time = time.time()
        reasoning_steps = reasoning_steps or self.reasoning_steps
        try:
            # Default to self.tools if tools argument is None
            if tools is None:
                tools = self.tools

            # Search for existing knowledge if any knowledge is provided
            if self._knowledge_sources and not self._knowledge_processed:
                self._ensure_knowledge_processed()
            
            if self.knowledge:
                search_results = self.knowledge.search(prompt, agent_id=self.agent_id)
                if search_results:
                    if isinstance(search_results, dict) and 'results' in search_results:
                        knowledge_content = "\n".join([result['memory'] for result in search_results['results']])
                    else:
                        knowledge_content = "\n".join(search_results)
                    prompt = f"{prompt}\n\nKnowledge: {knowledge_content}"

            if self._using_custom_llm:
                # Store chat history length for potential rollback
                chat_history_length = len(self.chat_history)
                
                # Normalize prompt content for consistent chat history storage
                normalized_content = prompt
                if isinstance(prompt, list):
                    # Extract text from multimodal prompts
                    normalized_content = next((item["text"] for item in prompt if item.get("type") == "text"), "")
                
                # Prevent duplicate messages
                if not (self.chat_history and 
                        self.chat_history[-1].get("role") == "user" and 
                        self.chat_history[-1].get("content") == normalized_content):
                    # Add user message to chat history BEFORE LLM call so handoffs can access it
                    self.chat_history.append({"role": "user", "content": normalized_content})
                
                try:
                    response_text = await self.llm_instance.get_response_async(
                        prompt=prompt,
                        system_prompt=self._build_system_prompt(tools),
                        chat_history=self.chat_history,
                        temperature=temperature,
                        tools=tools,
                        output_json=output_json,
                        output_pydantic=output_pydantic,
                        verbose=self.verbose,
                        markdown=self.markdown,
                        self_reflect=self.self_reflect,
                        max_reflect=self.max_reflect,
                        min_reflect=self.min_reflect,
                        console=self.console,
                        agent_name=self.name,
                        agent_role=self.role,
                        agent_tools=[t.__name__ if hasattr(t, '__name__') else str(t) for t in (tools if tools is not None else self.tools)],
                        task_name=task_name,
                        task_description=task_description,
                        task_id=task_id,
                        execute_tool_fn=self.execute_tool_async,
                        reasoning_steps=reasoning_steps
                    )

                    self.chat_history.append({"role": "assistant", "content": response_text})

                    if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
                        total_time = time.time() - start_time
                        logging.debug(f"Agent.achat completed in {total_time:.2f} seconds")
                    
                    # Apply guardrail validation for custom LLM response
                    try:
                        validated_response = self._apply_guardrail_with_retry(response_text, prompt, temperature, tools, task_name, task_description, task_id)
                        # Execute callback after validation
                        self._execute_callback_and_display(normalized_content, validated_response, time.time() - start_time, task_name, task_description, task_id)
                        return validated_response
                    except Exception as e:
                        logging.error(f"Agent {self.name}: Guardrail validation failed for custom LLM: {e}")
                        # Rollback chat history on guardrail failure
                        self.chat_history = self.chat_history[:chat_history_length]
                        return None
                except Exception as e:
                    # Rollback chat history if LLM call fails
                    self.chat_history = self.chat_history[:chat_history_length]
                    display_error(f"Error in LLM chat: {e}")
                    if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
                        total_time = time.time() - start_time
                        logging.debug(f"Agent.achat failed in {total_time:.2f} seconds: {str(e)}")
                    return None

            # For OpenAI client
            # Use the new _build_messages helper method
            messages, original_prompt = self._build_messages(prompt, temperature, output_json, output_pydantic)
            
            # Store chat history length for potential rollback
            chat_history_length = len(self.chat_history)
            
            # Normalize original_prompt for consistent chat history storage
            normalized_content = original_prompt
            if isinstance(original_prompt, list):
                # Extract text from multimodal prompts
                normalized_content = next((item["text"] for item in original_prompt if item.get("type") == "text"), "")
            
            # Prevent duplicate messages
            if not (self.chat_history and 
                    self.chat_history[-1].get("role") == "user" and 
                    self.chat_history[-1].get("content") == normalized_content):
                # Add user message to chat history BEFORE LLM call so handoffs can access it
                self.chat_history.append({"role": "user", "content": normalized_content})

            reflection_count = 0
            start_time = time.time()

            while True:
                try:
                    if self.verbose:
                        display_text = prompt
                        if isinstance(prompt, list):
                            display_text = next((item["text"] for item in prompt if item["type"] == "text"), "")
                        
                        if display_text and str(display_text).strip():
                            agent_tools = [t.__name__ if hasattr(t, '__name__') else str(t) for t in self.tools]
                            await adisplay_instruction(
                                f"Agent {self.name} is processing prompt: {display_text}",
                                console=self.console,
                                agent_name=self.name,
                                agent_role=self.role,
                                agent_tools=agent_tools
                            )

                    # Use the new _format_tools_for_completion helper method
                    formatted_tools = self._format_tools_for_completion(tools)
                    
                    # Check if OpenAI client is available
                    if self._openai_client is None:
                        error_msg = "OpenAI client is not initialized. Please provide OPENAI_API_KEY or use a custom LLM provider."
                        display_error(error_msg)
                        return None

                    # Make the API call based on the type of request
                    if tools:
                        response = await self._openai_client.async_client.chat.completions.create(
                            model=self.llm,
                            messages=messages,
                            temperature=temperature,
                            tools=formatted_tools,
                        )
                        result = await self._achat_completion(response, tools)
                        if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
                            total_time = time.time() - start_time
                            logging.debug(f"Agent.achat completed in {total_time:.2f} seconds")
                        # Execute callback after tool completion
                        self._execute_callback_and_display(original_prompt, result, time.time() - start_time, task_name, task_description, task_id)
                        return result
                    elif output_json or output_pydantic:
                        response = await self._openai_client.async_client.chat.completions.create(
                            model=self.llm,
                            messages=messages,
                            temperature=temperature,
                            response_format={"type": "json_object"}
                        )
                        response_text = response.choices[0].message.content
                        if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
                            total_time = time.time() - start_time
                            logging.debug(f"Agent.achat completed in {total_time:.2f} seconds")
                        # Execute callback after JSON/Pydantic completion
                        self._execute_callback_and_display(original_prompt, response_text, time.time() - start_time, task_name, task_description, task_id)
                        return response_text
                    else:
                        response = await self._openai_client.async_client.chat.completions.create(
                            model=self.llm,
                            messages=messages,
                            temperature=temperature
                        )
                        
                        response_text = response.choices[0].message.content
                        
                        # Handle self-reflection if enabled
                        if self.self_reflect:
                            reflection_count = 0
                            
                            while True:
                                reflection_prompt = f"""
Reflect on your previous response: '{response_text}'.
{self.reflect_prompt if self.reflect_prompt else "Identify any flaws, improvements, or actions."}
Provide a "satisfactory" status ('yes' or 'no').
Output MUST be JSON with 'reflection' and 'satisfactory'.
                                """
                                
                                # Add reflection prompt to messages
                                reflection_messages = messages + [
                                    {"role": "assistant", "content": response_text},
                                    {"role": "user", "content": reflection_prompt}
                                ]
                                
                                try:
                                    # Check if OpenAI client is available for self-reflection
                                    if self._openai_client is None:
                                        # For custom LLMs, self-reflection with structured output is not supported
                                        if self.verbose:
                                            display_self_reflection(f"Agent {self.name}: Self-reflection with structured output is not supported for custom LLM providers. Skipping reflection.", console=self.console)
                                        # Return the original response without reflection
                                        self.chat_history.append({"role": "user", "content": original_prompt})
                                        self.chat_history.append({"role": "assistant", "content": response_text})
                                        if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
                                            total_time = time.time() - start_time
                                            logging.debug(f"Agent.achat completed in {total_time:.2f} seconds")
                                        return response_text
                                    
                                    reflection_response = await self._openai_client.async_client.beta.chat.completions.parse(
                                        model=self.reflect_llm if self.reflect_llm else self.llm,
                                        messages=reflection_messages,
                                        temperature=temperature,
                                        response_format=ReflectionOutput
                                    )
                                    
                                    reflection_output = reflection_response.choices[0].message.parsed
                                    
                                    if self.verbose:
                                        display_self_reflection(f"Agent {self.name} self reflection (using {self.reflect_llm if self.reflect_llm else self.llm}): reflection='{reflection_output.reflection}' satisfactory='{reflection_output.satisfactory}'", console=self.console)
                                    
                                    # Only consider satisfactory after minimum reflections
                                    if reflection_output.satisfactory == "yes" and reflection_count >= self.min_reflect - 1:
                                        if self.verbose:
                                            display_self_reflection("Agent marked the response as satisfactory after meeting minimum reflections", console=self.console)
                                        break
                                    
                                    # Check if we've hit max reflections
                                    if reflection_count >= self.max_reflect - 1:
                                        if self.verbose:
                                            display_self_reflection("Maximum reflection count reached, returning current response", console=self.console)
                                        break
                                    
                                    # Regenerate response based on reflection
                                    regenerate_messages = reflection_messages + [
                                        {"role": "assistant", "content": f"Self Reflection: {reflection_output.reflection} Satisfactory?: {reflection_output.satisfactory}"},
                                        {"role": "user", "content": "Now regenerate your response using the reflection you made"}
                                    ]
                                    
                                    new_response = await self._openai_client.async_client.chat.completions.create(
                                        model=self.llm,
                                        messages=regenerate_messages,
                                        temperature=temperature
                                    )
                                    response_text = new_response.choices[0].message.content
                                    reflection_count += 1
                                    
                                except Exception as e:
                                    if self.verbose:
                                        display_error(f"Error in parsing self-reflection json {e}. Retrying", console=self.console)
                                    logging.error("Reflection parsing failed.", exc_info=True)
                                    reflection_count += 1
                                    if reflection_count >= self.max_reflect:
                                        break
                                    continue
                        
                        if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
                            total_time = time.time() - start_time
                            logging.debug(f"Agent.achat completed in {total_time:.2f} seconds")
                        
                        # Apply guardrail validation for OpenAI client response
                        try:
                            validated_response = self._apply_guardrail_with_retry(response_text, original_prompt, temperature, tools, task_name, task_description, task_id)
                            # Execute callback after validation
                            self._execute_callback_and_display(original_prompt, validated_response, time.time() - start_time, task_name, task_description, task_id)
                            return validated_response
                        except Exception as e:
                            logging.error(f"Agent {self.name}: Guardrail validation failed for OpenAI client: {e}")
                            # Rollback chat history on guardrail failure
                            self.chat_history = self.chat_history[:chat_history_length]
                            return None
                except Exception as e:
                    display_error(f"Error in chat completion: {e}")
                    if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
                        total_time = time.time() - start_time
                        logging.debug(f"Agent.achat failed in {total_time:.2f} seconds: {str(e)}")
                    return None
        except Exception as e:
            display_error(f"Error in achat: {e}")
            if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
                total_time = time.time() - start_time
                logging.debug(f"Agent.achat failed in {total_time:.2f} seconds: {str(e)}")
            return None

    async def _achat_completion(self, response, tools, reasoning_steps=False):
        """Async version of _chat_completion method"""
        try:
            message = response.choices[0].message
            if not hasattr(message, 'tool_calls') or not message.tool_calls:
                return message.content

            results = []
            for tool_call in message.tool_calls:
                try:
                    function_name = tool_call.function.name
                    # Parse JSON arguments safely 
                    try:
                        arguments = json.loads(tool_call.function.arguments)
                    except json.JSONDecodeError as json_error:
                        logging.error(f"Failed to parse tool arguments as JSON: {json_error}")
                        arguments = {}
                    
                    # Find the matching tool
                    tool = next((t for t in tools if t.__name__ == function_name), None)
                    if not tool:
                        display_error(f"Tool {function_name} not found")
                        continue
                    
                    # Check if the tool is async
                    if asyncio.iscoroutinefunction(tool):
                        result = await tool(**arguments)
                    else:
                        # Run sync function in executor to avoid blocking
                        loop = asyncio.get_event_loop()
                        result = await loop.run_in_executor(None, lambda: tool(**arguments))
                    
                    results.append(result)
                except Exception as e:
                    display_error(f"Error executing tool {function_name}: {e}")
                    results.append(None)

            # If we have results, format them into a response
            if results:
                formatted_results = "\n".join([str(r) for r in results if r is not None])
                if formatted_results:
                    messages = [
                        {"role": "system", "content": self.system_prompt},
                        {"role": "assistant", "content": "Here are the tool results:"},
                        {"role": "user", "content": formatted_results + "\nPlease process these results and provide a final response."}
                    ]
                    try:
                        final_response = await self._openai_client.async_client.chat.completions.create(
                            model=self.llm,
                            messages=messages,
                            temperature=1.0,
                            stream=True
                        )
                        full_response_text = ""
                        reasoning_content = ""
                        chunks = []
                        start_time = time.time()
                        
                        # Process stream without display_generating since streaming is active
                        async for chunk in final_response:
                            chunks.append(chunk)
                            if chunk.choices[0].delta.content:
                                full_response_text += chunk.choices[0].delta.content
                            
                            if reasoning_steps and hasattr(chunk.choices[0].delta, "reasoning_content"):
                                rc = chunk.choices[0].delta.reasoning_content
                                if rc:
                                    reasoning_content += rc
                        
                        self.console.print()
                        
                        final_response = process_stream_chunks(chunks)
                        # Return only reasoning content if reasoning_steps is True
                        if reasoning_steps and hasattr(final_response.choices[0].message, 'reasoning_content'):
                            return final_response.choices[0].message.reasoning_content
                        return final_response.choices[0].message.content if final_response else full_response_text

                    except Exception as e:
                        display_error(f"Error in final chat completion: {e}")
                        return formatted_results
                return formatted_results
            return None
        except Exception as e:
            display_error(f"Error in _achat_completion: {e}")
            return None

    async def arun(self, prompt: str, **kwargs):
        """Async alias for astart() method"""
        return await self.astart(prompt, **kwargs)

    async def astart(self, prompt: str, **kwargs):
        """Async version of start method"""
        return await self.achat(prompt, **kwargs)

    def run(self, prompt: str, **kwargs):
        """Alias for start() method"""
        return self.start(prompt, **kwargs)
    
    def _get_planning_agent(self):
        """Lazy load PlanningAgent for planning mode."""
        if self._planning_agent is None and self.planning:
            from ..planning import PlanningAgent
            self._planning_agent = PlanningAgent(
                llm=self.llm if hasattr(self, 'llm') else (self.llm_instance.model if hasattr(self, 'llm_instance') else "gpt-4o-mini"),
                tools=self.planning_tools,
                reasoning=self.planning_reasoning,
                verbose=1 if self.verbose else 0
            )
        return self._planning_agent
    
    def _start_with_planning(self, prompt: str, **kwargs):
        """Execute with planning mode - creates plan then executes each step."""
        from rich.console import Console
        from rich.panel import Panel
        from rich.markdown import Markdown
        
        console = Console()
        
        # Step 1: Create the plan
        console.print("\n[bold blue]📋 PLANNING PHASE[/bold blue]")
        console.print("[dim]Creating implementation plan...[/dim]\n")
        
        planner = self._get_planning_agent()
        plan = planner.create_plan_sync(request=prompt, agents=[self])
        
        if not plan or not plan.steps:
            console.print("[yellow]⚠️ Planning failed, falling back to direct execution[/yellow]")
            return self.chat(prompt, **kwargs)
        
        # Display the plan
        console.print(Panel(
            Markdown(plan.to_markdown()),
            title="[bold green]Generated Plan[/bold green]",
            border_style="green"
        ))
        
        # Step 2: Execute each step
        console.print("\n[bold blue]🚀 EXECUTION PHASE[/bold blue]\n")
        
        results = []
        context = ""
        
        for i, step in enumerate(plan.steps):
            progress = (i + 1) / len(plan.steps)
            bar_length = 30
            filled = int(bar_length * progress)
            bar = "█" * filled + "░" * (bar_length - filled)
            
            console.print(f"[dim]Progress: [{bar}] {progress * 100:.0f}%[/dim]")
            console.print(f"\n[bold]📌 Step {i + 1}/{len(plan.steps)}:[/bold] {step.description[:60]}...")
            
            # Build prompt with context from previous steps
            step_prompt = step.description
            if context:
                step_prompt = f"{step.description}\n\nContext from previous steps:\n{context}"
            
            # Execute the step
            result = self.chat(step_prompt, **kwargs)
            results.append(result)
            
            # Update context for next step
            context += f"\n\nStep {i + 1} result: {result[:500] if result else 'No result'}"
            
            console.print(f"   [green]✅ Completed[/green]")
        
        console.print(f"\n[bold green]🎉 EXECUTION COMPLETE[/bold green]")
        console.print(f"[dim]Progress: [{'█' * bar_length}] 100%[/dim]")
        console.print(f"Completed {len(plan.steps)}/{len(plan.steps)} steps!\n")
        
        # Return the final result
        return results[-1] if results else None

    def start(self, prompt: str, **kwargs):
        """Start the agent with a prompt. This is a convenience method that wraps chat()."""
        # Check if planning mode is enabled
        if self.planning:
            return self._start_with_planning(prompt, **kwargs)
        
        # Check if streaming is enabled (either from kwargs or agent's stream attribute)
        stream_enabled = kwargs.get('stream', getattr(self, 'stream', False))
        
        if stream_enabled:
            # Return a generator for streaming response
            return self._start_stream(prompt, **kwargs)
        else:
            # Return regular chat response for backward compatibility
            # Explicitly pass the resolved stream parameter to avoid chat() method default
            kwargs['stream'] = stream_enabled
            return self.chat(prompt, **kwargs)

    def _start_stream(self, prompt: str, **kwargs) -> Generator[str, None, None]:
        """Stream generator for real-time response chunks."""
        try:
            # Reset the final display flag for each new conversation
            self._final_display_shown = False
            
            # Temporarily disable verbose mode to prevent console output conflicts during streaming
            original_verbose = self.verbose
            self.verbose = False
            
            # For custom LLM path, use the new get_response_stream generator
            if self._using_custom_llm:
                # Handle knowledge search
                actual_prompt = prompt
                if self._knowledge_sources and not self._knowledge_processed:
                    self._ensure_knowledge_processed()
                
                if self.knowledge:
                    search_results = self.knowledge.search(prompt, agent_id=self.agent_id)
                    if search_results:
                        if isinstance(search_results, dict) and 'results' in search_results:
                            knowledge_content = "\n".join([result['memory'] for result in search_results['results']])
                        else:
                            knowledge_content = "\n".join(search_results)
                        actual_prompt = f"{prompt}\n\nKnowledge: {knowledge_content}"
                
                # Handle tools properly
                tools = kwargs.get('tools', self.tools)
                if tools is None or (isinstance(tools, list) and len(tools) == 0):
                    tool_param = self.tools
                else:
                    tool_param = tools
                
                # Convert MCP tools if needed
                if tool_param is not None:
                    from ..mcp.mcp import MCP
                    if isinstance(tool_param, MCP) and hasattr(tool_param, 'to_openai_tool'):
                        openai_tool = tool_param.to_openai_tool()
                        if openai_tool:
                            if isinstance(openai_tool, list):
                                tool_param = openai_tool
                            else:
                                tool_param = [openai_tool]
                
                # Store chat history length for potential rollback
                chat_history_length = len(self.chat_history)
                
                # Normalize prompt content for chat history
                normalized_content = actual_prompt
                if isinstance(actual_prompt, list):
                    normalized_content = next((item["text"] for item in actual_prompt if item.get("type") == "text"), "")
                
                # Prevent duplicate messages in chat history
                if not (self.chat_history and 
                        self.chat_history[-1].get("role") == "user" and 
                        self.chat_history[-1].get("content") == normalized_content):
                    self.chat_history.append({"role": "user", "content": normalized_content})
                
                try:
                    # Use the new streaming generator from LLM class
                    response_content = ""
                    for chunk in self.llm_instance.get_response_stream(
                        prompt=actual_prompt,
                        system_prompt=self._build_system_prompt(tool_param),
                        chat_history=self.chat_history,
                        temperature=kwargs.get('temperature', 1.0),
                        tools=tool_param,
                        output_json=kwargs.get('output_json'),
                        output_pydantic=kwargs.get('output_pydantic'),
                        verbose=False,  # Keep verbose false for streaming
                        markdown=self.markdown,
                        agent_name=self.name,
                        agent_role=self.role,
                        agent_tools=[t.__name__ if hasattr(t, '__name__') else str(t) for t in (tool_param or [])],
                        task_name=kwargs.get('task_name'),
                        task_description=kwargs.get('task_description'),
                        task_id=kwargs.get('task_id'),
                        execute_tool_fn=self.execute_tool
                    ):
                        response_content += chunk
                        yield chunk
                    
                    # Add complete response to chat history
                    if response_content:
                        self.chat_history.append({"role": "assistant", "content": response_content})
                        
                except Exception as e:
                    # Rollback chat history on error
                    self.chat_history = self.chat_history[:chat_history_length]
                    logging.error(f"Custom LLM streaming error: {e}")
                    raise
                    
            else:
                # For OpenAI-style models, implement proper streaming without display
                # Handle knowledge search
                actual_prompt = prompt
                if self._knowledge_sources and not self._knowledge_processed:
                    self._ensure_knowledge_processed()
                
                if self.knowledge:
                    search_results = self.knowledge.search(prompt, agent_id=self.agent_id)
                    if search_results:
                        if isinstance(search_results, dict) and 'results' in search_results:
                            knowledge_content = "\n".join([result['memory'] for result in search_results['results']])
                        else:
                            knowledge_content = "\n".join(search_results)
                        actual_prompt = f"{prompt}\n\nKnowledge: {knowledge_content}"
                
                # Handle tools properly
                tools = kwargs.get('tools', self.tools)
                if tools is None or (isinstance(tools, list) and len(tools) == 0):
                    tool_param = self.tools
                else:
                    tool_param = tools
                
                # Build messages using the helper method
                messages, original_prompt = self._build_messages(actual_prompt, kwargs.get('temperature', 1.0), 
                                                               kwargs.get('output_json'), kwargs.get('output_pydantic'))
                
                # Store chat history length for potential rollback
                chat_history_length = len(self.chat_history)
                
                # Normalize original_prompt for consistent chat history storage
                normalized_content = original_prompt
                if isinstance(original_prompt, list):
                    normalized_content = next((item["text"] for item in original_prompt if item.get("type") == "text"), "")
                
                # Prevent duplicate messages in chat history
                if not (self.chat_history and 
                        self.chat_history[-1].get("role") == "user" and 
                        self.chat_history[-1].get("content") == normalized_content):
                    self.chat_history.append({"role": "user", "content": normalized_content})
                
                try:
                    # Check if OpenAI client is available
                    if self._openai_client is None:
                        raise ValueError("OpenAI client is not initialized. Please provide OPENAI_API_KEY or use a custom LLM provider.")
                    
                    # Format tools for OpenAI
                    formatted_tools = self._format_tools_for_completion(tool_param)
                    
                    # Create streaming completion directly without display function
                    completion_args = {
                        "model": self.llm,
                        "messages": messages,
                        "temperature": kwargs.get('temperature', 1.0),
                        "stream": True
                    }
                    if formatted_tools:
                        completion_args["tools"] = formatted_tools
                    
                    completion = self._openai_client.sync_client.chat.completions.create(**completion_args)
                    
                    # Stream the response chunks without display
                    response_text = ""
                    tool_calls_data = []
                    
                    for chunk in completion:
                        delta = chunk.choices[0].delta
                        
                        # Handle text content
                        if delta.content is not None:
                            chunk_content = delta.content
                            response_text += chunk_content
                            yield chunk_content
                        
                        # Handle tool calls (accumulate but don't yield as chunks)
                        if hasattr(delta, 'tool_calls') and delta.tool_calls:
                            for tool_call_delta in delta.tool_calls:
                                # Extend tool_calls_data list to accommodate the tool call index
                                while len(tool_calls_data) <= tool_call_delta.index:
                                    tool_calls_data.append({'id': '', 'function': {'name': '', 'arguments': ''}})
                                
                                # Accumulate tool call data
                                if tool_call_delta.id:
                                    tool_calls_data[tool_call_delta.index]['id'] = tool_call_delta.id
                                if tool_call_delta.function.name:
                                    tool_calls_data[tool_call_delta.index]['function']['name'] = tool_call_delta.function.name
                                if tool_call_delta.function.arguments:
                                    tool_calls_data[tool_call_delta.index]['function']['arguments'] += tool_call_delta.function.arguments
                    
                    # Handle any tool calls that were accumulated
                    if tool_calls_data:
                        # Add assistant message with tool calls to chat history
                        assistant_message = {"role": "assistant", "content": response_text}
                        if tool_calls_data:
                            assistant_message["tool_calls"] = [
                                {
                                    "id": tc['id'],
                                    "type": "function", 
                                    "function": tc['function']
                                } for tc in tool_calls_data if tc['id']
                            ]
                        self.chat_history.append(assistant_message)
                        
                        # Execute tool calls and add results to chat history
                        for tool_call in tool_calls_data:
                            if tool_call['id'] and tool_call['function']['name']:
                                try:
                                    # Parse JSON arguments safely 
                                    try:
                                        parsed_args = json.loads(tool_call['function']['arguments']) if tool_call['function']['arguments'] else {}
                                    except json.JSONDecodeError as json_error:
                                        logging.error(f"Failed to parse tool arguments as JSON: {json_error}")
                                        parsed_args = {}
                                    
                                    tool_result = self.execute_tool(
                                        tool_call['function']['name'], 
                                        parsed_args
                                    )
                                    # Add tool result to chat history
                                    self.chat_history.append({
                                        "role": "tool",
                                        "tool_call_id": tool_call['id'],
                                        "content": str(tool_result)
                                    })
                                except Exception as tool_error:
                                    logging.error(f"Tool execution error in streaming: {tool_error}")
                                    # Add error result to chat history
                                    self.chat_history.append({
                                        "role": "tool", 
                                        "tool_call_id": tool_call['id'],
                                        "content": f"Error: {str(tool_error)}"
                                    })
                    else:
                        # Add complete response to chat history (text-only response)
                        if response_text:
                            self.chat_history.append({"role": "assistant", "content": response_text})
                        
                except Exception as e:
                    # Rollback chat history on error
                    self.chat_history = self.chat_history[:chat_history_length]
                    logging.error(f"OpenAI streaming error: {e}")
                    # Fall back to simulated streaming
                    response = self.chat(prompt, **kwargs)
                    if response:
                        words = str(response).split()
                        chunk_size = max(1, len(words) // 20)
                        for i in range(0, len(words), chunk_size):
                            chunk_words = words[i:i + chunk_size]
                            chunk = ' '.join(chunk_words)
                            if i + chunk_size < len(words):
                                chunk += ' '
                            yield chunk
            
            # Restore original verbose mode
            self.verbose = original_verbose
                    
        except Exception as e:
            # Restore verbose mode on any error
            self.verbose = original_verbose
            # Graceful fallback to non-streaming if streaming fails
            logging.warning(f"Streaming failed, falling back to regular response: {e}")
            response = self.chat(prompt, **kwargs)
            if response:
                yield response

    def execute(self, task, context=None):
        """Execute a task synchronously - backward compatibility method"""
        if hasattr(task, 'description'):
            prompt = task.description
        elif isinstance(task, str):
            prompt = task
        else:
            prompt = str(task)
        return self.chat(prompt)

    async def aexecute(self, task, context=None):
        """Execute a task asynchronously - backward compatibility method"""
        if hasattr(task, 'description'):
            prompt = task.description
        elif isinstance(task, str):
            prompt = task
        else:
            prompt = str(task)
        # Extract task info if available
        task_name = getattr(task, 'name', None)
        task_description = getattr(task, 'description', None)
        task_id = getattr(task, 'id', None)
        return await self.achat(prompt, task_name=task_name, task_description=task_description, task_id=task_id)

    async def execute_tool_async(self, function_name: str, arguments: Dict[str, Any]) -> Any:
        """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:
                if (callable(tool) and getattr(tool, '__name__', '') == function_name):
                    func = tool
                    break
            
            if func is None:
                logging.error(f"Function {function_name} not found in tools")
                return {"error": f"Function {function_name} not found in tools"}

            try:
                if inspect.iscoroutinefunction(func):
                    logging.debug(f"Executing async function: {function_name}")
                    result = await func(**arguments)
                else:
                    logging.debug(f"Executing sync function in executor: {function_name}")
                    loop = asyncio.get_event_loop()
                    result = await loop.run_in_executor(None, lambda: func(**arguments))
                
                # Ensure result is JSON serializable
                logging.debug(f"Raw result from tool: {result}")
                if result is None:
                    return {"result": None}
                try:
                    json.dumps(result)  # Test serialization
                    return result
                except TypeError:
                    logging.warning(f"Result not JSON serializable, converting to string: {result}")
                    return {"result": str(result)}

            except Exception as e:
                logging.error(f"Error executing {function_name}: {str(e)}", exc_info=True)
                return {"error": f"Error executing {function_name}: {str(e)}"}

        except Exception as e:
            logging.error(f"Error in execute_tool_async: {str(e)}", exc_info=True)
            return {"error": f"Error in execute_tool_async: {str(e)}"}

    def launch(self, path: str = '/', port: int = 8000, host: str = '0.0.0.0', debug: bool = False, protocol: str = "http"):
        """
        Launch the agent as an HTTP API endpoint or an MCP server.
        
        Args:
            path: API endpoint path (default: '/') for HTTP, or base path for MCP.
            port: Server port (default: 8000)
            host: Server host (default: '0.0.0.0')
            debug: Enable debug mode for uvicorn (default: False)
            protocol: "http" to launch as FastAPI, "mcp" to launch as MCP server.
            
        Returns:
            None
        """
        if protocol == "http":
            global _server_started, _registered_agents, _shared_apps
            
            # Try to import FastAPI dependencies - lazy loading
            try:
                import uvicorn
                from fastapi import FastAPI, HTTPException, Request
                from fastapi.responses import JSONResponse
                from pydantic import BaseModel
                import threading
                import time
                import asyncio
                
                # Define the request model here since we need pydantic
                class AgentQuery(BaseModel):
                    query: str
                    
            except ImportError as e:
                # Check which specific module is missing
                missing_module = str(e).split("No module named '")[-1].rstrip("'")
                display_error(f"Missing dependency: {missing_module}. Required for launch() method with HTTP mode.")
                logging.error(f"Missing dependency: {missing_module}. Required for launch() method with HTTP mode.")
                print(f"\nTo add API capabilities, install the required dependencies:")
                print(f"pip install {missing_module}")
                print("\nOr install all API dependencies with:")
                print("pip install 'praisonaiagents[api]'")
                return None
                
            # Initialize port-specific collections if needed
            if port not in _registered_agents:
                _registered_agents[port] = {}
                
            # Initialize shared FastAPI app if not already created for this port
            if _shared_apps.get(port) is None:
                _shared_apps[port] = FastAPI(
                    title=f"PraisonAI Agents API (Port {port})",
                    description="API for interacting with PraisonAI Agents"
                )
                
                # Add a root endpoint with a welcome message
                @_shared_apps[port].get("/")
                async def root():
                    return {
                        "message": f"Welcome to PraisonAI Agents API on port {port}. See /docs for usage.",
                        "endpoints": list(_registered_agents[port].keys())
                    }
                
                # Add healthcheck endpoint
                @_shared_apps[port].get("/health")
                async def healthcheck():
                    return {
                        "status": "ok", 
                        "endpoints": list(_registered_agents[port].keys())
                    }
            
            # Normalize path to ensure it starts with /
            if not path.startswith('/'):
                path = f'/{path}'
                
            # Check if path is already registered for this port
            if path in _registered_agents[port]:
                logging.warning(f"Path '{path}' is already registered on port {port}. Please use a different path.")
                print(f"⚠️ Warning: Path '{path}' is already registered on port {port}.")
                # Use a modified path to avoid conflicts
                original_path = path
                path = f"{path}_{self.agent_id[:6]}"
                logging.warning(f"Using '{path}' instead of '{original_path}'")
                print(f"🔄 Using '{path}' instead")
            
            # Register the agent to this path
            _registered_agents[port][path] = self.agent_id
            
            # Define the endpoint handler
            @_shared_apps[port].post(path)
            async def handle_agent_query(request: Request, query_data: Optional[AgentQuery] = None):
                # Handle both direct JSON with query field and form data
                if query_data is None:
                    try:
                        request_data = await request.json()
                        if "query" not in request_data:
                            raise HTTPException(status_code=400, detail="Missing 'query' field in request")
                        query = request_data["query"]
                    except:
                        # Fallback to form data or query params
                        form_data = await request.form()
                        if "query" in form_data:
                            query = form_data["query"]
                        else:
                            raise HTTPException(status_code=400, detail="Missing 'query' field in request")
                else:
                    query = query_data.query
                    
                try:
                    # Use async version if available, otherwise use sync version
                    if asyncio.iscoroutinefunction(self.chat):
                        response = await self.achat(query, task_name=None, task_description=None, task_id=None)
                    else:
                        # Run sync function in a thread to avoid blocking
                        loop = asyncio.get_event_loop()
                        response = await loop.run_in_executor(None, lambda p=query: self.chat(p))
                    
                    return {"response": response}
                except Exception as e:
                    logging.error(f"Error processing query: {str(e)}", exc_info=True)
                    return JSONResponse(
                        status_code=500,
                        content={"error": f"Error processing query: {str(e)}"}
                    )
            
            print(f"🚀 Agent '{self.name}' available at http://{host}:{port}")
            
            # Start the server if it's not already running for this port
            if not _server_started.get(port, False):
                # Mark the server as started first to prevent duplicate starts
                _server_started[port] = True
                
                # Start the server in a separate thread
                def run_server():
                    try:
                        print(f"✅ FastAPI server started at http://{host}:{port}")
                        print(f"📚 API documentation available at http://{host}:{port}/docs")
                        print(f"🔌 Available endpoints: {', '.join(list(_registered_agents[port].keys()))}")
                        uvicorn.run(_shared_apps[port], host=host, port=port, log_level="debug" if debug else "info")
                    except Exception as e:
                        logging.error(f"Error starting server: {str(e)}", exc_info=True)
                        print(f"❌ Error starting server: {str(e)}")
                
                # Run server in a background thread
                server_thread = threading.Thread(target=run_server, daemon=True)
                server_thread.start()
                
                # Wait for a moment to allow the server to start and register endpoints
                time.sleep(0.5)
            else:
                # If server is already running, wait a moment to make sure the endpoint is registered
                time.sleep(0.1)
                print(f"🔌 Available endpoints on port {port}: {', '.join(list(_registered_agents[port].keys()))}")
            
            # Get the stack frame to check if this is the last launch() call in the script
            import inspect
            stack = inspect.stack()
            
            # If this is called from a Python script (not interactive), try to detect if it's the last launch call
            if len(stack) > 1 and stack[1].filename.endswith('.py'):
                caller_frame = stack[1]
                caller_line = caller_frame.lineno
                
                try:
                    # Read the file to check if there are more launch calls after this one
                    with open(caller_frame.filename, 'r') as f:
                        lines = f.readlines()
                    
                    # Check if there are more launch() calls after the current line
                    has_more_launches = False
                    for line_content in lines[caller_line:]: # renamed line to line_content
                        if '.launch(' in line_content and not line_content.strip().startswith('#'):
                            has_more_launches = True
                            break
                    
                    # If this is the last launch call, block the main thread
                    if not has_more_launches:
                        try:
                            print("\nAll agents registered for HTTP mode. Press Ctrl+C to stop the servers.")
                            while True:
                                time.sleep(1)
                        except KeyboardInterrupt:
                            print("\nServers stopped")
                except Exception as e:
                    # If something goes wrong with detection, block anyway to be safe
                    logging.error(f"Error in launch detection: {e}")
                    try:
                        print("\nKeeping HTTP servers alive. Press Ctrl+C to stop.")
                        while True:
                            time.sleep(1)
                    except KeyboardInterrupt:
                        print("\nServers stopped")
            return None
            
        elif protocol == "mcp":
            try:
                import uvicorn
                from mcp.server.fastmcp import FastMCP
                from mcp.server.sse import SseServerTransport
                from starlette.applications import Starlette
                from starlette.requests import Request
                from starlette.routing import Mount, Route
                from mcp.server import Server as MCPServer # Alias to avoid conflict
                import threading
                import time
                import inspect
                import asyncio  # Import asyncio in the MCP scope
                # logging is already imported at the module level
                
            except ImportError as e:
                missing_module = str(e).split("No module named '")[-1].rstrip("'")
                display_error(f"Missing dependency: {missing_module}. Required for launch() method with MCP mode.")
                logging.error(f"Missing dependency: {missing_module}. Required for launch() method with MCP mode.")
                print(f"\nTo add MCP capabilities, install the required dependencies:")
                print(f"pip install {missing_module} mcp praison-mcp starlette uvicorn") # Added mcp, praison-mcp, starlette, uvicorn
                print("\nOr install all MCP dependencies with relevant packages.")
                return None

            mcp_server_instance_name = f"{self.name}_mcp_server" if self.name else "agent_mcp_server"
            mcp = FastMCP(mcp_server_instance_name)

            # Determine the MCP tool name based on self.name
            actual_mcp_tool_name = f"execute_{self.name.lower().replace(' ', '_').replace('-', '_')}_task" if self.name \
                else "execute_task"

            @mcp.tool(name=actual_mcp_tool_name)
            async def execute_agent_task(prompt: str) -> str:
                """Executes the agent's primary task with the given prompt."""
                logging.info(f"MCP tool '{actual_mcp_tool_name}' called with prompt: {prompt}")
                try:
                    # Ensure self.achat is used as it's the async version and pass its tools
                    if hasattr(self, 'achat') and asyncio.iscoroutinefunction(self.achat):
                        response = await self.achat(prompt, tools=self.tools, task_name=None, task_description=None, task_id=None)
                    elif hasattr(self, 'chat'): # Fallback for synchronous chat
                        loop = asyncio.get_event_loop()
                        response = await loop.run_in_executor(None, lambda p=prompt: self.chat(p, tools=self.tools))
                    else:
                        logging.error(f"Agent {self.name} has no suitable chat or achat method for MCP tool.")
                        return f"Error: Agent {self.name} misconfigured for MCP."
                    return response if response is not None else "Agent returned no response."
                except Exception as e:
                    logging.error(f"Error in MCP tool '{actual_mcp_tool_name}': {e}", exc_info=True)
                    return f"Error executing task: {str(e)}"

            # Normalize base_path for MCP routes
            base_path = path.rstrip('/')
            sse_path = f"{base_path}/sse"
            messages_path_prefix = f"{base_path}/messages" # Prefix for message posting
            
            # Ensure messages_path ends with a slash for Mount
            if not messages_path_prefix.endswith('/'):
                messages_path_prefix += '/'


            sse_transport = SseServerTransport(messages_path_prefix) # Pass the full prefix

            async def handle_sse_connection(request: Request) -> None:
                logging.debug(f"SSE connection request received from {request.client} for path {request.url.path}")
                async with sse_transport.connect_sse(
                        request.scope,
                        request.receive,
                        request._send,  # noqa: SLF001
                ) as (read_stream, write_stream):
                    await mcp._mcp_server.run( # Use the underlying server from FastMCP
                        read_stream,
                        write_stream,
                        mcp._mcp_server.create_initialization_options(),
                    )
            
            starlette_app = Starlette(
                debug=debug,
                routes=[
                    Route(sse_path, endpoint=handle_sse_connection),
                    Mount(messages_path_prefix, app=sse_transport.handle_post_message),
                ],
            )

            print(f"🚀 Agent '{self.name}' MCP server starting on http://{host}:{port}")
            print(f"📡 MCP SSE endpoint available at {sse_path}")
            print(f"📢 MCP messages post to {messages_path_prefix}")
            # Instead of trying to extract tool names, hardcode the known tool name
            tool_names = [actual_mcp_tool_name]  # Use the determined dynamic tool name
            print(f"🛠️ Available MCP tools: {', '.join(tool_names)}")

            # Uvicorn server running logic (similar to HTTP mode but standalone for MCP)
            def run_mcp_server():
                try:
                    uvicorn.run(starlette_app, host=host, port=port, log_level="debug" if debug else "info")
                except Exception as e:
                    logging.error(f"Error starting MCP server: {str(e)}", exc_info=True)
                    print(f"❌ Error starting MCP server: {str(e)}")

            server_thread = threading.Thread(target=run_mcp_server, daemon=True)
            server_thread.start()
            time.sleep(0.5) # Allow server to start

            # Blocking logic for MCP mode
            import inspect # Already imported but good for clarity
            stack = inspect.stack()
            if len(stack) > 1 and stack[1].filename.endswith('.py'):
                caller_frame = stack[1]
                caller_line = caller_frame.lineno
                try:
                    with open(caller_frame.filename, 'r') as f:
                        lines = f.readlines()
                    has_more_launches = False
                    for line_content in lines[caller_line:]: # renamed line to line_content
                        if '.launch(' in line_content and not line_content.strip().startswith('#'):
                            has_more_launches = True
                            break
                    if not has_more_launches:
                        try:
                            print("\nAgent MCP server running. Press Ctrl+C to stop.")
                            while True:
                                time.sleep(1)
                        except KeyboardInterrupt:
                            print("\nMCP Server stopped")
                except Exception as e:
                    logging.error(f"Error in MCP launch detection: {e}")
                    try:
                        print("\nKeeping MCP server alive. Press Ctrl+C to stop.")
                        while True:
                            time.sleep(1)
                    except KeyboardInterrupt:
                        print("\nMCP Server stopped")
            return None
        else:
            display_error(f"Invalid protocol: {protocol}. Choose 'http' or 'mcp'.")
            return None 